diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..ed4671ebcae9651fcaf9c2c12cf57fae0e999499 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +assets/dataset_analysis.png filter=lfs diff=lfs merge=lfs -text +assets/dataset_overview.png filter=lfs diff=lfs merge=lfs -text +assets/hardware_setup.png filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..c359d1decacfd23401f554bdb2dc657e14a48567 --- /dev/null +++ b/.gitignore @@ -0,0 +1,246 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class +wandb/ + +# C extensions + + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +*.swp +*.swo +gear_working_dir/ +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ +/groot/control/wbc_checkpoints/ +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +.vscode/ +.cursor/ +.cursor* +.claude/ +CLAUDE.md + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# UV +# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +#uv.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/latest/usage/project/#working-with-version-control +.pdm.toml +.pdm-python +.pdm-build/ + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ +# IDE +.idea/ +.vscode/ + +# log +outputs/ +logs/ +*logs_rl* +!external_dependencies/OpenHomie/HomieRL/legged_gym/logs/exported/policies/*.onnx + +# Ruff stuff: +.ruff_cache/ + +# PyPI configuration file +.pypirc + +outputs/ +osmo/ +source/ +groot/rl/data/ + +.DS_Store +sim_dependencies/ +.venv/* + + +# Mujoco +MUJOCO_LOG.TXT +playground/ + +*.code-workspace +batch_*.sh +logs_*/ +runs/ +out/ +recordings/ +pyrightconfig.json +*.npz +*.nbize.py + +# Git worktrees +/wt + +# Gear working directory +gear_working_dir/ + +episode_data/ + +# third party packages +third_parties/ + +# redis related +*.rdb +download_model.py +.dockerignore + +# Sysid files +plots/ +sysid_data/ +isaac_data/ + +# External dependencies +external_dependencies/ws_lidar_slam_ros2/build +external_dependencies/ws_lidar_slam_ros2/install +external_dependencies/ws_lidar_slam_ros2/log +external_dependencies/ws_slam/ +external_dependencies/Livox-SDK2/ +external_dependencies/mfm +# external_dependencies/genmo +external_dependencies/whole_body_tracking +inputs/ +outputs/ + +groot/dexmg/grootrobosuite/docs/images/ +external_dependencies/mfm/ +dump.rdb + +# Open loop eval plots +offline_open_loop_eval_plots/ + +# yam eval data +groot/control/envs/yam/data/eval/* +/models +/data + +checkpoints/ +video_pred_output*/ \ No newline at end of file diff --git a/COPYRIGHT b/COPYRIGHT new file mode 100644 index 0000000000000000000000000000000000000000..dd908a2529902a5e576649f2a9064c34a8b048d4 --- /dev/null +++ b/COPYRIGHT @@ -0,0 +1,4 @@ +Copyright (c) 2025 NVIDIA Corporation. All rights reserved. + +Licensed under the Apache License, Version 2.0. +See LICENSE for the full license text. diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..ad169f672b29881f9bf740ce4d91aef8e9021fdb --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 The Regents of the University of California + +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/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..256c08dc85825402d9d3e95b6aa1fa958a7dfa7e --- /dev/null +++ b/README.md @@ -0,0 +1,194 @@ +--- +license: mit +language: +- en +task_categories: +- robotics +tags: +- LeRobot +- robotics +- manipulation +- tactile +- bimanual +- dexterous-manipulation +pretty_name: T-Rex Dataset +size_categories: +- 1M + +
+ One episode from each of 20 motor primitives (head-camera view, cropped to the workspace), each with a different object. +

+ +

+ +
+ Teleoperation setup: Manus gloves + VIVE trackers drive the bimanual Dexmate Vega-1 with two Sharpa Wave hands; observations come from a head-mounted ZED X Mini camera and two wide-view ZED X One S wrist cameras. +

+ +## At a glance + +| | | +|---|---| +| Robot | Dexmate Vega-1 dual-arm (7 actuated joints per arm) + 2× Sharpa Wave dexterous hands (5 fingertip tactile sensors each) | +| Modalities | 3 RGB cameras · 10 raw tactile + 10 deformation tactile videos · 6-axis fingertip wrenches · joint states/targets | +| Episodes | 5,464 | +| Frames | 5,473,459 (~50 hours @ 30 fps) | +| Tasks | 5,370 language-annotated trajectories · 22 motor primitives · 207 objects | +| Format | LeRobotDataset v3.0 (`codebase_version: v3.0`) | + +## Composition + +

+ +
+ Object categories, episodes per motor primitive, and per-object episode counts. +

+ +## Collection setup + +The full teleoperation/hardware stack used to collect the dataset is open-sourced in +[`hardware_code/`](https://github.com/ZhuoyangLiu2005/T-Rex/tree/main/hardware_code) in the T-Rex repo. + +- **Robot.** The Dexmate Vega-1 is a dual-arm mobile robot with 7 actuated joints per arm, here + equipped with two Sharpa Wave dexterous hands. During collection the wheels, torso, and head + joints are fixed; only the 14 arm joints and the two hands are actuated. +- **Cameras.** A head-mounted ZED X Mini stereo camera (its left monocular RGB stream is recorded) + plus two wide-view ZED X One S monocular RGB cameras mounted on the wrists, posed so the head + camera covers the full reachable workspace while the wrist cameras keep the fingers visible + without significant palm occlusion. All three streams are recorded at 640×360. +- **Tactile.** Each hand carries five fingertip tactile sensors. Per sensor, the dataset stores the + raw sensor image and the estimated deformation map (both as video), and the estimated 6-axis net + wrench (`observation.tactile_force`). +- **Teleoperation.** Manus gloves capture fingertip positions relative to the hand base, retargeted + to the Sharpa Wave hands with the manufacturer's differential-inverse-kinematics package + (Pinocchio + CasADi). Two VIVE trackers provide SE(3) wrist poses, converted to arm joint + commands via differential inverse kinematics ([Pink](https://github.com/stephane-caron/pink)), + low-pass filtered, and tracked by the manufacturer's low-level cascade PID controller. A 30 Hz + high-level thread records observations and joint-space targets while asynchronously updating a + 300 Hz low-level control thread — the dataset's 30 fps matches the high-level loop, and `action` + holds its 30 Hz joint-space targets. + +## Per-frame features + +| feature | shape | description | +|---|---|---| +| `observation.state` | `(58,)` | joint **positions**, laid out `[L_arm 7 \| L_hand 22 \| R_arm 7 \| R_hand 22]` | +| `action` | `(58,)` | **target** joint positions (same layout) | +| `observation.tactile_force` | `(60,)` | per-fingertip 6-axis wrench: `(left, right) × (thumb…pinky) × (Fx, Fy, Fz, Mx, My, Mz)` | +| `observation.images.{head_left, left_wrist, right_wrist}` | `360×640×3` | scene + wrist RGB cameras | +| `observation.images.tactile_{left,right}_raw_{finger}` | `240×320` (grayscale) | raw tactile sensor images (10 = 2 hands × 5 fingers) | +| `observation.images.tactile_{left,right}_deform_{finger}` | `240×240` (grayscale) | tactile deformation fields (10) | + +Finger order is `thumb, index, middle, ring, pinky`. Full per-dimension joint names are in +`meta/info.json` (`features[*].names`). + +## Per-episode metadata + +`meta/episodes/*.parquet` carries language and task labels per episode, in addition to the standard +LeRobot fields (`episode_index`, `tasks`, `length`, per-feature stats, video pointers): + +| field | description | +|---|---| +| `caption` | human-verified natural-language instruction (5,370 unique) | +| `motor_primitive` | one of 22 primitives (`reach`, `lift_and_place`, …) | +| `object` | canonical object name (207 unique) | +| `target` | canonical target/receptacle (only set for `lift_and_place`; null otherwise) | + +## Tactile video encoding (read before decoding) + +The raw and deformation tactile videos are stored **losslessly** (`libx264 -qp 0`) because their +pixel values are physically meaningful (raw sensor images and deformation maps). They are grayscale — +the signal lives entirely in the **luma (Y) plane** — and use **full-range** `yuvj420p`, so values +span the full `0–255` with no range conversion. Decode the luma plane to recover the original `uint8` +images exactly, e.g. `frame.to_ndarray(format="gray")` in PyAV. + +> **No web thumbnails for tactile.** Lossless H.264 forces the *High 4:4:4 Predictive* profile, +> which most browsers and the Hugging Face preview cannot decode — so the tactile videos do not show +> thumbnails on the dataset page. This is expected; decode them locally (ffmpeg / PyAV / torchcodec). +> The RGB cameras use standard `yuv420p` (limited range, BT.709) and preview normally. + +## Usage + +### With LeRobot + +```python +from lerobot.datasets.lerobot_dataset import LeRobotDataset + +ds = LeRobotDataset("zekaiwang/trex_dataset") +frame = ds[0] # dict of tensors: observation.state, action, observation.tactile_force, images... +``` + +### Stream individual episodes (no full download) + +The [**T-Rex Quick Start**](https://github.com/ZhuoyangLiu2005/T-Rex/tree/main/dataset_quickstart) +companion repo browses, inspects, and replays single episodes without downloading the full dataset, +and includes a notebook you can +[open directly in Colab](https://colab.research.google.com/github/ZhuoyangLiu2005/T-Rex/blob/main/dataset_quickstart/quickstart.ipynb). + +## Dataset viewer + +The Hugging Face table viewer is manually configured (the `configs` block above) with two views: + +- **`frames`** (default) — the raw per-frame `observation.state`, `action`, and + `observation.tactile_force` arrays (plus index columns). This is the actual data, and what + `datasets.load_dataset("zekaiwang/trex_dataset")` returns by default. +- **`episodes`** — one row per episode with its language `caption`, `motor_primitive`, `object`, and + `target`: a readable table of contents. Switch to it with the config dropdown, or load it with + `datasets.load_dataset("zekaiwang/trex_dataset", "episodes")`. + +Videos are stored as separate `.mp4` files referenced by timestamp pointers, so they are not shown in +the table viewer — load them via LeRobot or the quick-start tools. For interactive video playback, use +the LeRobot dataset visualizer. + +## Layout + +``` +data/chunk-000/file-*.parquet per-frame state / action / tactile_force (+ index columns) +videos//chunk-*/file-*.mp4 23 video streams (3 RGB + 20 tactile) +episodes_preview.parquet curated per-episode labels for the web viewer (see "Dataset viewer") +meta/info.json features, shapes, fps, codebase_version +meta/episodes/*.parquet per-episode metadata + stats + video pointers +meta/tasks.parquet task (caption) table +meta/stats.json global feature statistics +``` + +## Citation + +If you find the T-Rex Dataset useful, please cite: + +```bibtex +@misc{trex2026, + title={T-Rex: Tactile-Reactive Dexterous Manipulation}, + author={Dantong Niu and Zhuoyang Liu and Zekai Wang and Boning Shao and Zhao-Heng Yin and Anirudh Pai and Yuvan Sharma and Stefano Saravalle and Ruijie Zheng and Jing Wang and Ryan Punamiya and Mengda Xu and Yuqi Xie and Yunfan Jiang and Letian Fu and Konstantinos Kallidromitis and Matteo Gioia and Junyi Zhang and Jiaxin Ge and Haiwen Feng and Fabio Galasso and Wei Zhan and David M. Chan and Yutong Bai and Roei Herzig and Jiahui Lei and Fei-Fei Li and Ken Goldberg and Jitendra Malik and Pieter Abbeel and Yuke Zhu and Danfei Xu and Jim Fan and Trevor Darrell}, + year={2026}, + eprint={2606.17055}, + archivePrefix={arXiv}, + primaryClass={cs.RO}, + url={https://arxiv.org/abs/2606.17055}, +} +``` + +## License + +Released under the **MIT License** © 2026 The Regents of the University of California. See `LICENSE`. diff --git a/assets/dataset_analysis.png b/assets/dataset_analysis.png new file mode 100644 index 0000000000000000000000000000000000000000..dd8c29586336bfdec29a0567f549dd6976eabc17 --- /dev/null +++ b/assets/dataset_analysis.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:725cd8b39919bd84df67a0e487d18995910ee328aafe3eed10f3b058c8933c27 +size 360788 diff --git a/assets/dataset_overview.png b/assets/dataset_overview.png new file mode 100644 index 0000000000000000000000000000000000000000..4c048a674d8205bb51d68ff5833fadce1e75cbac --- /dev/null +++ b/assets/dataset_overview.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9b4f28523cbdad67c4b7779ecbfb3394c9f21db9ac6b42e9579a077b27bbdd06 +size 3806789 diff --git a/assets/hardware_setup.png b/assets/hardware_setup.png new file mode 100644 index 0000000000000000000000000000000000000000..2c0d38ef0a9582ba349147ebdd6ff0170f80590a --- /dev/null +++ b/assets/hardware_setup.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:06579da8a21d7a860c70a5bebc3350e3fac225f3e358f23ec8e2be0f3da4b79a +size 1323388 diff --git a/assets/trex_camera_calib.json b/assets/trex_camera_calib.json new file mode 100644 index 0000000000000000000000000000000000000000..ac06931513bcb460d54fa488ec4ba818cdfa4f8c --- /dev/null +++ b/assets/trex_camera_calib.json @@ -0,0 +1,51 @@ +{ + "native_hw": [360, 640], + "video_hw": [180, 320], + "notes": "Default T-Rex camera calibration. Head uses URDF zed_left_camera FK; wrist uses L_ee/R_ee + T_ee_to_cam. Tune wrist mounts if projections look misaligned.", + "views": { + "head_left": { + "camera_model": "zed_x_mini_left", + "extrinsic_source": "fk_zed_left_camera", + "K_native": [ + [349.5075, 0.0, 313.375], + [0.0, 349.5075, 180.47075], + [0.0, 0.0, 1.0] + ], + "dist": [-0.173733, 0.0272837, -6.10446e-05, 0.0, 0.0] + }, + "left_wrist": { + "camera_model": "zed_one", + "extrinsic_source": "fk_L_ee_times_T_ee_to_cam", + "ee_frame": "L_ee", + "T_ee_to_cam": [ + [0.0, 0.0, 1.0, 0.08], + [-1.0, 0.0, 0.0, 0.0], + [0.0, -1.0, 0.0, 0.02], + [0.0, 0.0, 0.0, 1.0] + ], + "K_native": [ + [350.0, 0.0, 320.0], + [0.0, 350.0, 180.0], + [0.0, 0.0, 1.0] + ], + "dist": [0.0, 0.0, 0.0, 0.0, 0.0] + }, + "right_wrist": { + "camera_model": "zed_one", + "extrinsic_source": "fk_R_ee_times_T_ee_to_cam", + "ee_frame": "R_ee", + "T_ee_to_cam": [ + [0.0, 0.0, 1.0, 0.08], + [-1.0, 0.0, 0.0, 0.0], + [0.0, -1.0, 0.0, 0.02], + [0.0, 0.0, 0.0, 1.0] + ], + "K_native": [ + [350.0, 0.0, 320.0], + [0.0, 350.0, 180.0], + [0.0, 0.0, 1.0] + ], + "dist": [0.0, 0.0, 0.0, 0.0, 0.0] + } + } +} diff --git a/docs/DATASET_TO_GEAR_AND_TRAIN.md b/docs/DATASET_TO_GEAR_AND_TRAIN.md new file mode 100644 index 0000000000000000000000000000000000000000..3b9c462ce8723d831a49548b166639c5ee91af2f --- /dev/null +++ b/docs/DATASET_TO_GEAR_AND_TRAIN.md @@ -0,0 +1,471 @@ +# Adding a New Embodiment to DreamZero + +How to take a LeRobot v2 dataset for a new robot, convert it to GEAR format, define its modality config, and train a DreamZero policy. + +Throughout this guide, replace `` with your robot's name (e.g. `myrobot`, `franka`, `aloha`). + +--- + +## Overview + +``` +Step 1 Convert LeRobot v2 dataset → GEAR metadata +Step 2 Register the embodiment tag +Step 3 Add modality config + transforms to base YAML +Step 4 Create a dataset YAML +Step 5 Create a training script +Step 6 Train +``` + +--- + +## Step 1: Convert Dataset to GEAR Format + +The converter reads a LeRobot v2 dataset and generates the metadata files DreamZero needs. It does **not** modify your parquet files or videos — it only writes to `meta/`. + +### Expected input structure + +``` +your_dataset/ +├── data/ +│ └── chunk-000/ +│ ├── episode_000000.parquet +│ └── ... +├── videos/ +│ └── chunk-000/ +│ ├── observation.images.cam0/ +│ │ ├── episode_000000.mp4 +│ │ └── ... +│ └── observation.images.cam1/ +│ └── ... +└── meta/ + └── info.json # must contain: features, total_episodes, fps +``` + +### Run the converter + +```bash +python scripts/data/convert_lerobot_to_gear.py \ + --dataset-path /path/to/your_dataset \ + --embodiment-tag \ + --state-keys '{"joint_pos": [0, 6], "gripper_pos": [6, 7]}' \ + --action-keys '{"joint_pos": [0, 6], "gripper_pos": [6, 7]}' \ + --relative-action-keys joint_pos gripper_pos \ + --task-key annotation.task +``` + +`--state-keys` and `--action-keys` tell the converter how to split a packed vector column into named sub-keys. The JSON maps sub-key name → `[start_index, end_index]`. Omit these flags to let the converter auto-detect. + +### Arguments + +| Argument | Default | Description | +|---|---|---| +| `--dataset-path` | *(required)* | Path to the LeRobot v2 dataset | +| `--output-path` | *(in-place)* | Write to a different directory instead of in-place | +| `--embodiment-tag` | `xdof` | Tag for `meta/embodiment.json`; must match the key you use in Step 3 | +| `--state-keys` | *(auto)* | JSON: sub-key name → `[start, end]` index range | +| `--action-keys` | *(auto)* | JSON: sub-key name → `[start, end]` index range | +| `--relative-action-keys` | *(none)* | Sub-key names to compute relative action stats for | +| `--task-key` | *(auto)* | Column name for language/task annotations | +| `--fps` | *(from info.json)* | Override dataset FPS | +| `--action-horizon` | `24` | Horizon for relative stats computation | +| `--force` | `false` | Overwrite existing metadata files | + +### Generated files + +The converter creates these under `meta/`: + +| File | Contents | +|---|---| +| `modality.json` | Maps state, action, video, and annotation keys with index ranges and dtypes | +| `embodiment.json` | `{"embodiment_tag": ""}` | +| `stats.json` | Per-feature statistics (mean, std, min, max, q01, q99) | +| `relative_stats_dreamzero.json` | Relative action statistics (action − reference state) | +| `tasks.jsonl` | Unique task descriptions | +| `episodes.jsonl` | Per-episode metadata (index, tasks, length) | + +--- + +## Step 2: Register the Embodiment Tag + +1. Add to the enum in `groot/vla/data/schema/embodiment_tags.py`: + +```python +class EmbodimentTag(str, Enum): + ... + MY_ROBOT = "" +``` + +2. Add to `VALID_EMBODIMENT_TAGS` in `scripts/data/convert_lerobot_to_gear.py` (if you want the converter to accept the tag without `--force`): + +```python +VALID_EMBODIMENT_TAGS = [ + ..., + "", +] +``` + +--- + +## Step 3: Add Modality Config and Transforms + +Edit `groot/vla/configs/data/dreamzero/base_48_wan_fine_aug_relative.yaml`. + +### 3a. Understanding modality + +**Modality** connects your dataset columns to the training pipeline. The `modality.json` from Step 1 contains entries like: + +```json +{ + "state": { + "joint_pos": {"original_key": "observation.state", "start": 0, "end": 6}, + "gripper_pos": {"original_key": "observation.state", "start": 6, "end": 7} + }, + "action": { + "joint_pos": {"original_key": "action", "start": 0, "end": 6}, + "gripper_pos": {"original_key": "action", "start": 6, "end": 7} + }, + "video": { + "cam0": {"original_key": "observation.images.cam0"} + }, + "annotation": { + "task": {"original_key": "annotation.task"} + } +} +``` + +The YAML config must reference **exactly these key names** with type prefixes: + +| Modality | YAML key format | Example | +|---|---|---| +| State | `state.` | `state.joint_pos` | +| Action | `action.` | `action.joint_pos` | +| Video | `video.` | `video.cam0` | +| Language | `annotation.` | `annotation.task` | + +If the YAML keys don't match `modality.json`, training will fail with missing-key errors. + +### 3b. Add `modality_config_` + +```yaml +modality_config_: + video: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] + eval_delta_indices: [0] + modality_keys: # one entry per camera, matching modality.json + - video.cam0 + - video.cam1 + - video.cam2 + state: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0] + modality_keys: # matching modality.json state keys + - state.joint_pos + - state.gripper_pos + action: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23] + modality_keys: # matching modality.json action keys + - action.joint_pos + - action.gripper_pos + language: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0] + modality_keys: + - annotation.task +``` + +**`delta_indices` explained:** + +- **Video** — frame offsets to sample (25 entries = 25 frames from the trajectory). +- **State / Language** — `[0]` = current timestep only. +- **Action** — future offsets (24 entries = 24-step action chunk). + +Adjust these to match your `num_frames` and `action_horizon` training settings. + +### 3c. Add `transform_` + +```yaml +transform_: + _target_: groot.vla.data.transform.ComposedModalityTransform + transforms: + # Video + - <<: *totensor_cfg + apply_to: ${modality_config_.video.modality_keys} + - <<: *crop_cfg + apply_to: ${modality_config_.video.modality_keys} + - <<: *resize_cfg + apply_to: ${modality_config_.video.modality_keys} + - <<: *color_jitter_cfg + apply_to: ${modality_config_.video.modality_keys} + - <<: *to_numpy_cfg + apply_to: ${modality_config_.video.modality_keys} + + # State + - _target_: groot.vla.data.transform.StateActionToTensor + apply_to: ${modality_config_.state.modality_keys} + - _target_: groot.vla.data.transform.StateActionTransform + apply_to: ${modality_config_.state.modality_keys} + normalization_modes: + state.joint_pos: q99 # every state key needs a normalization mode + state.gripper_pos: q99 + + # Action + - _target_: groot.vla.data.transform.StateActionToTensor + apply_to: ${modality_config_.action.modality_keys} + - _target_: groot.vla.data.transform.StateActionTransform + apply_to: ${modality_config_.action.modality_keys} + normalization_modes: + action.joint_pos: q99 # every action key needs a normalization mode + action.gripper_pos: q99 + + # Concat + - _target_: groot.vla.data.transform.ConcatTransform + video_concat_order: ${modality_config_.video.modality_keys} + state_concat_order: ${modality_config_.state.modality_keys} + action_concat_order: ${modality_config_.action.modality_keys} + + # Model-specific (required, don't change) + - ${model_specific_transform} +``` + +Every state and action key **must** appear in `normalization_modes`. The strategy is typically `q99`. + +### 3d. Register in the global maps + +Add your embodiment to each of these four maps (at the bottom of the base YAML): + +```yaml +modality_configs: + ... + : ${modality_config_} + +transforms: + ... + : ${transform_} + +metadata_versions: + ... + : '0221' + +fps: + ... + : 30 # set to your dataset's FPS +``` + +--- + +## Step 4: Create a Dataset YAML + +Create `groot/vla/configs/data/dreamzero/_relative.yaml`: + +```yaml +# @package _global_ + +defaults: + - dreamzero/base_48_wan_fine_aug_relative + - _self_ + +max_state_dim: 64 +use_global_metadata: false +relative_action: true +relative_action_per_horizon: false +relative_action_keys: + - joint_pos # sub-key names (without state./action. prefix) + - gripper_pos # that should use relative actions +max_chunk_size: 5 +dataset_shard_sampling_rate: 0.1 +mixture_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotMixtureDataset.from_mixture_spec +single_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotSubLangSingleActionChunkDatasetDROID + +_data_root: ??? # set via CLI or env var + +train_dataset: + _target_: ${mixture_dataset_cls} + _convert_: object + mixture_spec: + - dataset_path: + : # must match key in modality_configs/transforms + - ${_data_root} + dataset_weight: 1.0 + distribute_weights: true + + dataset_class: ${single_dataset_cls} + all_modality_configs: ${modality_configs} + all_transforms: ${transforms} + metadata_versions: ${metadata_versions} + fps: ${fps} + dataset_kwargs: + video_backend: decord + use_global_metadata: ${use_global_metadata} + max_chunk_size: ${max_chunk_size} + relative_action: ${relative_action} + relative_action_keys: ${relative_action_keys} + relative_action_per_horizon: ${relative_action_per_horizon} + mixture_kwargs: + training: true + balance_dataset_weights: false + seed: 42 + shard_sampling_rate: ${dataset_shard_sampling_rate} +``` + +The critical things to get right: + +- `` in `mixture_spec.dataset_path` must match the key in `modality_configs` and `transforms`. +- `relative_action_keys` lists the sub-key names (without `state.`/`action.` prefix) that exist in **both** state and action modalities. + +--- + +## Step 5: Create a Training Script + +Create `scripts/train/_training.sh`: + +```bash +#!/bin/bash +export HYDRA_FULL_ERROR=1 + +# ============ CONFIGURATION ============ +DATA_ROOT=${DATA_ROOT:?"Set DATA_ROOT to your GEAR-converted dataset"} +OUTPUT_DIR=${OUTPUT_DIR:-"./checkpoints/dreamzero__lora"} + +if [ -z "${NUM_GPUS:-}" ]; then + NUM_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) +fi +NUM_GPUS=${NUM_GPUS:-8} + +WAN_CKPT_DIR=${WAN_CKPT_DIR:-"./checkpoints/Wan2.1-I2V-14B-480P"} +TOKENIZER_DIR=${TOKENIZER_DIR:-"./checkpoints/umt5-xxl"} +# ======================================= + +# Auto-download weights if missing +if [ ! -d "$WAN_CKPT_DIR" ] || [ -z "$(ls -A "$WAN_CKPT_DIR" 2>/dev/null)" ]; then + huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir "$WAN_CKPT_DIR" +fi +if [ ! -d "$TOKENIZER_DIR" ] || [ -z "$(ls -A "$TOKENIZER_DIR" 2>/dev/null)" ]; then + huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR" +fi + +if [ ! -d "$DATA_ROOT" ]; then + echo "ERROR: Dataset not found at $DATA_ROOT" + exit 1 +fi +if [ ! -f "$DATA_ROOT/meta/embodiment.json" ]; then + echo "ERROR: meta/embodiment.json missing — run convert_lerobot_to_gear.py first" + exit 1 +fi + +torchrun --nproc_per_node $NUM_GPUS --standalone \ + groot/vla/experiment/experiment.py \ + report_to=wandb \ + data=dreamzero/_relative \ + wandb_project=dreamzero \ + train_architecture=lora \ + num_frames=33 \ + action_horizon=24 \ + num_views=3 \ + model=dreamzero/vla \ + model/dreamzero/action_head=wan_flow_matching_action_tf \ + model/dreamzero/transform=dreamzero_cotrain \ + num_frame_per_block=2 \ + num_action_per_block=24 \ + num_state_per_block=1 \ + seed=42 \ + training_args.learning_rate=1e-5 \ + training_args.deepspeed="groot/vla/configs/deepspeed/zero2.json" \ + save_steps=10000 \ + training_args.warmup_ratio=0.05 \ + output_dir=$OUTPUT_DIR \ + per_device_train_batch_size=4 \ + max_steps=100000 \ + weight_decay=1e-5 \ + save_total_limit=10 \ + upload_checkpoints=false \ + bf16=true \ + tf32=true \ + eval_bf16=true \ + dataloader_pin_memory=false \ + dataloader_num_workers=1 \ + image_resolution_width=320 \ + image_resolution_height=176 \ + save_lora_only=true \ + max_chunk_size=4 \ + frame_seqlen=880 \ + save_strategy=steps \ + _data_root=$DATA_ROOT \ + dit_version=$WAN_CKPT_DIR \ + text_encoder_pretrained_path=$WAN_CKPT_DIR/models_t5_umt5-xxl-enc-bf16.pth \ + image_encoder_pretrained_path=$WAN_CKPT_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth \ + vae_pretrained_path=$WAN_CKPT_DIR/Wan2.1_VAE.pth \ + tokenizer_path=$TOKENIZER_DIR \ + pretrained_model_path=./checkpoints/DreamZero-AgiBot \ + ++action_head_cfg.config.skip_component_loading=true \ + ++action_head_cfg.config.defer_lora_injection=true +``` + +### Key parameters to adjust per embodiment + +| Parameter | Default | When to change | +|---|---|---| +| `num_views` | `3` | Number of cameras your robot has | +| `action_horizon` | `24` | Must match the number of action `delta_indices` | +| `num_frames` | `33` | Must be `len(video delta_indices) + num_frame_per_block * (blocks - 1)` | +| `image_resolution_width` | `320` | Match your camera resolution (or desired resize) | +| `image_resolution_height` | `176` | Match your camera resolution (or desired resize) | +| `max_steps` | `100000` | Scale with dataset size | +| `per_device_train_batch_size` | `4` | Adjust for GPU memory | + +--- + +## Step 6: Train + +### Download the pretrained checkpoint + +The training scripts load from a pretrained DreamZero checkpoint for LoRA fine-tuning. Download [DreamZero-AgiBot](https://huggingface.co/GEAR-Dreams/DreamZero-AgiBot) (~45GB) to `./checkpoints/DreamZero-AgiBot`: + +```bash +git clone https://huggingface.co/GEAR-Dreams/DreamZero-AgiBot ./checkpoints/DreamZero-AgiBot +``` + +Or with the Hugging Face CLI: + +```bash +hf download GEAR-Dreams/DreamZero-AgiBot --repo-type model --local-dir ./checkpoints/DreamZero-AgiBot +``` + +### Launch training + +```bash +DATA_ROOT=/path/to/your_dataset bash scripts/train/_training.sh + +# With overrides: +DATA_ROOT=/path/to/your_dataset OUTPUT_DIR=./checkpoints/run1 NUM_GPUS=4 \ + bash scripts/train/_training.sh +``` + +--- + +## Pre-Training Checklist + +- [ ] `meta/embodiment.json` exists and has the correct tag +- [ ] `meta/modality.json` state/action/video/annotation keys are populated +- [ ] `meta/stats.json` and `meta/relative_stats_dreamzero.json` exist +- [ ] `meta/tasks.jsonl` and `meta/episodes.jsonl` exist +- [ ] Embodiment tag in `embodiment.json` matches the key in `modality_configs` / `transforms` / `metadata_versions` / `fps` +- [ ] YAML `modality_keys` match `modality.json` keys exactly (with `state.`/`action.`/`video.`/`annotation.` prefix) +- [ ] Every state and action key appears in `normalization_modes` in the transform block +- [ ] `relative_action_keys` are sub-key names that exist in both state and action +- [ ] Wan2.1-I2V-14B-480P and umt5-xxl weights are available +- [ ] DreamZero-AgiBot checkpoint is downloaded to `./checkpoints/DreamZero-AgiBot` + +--- + +## Quick Reference: Existing Embodiments + +| Embodiment | Data Config | Layout | +|---|---|---| +| `oxe_droid` | `droid_relative.yaml` | 3 cameras, joint_position + gripper_position | +| `agibot` | `agibot_relative.yaml` | 3 cameras, 6 state keys, 7 action keys | +| `yam` | `yam_relative.yaml` | 3 cameras (top/left/right), bimanual left/right joint_pos + gripper_pos | + +Use these as concrete examples when building your own config. diff --git a/docs/DROID_CONVERSION.md b/docs/DROID_CONVERSION.md new file mode 100644 index 0000000000000000000000000000000000000000..e2f316dbc44ceaa0085355d8d5933ab1d5307764 --- /dev/null +++ b/docs/DROID_CONVERSION.md @@ -0,0 +1,62 @@ +# Converting DROID from Scratch + +If you want to reproduce the DreamZero DROID dataset conversion yourself (or modify the filtering), follow the steps below. This requires the raw DROID 1.0.1 dataset in RLDS format and the idle filter ranges JSON. + +> **Most users should skip this** and simply download the preprocessed dataset: +> ```bash +> huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir ./data/droid_lerobot +> ``` + +## Step 1: Install conversion dependencies + +```bash +pip install tensorflow tensorflow-datasets polars av +``` + +## Step 2: Download the raw DROID 1.0.1 dataset + +This requires `gsutil` ([Google Cloud CLI](https://cloud.google.com/storage/docs/gsutil_install)). The full dataset is ~1.7TB. + +```bash +gsutil -m cp -r gs://gresearch/robotics/droid/1.0.1 ./data/droid/1.0.1 +``` + +> **Important:** Use version 1.0.1, not 1.0.0. Version 1.0.1 contains the complete set of language annotations (~75k episodes). + +## Step 3: Download the idle filter ranges + +This JSON file maps each episode to the frame ranges that should be kept (non-idle frames). It was originally computed by [Physical Intelligence](https://github.com/Physical-Intelligence/openpi) for training pi0-DROID models. + +```bash +gsutil cp gs://openpi-assets/droid/droid_sample_ranges_v1_0_1.json ./data/keep_ranges.json +``` + +## Step 4: Run the conversion + +```bash +python scripts/data/convert_droid.py \ + ./data/droid/1.0.1 \ + ./data/droid_lerobot \ + --keep-ranges-path ./data/keep_ranges.json \ + --filter-failed \ + -n 16 +``` + +For a quick test with a small subset: +```bash +python scripts/data/convert_droid.py \ + ./data/droid/1.0.1 \ + ./data/droid_lerobot_test \ + --keep-ranges-path ./data/keep_ranges.json \ + --filter-failed \ + --first-n 5 \ + -n 4 +``` + +## Script reference + +See [`scripts/data/convert_droid.py`](scripts/data/convert_droid.py) for full usage: + +``` +python scripts/data/convert_droid.py --help +``` \ No newline at end of file diff --git a/docs/TREX_TRACK_FORCE_MODEL.md b/docs/TREX_TRACK_FORCE_MODEL.md new file mode 100644 index 0000000000000000000000000000000000000000..ba4caa936e7dfe37f14ef85bcbabe5977625844a --- /dev/null +++ b/docs/TREX_TRACK_FORCE_MODEL.md @@ -0,0 +1,301 @@ +# T-Rex Track-Force：16-step、20 Hz 动作与 5 Hz 触觉模型 + +本文描述独立的 `trex_track_force` 模型、数据契约、embedding、注意力、两阶段 +flow matching、训练和在线执行。它不会修改原 DreamZero 的模型调用路径。 + +## 1. 固定时序契约 + +- 每个 action chunk 含 **16 个动作**，动作频率为 **20 Hz**。 +- 一个 chunk 覆盖 `16 / 20 = 0.8 s`；动作时间戳跨度为 `0.75 s`。 +- 触觉频率为 **5 Hz**，因此每 **4 个动作步**刷新一次。 +- 一个 chunk 内的触觉刷新 offset 固定为 `[0, 4, 8, 12]`。 +- 每次刷新只重新去噪尚未执行的 action suffix；已经执行的前缀保持不变。 +- `max_chunk_size=4` 保留 DreamZero 原有的 autoregressive block memory。 + 这里的 `max_chunk_size` 是跨 block 的记忆窗口，不是单个 action chunk 的长度。 +- 一个训练 sample 仍由 4 个 autoregressive blocks 组成，因此共有 64 个动作； + 只有每 block 的 action 长度从 24 改为 16。 +- 三路视频保持 DreamZero 的 33 帧布局：1 帧 clean conditioning observation， + 后接每 block 8 帧、共 32 帧 10 Hz future-video targets。 +- 每个 block 都有 16 帧过去 track 和 16 步 target track；target 的第 0 步就是 + 当前 anchor 帧，与 OpenPI 的 track window 契约一致。 + +## 2. 总体架构 + +```mermaid +flowchart LR + V["三视角 RGB\nhead + left wrist + right wrist"] --> GRID["三视角拼图\n160 × 320"] + GRID --> VAE["Wan2.2 VAE38\n48-channel latent"] + GRID --> CLIP["Wan CLIP\n图像条件"] + TXT["任务文本"] --> T5["UMT5\n文本条件"] + + TRK["250 点 track\npast clean + future noisy"] --> TE["TrackEncoder\n几何 + 可见性 + 身份 embedding"] + ACT["16 × 62D delta-base action\n补零到 64D"] --> AE["Action encoder"] + ST["62D 当前 EEF/hand state"] --> SE["State encoder"] + + VAE --> WAN["CausalWanTrackForceModel\nWan2.2-TI2V-5B · 30 blocks"] + CLIP --> WAN + T5 --> WAN + TE --> WAN + AE --> WAN + SE --> WAN + + WAN --> COARSE["τ: 1 → 0.4\n6-step coarse action"] + WAN --> TFLOW["未来 track flow"] + WAN --> VFLOW["未来 video flow"] + WAN --> MEM["τ=0.4 coarse memory"] + + RAW["10 fingers × 6D force\n16-sample history"] --> VQ["Per-finger force VQ-VAE\n64 codes · 256D"] + FRESH["当前 10 × 6D force"] --> FT["Force-only Transformer\n6 layers · width 768"] + VQ --> FT + MEM --> FT + COARSE --> FT + FT --> FINE["τ: 0.4 → 0\n4-step tactile action suffix"] +``` + +粗阶段负责从视觉、状态、语言和 track 中得到动作的大尺度结构；独立的 +force-only transformer 只使用力信号和粗阶段 memory 完成剩余去噪，使触觉 +负责接触后的细粒度修正。 + +## 3. 62DoF 动作空间 + +每只手臂占 31 维，两侧共 62 维： + +- 手腕/末端位姿：`xyz + rotation-6D`，共 9 维。 +- 手部关节目标：22 维。 + +训练动作不是世界坐标绝对位姿。loader 以 chunk 起始状态为 reference，把每个 +目标手腕位姿转换为 **delta-base**： + +- 平移在 chunk 起始手腕坐标系中表示。 +- 旋转为 `R_reference^-1 × R_target`，再编码为 rotation-6D。 +- 手部 22 维仍是绝对目标。 + +模型内部把 62 个物理维补零为 64 维；采样的两个 padding 维在初始化和每个 +Euler step 后都会重新置零。输出先按 relative-action 统计量反归一化，再用 +chunk 起始 state 恢复绝对 EEF 位姿。 + +## 4. Track 数据和点身份 + +250 个点具有固定、可验证的顺序： + +- `0:50`：头部视角左手及手臂。 +- `50:100`：头部视角右手及手臂。 +- `100:125`：左腕视角 5×5 背景点。 +- `125:175`：左腕视角手掌点。 +- `175:200`：右腕视角 5×5 背景点。 +- `200:250`：右腕视角手掌点。 + +SAM2 只在 episode 第 0 帧根据固定 prompt 产生手/臂 mask；点从 mask 中采样， +随后由 CoTracker 跟踪整段视频。保存字段为： + +- `observation.track_xy`：`[250, 2]`，坐标归一化到 `[0, 1]`。 +- `observation.track_visibility`：`[250]`。 +- metadata 中同时记录 view、hand、role、point index 和各 segment 边界。 + +## 5. Embedding 设计 + +### 5.1 Video、语言和图像条件 + +- 三视角拼入 2×2 canvas 的三个有效格，再统一缩放到 `160×320`。 +- Wan2.2 VAE38 生成 48-channel latent。 +- `patch_size=[1,2,2]`，每 latent frame 形成 50 个 video tokens。 +- video token 使用 Wan 原生 3D RoPE，编码时间、高度和宽度。 +- UMT5 输出 4096D 文本条件。 +- Wan CLIP 输出 1280D 图像条件，并经 `img_emb` 投影。 +- CLIP/T5 cross-attention 只作用于 observation/action query，不向 track query + 泄漏视觉或语言信息。 + +### 5.2 Action 和 state + +- action：每个 62D 动作补零到 64D，经线性层投影到 Wan hidden dim 3072。 +- 一个 block 有 16 个 action tokens；位置由 1D action RoPE 编码。 +- state：62D 当前状态补到 64D，经线性层得到一个 state token。 +- flow 时间 `τ` 经 sinusoidal embedding 和 MLP 后参与 Wan modulation。 + +### 5.3 TrackEncoder + +每个点分别产生一个 `past token` 和一个 `future token`。时间序列输入特征为： + +```text +[x, y, visibility, Δx, Δy] +``` + +不可见时刻的坐标和 motion 先清零，随后完整的 `16×5` 时序按固定顺序展平，再经 +`Linear → SiLU → Linear` 投影。不能在时间维求平均：future flow 的每个时刻含有 +独立噪声，平均会丢失“哪一个噪声属于哪一个 timestep”，使 16-step 重建不可解。 + +最终每个点 token 是下列 embedding 的和，再经 LayerNorm： + +```text +trajectory ++ view(head / left_wrist / right_wrist) ++ hand(none / left / right) ++ role(head_hand / wrist_background / wrist_hand) ++ canonical point id(0...249) ++ autoregressive block id ++ temporal role(past / future) +``` + +左右手不只通过点序号区分，还具有显式 hand embedding。背景点使用 +`hand=none` 和独立 role embedding。 + +与 OpenPI 一致，track target 始终是 `[0,1]` 内的绝对归一化 XY；target window +从当前帧开始，因此第 0 步等于 GT anchor。`TrackDecoder` 对 250 个 +future-track hidden token 分别执行 `LayerNorm → Linear(16×2)`，预测 CFM +velocity `noise-clean_xy`。Euler 积分后的结果已经是绝对坐标，不做空间 +`cumsum`，也不再额外加 anchor。 + +### 5.4 Force-only VQ-VAE + +输入只包含 force/wrench，不包含 deformation map： + +```text +[batch, history=16, fingers=10, wrench=6] +``` + +- 两只手共享同一套时序卷积 encoder/decoder。 +- 每只手内部加入 5 个 finger identity embeddings。 +- 输出每根手指一个 256D latent，共 10 个 tactile history tokens。 +- EMA codebook 大小为 64，带 commitment loss、perplexity 统计和 dead-code + revival。 +- 左右手通过 side embedding 区分，手指位置通过 finger embedding 区分。 +- 训练可直接输入 raw 16-step history；推理也可输入预计算的 10 个离散 codes。 +- episode 前缀缺失的历史填为归一化中性值，并用 validity mask 从 VQ 重建 + loss 中排除 padding 步。 + +### 5.5 Force-only Transformer + +输入 token 包括： + +- 16 个 noisy action tokens。 +- 10 个当前 force tokens。 +- 10 个 VQ history tokens。 +- 从 Wan `τ=0.4` hidden state 提取的 coarse memory tokens。 + +action token 叠加 action position、5Hz force slot、token type、当前 refresh +offset 和 `τ` embedding。force/history token叠加 finger position、token type 和 +refresh offset。独立 transformer 为 6 层、12 heads、hidden dim 768。 + +## 6. 非对称 attention 契约 + +每个 autoregressive block 的 packed 顺序为： + +```text +[obs, action, state, track_past, track_future] +``` + +其可见性严格为： + +- observation query 和 action query 在同一 block 内互相可见。 +- observation/action query 可读取同 block 的 state、past track 和 noisy future + track，也可读取有限 AR 窗口内的历史 obs/action/track。 +- state query 只读取自己的 state token，避免把多模态信息反向带给 track。 +- past-track query 读取历史 block 的 track 和当前 past track，但不能读取当前 + noisy future track。 +- future-track query 只读取当前及历史 track。 +- **任何 track query 都不能读取 observation、action、state、CLIP 或文本。** +- 任何 query 都不能读取未来 block。 + +因此实现了要求的方向性：`obs/action → track` 表示 obs/action 可以把 track +作为条件；反方向被 mask 禁止。 + +## 7. 两阶段 flow matching + +使用线性插值： + +```text +x_τ = τ · noise + (1 - τ) · clean +target flow = noise - clean +x_next = x_τ + (τ_next - τ) · predicted_flow +``` + +### 粗阶段：Wan，`τ=1 → 0.4` + +- 训练时 action/Wan expert 遵循原始 T-Rex，在完整 `(0,1]` 上按 + `Beta(1.5,1.0)` 采样；其均值约为 `0.6`，但这不是 split timestep。 +- Wan 联合预测 action flow、future-track flow 和 future-video flow。 +- 推理按总计 10 个、`Δτ=-0.1` 的 Euler 网格运行前 6 步，到达 `τ=0.4`。 +- 在精确的 `τ=0.4` 再执行一次 Wan，生成给触觉 transformer 使用的 detached + coarse memory。 + +### 精阶段：触觉，`τ=0.4 → 0` + +- 训练时 `τ_tactile = 0.4 × Beta(1.5,1.0)`，覆盖 `(0,0.4]`。 +- 每个训练 chunk 对 offset `0/4/8/12` 分别提供对应的 force 与 16-sample + history。 +- offset 为 `k` 时，loss 只覆盖 action `k:16`。 +- 推理运行剩余 4 个 Euler updates；每次 5Hz 刷新只写入未执行 suffix。 +- runtime controller 保存之前已经发出的 prefix，确保后续触觉刷新不能改写历史 + command。 + +## 8. 训练 loss + +总 loss 由以下部分加权求和： + +- `dynamics_loss`：未来视频 latent flow MSE。 +- `action_loss`：粗阶段 62 个物理动作维的 flow MSE。 +- `track_loss`：可见性 mask 后的绝对 XY CFM velocity MSE。 +- `force_loss`：按 refresh offset mask 后的动作 suffix flow MSE。 +- `vq_loss`：仅在有效历史步计算的 force history reconstruction loss。 +- `commitment_loss`：VQ commitment loss，默认系数 0.25。 + +训练使用 Wan2.2-TI2V-5B 共享权重和 LoRA。LoRA 注入 +`q,k,v,o,k_img,v_img,ffn.0,ffn.2`；新建的 action/state/track/force、video token +投影和 decoder 完整训练。checkpoint 保存 LoRA、新模块参数以及 VQ EMA buffers。 + +## 9. 在线执行 + +1. 收集三视角 conditioning observation（在线接口也可编码多帧历史）、当前 + 62D state、16 帧 past track、当前 force 和 16-sample force history。 +2. Wan 对 action、future track 和 future video 运行 6 步，从 `τ=1` 到 + `τ=0.4`。 +3. offset 0 的触觉运行剩余 4 步，把完整 16-step action 去噪到 `τ=0`。 +4. 以 20 Hz 执行动作。 +5. 执行 4、8、12 步后，各接收一次新的 5Hz 触觉；每次从缓存的 coarse state + 重新去噪剩余 suffix，同时保留已执行 prefix。 +6. 输出 delta-base 动作反归一化并恢复为绝对手腕 `xyz + rotation-6D`。 + +## 10. 入口 + +训练： + +```bash +cd /scratch1/home/zhicao/dreamzero +bash scripts/train/trex_track_force_training_wan22.sh +``` + +训练脚本直接使用已经构建完成的 `data/trex_small`，启动前只执行 schema +完整性校验，不会再次运行 SAM2/CoTracker。 + +日志、checkpoint 与原 T-Rex 训练共用同一套 Trainer 回调。每隔 +`wandb_video_reconstruction_steps` 个 global step 会分别保存两个监控视频： +不带标记的原始重建视频位于 +`OUTPUT_DIR/eval_videos/train_step_XXXXXX.mp4`，20Hz 预测 track motion +叠加到 10Hz head/left-wrist/right-wrist 三面板后的版本位于 +`OUTPUT_DIR/eval_track_videos/train_step_XXXXXX.mp4`。原始预测与目标 +track 同时写入 `OUTPUT_DIR/eval_tracks/train_step_XXXXXX.npz`。两个视频 +分别记录到 W&B `eval/predicted_video` 和 `eval/predicted_track_video`； +`dynamics/action/track/force/VQ/commitment` loss 和触觉 codebook +perplexity/active-code 指标同时写入 W&B 与 `OUTPUT_DIR/loss_log.jsonl`。 +两个视频的首帧都是精确 GT conditioning frame；track target/prediction 的第 0 +步在模型链路内就是当前 GT anchor。可视化不再对整段预测做事后平移。 + +训练可视化默认使用 +`wandb_video_reconstruction_inference_steps=1`：先运行一次 coarse WAN +Euler update，再在 `tau=0.4` 用一次边界预测恢复 clean video/track；不会直接 +解码半噪声状态，并跳过不会改变这两个输出的 action-only 触觉精修。 +这不会改变训练或正式在线推理的 6+4 两阶段 schedule；如需更高质量的监控视频， +可以单独提高该值，但耗时近似按 inference steps 线性增长。 + +离线 NPZ 推理： + +```bash +python scripts/eval/trex_track_force_inference.py \ + --checkpoint checkpoints/trex_track_force_wan22_lora/checkpoint-8000 \ + --dataset-root data/trex_small \ + --input sample_input.npz \ + --output prediction.npz +``` + +核心实现位于 `groot/vla/model/trex_track_force/`；独立 Hydra 配置为 +`model=trex_track_force/vla` 和 `data=dreamzero/trex_track_force_wan22`。 diff --git a/docs/WAN22_BACKBONE.md b/docs/WAN22_BACKBONE.md new file mode 100644 index 0000000000000000000000000000000000000000..7e9bdefa9c90414496d71c35d00297256dad46b1 --- /dev/null +++ b/docs/WAN22_BACKBONE.md @@ -0,0 +1,167 @@ +# Training DreamZero with Wan2.2-TI2V-5B Backbone + +This guide explains how to train DreamZero on the DROID dataset using **Wan2.2-TI2V-5B** as the backbone instead of the default Wan2.1-I2V-14B. + +## Architecture Differences + +| Component | Wan2.1-I2V-14B | Wan2.2-TI2V-5B | +|-----------|-----------------|----------------| +| DiT dim | 5120 | 3072 | +| DiT layers | 32 | 30 | +| DiT heads | 16 | 24 | +| FFN dim | 13824 | 14336 | +| VAE latent channels | 16 | 48 | +| VAE spatial stride | 8× | 16× | +| Model type | i2v | ti2v | + +**FFN** = Feed-Forward Network: the two-layer MLP in each transformer block (Linear → GELU → Linear). FFN dim is the intermediate hidden size (e.g. 14336 for 5B). + +DreamZero uses a **CausalWanModel** wrapper that extends the base Wan architecture with **action/state registers** for robot policy learning. The same `CausalWanModel` class supports both Wan2.1 and Wan2.2 backbones via configuration—no new class is required. The config switches the architecture parameters (dim, in_dim, out_dim, etc.) and uses `WanVideoVAE38` for the 48-channel Wan2.2 VAE. + +**What action/state registers do:** The DiT sees a single sequence `[video_tokens | action_register]` where the action register is encoded action and state features (one chunk per block). All tokens share the same transformer (with causal masking and RoPE). The model learns to predict **video noise** (dynamics) and **action noise** (policy): the action-register slice is decoded by `action_decoder` to produce action noise predictions. So the model is conditioned on current state and (noisy) actions and learns to denoise both video and actions for closed-loop policy learning. + +## Causal masking, RoPE, and sequence layout + +### Causal masking + +In attention, **causal masking** means each position can only attend to **past and current** positions (no future). So token at index `i` can see keys at indices `j ≤ i`. That keeps the model autoregressive: it never uses future video frames or future actions when predicting the current step. In CausalWanModel the masking is **blockwise**: the first frame attends to itself; each later block of video frames can attend to the first frame plus previous (and optionally current) blocks. Action and state tokens have their own causal pattern so each action chunk only sees past video and past actions/state. This matches policy learning where you condition on observed history and predict the next action chunk. + +### RoPE (Rotary Position Embeddings) + +**RoPE** encodes position by rotating query and key vectors in a complex plane with position-dependent angles. Unlike adding a position vector, RoPE makes attention scores depend on the *relative* position of query and key, which generalizes better to longer sequences. In CausalWanModel: + +- **Video tokens** use **3D RoPE**: separate frequency components for frame index (time), height, and width of the patch grid. So each token knows its (t, h, w) in the video. +- **Action and state tokens** use **1D RoPE**: a single position index along the sequence (frame/block index). So the model knows the temporal order of action chunks and state. + +Freqs are built in `_create_freqs()` from the patch grid size (F, H, W) and concatenated with separate 1D freqs for the action register. + +### Tokens, blocks, and chunks + +- **Token**: The smallest unit the transformer sees. After **patch_embedding** (stride 1×2×2 on the latent), one frame yields a 2D grid of tokens; the total per frame is **frame_seqlen** (e.g. 50 for 160×320). So one **token** = one patch (e.g. 1×2×2 in latent space). + +- **Block (image block)**: A group of consecutive **frames**, not tokens. **num_frame_per_block** (e.g. 2) frames form one “image block.” So with 33 frames you get multiple blocks. **num_image_blocks** = `(num_frames - 1) // num_frame_per_block`. Blocks are used for blockwise causal attention and to align video with action/state. + +- **Chunk**: In policy terms, an **action chunk** is the sequence of actions the policy outputs for one block (e.g. **num_action_per_block** = 24 actions per block). The **action register** in the DiT has one chunk per image block: for each block there are `num_action_per_block` action tokens and `num_state_per_block` state tokens. So the register length is `num_image_blocks * (num_action_per_block + num_state_per_block)`. “Chunk” and “block” are often used together: one video block corresponds to one action chunk (and one state token) in the register. + +Summary: **tokens** = patch-level units (50 per frame); **blocks** = groups of frames (e.g. 2 frames per block); **chunks** = per-block action (and state) outputs that are packed into the action register. + +## Inference: blocks, chunks, and closed-loop + +### How blocks and chunks are used when predicting actions + +At **inference**, the model predicts **one action chunk** per call, conditioned on **one block** of video (and current state): + +1. **Input**: A short video of the current block — e.g. `num_frame_per_block` frames (e.g. 2) — plus current **state** and (during the denoising loop) **noisy actions** for the chunk being predicted. The first time in a trajectory, the “context” is the first frame (and optionally a warm-up pass with no action to fill the KV cache). + +2. **DiT input**: The sequence is `[video_tokens for this block | action_register]`. The action register holds encoded **noisy** actions and **state** for this block only (one chunk). So the DiT sees: “this block of video + this chunk’s noisy actions and state.” + +3. **KV cache**: To keep inference causal and efficient, the model uses a **KV cache** over previous blocks. So for the *next* block, the cache already contains keys/values for earlier frames; the DiT only runs on the **new** block’s tokens plus the new action register. `current_start_frame` tells the DiT which block we’re on so RoPE and cache indexing are correct. + +4. **Output**: The DiT predicts **video noise** and **action noise**. The action noise is decoded by `action_decoder` into a prediction for the **current chunk**. The scheduler then updates the noisy action toward clean; after `num_inference_steps` denoising steps you get **one clean action chunk** (e.g. 24 actions). + +So: **one block of frames** (and state) in → **one action chunk** out. Blocks and chunks are aligned: one image block ↔ one action chunk in the register. + +### How DreamZero does closed-loop inference + +Closed-loop execution reuses the same block/chunk logic in a loop: + +1. **Observe**: Robot has current observation (e.g. image history + state). The policy is called with this observation (e.g. via `lazy_joint_video_action` or `lazy_joint_video_action_causal`). + +2. **Predict**: The action head runs the diffusion loop for the **current block**: it encodes the observed frames to latent, runs the DiT (with KV cache and `current_start_frame`) for each denoising step, and returns one denoised **action chunk** (e.g. 24 actions). + +3. **Execute**: The robot **executes** that chunk (e.g. 24 steps at 5 Hz → ~4.8 s). No new model call during execution. + +4. **Repeat**: After execution, new observation is available. The policy is called again with the new video (e.g. last N frames). If the task/language is unchanged, `current_start_frame` is incremented by `num_frame_per_block` and the KV cache is reused; the DiT only processes the **new** block and predicts the **next** action chunk. If the task or language changes (or the cache is full), the cache and `current_start_frame` are reset. + +So closed-loop = **repeated “one block in → one chunk out”** with KV cache across steps so the model never re-processes past frames. + +### What changes when you swap the backbone to 5B + +The **inference algorithm and API stay the same** for 14B vs 5B: + +- Same **block/chunk layout**: `num_frame_per_block`, `num_action_per_block`, `num_state_per_block` (and thus one block → one chunk) are defined by config and data; they do not depend on which backbone (14B vs 5B) you use. +- Same **closed-loop flow**: `lazy_joint_video_action`, KV cache, `current_start_frame`, and the denoising loop are in the **action head** and are shared. The policy still calls the same methods (`get_action`, `lazy_joint_video_action`, etc.). +- Same **backbone role**: The backbone only produces conditioning (e.g. text embeddings). The action head owns the DiT, VAE, and action/state encoders. So “swapping to 5B” means swapping the **action head config** (and checkpoints) to the Wan22 5B DiT + VAE38 + 160×320; the high-level inference path (backbone → action_head → one chunk) is unchanged. + +What **does** change with 5B: + +- **DiT size and layout**: 5B uses a smaller DiT (dim 3072, 30 layers, 24 heads), **frame_seqlen = 50** (for 160×320), and **no** first-frame latent concat (`concat_first_frame_latent=False`). First frame is conditioned via **CLIP** in the context, not as extra channel in the latent. +- **VAE and resolution**: 5B uses **WanVideoVAE38** (48 channels, 16× spatial) and **160×320** video. So latent is 10×20; tokens per frame = 50. +- **Conditioning**: 5B uses CLIP image embedding for the first frame in the context; 14B can concatenate the first-frame latent to the DiT input. The action head handles this inside the same `_forward_inference` / `_forward_blocks`; no change to the external inference API. + +So: **blocks and chunks** are used the same way for predicting actions at inference; **closed-loop** is the same loop of “observe → predict one chunk → execute → repeat” with KV cache; **swapping to 5B** keeps that flow and only changes the internal model (DiT/VAE) and resolution/conditioning. + +## Prerequisites + +1. **Wan2.2-TI2V-5B** weights: + ```bash + huggingface-cli download Wan-AI/Wan2.2-TI2V-5B --local-dir ./checkpoints/Wan2.2-TI2V-5B + ``` + Or clone from [Wan2.2 GitHub](https://github.com/Wan-Video/Wan2.2) and follow their download instructions. + +2. **Image encoder (CLIP)**: Wan2.2-TI2V-5B does not include the CLIP image encoder. Use the one from Wan2.1: + ```bash + huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir ./checkpoints/Wan2.1-I2V-14B-480P + ``` + Only `models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth` is needed. + +3. **DROID dataset** in LeRobot format: + ```bash + huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir ./data/droid_lerobot + ``` + +## Quick Start + +```bash +# Set paths (optional - defaults shown) +export WAN22_CKPT_DIR=./checkpoints/Wan2.2-TI2V-5B +export IMAGE_ENCODER_DIR=./checkpoints/Wan2.1-I2V-14B-480P # for CLIP only +export DROID_DATA_ROOT=./data/droid_lerobot + +# Run training +bash scripts/train/droid_training_wan22.sh +``` + +## Configuration Details + +The Wan2.2 config (`wan_flow_matching_action_tf_wan22.yaml`) overrides: + +- **model/dreamzero/action_head**: `wan_flow_matching_action_tf_wan22` +- **diffusion_model_cfg**: Wan2.2 architecture (dim=3072, in_dim=48, out_dim=48, etc.) +- **vae_cfg**: `WanVideoVAE38` (48-channel Wan2.2 VAE) +- **frame_seqlen**: 50 (patch output per frame) +- **target_video_height / target_video_width**: 160 and 320 so latent spatial size is **even** (10×20 after VAE38 16×), avoiding a dynamics-loss crop. Previously 176×320 gave latent 11×20 (odd height); we use **160×320** (H×W) so both latent dimensions are even after the DiT’s stride-(1,2,2) patch embedding. + +For other resolutions, `frame_seqlen` must match patch output per frame; use H and W divisible by 32 for even latent: +- 160×320 (H×W): latent 10×20 → 50 +- 176×320: latent 11×20 → 50 (odd H; loss uses crop) +- 640×352: 220 + +## Using with Custom Training Scripts + +To use Wan2.2 in your own training script, add: + +```bash +model/dreamzero/action_head=wan_flow_matching_action_tf_wan22 \ +dit_version=$WAN22_CKPT_DIR \ +text_encoder_pretrained_path=$WAN22_CKPT_DIR/models_t5_umt5-xxl-enc-bf16.pth \ +image_encoder_pretrained_path=$IMAGE_ENCODER_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth \ +vae_pretrained_path=$WAN22_CKPT_DIR/Wan2.2_VAE.pth +``` +(Do not pass `frame_seqlen`; the Wan22 config uses 50.) + +## File Layout + +``` +dreamzero/ +├── groot/vla/configs/model/dreamzero/action_head/ +│ ├── wan_flow_matching_action_tf.yaml # Wan2.1 (default) +│ └── wan_flow_matching_action_tf_wan22.yaml # Wan2.2-TI2V-5B +├── scripts/train/ +│ ├── droid_training.sh # Wan2.1 backbone +│ └── droid_training_wan22.sh # Wan2.2 backbone +└── docs/ + └── WAN22_BACKBONE.md # This file +``` + +The action head (`wan_flow_matching_action_tf.py`) automatically detects Wan2.2 vs Wan2.1 based on `in_dim` (48 vs 16) and `vae.z_dim` (48 vs 16), and loads the correct checkpoint files from the appropriate HuggingFace repos when local paths are not found. diff --git a/eval_utils/policy_client.py b/eval_utils/policy_client.py new file mode 100644 index 0000000000000000000000000000000000000000..5511ff3afdad382e02fd04fc72f018e3fcfbb80c --- /dev/null +++ b/eval_utils/policy_client.py @@ -0,0 +1,93 @@ +"""Client for communicating with a policy server. + +Adapted from https://github.com/robo-arena/roboarena/ + +""" + +import logging +import time +from typing import Dict, Tuple + +import websockets.sync.client +from typing_extensions import override + +from openpi_client.base_policy import BasePolicy +from openpi_client import msgpack_numpy + +# The websockets library by default sends a ping every 20 seconds and +# expects a pong response within 20 seconds. However, the sever may not +# send a pong response immediately if it is busy processing a request. +# Increase the ping interval and timeout so that the client can wait +# for a longer time before closing the connection. +PING_INTERVAL_SECS = 60 +PING_TIMEOUT_SECS = 600 + +class WebsocketClientPolicy(BasePolicy): + """Implements the Policy interface by communicating with a server over websocket. + + See WebsocketPolicyServer for a corresponding server implementation. + """ + + def __init__(self, host: str = "0.0.0.0", port: int = 8000) -> None: + self._uri = f"ws://{host}:{port}" + self._packer = msgpack_numpy.Packer() + self._ws, self._server_metadata = self._wait_for_server() + + def get_server_metadata(self) -> Dict: + return self._server_metadata + + def _wait_for_server(self) -> Tuple[websockets.sync.client.ClientConnection, Dict]: + logging.info(f"Waiting for server at {self._uri}...") + try: + conn = websockets.sync.client.connect( + self._uri, + compression=None, + max_size=None, + ping_interval=PING_INTERVAL_SECS, + ping_timeout=PING_TIMEOUT_SECS, + ) + metadata = msgpack_numpy.unpackb(conn.recv()) + return conn, metadata + except: + logging.info("Connection to server with ws:// failed. Trying wss:// ...") + + self._uri = "wss://" + self._uri.split("//")[1] + conn = websockets.sync.client.connect( + self._uri, + compression=None, + max_size=None, + ping_interval=PING_INTERVAL_SECS, + ping_timeout=PING_TIMEOUT_SECS, + ) + metadata = msgpack_numpy.unpackb(conn.recv()) + return conn, metadata + + @override + def infer(self, obs: Dict) -> Dict: # noqa: UP006 + # Notify server that we're calling the infer endpoint (as opposed to the reset endpoint) + obs["endpoint"] = "infer" + + data = self._packer.pack(obs) + self._ws.send(data) + response = self._ws.recv() + if isinstance(response, str): + # we're expecting bytes; if the server sends a string, it's an error. + raise RuntimeError(f"Error in inference server:\n{response}") + return msgpack_numpy.unpackb(response) + + @override + def reset(self, reset_info: Dict) -> None: + # Notify server that we're calling the reset endpoint (as opposed to the infer endpoint) + reset_info["endpoint"] = "reset" + + data = self._packer.pack(reset_info) + self._ws.send(data) + response = self._ws.recv() + return response + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + client = WebsocketClientPolicy() + actions = client.infer({}) + print(f"Actions received: {actions}") + client.reset({}) \ No newline at end of file diff --git a/eval_utils/policy_server.py b/eval_utils/policy_server.py new file mode 100644 index 0000000000000000000000000000000000000000..9b34bb3bd66064f9d6e83707c3181d7bc3c15624 --- /dev/null +++ b/eval_utils/policy_server.py @@ -0,0 +1,130 @@ +"""Server for serving a policy over websockets. + +Adapted from https://github.com/robo-arena/roboarena/ + +""" + + +import asyncio +import dataclasses +import logging +import traceback + +from openpi_client.base_policy import BasePolicy +from openpi_client import msgpack_numpy +import websockets.asyncio.server +import websockets.frames + + +@dataclasses.dataclass +class PolicyServerConfig: + # Resolution that images get resized to client-side, None means no resizing. + # It's beneficial to resize images to the desired resolution client-side for faster communication. + image_resolution: tuple[int, int] | None = (224, 224) + # Whether or not wrist camera image(s) should be sent. + needs_wrist_camera: bool = True + # Number of external cameras to send. + n_external_cameras: int = 1 # can be in [0, 1, 2] + # Whether or not stereo camera image(s) should be sent. + needs_stereo_camera: bool = False + # Whether or not the unique eval session id should be sent (e.g. for policies that want to keep track of history). + needs_session_id: bool = False + # Which action space to use. + action_space: str = "joint_position" # can be in ["joint_position", "joint_velocity", "cartesian_position", "cartesian_velocity"] + + +class WebsocketPolicyServer: + """ + Serves a policy using the websocket protocol. + + Interface: + Observation: + - observation/wrist_image_left: (H, W, 3) if needs_wrist_camera is True + - observation/wrist_image_right: (H, W, 3) if needs_wrist_camera is True and needs_stereo_camera is True + - observation/exterior_image_{i}_left: (H, W, 3) if n_external_cameras >= 1 + - observation/exterior_image_{i}_right: (H, W, 3) if needs_stereo_camera is True + - session_id: (1,) if needs_session_id is True + - observation/joint_position: (7,) + - observation/cartesian_position: (6,) + - observation/gripper_position: (1,) + - prompt: str, the natural language task instruction for the policy + + Action: + - action: (N, 8,) or (N, 7,): either 7 movement actions (for joint action spaces) or 6 (for cartesian) plus one dimension for gripper position + --> all N actions will get executed on the robot before the server is queried again + + """ + + def __init__( + self, + policy: BasePolicy, + server_config: PolicyServerConfig, + host: str = "0.0.0.0", + port: int = 8000, + ) -> None: + self._policy = policy + self._server_config = server_config + self._host = host + self._port = port + 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: + await server.serve_forever() + + async def _handler(self, websocket: websockets.asyncio.server.ServerConnection): + logging.info(f"Connection from {websocket.remote_address} opened") + packer = msgpack_numpy.Packer() + + # Send server config to client to configure what gets sent to server. + await websocket.send(packer.pack(dataclasses.asdict(self._server_config))) + + while True: + try: + obs = msgpack_numpy.unpackb(await websocket.recv()) + + endpoint = obs["endpoint"] + del obs["endpoint"] + if endpoint == "reset": + self._policy.reset(obs) + to_return = "reset successful" + else: + action = self._policy.infer(obs) + to_return = packer.pack(action) + await websocket.send(to_return) + 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 + + +if __name__ == "__main__": + import numpy as np + + class DummyPolicy(BasePolicy): + def infer(self, obs): + return np.zeros((1, 8), dtype=np.float32) + + def reset(self, reset_info): + pass + + logging.basicConfig(level=logging.INFO) + policy = DummyPolicy() + server = WebsocketPolicyServer(policy, PolicyServerConfig()) + server.serve_forever() + \ No newline at end of file diff --git a/eval_utils/run_sim_eval.py b/eval_utils/run_sim_eval.py new file mode 100644 index 0000000000000000000000000000000000000000..cb194e77113a09db0e8756cf27a5bc31e0e04e51 --- /dev/null +++ b/eval_utils/run_sim_eval.py @@ -0,0 +1,219 @@ +""" +Example script for running 10 rollouts of a DROID policy on the example environment. + +Usage: + +First, make sure you download the simulation assets and unpack them into the root directory of this package. + +Then, in a separate terminal, launch the policy server on localhost:8000 +-- make sure to set XLA_PYTHON_CLIENT_MEM_FRACTION to avoid JAX hogging all the GPU memory. + +For example, to launch a pi0-FAST-DROID policy (with joint position control), +run the command below in a separate terminal from the openpi "karl/droid_policies" branch: + +XLA_PYTHON_CLIENT_MEM_FRACTION=0.5 uv run scripts/serve_policy.py policy:checkpoint --policy.config=pi0_fast_droid_jointpos --policy.dir=s3://openpi-assets-simeval/pi0_fast_droid_jointpos + +Finally, run the evaluation script: + +python run_eval.py --episodes 10 --headless +""" + +import uuid + +import tyro +import argparse +import gymnasium as gym +import torch +import cv2 +import mediapy +import numpy as np +from datetime import datetime +from pathlib import Path +from PIL import Image +from tqdm import tqdm + +from openpi_client import image_tools +from sim_evals.inference.abstract_client import InferenceClient +from policy_client import WebsocketClientPolicy + + +class DreamZeroJointPosClient(InferenceClient): + def __init__(self, + remote_host:str = "localhost", + remote_port:int = 6000, + open_loop_horizon:int = 8, + ) -> None: + self.client = WebsocketClientPolicy(remote_host, remote_port) + self.open_loop_horizon = open_loop_horizon + self.actions_from_chunk_completed = 0 + self.pred_action_chunk = None + self.session_id = str(uuid.uuid4()) + + def visualize(self, request: dict): + """ + Return the camera views how the model sees it + """ + curr_obs = self._extract_observation(request) + right_img = image_tools.resize_with_pad(curr_obs["right_image"], 224, 224) + wrist_img = image_tools.resize_with_pad(curr_obs["wrist_image"], 224, 224) + left_img = image_tools.resize_with_pad(curr_obs["left_image"], 224, 224) + combined = np.concatenate([right_img, wrist_img, left_img], axis=1) + return combined + + def reset(self): + self.actions_from_chunk_completed = 0 + self.pred_action_chunk = None + self.session_id = str(uuid.uuid4()) + + def infer(self, obs: dict, instruction: str) -> dict: + """ + Infer the next action from the policy in a server-client setup + """ + curr_obs = self._extract_observation(obs) + if ( + self.actions_from_chunk_completed == 0 + or self.actions_from_chunk_completed >= self.open_loop_horizon + ): + self.actions_from_chunk_completed = 0 + request_data = { + "observation/exterior_image_0_left": image_tools.resize_with_pad(curr_obs["right_image"], 180, 320), + "observation/exterior_image_1_left": image_tools.resize_with_pad(curr_obs["left_image"], 180, 320), + "observation/wrist_image_left": image_tools.resize_with_pad(curr_obs["wrist_image"], 180, 320), + "observation/joint_position": curr_obs["joint_position"].astype(np.float64), + "observation/cartesian_position": np.zeros((6,), dtype=np.float64), # dummy cartesian position + "observation/gripper_position": curr_obs["gripper_position"].astype(np.float64), + "prompt": instruction, + "session_id": self.session_id, + } + for k, v in request_data.items(): + print(f"{k}: {v.shape if not isinstance(v, str) else v}") + + result = self.client.infer(request_data) + actions = result["actions"] if isinstance(result, dict) else result + assert len(actions.shape) == 2, f"Expected 2D array, got shape {actions.shape}" + assert actions.shape[-1] == 8, f"Expected 8 action dimensions (7 joints + 1 gripper), got {actions.shape[-1]}" + self.pred_action_chunk = actions + + + action = self.pred_action_chunk[self.actions_from_chunk_completed] + self.actions_from_chunk_completed += 1 + + # binarize gripper action + if action[-1].item() > 0.5: + action = np.concatenate([action[:-1], np.ones((1,))]) + else: + action = np.concatenate([action[:-1], np.zeros((1,))]) + + img1 = image_tools.resize_with_pad(curr_obs["right_image"], 224, 224) + img2 = image_tools.resize_with_pad(curr_obs["wrist_image"], 224, 224) + img3 = image_tools.resize_with_pad(curr_obs["left_image"], 224, 224) + both = np.concatenate([img1, img2, img3], axis=1) + + return {"action": action, "viz": both} + + def _extract_observation(self, obs_dict, *, save_to_disk=False): + # Assign images + right_image = obs_dict["policy"]["external_cam"][0].clone().detach().cpu().numpy() + left_image = obs_dict["policy"]["external_cam_2"][0].clone().detach().cpu().numpy() + wrist_image = obs_dict["policy"]["wrist_cam"][0].clone().detach().cpu().numpy() + + # Capture proprioceptive state + robot_state = obs_dict["policy"] + joint_position = robot_state["arm_joint_pos"].clone().detach().cpu().numpy() + gripper_position = robot_state["gripper_pos"].clone().detach().cpu().numpy() + + if save_to_disk: + combined_image = np.concatenate([right_image, wrist_image], axis=1) + combined_image = Image.fromarray(combined_image) + combined_image.save("robot_camera_views.png") + + return { + "right_image": right_image, + "left_image": left_image, + "wrist_image": wrist_image, + "joint_position": joint_position, + "gripper_position": gripper_position, + } + + + + +def main( + episodes: int = 10, + scene: int = 1, + headless: bool = True, + host: str = "localhost", + port: int = 6000, + ): + # launch omniverse app with arguments (inside function to prevent overriding tyro) + from isaaclab.app import AppLauncher + parser = argparse.ArgumentParser(description="Tutorial on creating an empty stage.") + AppLauncher.add_app_launcher_args(parser) + args_cli, _ = parser.parse_known_args() + args_cli.enable_cameras = True + args_cli.headless = headless + app_launcher = AppLauncher(args_cli) + simulation_app = app_launcher.app + + # All IsaacLab dependent modules should be imported after the app is launched + import sim_evals.environments # noqa: F401 + from isaaclab_tasks.utils import parse_env_cfg + + + # Initialize the env + env_cfg = parse_env_cfg( + "DROID", + device=args_cli.device, + num_envs=1, + use_fabric=True, + ) + instruction = None + match scene: + case 1: + instruction = "put the cube in the bowl" + case 2: + instruction = "pick up the can and put it in the mug" + case 3: + instruction = "put the banana in the bin" + case _: + raise ValueError(f"Scene {scene} not supported") + + env_cfg.set_scene(scene) + env = gym.make("DROID", cfg=env_cfg) + + obs, _ = env.reset() + obs, _ = env.reset() # need second render cycle to get correctly loaded materials + client = DreamZeroJointPosClient(remote_host=host, remote_port=port) + + + video_dir = Path("runs") / datetime.now().strftime("%Y-%m-%d") / datetime.now().strftime("%H-%M-%S") + video_dir.mkdir(parents=True, exist_ok=True) + video = [] + ep = 0 + max_steps = env.env.max_episode_length + with torch.no_grad(): + for ep in range(episodes): + for _ in tqdm(range(max_steps), desc=f"Episode {ep+1}/{episodes}"): + ret = client.infer(obs, instruction) + if not headless: + cv2.imshow("Right Camera", cv2.cvtColor(ret["viz"], cv2.COLOR_RGB2BGR)) + cv2.waitKey(1) + video.append(ret["viz"]) + action = torch.tensor(ret["action"])[None] + obs, _, term, trunc, _ = env.step(action) + if term or trunc: + break + + client.reset() + mediapy.write_video( + video_dir / f"episode_{ep}.mp4", + video, + fps=15, + ) + video = [] + + env.close() + simulation_app.close() + +if __name__ == "__main__": + args = tyro.cli(main) diff --git a/eval_utils/serve_dreamzero_wan22.py b/eval_utils/serve_dreamzero_wan22.py new file mode 100644 index 0000000000000000000000000000000000000000..725f652ebadecef6ed24b01a20a052e44cb19871 --- /dev/null +++ b/eval_utils/serve_dreamzero_wan22.py @@ -0,0 +1,391 @@ +""" +Serve the DreamZero 5B implementation (Wan2.2-TI2V-5B) over the websocket policy server. + +This is the 5B model: Wan2.2 diffusion backbone, 48-channel VAE38, frame_seqlen=50 (160×320 +latent 10×20). Inference is causal with KV caching: first request in a session uses 1 frame +and warms the cache; subsequent requests use FRAMES_PER_CHUNK=4 frames and append to the cache. +On session_id change (or explicit reset), buffers and action_head.current_start_frame are cleared. + +The checkpoint at model_path should be DreamZero with Wan22 5B (model/dreamzero/action_head= +wan_flow_matching_action_tf_wan22, data droid_relative_wan22 → 160×320). GrootSimPolicy loads +that checkpoint and runs inference; it is the correct policy class for DreamZero. + +Usage (single GPU): + + torchrun --nproc_per_node=1 eval_utils/serve_dreamzero_wan22.py --model_path ./checkpoints/dreamzero_droid_wan22_smoke --port 8000 + + # Or single process: + python eval_utils/serve_dreamzero_wan22.py --model_path ./checkpoints/dreamzero_droid_wan22_smoke --port 8000 + +Client: send observations per PolicyServerConfig (policy_server.py). Video is resized to the +checkpoint's expected resolution (e.g. 180×320) so the eval transform accepts it; the 5B action +head resizes to 160×320 internally. Override with --image_height/--image_width if needed. +Response is an action chunk (N, 8). Use session_id for episode boundaries. +""" + +import datetime +import logging +import os +import sys + +import imageio + +logger = logging.getLogger(__name__) + +import cv2 +import numpy as np +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import init_device_mesh +import tyro + +# Avoid FailOnRecompileLimitHit when serving: the flow scheduler's torch.compile'd +# multistep_uni_p_bh_update recompiles under varying shapes/inputs (e.g. batch size, +# step_index, order). Increase limits so the server doesn't hit the default cap. +_dynamo = torch._dynamo.config +if hasattr(_dynamo, "cache_size_limit"): + _dynamo.cache_size_limit = 1000 +if hasattr(_dynamo, "recompile_limit"): + _dynamo.recompile_limit = 800 +if hasattr(_dynamo, "accumulated_cache_size_limit"): + _dynamo.accumulated_cache_size_limit = 1000 +if hasattr(_dynamo, "accumulated_recompile_limit"): + _dynamo.accumulated_recompile_limit = 2000 +from pathlib import Path +from tianshou.data import Batch + +# Add repo root for imports +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from openpi_client.base_policy import BasePolicy + +from eval_utils.policy_server import WebsocketPolicyServer, PolicyServerConfig +from groot.vla.model.n1_5.sim_policy import GrootSimPolicy +from groot.vla.data.schema import EmbodimentTag +from groot.vla.data.transform import ComposedModalityTransform + + +# DreamZero Wan 5B is trained with 160×320 (droid_relative_wan22). Fallback if we cannot read from policy. +DEFAULT_IMAGE_HEIGHT = 160 +DEFAULT_IMAGE_WIDTH = 320 +FRAMES_PER_CHUNK = 4 # matches 5B num_frame_per_block for causal chunked inference + + +def _get_expected_video_resolution(policy: GrootSimPolicy) -> tuple[int, int]: + """Get (height, width) the policy's eval_transform expects for video (from checkpoint + metadata). Resolution in metadata is (width, height); we return (height, width) for resize. + DreamZero Wan 5B (droid_relative_wan22) uses 160×320; other configs may use e.g. 180×320. + """ + eval_transform = getattr(policy, "eval_transform", None) + if eval_transform is None: + return (DEFAULT_IMAGE_HEIGHT, DEFAULT_IMAGE_WIDTH) + if not isinstance(eval_transform, ComposedModalityTransform): + return (DEFAULT_IMAGE_HEIGHT, DEFAULT_IMAGE_WIDTH) + for t in eval_transform.transforms: + if hasattr(t, "original_resolutions") and getattr(t, "original_resolutions", None): + res = t.original_resolutions + if res: + # original_resolutions values are (width, height) + w, h = next(iter(res.values())) + return (int(h), int(w)) + return (DEFAULT_IMAGE_HEIGHT, DEFAULT_IMAGE_WIDTH) + + +def _resize_frames_to_resolution(frames: np.ndarray, target_h: int, target_w: int) -> np.ndarray: + """Resize video frames to (target_h, target_w). Accepts (H,W,C) or (T,H,W,C).""" + if frames.ndim == 3: + if (frames.shape[0], frames.shape[1]) != (target_h, target_w): + frames = cv2.resize(frames, (target_w, target_h), interpolation=cv2.INTER_LINEAR) + return frames + out = np.stack( + [cv2.resize(f, (target_w, target_h), interpolation=cv2.INTER_LINEAR) for f in frames], + axis=0, + ) + return out + + +def _maybe_init_distributed(): + """Initialize process group for single-GPU or multi-GPU. Required by GrootSimPolicy.""" + if dist.is_initialized(): + return + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", "29500") + dist.init_process_group(backend="nccl", rank=0, world_size=1) + torch.cuda.set_device(0) + + +# Modality key mappings: client observation keys -> model input keys per embodiment. +# Client sends: observation/exterior_image_0_left, exterior_image_1_left, wrist_image_left. +VIDEO_KEY_MAPPING = { + "oxe_droid": { + "observation/exterior_image_0_left": "video.exterior_image_1_left", + "observation/exterior_image_1_left": "video.exterior_image_2_left", + "observation/wrist_image_left": "video.wrist_image_left", + }, +} +STATE_KEY_MAPPING = { + "oxe_droid": ("state.joint_position", "state.gripper_position"), +} +LANGUAGE_KEY_MAPPING = { + "oxe_droid": "annotation.language.action_text", +} + + +class DreamZeroWan225BPolicy(BasePolicy): + """ + Wraps GrootSimPolicy for the DreamZero 5B implementation (Wan2.2-TI2V-5B). + + Converts roboarena observation/action format to DROID/Batch. Video is resized to the + resolution expected by the policy's eval_transform (from checkpoint metadata) so + VideoToTensor validation passes. The 5B action head then resizes to 160×320 internally. + First call in a session uses 1 frame; later calls use 4 frames (FRAMES_PER_CHUNK). + Session reset clears frame buffers and action_head.current_start_frame. + """ + + def __init__( + self, + groot_policy: GrootSimPolicy, + image_height: int, + image_width: int, + embodiment_tag: str = "oxe_droid", + save_video_pred: bool = False, + video_output_dir: str = "./video_pred_output", + ): + super().__init__() + self._policy = groot_policy + self._image_height = image_height + self._image_width = image_width + self._embodiment_tag = ( + embodiment_tag if embodiment_tag in VIDEO_KEY_MAPPING else "oxe_droid" + ) + video_keys = list(VIDEO_KEY_MAPPING[self._embodiment_tag].values()) + self._frame_buffers = {k: [] for k in video_keys} + self._is_first_call = True + self._current_session_id = None + self._save_video_pred = save_video_pred + self._video_output_dir = video_output_dir + self._video_pred_latents: list[torch.Tensor] = [] + self._current_prompt: str = "" + + def _convert_observation(self, obs: dict) -> dict: + """Convert roboarena observation format to model Batch format. + Incoming frames are resized to the policy's expected (height, width) so + eval_transform's VideoToTensor check passes. + """ + image_key_mapping = VIDEO_KEY_MAPPING[self._embodiment_tag] + for roboarena_key, model_key in image_key_mapping.items(): + if roboarena_key in obs: + data = obs[roboarena_key] + if isinstance(data, np.ndarray): + data = _resize_frames_to_resolution( + data, self._image_height, self._image_width + ) + if data.ndim == 4: + self._frame_buffers[model_key].extend(list(data)) + else: + self._frame_buffers[model_key].append(data) + + num_frames = 1 if self._is_first_call else FRAMES_PER_CHUNK + converted = {} + for model_key, buffer in self._frame_buffers.items(): + if len(buffer) > 0: + if len(buffer) >= num_frames: + frames_to_use = buffer[-num_frames:] + else: + frames_to_use = buffer.copy() + while len(frames_to_use) < num_frames: + frames_to_use.insert(0, buffer[0]) + video = np.stack(frames_to_use, axis=0) + converted[model_key] = video + + state_joint_key, state_gripper_key = STATE_KEY_MAPPING[self._embodiment_tag] + if "observation/joint_position" in obs: + joint_pos = np.asarray(obs["observation/joint_position"]) + if joint_pos.ndim == 1: + joint_pos = joint_pos.reshape(1, -1) + converted[state_joint_key] = joint_pos.astype(np.float64) + else: + converted[state_joint_key] = np.zeros((1, 7), dtype=np.float64) + + if "observation/gripper_position" in obs: + gripper_pos = np.asarray(obs["observation/gripper_position"]) + if gripper_pos.ndim == 1: + gripper_pos = gripper_pos.reshape(1, -1) + converted[state_gripper_key] = gripper_pos.astype(np.float64) + else: + converted[state_gripper_key] = np.zeros((1,1), dtype=np.float64) + + text_prompt = obs.get("prompt", "") + logger.info("Text prompt: %s", text_prompt) + if text_prompt: + self._current_prompt = text_prompt + lang_key = LANGUAGE_KEY_MAPPING[self._embodiment_tag] + converted[lang_key] = text_prompt + return converted + + def _convert_action(self, action_dict: dict) -> np.ndarray: + """Convert model action dict to (N, 8) array (7 joint + 1 gripper).""" + joint_action = None + gripper_action = None + for key, value in action_dict.items(): + if ("joint_position" in key or "joint_pos" in key) and "gripper" not in key: + joint_action = value + elif "gripper_position" in key or "gripper" in key: + gripper_action = value + if joint_action is None: + return np.zeros((1, 8), dtype=np.float32) + if isinstance(joint_action, torch.Tensor): + joint_action = joint_action.cpu().numpy() + if joint_action.ndim == 1: + joint_action = joint_action.reshape(1, -1) + N = joint_action.shape[0] + if gripper_action is not None: + if isinstance(gripper_action, torch.Tensor): + gripper_action = gripper_action.cpu().numpy() + if gripper_action.ndim == 1: + gripper_action = gripper_action.reshape(-1, 1) + if gripper_action.shape[-1] > 1: + gripper_action = gripper_action[..., :1] + else: + gripper_action = np.zeros((N, 1), dtype=np.float32) + return np.concatenate([joint_action, gripper_action], axis=-1).astype(np.float32) + + def infer(self, obs: dict) -> np.ndarray: + session_id = obs.get("session_id") + if session_id is not None and session_id != self._current_session_id: + if self._current_session_id is not None: + self.reset({}) + self._current_session_id = session_id + + converted_obs = self._convert_observation(obs) + batch = Batch(obs=converted_obs) + with torch.no_grad(): + result_batch, video_pred = self._policy.lazy_joint_forward_causal(batch) + if self._save_video_pred and video_pred is not None: + self._video_pred_latents.append(video_pred.detach()) + action_dict = {} + action_chunk_dict = result_batch.act + for k in dir(action_chunk_dict): + if k.startswith("action."): + action_dict[k] = getattr(action_chunk_dict, k) + action = self._convert_action(action_dict) + if self._is_first_call: + self._is_first_call = False + return action + + def _save_predicted_video(self) -> None: + """Decode accumulated video prediction latents through the VAE and save as mp4.""" + if not self._video_pred_latents: + return + try: + from einops import rearrange + + action_head = self._policy.trained_model.action_head + latents = torch.cat(self._video_pred_latents, dim=2) + with torch.no_grad(): + frames = action_head.vae.decode( + latents, + tiled=action_head.tiled, + tile_size=(action_head.tile_size_height, action_head.tile_size_width), + tile_stride=(action_head.tile_stride_height, action_head.tile_stride_width), + ) + frames = rearrange(frames, "B C T H W -> B T H W C")[0] + frames = ((frames.float() + 1) * 127.5).clip(0, 255).cpu().numpy().astype(np.uint8) + + os.makedirs(self._video_output_dir, exist_ok=True) + timestamp = datetime.datetime.now().strftime("%m_%d_%H_%M_%S") + n_latent_frames = latents.shape[2] + existing = [f for f in os.listdir(self._video_output_dir) if f.endswith(".mp4")] + safe_prompt = self._current_prompt.replace(" ", "_") + safe_prompt = "".join(c for c in safe_prompt if c.isalnum() or c in "_-.") + if len(safe_prompt) > 80: + safe_prompt = safe_prompt[:80] + if not safe_prompt: + safe_prompt = "no_prompt" + output_path = os.path.join( + self._video_output_dir, + f"{len(existing):06}_{safe_prompt}_{timestamp}.mp4", + ) + imageio.mimsave(output_path, list(frames), fps=5, codec="libx264") + logger.info("Saved video prediction (%d frames) to %s", len(frames), output_path) + except Exception as e: + logger.warning("Failed to save video prediction: %s", e) + + def reset(self, reset_info: dict) -> None: + if self._save_video_pred: + self._save_predicted_video() + self._video_pred_latents.clear() + self._current_prompt = "" + for key in self._frame_buffers: + self._frame_buffers[key] = [] + self._is_first_call = True + self._current_session_id = None + if hasattr(self._policy.trained_model, "action_head") and hasattr( + self._policy.trained_model.action_head, "current_start_frame" + ): + self._policy.trained_model.action_head.current_start_frame = 0 + + +def main( + model_path: str = "./checkpoints/dreamzero_droid_wan22_smoke", + embodiment_tag: str = "oxe_droid", + tokenizer_path: str | None = None, + port: int = 8000, + host: str = "0.0.0.0", + image_height: int | None = None, + image_width: int | None = None, + save_video_pred: bool = False, + video_output_dir: str = "./video_pred_output", +) -> None: + logging.basicConfig(level=logging.INFO, force=True) + + _maybe_init_distributed() + device_mesh = init_device_mesh("cuda", mesh_shape=(1,), mesh_dim_names=("ip",)) + + logger.info("Loading DreamZero Wan22 policy from %s (embodiment=%s)", model_path, embodiment_tag) + checkpoint_name = os.path.basename(model_path.rstrip("/")) + video_output_dir = os.path.join(video_output_dir, checkpoint_name) + policy = GrootSimPolicy( + embodiment_tag=EmbodimentTag(embodiment_tag), + model_path=model_path, + tokenizer_path_override=tokenizer_path, + device="cuda" if torch.cuda.is_available() else "cpu", + device_mesh=device_mesh, + ) + if image_height is not None and image_width is not None: + h, w = image_height, image_width + logger.info("Using CLI video resolution: %dx%d", h, w) + else: + h, w = _get_expected_video_resolution(policy) + logger.info("Using checkpoint video resolution: %dx%d (HxW)", h, w) + wrapper = DreamZeroWan225BPolicy( + groot_policy=policy, + image_height=h, + image_width=w, + embodiment_tag=embodiment_tag, + save_video_pred=save_video_pred, + video_output_dir=video_output_dir, + ) + + server_config = PolicyServerConfig( + image_resolution=(h, w), + needs_wrist_camera=True, + n_external_cameras=2, + needs_stereo_camera=False, + needs_session_id=True, + action_space="joint_position", + ) + logger.info("Starting WebsocketPolicyServer on %s:%d (DreamZero 5B, %dx%d)", host, port, h, w) + server = WebsocketPolicyServer( + policy=wrapper, + server_config=server_config, + host=host, + port=port, + ) + server.serve_forever() + + +if __name__ == "__main__": + tyro.cli(main) diff --git a/groot/__init__.py b/groot/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/control/__init__.py b/groot/control/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/control/tensorrt_utils.py b/groot/control/tensorrt_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..74e25eb85698d6e5736d3a8708fe4882aa39db87 --- /dev/null +++ b/groot/control/tensorrt_utils.py @@ -0,0 +1,852 @@ +import torch +import os +import subprocess +import tensorrt as trt +import sys +import atexit +import ctypes +import modelopt.torch.quantization as mtq +from typing import Dict, List, Tuple +import shutil + +import numpy as np +import torch + + +FP8_DEFAULT_CONFIG = { + "quant_cfg": { + "*weight_quantizer": {"num_bits": (4, 3), "axis": None}, + "*input_quantizer": {"num_bits": (4, 3), "axis": None}, + "*output_quantizer": {"enable": False}, + "*[qkv]_bmm_quantizer": {"num_bits": (4, 3), "axis": None}, + "*softmax_quantizer": { + "num_bits": (4, 3), + "axis": None, + }, + "default": {"enable": False}, + }, + "algorithm": "max", +} + +NVFP4_DEFAULT_CONFIG = { + "quant_cfg": { + "*weight_quantizer": { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + "axis": None, + "enable": True, + }, + "*input_quantizer": { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + "axis": None, + "enable": True, + }, + "*output_quantizer": {"enable": False}, + "*[qkv]_bmm_quantizer": {"num_bits": (4, 3), "axis": None}, + "*softmax_quantizer": { + "num_bits": (4, 3), + "axis": None, + }, + "default": {"enable": False}, + }, + "algorithm": "max", +} + + + +def wan_quantize( + policy, + quantization_config, + model_type, + forward_loop, +): + """Quantize the VLA model using ModelOpt - simplified to use calc_mse_for_single_trajectory.""" + + # Configure quantization - disable problematic layers + if "quant_cfg" in quantization_config: + quantization_config["quant_cfg"]["*patch_embedding*"] = {"enable": False} + # if model_type == "14B" or model_type == "ar_14B": + # # Workaround: until we understand the issue https://nvbugspro.nvidia.com/bug/5612316 + # quantization_config["quant_cfg"]["*.self_attn.o.*"] = {"enable": False} + # quantization_config["quant_cfg"]["*.cross_attn.o.*"] = {"enable": False} + + policy.trained_model.action_head.model = mtq.quantize( + policy.trained_model.action_head.model, quantization_config, forward_loop=forward_loop + ) + mtq.print_quant_summary(policy.trained_model.action_head.model) + + return + + +def wan_trt_quantize_and_load_engine( + policy, + cfg, + onnx_path, + engine_path, + model_type, + forward_loop, +): + if ( + os.path.exists(os.path.dirname(engine_path)) + and cfg.inference_mode == "trt_build" + ): + shutil.rmtree(os.path.dirname(engine_path)) + + quantization_config = None + if cfg.quantize_dtype == "fp8": + quantization_config = FP8_DEFAULT_CONFIG.copy() + elif cfg.quantize_dtype == "nvfp4": + quantization_config = NVFP4_DEFAULT_CONFIG.copy() + else: + print(f"Quantization type {cfg.quantize_dtype} not supported. Skipping quantization.") + + if quantization_config is not None and cfg.inference_mode == "trt_build": + #policy.trained_model.action_head.model.to(torch.float16) + wan_quantize( + policy, + quantization_config, + model_type=model_type, + forward_loop=forward_loop, + ) + + if cfg.inference_mode == "trt_build": + policy.trained_model.action_head.model.to(torch.float16) + + print("Export model:", policy.trained_model.action_head.model) + + test_inputs = create_wan_test_inputs(policy, device="cuda", model_type=model_type) + min_shape = None + max_shape = None + opt_shape = None + + if model_type == "ar_14B": + + policy.trained_model.action_head.model.forward = policy.trained_model.action_head.model._forward_inference_trt + dynamic_axes = { + "kv_cache_packed": {3: "kv_cache_len"}, + } + min_shape = "kv_cache_packed:40x2x1x880x40x128" + max_shape = "kv_cache_packed:40x2x1x8800x40x128" + opt_shape = "kv_cache_packed:40x2x1x7920x40x128" + elif model_type == "ar_14B_droid": + policy.trained_model.action_head.model.forward = policy.trained_model.action_head.model._forward_inference_trt + dynamic_axes = { + "kv_cache_packed": {3: "kv_cache_len"}, + } + min_shape = "kv_cache_packed:40x2x1x880x40x128" + max_shape = "kv_cache_packed:40x2x1x8800x40x128" + opt_shape = "kv_cache_packed:40x2x1x7920x40x128" + elif model_type == "ar_5B_n6": + policy.trained_model.action_head.model.forward = policy.trained_model.action_head.model._forward_inference_trt + dynamic_axes = { + "kv_cache_packed": {3: "kv_cache_len"}, + } + min_shape = "kv_cache_packed:30x2x1x220x24x128" + max_shape = "kv_cache_packed:30x2x1x3080x24x128" + opt_shape = "kv_cache_packed:30x2x1x2860x24x128" + else: + dynamic_axes = None + + if cfg.quantize_dtype == "nvfp4": + export_to_onnx_fp4(policy.trained_model.action_head.model, test_inputs, onnx_path, dynamic_axes=dynamic_axes) + else: + export_to_onnx( + policy.trained_model.action_head.model, + test_inputs, + onnx_path, + model_type=model_type, + quantization_mode=cfg.quantize_dtype, + dynamic_axes=dynamic_axes, + ) + + build_tensorrt_engine(onnx_path, engine_path, min_shape, max_shape, opt_shape) + + trt_wan_model = load_tensorrt_engine(engine_path, model_type=model_type) + policy.trained_model.action_head.model = trt_wan_model + +def export_to_onnx_fp4(model, test_inputs, onnx_save_path, dynamic_axes=None): + from modelopt.torch._deploy.utils.torch_onnx import OnnxBytes + from modelopt.torch._deploy.utils.torch_onnx import get_onnx_bytes_and_metadata + + print("exporting to onnx fp4") + try: + onnx_bytes, _ = get_onnx_bytes_and_metadata(model=model, dummy_input=test_inputs, dynamic_axes=dynamic_axes) + onnx_model = OnnxBytes.from_bytes(onnx_bytes) + except Exception as e: + print(f"Error exporting model to ONNX: {e}") + return + save_dir = os.path.dirname(os.path.abspath(onnx_save_path)) + os.makedirs(save_dir, exist_ok=True) + for filename, file_bytes in onnx_model.onnx_model.items(): + file_path = os.path.join(save_dir, filename) + with open(file_path, "wb") as f: + f.write(file_bytes) + print(f"exported onnx to {file_path}") + + +def export_to_onnx( + pytorch_model, + test_inputs, + onnx_path="tensorrt/wan_model.onnx", + model_type="5B", + quantization_mode="fp8", + dynamic_axes=None, +): + # + if model_type == "5B": + return export_to_onnx_5B(pytorch_model, test_inputs, onnx_path, dynamic_axes) + elif model_type == "14B": + return export_to_onnx_14B(pytorch_model, test_inputs, onnx_path, dynamic_axes) + elif model_type == "ar_14B" or model_type == "ar_14B_droid": + return export_to_onnx_ar_14B(pytorch_model, test_inputs, onnx_path, dynamic_axes) + else: + raise ValueError(f"Model type {model_type} not supported") + + +def export_to_onnx_ar_14B(pytorch_model, test_inputs, onnx_path="tensorrt/wan_model.onnx", dynamic_axes=None): + """Export PyTorch model to ONNX""" + print("Exporting AR 14B model to ONNX...", onnx_path) + + # Create directory if it doesn't exist + os.makedirs(os.path.dirname(onnx_path), exist_ok=True) + pytorch_model.eval() + pytorch_model.to(torch.float16) + + input_names = [ + "x", + "timestep", + "context", + "kv_cache_packed", + "y", + "clip_feature", + "action", + "timestep_action", + "state", + ] + output_names = ["video_noise_pred", "action_noise_pred"] + + try: + with torch.no_grad(): + torch.onnx.export( + pytorch_model, + test_inputs, + onnx_path, + export_params=True, + opset_version=20, + do_constant_folding=True, + input_names=input_names, + output_names=output_names, + dynamic_axes=dynamic_axes, + ) + print(f" ONNX model exported to: {onnx_path}") + return onnx_path + + except Exception as e: + import traceback + print(f" ERROR: ONNX export failed. Exception type: {type(e)}") + print("Traceback:") + traceback.print_exc() + return None + + + +def export_to_onnx_5B(pytorch_model, test_inputs, onnx_path="tensorrt/wan_model.onnx"): + """Export PyTorch model to ONNX""" + print("Exporting model to ONNX...") + + # Create directory if it doesn't exist + os.makedirs(os.path.dirname(onnx_path), exist_ok=True) + pytorch_model.eval() + pytorch_model.to(torch.float16) + + x, action, timestep, context, state, embodiment_id = test_inputs + + # Define input names for better ONNX graph + input_names = ["x", "action", "timestep", "context", "state", "embodiment_id"] + output_names = ["video_noise_pred", "action_noise_pred"] + + try: + with torch.no_grad(): + torch.onnx.export( + pytorch_model, + (x, action, timestep, context, state, embodiment_id), + onnx_path, + export_params=True, + opset_version=20, + do_constant_folding=True, + input_names=input_names, + output_names=output_names, + ) + print(f" ONNX model exported to: {onnx_path}") + return onnx_path + + except Exception as e: + import traceback + print(f" ERROR: ONNX export failed. Exception type: {type(e)}") + print("Traceback:") + traceback.print_exc() + return None + + +def export_to_onnx_14B(pytorch_model, test_inputs, onnx_path="tensorrt/wan_model.onnx"): + """Export PyTorch model to ONNX""" + print("Exporting model to ONNX...") + + # Create directory if it doesn't exist + os.makedirs(os.path.dirname(onnx_path), exist_ok=True) + pytorch_model.eval() + pytorch_model.to(torch.float16) + + x, action, timestep, context, state, embodiment_id, clip_feature, y = test_inputs + + # Define input names for better ONNX graph + input_names = [ + "x", + "action", + "timestep", + "context", + "state", + "embodiment_id", + "clip_feature", + "y", + ] + output_names = ["video_noise_pred", "action_noise_pred"] + + try: + with torch.no_grad(): + torch.onnx.export( + pytorch_model, + (x, action, timestep, context, state, embodiment_id, clip_feature, y), + onnx_path, + export_params=True, + opset_version=20, + do_constant_folding=True, + input_names=input_names, + output_names=output_names, + ) + print(f" ONNX model exported to: {onnx_path}") + return onnx_path + + except Exception as e: + import traceback + print(f" ERROR: ONNX export failed. Exception type: {type(e)}") + print("Traceback:") + traceback.print_exc() + return None + + +def build_tensorrt_engine(onnx_path, engine_path="tensorrt/wan_model.trt", min_shape=None, max_shape=None, opt_shape=None): + """Build TensorRT engine from ONNX using trtexec""" + print("Building TensorRT engine with trtexec...") + + if not os.path.exists(onnx_path): + print(f" ERROR: ONNX file not found: {onnx_path}") + return None + + # Create directory if it doesn't exist + os.makedirs(os.path.dirname(engine_path), exist_ok=True) + + # Build engine using trtexec (much faster than torch_tensorrt) + trtexec_bin = shutil.which("trtexec") or "/opt/tensorrt/bin/trtexec" + cmd = [ + trtexec_bin, + f"--onnx={onnx_path}", + f"--saveEngine={engine_path}", + "--fp8", + "--fp16", + "--bf16", + "--separateProfileRun", + "--profilingVerbosity=detailed", + "--memPoolSize=workspace:65536", + "--dumpProfile", + "--dumpLayerInfo", + "--useCudaGraph", + "--verbose", + ] + + if min_shape is not None: + cmd.append(f"--minShapes={min_shape}") + if max_shape is not None: + cmd.append(f"--maxShapes={max_shape}") + if opt_shape is not None: + cmd.append(f"--optShapes={opt_shape}") + + # Create log file for trtexec output + log_file = engine_path.replace(".trt", "_build.log") + + try: + print(f" Running: {' '.join(cmd)}") + print(f" Logging output to: {log_file}") + + with open(log_file, "w") as f: + result = subprocess.run(cmd, stdout=f, stderr=subprocess.STDOUT, text=True, timeout=600) + + if result.returncode == 0: + print(f" TensorRT engine built successfully: {engine_path}") + print(f" Build log saved to: {log_file}") + return engine_path + else: + print(f" ERROR: trtexec failed with return code {result.returncode}") + print(f" Check build log for details: {log_file}") + # Print last few lines of log file for immediate feedback + try: + with open(log_file, "r") as f: + lines = f.readlines() + if lines: + print(" Last few lines from build log:") + for line in lines[-10:]: # Show last 10 lines + print(f" {line.rstrip()}") + except: + pass + return None + + except subprocess.TimeoutExpired: + print(" ERROR: trtexec timed out after 5 minutes") + print(f" Partial build log saved to: {log_file}") + return None + except Exception as e: + print(f" ERROR: Failed to run trtexec: {e}") + return None + + +def torch_type(trt_type): + mapping = { + trt.float32: torch.float32, # Added missing FLOAT mapping + trt.float16: torch.float16, + trt.bfloat16: torch.bfloat16, + trt.int8: torch.int8, + trt.int32: torch.int32, + trt.bool: torch.bool, + trt.uint8: torch.uint8, + trt.int64: torch.int64, + } + if trt_type in mapping: + return mapping[trt_type] + + raise TypeError( + f"Could not resolve TensorRT datatype to an equivalent torch datatype. {trt_type}" + ) + + +class Engine(object): + def __init__(self, file, plugins=[]): + super().__init__() + + self.logger = trt.Logger(trt.Logger.ERROR) + trt.init_libnvinfer_plugins(self.logger, "") + + self.plugins = [ctypes.CDLL(plugin, ctypes.RTLD_GLOBAL) for plugin in plugins] + self.file = file + self.load(file) + + def destroy(self): + del self.execution_context + del self.handle + + atexit.register(destroy, self) + self.print() + + def print(self): + + print("============= TRT Engine Detail =============") + print(f"Engine file: {self.file}") + print(f"Inputs: {len(self.in_meta)}") + for ib, item in enumerate(self.in_meta): + tensor_name, shape, dtype = item[:3] + print(f" {ib}. {tensor_name}: {'x'.join(map(str, shape))} [{dtype}]") + + print(f"Outputs: {len(self.out_meta)}") + for ib, item in enumerate(self.out_meta): + tensor_name, shape, dtype = item[:3] + print(f" {ib}. {tensor_name}: {'x'.join(map(str, shape))} [{dtype}]") + print("=============================================") + + def load(self, file): + runtime = trt.Runtime(self.logger) + + with open(file, "rb") as f: + self.handle = runtime.deserialize_cuda_engine(f.read()) + assert ( + self.handle is not None + ), f"Failed to deserialize the cuda engine from file: {file}" + + self.execution_context = self.handle.create_execution_context() + self.meta, self.in_meta, self.out_meta = [], [], [] + for tensor_name in self.handle: + shape = self.handle.get_tensor_shape(tensor_name) + print(f"Tensor name: {tensor_name}, shape: {shape}") + dtype = torch_type(self.handle.get_tensor_dtype(tensor_name)) + if self.handle.get_tensor_mode(tensor_name) == trt.TensorIOMode.INPUT: + self.in_meta.append([tensor_name, shape, dtype]) + else: + self.out_meta.append([tensor_name, shape, dtype]) + + def __call__(self, *args, **inputs): + return self.forward(*args, **inputs) + + def set_runtime_tensor_shape(self, name, shape): + self.execution_context.set_input_shape(name, shape) + + def forward(self, *args, **kwargs): + return_list = kwargs.pop("return_list", False) + reference_tensors = [] + stream = torch.cuda.current_stream() + for iarg, x in enumerate(args): + name, shape, dtype = self.in_meta[iarg] + runtime_shape = self.execution_context.get_tensor_shape(name) + assert isinstance(x, torch.Tensor), f"Unsupported tensor type: {type(x)}" + assert runtime_shape == x.shape, f"Invalid input shape: {runtime_shape} != {x.shape}" + assert ( + dtype == x.dtype + ), f"Invalid tensor dtype, excepted dtype is {dtype}, but got {x.dtype}" + assert x.is_cuda, f"Invalid tensor device, excepted device is cuda, but got {x.device}" + x = x.cuda().contiguous() + self.execution_context.set_tensor_address(name, x.data_ptr()) + reference_tensors.append(x) + + for name, shape, dtype in self.in_meta: + if name not in kwargs: + continue + + runtime_shape = self.execution_context.get_tensor_shape(name) + x = kwargs[name] + assert isinstance(x, torch.Tensor), f"Unsupported tensor[{name}] type: {type(x)}" + assert ( + runtime_shape == x.shape + ), f"Invalid input[{name}] shape: {x.shape}, but the expected shape is: {runtime_shape}" + assert ( + dtype == x.dtype + ), f"Invalid tensor[{name}] dtype, expected dtype is {dtype}, but got {x.dtype}" + assert ( + x.is_cuda + ), f"Invalid tensor[{name}] device, expected device is cuda, but got {x.device}" + x = x.cuda().contiguous() + self.execution_context.set_tensor_address(name, x.data_ptr()) + reference_tensors.append(x) + + for item in self.out_meta: + name = item[0] + runtime_shape = self.execution_context.get_tensor_shape(name) + output_tensor = torch.zeros( + *runtime_shape, dtype=item[2], device=reference_tensors[0].device + ) + self.execution_context.set_tensor_address(name, output_tensor.data_ptr()) + reference_tensors.append(output_tensor) + + self.execution_context.execute_async_v3(stream.cuda_stream) + stream.synchronize() + assert len(reference_tensors) == len(self.in_meta) + len( + self.out_meta + ), f"Invalid input tensors. The expected I/O tensors are {len(self.in_meta) + len(self.out_meta)}, but got {len(reference_tensors)}" + + if return_list: + return [ + reference_tensors[len(self.in_meta) + i] for i, item in enumerate(self.out_meta) + ] + else: + return { + item[0]: reference_tensors[len(self.in_meta) + i] + for i, item in enumerate(self.out_meta) + } + + +class WanTrtModel5B(torch.nn.Module): + def __init__(self, eng_path: str): + super().__init__() + self.engine = Engine(eng_path) + + def forward( + self, + x: torch.Tensor, + action: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + state: torch.Tensor, + embodiment_id: torch.Tensor, + ): + + self.engine.set_runtime_tensor_shape("x", x.shape) + self.engine.set_runtime_tensor_shape("action", action.shape) + self.engine.set_runtime_tensor_shape("context", context.shape) + self.engine.set_runtime_tensor_shape("state", state.shape) + + output = self.engine( + x=x.to(torch.float16), + action=action.to(torch.float16), + timestep=timestep.to(torch.float16), + context=context.to(torch.float16), + state=state.to(torch.float16), + embodiment_id=embodiment_id.to(torch.int32), + ) + if "out.0" in output: # for nvfp4 model export through modelopt + return output["out.0"].to(torch.bfloat16).contiguous(), output["out.1"].to(torch.bfloat16).contiguous() + else: + return output["video_noise_pred"].to(torch.bfloat16).contiguous(), output["action_noise_pred"].to(torch.bfloat16).contiguous() + + +class WanTrtModel14B(torch.nn.Module): + def __init__(self, eng_path: str): + super().__init__() + self.engine = Engine(eng_path) + + def forward( + self, + x: torch.Tensor, + action: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + state: torch.Tensor, + embodiment_id: torch.Tensor, + clip_feature: torch.Tensor, + y: torch.Tensor, + ): + + self.engine.set_runtime_tensor_shape("x", x.shape) + self.engine.set_runtime_tensor_shape("action", action.shape) + self.engine.set_runtime_tensor_shape("context", context.shape) + self.engine.set_runtime_tensor_shape("state", state.shape) + self.engine.set_runtime_tensor_shape("clip_feature", clip_feature.shape) + self.engine.set_runtime_tensor_shape("y", y.shape) + + output = self.engine( + x=x.to(torch.float16), + action=action.to(torch.float16), + timestep=timestep.to(torch.float16), + context=context.to(torch.float16), + state=state.to(torch.float16), + embodiment_id=embodiment_id.to(torch.int32), + clip_feature=clip_feature.to(torch.float16), + y=y.to(torch.float16), + ) + if "out.0" in output: # for nvfp4 model export through modelopt + return output["out.0"].to(torch.bfloat16).contiguous(), output["out.1"].to(torch.bfloat16).contiguous() + else: + return output["video_noise_pred"].to(torch.bfloat16).contiguous(), output["action_noise_pred"].to(torch.bfloat16).contiguous() + + +class WanTrtModelAr5B(torch.nn.Module): + """TRT wrapper for ar_5B_n6 model type - uses kv_cache but no clip_feature.""" + def __init__(self, eng_path: str): + super().__init__() + self.engine = Engine(eng_path) + + def forward( + self, + x, + timestep, + context, + kv_cache: list[torch.Tensor], + y=None, + action=None, + timestep_action=None, + state=None, + ): + + kv_cache_packed = torch.stack(kv_cache, dim=0) + + self.engine.set_runtime_tensor_shape("x", x.shape) + self.engine.set_runtime_tensor_shape("timestep", timestep.shape) + self.engine.set_runtime_tensor_shape("context", context.shape) + self.engine.set_runtime_tensor_shape("kv_cache_packed", kv_cache_packed.shape) + # self.engine.set_runtime_tensor_shape("y", y.shape) + self.engine.set_runtime_tensor_shape("action", action.shape) + self.engine.set_runtime_tensor_shape("timestep_action", timestep_action.shape) + self.engine.set_runtime_tensor_shape("state", state.shape) + + + output = self.engine( + x.to(torch.float16), + timestep.to(torch.float16), + context.to(torch.float16), + kv_cache_packed.to(torch.float16), + # y.to(torch.float16), + action.to(torch.float16), + timestep_action.to(torch.float16), + state.to(torch.float16), + ) + + if "out.0" in output: # for nvfp4 model export through modelopt + return output["out.0"].to(torch.bfloat16).contiguous(), output["out.1"].to(torch.bfloat16).contiguous() + else: + return output["video_noise_pred"].to(torch.bfloat16).contiguous(), output["action_noise_pred"].to(torch.bfloat16).contiguous() + + +class WanTrtModelAr14B(torch.nn.Module): + def __init__(self, eng_path: str): + super().__init__() + self.engine = Engine(eng_path) + + def forward( + self, + x, + timestep, + context, + kv_cache: list[torch.Tensor], + y=None, + clip_feature=None, + action=None, + timestep_action=None, + state=None, + ): + + kv_cache_packed = torch.stack(kv_cache, dim=0) + + self.engine.set_runtime_tensor_shape("x", x.shape) + self.engine.set_runtime_tensor_shape("timestep", timestep.shape) + self.engine.set_runtime_tensor_shape("context", context.shape) + self.engine.set_runtime_tensor_shape("kv_cache_packed", kv_cache_packed.shape) + self.engine.set_runtime_tensor_shape("y", y.shape) + self.engine.set_runtime_tensor_shape("clip_feature", clip_feature.shape) + self.engine.set_runtime_tensor_shape("action", action.shape) + self.engine.set_runtime_tensor_shape("timestep_action", timestep_action.shape) + self.engine.set_runtime_tensor_shape("state", state.shape) + + + output = self.engine( + x.to(torch.float16), + timestep.to(torch.float16), + context.to(torch.float16), + kv_cache_packed.to(torch.float16), + y.to(torch.float16), + clip_feature.to(torch.float16), + action.to(torch.float16), + timestep_action.to(torch.float16), + state.to(torch.float16), + ) + + if "out.0" in output: # for nvfp4 model export through modelopt + return output["out.0"].to(torch.bfloat16).contiguous(), output["out.1"].to(torch.bfloat16).contiguous() + else: + return output["video_noise_pred"].to(torch.bfloat16).contiguous(), output["action_noise_pred"].to(torch.bfloat16).contiguous() + +def load_tensorrt_engine(engine_path="tensorrt/wan_model.trt", model_type="5B"): + """Load TensorRT engine""" + if model_type == "5B": + trt_inference = WanTrtModel5B(engine_path) + elif model_type == "ar_5B_n6" or model_type == "ar_5B": + trt_inference = WanTrtModelAr5B(engine_path) + elif model_type == "14B": + trt_inference = WanTrtModel14B(engine_path) + elif model_type == "ar_14B" or model_type == "ar_14B_droid": + trt_inference = WanTrtModelAr14B(engine_path) + else: + raise ValueError(f"Model type {model_type} not supported") + return trt_inference + + +def create_wan_test_inputs(policy, device="cuda", model_type="5B"): + # Get dtype from model parameters + dtype = torch.float16 + + # Use hardcoded dimensions from the original working version of the script + if model_type == "5B": + x = torch.randn(1, 48, 13, 22, 40, dtype=dtype, device=device) + action = torch.randn(1, 48, 32, dtype=dtype, device=device) + timestep = torch.randn(1, dtype=dtype, device=device) + context = torch.randn(1, 512, 4096, dtype=dtype, device=device) + state = torch.randn(1, 1, 64, dtype=dtype, device=device) + embodiment_id = torch.zeros(1, dtype=torch.int32, device=device) + timestep_action = torch.randn(1, 48, dtype=dtype, device=device) + seq_len = torch.tensor(440, dtype=torch.int32, device=device) + return x, action, timestep, context, state, embodiment_id, timestep_action, seq_len + elif model_type == "ar_5B_n6": + # ar_5B_n6 uses _forward_inference_trt which requires kv_cache_packed + # Shape from dynamic_axes: kv_cache_packed:30x2x1x220x24x128 + # Note: 5B model doesn't use clip_feature (unlike 14B), but still needs y + x = torch.randn(1, 48, 2, 22, 40, dtype=dtype, device=device) + timestep = torch.randn(1, 2, dtype=dtype, device=device) + context = torch.randn(1, 512, 4096, dtype=dtype, device=device) + # y = torch.randn(1, 52, 2, 22, 40, dtype=dtype, device=device) # y is required by _forward_inference_trt + action = torch.randn(1, 48, 32, dtype=dtype, device=device) + timestep_action = torch.randn(1, 48, dtype=dtype, device=device) + state = torch.randn(1, 1, 64, dtype=dtype, device=device) + + num_heads = 24 + head_dim = 128 + num_layers = 30 + B = 1 + + kv_cache = [] + for _ in range(num_layers): + kv_cache.append( + torch.zeros([2, B, 13*220, num_heads, head_dim], dtype=dtype, device=device) + ) + + kv_cache_packed = torch.stack(kv_cache, dim=0) + # Return order matches _forward_inference_trt signature: x, timestep, context, kv_cache_packed, y, action, timestep_action, state + return (x, timestep, context, kv_cache_packed, action, timestep_action, state) + elif model_type == "14B": + x = torch.randn(1, 16, 13, 44, 80, dtype=dtype, device=device) + action = torch.randn(1, 48, 32, dtype=dtype, device=device) + timestep = torch.randn(1, dtype=dtype, device=device) + context = torch.randn(1, 512, 4096, dtype=dtype, device=device) + state = torch.randn(1, 1, 64, dtype=dtype, device=device) + embodiment_id = torch.zeros(1, dtype=torch.int32, device=device) + clip_feature = torch.randn(1, 257, 1280, dtype=dtype, device=device) + y = torch.randn(1, 20, 13, 44, 80, dtype=dtype, device=device) + return x, action, timestep, context, state, embodiment_id, clip_feature, y + elif model_type == "ar_14B": + clip_feature = torch.randn(1, 257, 1280, dtype=dtype, device=device) + y = torch.randn(1, 20, 2, 44, 80, dtype=dtype, device=device) + timestep_action = torch.randn(1, 48, dtype=dtype, device=device) + x = torch.randn(1, 16, 2, 44, 80, dtype=dtype, device=device) + timestep = torch.randn(1, 2, dtype=dtype, device=device) + context = torch.randn(1, 512, 4096, dtype=dtype, device=device) + seq_len = torch.tensor(1760, dtype=torch.int32, device=device) + action = torch.randn(1, 48, 32, dtype=dtype, device=device) + state = torch.randn(1, 1, 64, dtype=dtype, device=device) + embodiment_id = torch.zeros(1, dtype=torch.int32, device=device) + + num_heads = 40 + head_dim = 5120 // num_heads + num_layers = 40 + B = 1 + + kv_cache = [] + for _ in range(num_layers): + kv_cache.append( + torch.zeros([2, B, 9*880, num_heads, head_dim], dtype=dtype, device=device) + ) + + crossattn_k_cache = [] + for _ in range(num_layers): + crossattn_k_cache.append( + torch.zeros([2, B, 9*880, num_heads, head_dim], dtype=dtype, device=device) + ) + kv_cache_packed = torch.stack(kv_cache, dim=0) + crossattn_packed = torch.stack(crossattn_k_cache, dim=0) + return (x, timestep, context, kv_cache_packed, y, clip_feature, action, timestep_action, state) + elif model_type == "ar_14B_droid": + clip_feature = torch.randn(1, 257, 1280, dtype=dtype, device=device) + y = torch.randn(1, 20, 2, 44, 80, dtype=dtype, device=device) + timestep_action = torch.randn(1, 24, dtype=dtype, device=device) + x = torch.randn(1, 16, 2, 44, 80, dtype=dtype, device=device) + timestep = torch.randn(1, 2, dtype=dtype, device=device) + context = torch.randn(1, 512, 4096, dtype=dtype, device=device) + seq_len = torch.tensor(1760, dtype=torch.int32, device=device) + action = torch.randn(1, 24, 32, dtype=dtype, device=device) + state = torch.randn(1, 1, 64, dtype=dtype, device=device) + embodiment_id = torch.zeros(1, dtype=torch.int32, device=device) + + num_heads = 40 + head_dim = 5120 // num_heads + num_layers = 40 + B = 1 + + kv_cache = [] + for _ in range(num_layers): + kv_cache.append( + torch.zeros([2, B, 9*880, num_heads, head_dim], dtype=dtype, device=device) + ) + + crossattn_k_cache = [] + for _ in range(num_layers): + crossattn_k_cache.append( + torch.zeros([2, B, 9*880, num_heads, head_dim], dtype=dtype, device=device) + ) + kv_cache_packed = torch.stack(kv_cache, dim=0) + crossattn_packed = torch.stack(crossattn_k_cache, dim=0) + return (x, timestep, context, kv_cache_packed, y, clip_feature, action, timestep_action, state) + + diff --git a/groot/vla/__init__.py b/groot/vla/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/common/__init__.py b/groot/vla/common/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/groot/vla/common/__init__.py @@ -0,0 +1 @@ + diff --git a/groot/vla/common/utils/__init__.py b/groot/vla/common/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..762cc7c5bd7c6231b10f1155ed1f3c83ae933ac0 --- /dev/null +++ b/groot/vla/common/utils/__init__.py @@ -0,0 +1,3 @@ +from .data_structure import * # noqa: F403 +from .io import * # noqa: F403 +from .misc import * # noqa: F403 diff --git a/groot/vla/common/utils/data_structure/__init__.py b/groot/vla/common/utils/data_structure/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c3e40684ae7be6d8a216204a4d89b6c5352a72f2 --- /dev/null +++ b/groot/vla/common/utils/data_structure/__init__.py @@ -0,0 +1,2 @@ +from .shape_utils import * # noqa: F403 +from .tree_utils import * # noqa: F403 diff --git a/groot/vla/common/utils/data_structure/shape_utils.py b/groot/vla/common/utils/data_structure/shape_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..343524c5db1247a6c27f8322b983a5058ddaa17a --- /dev/null +++ b/groot/vla/common/utils/data_structure/shape_utils.py @@ -0,0 +1,283 @@ +""" +Shape inference methods +""" + +from functools import partial +import math +from typing import List, Tuple, Union +import warnings + +import numpy as np +import torch + +# fmt: off +__all__ = [ + "shape_convnd", + "shape_conv1d", "shape_conv2d", "shape_conv3d", + "shape_transpose_convnd", + "shape_transpose_conv1d", "shape_transpose_conv2d", "shape_transpose_conv3d", + "shape_poolnd", + "shape_maxpool1d", "shape_maxpool2d", "shape_maxpool3d", + "shape_avgpool1d", "shape_avgpool2d", "shape_avgpool3d", + "shape_slice", + "check_shape" +] +# fmt: on + + +def _get_shape(x): + "single object" + if isinstance(x, np.ndarray): + return tuple(x.shape) + else: + return tuple(x.size()) + + +def _expands(dim, *xs): + "repeat vars like kernel and stride to match dim" + + def _expand(x): + if isinstance(x, int): + return (x,) * dim + else: + assert len(x) == dim + return x + + return map(lambda x: _expand(x), xs) + + +_HELPER_TENSOR = torch.zeros((1,)) + + +def shape_slice(input_shape, slice): + """ + Credit to Adam Paszke for the trick. Shape inference without instantiating + an actual tensor. + The key is that `.expand()` does not actually allocate memory + Still needs to allocate a one-element HELPER_TENSOR. + """ + shape = _HELPER_TENSOR.expand(*input_shape)[slice] + if hasattr(shape, "size"): + return tuple(shape.size()) + return (1,) + + +class ShapeSlice: + """ + shape_slice inference with easy []-operator + """ + + def __init__(self, input_shape): + self.input_shape = input_shape + + def __getitem__(self, slice): + return shape_slice(self.input_shape, slice) + + +def check_shape( + value: Union[Tuple, List, torch.Tensor, np.ndarray], + expected: Union[Tuple, List, torch.Tensor, np.ndarray], + err_msg="", + mode="raise", +): + """ + Args: + value: np array or torch Tensor + expected: + - list[int], tuple[int]: if any value is None, will match any dim + - np array or torch Tensor: must have the same dimensions + mode: + - "raise": raise ValueError, shape mismatch + - "return": returns True if shape matches, otherwise False + - "warning": warnings.warn + """ + assert mode in ["raise", "return", "warning"] + if torch.is_tensor(value): + actual_shape = value.size() + elif hasattr(value, "shape"): + actual_shape = value.shape + else: + assert isinstance(value, (list, tuple)) + actual_shape = value + assert all( + isinstance(s, int) for s in actual_shape + ), f"actual shape: {actual_shape} is not a list of ints" + + if torch.is_tensor(expected): + expected_shape = expected.size() + elif hasattr(expected, "shape"): + expected_shape = expected.shape + else: + assert isinstance(expected, (list, tuple)) + expected_shape = expected + + err_msg = f" for {err_msg}" if err_msg else "" + + if len(actual_shape) != len(expected_shape): + err_msg = ( + f"Dimension mismatch{err_msg}: actual shape {actual_shape} " + f"!= expected shape {expected_shape}." + ) + if mode == "raise": + raise ValueError(err_msg) + elif mode == "warning": + warnings.warn(err_msg) + return False + + for s_a, s_e in zip(actual_shape, expected_shape): + if s_e is not None and s_a != s_e: + err_msg = ( + f"Shape mismatch{err_msg}: actual shape {actual_shape} " + f"!= expected shape {expected_shape}." + ) + if mode == "raise": + raise ValueError(err_msg) + elif mode == "warning": + warnings.warn(err_msg) + return False + return True + + +def shape_convnd( + dim, + input_shape, + out_channels, + kernel_size, + stride=1, + padding=0, + dilation=1, + has_batch=False, +): + """ + http://pytorch.org/docs/nn.html#conv1d + http://pytorch.org/docs/nn.html#conv2d + http://pytorch.org/docs/nn.html#conv3d + + Args: + dim: supports 1D to 3D + input_shape: + - 1D: [channel, length] + - 2D: [channel, height, width] + - 3D: [channel, depth, height, width] + has_batch: whether the first dim is batch size or not + """ + if has_batch: + assert ( + len(input_shape) == dim + 2 + ), "input shape with batch should be {}-dimensional".format(dim + 2) + else: + assert ( + len(input_shape) == dim + 1 + ), "input shape without batch should be {}-dimensional".format(dim + 1) + if stride is None: + # for pooling convention in PyTorch + stride = kernel_size + kernel_size, stride, padding, dilation = _expands(dim, kernel_size, stride, padding, dilation) + if has_batch: + batch = input_shape[0] + input_shape = input_shape[1:] + else: + batch = None + _, *img = input_shape + new_img_shape = [ + math.floor( + (img[i] + 2 * padding[i] - dilation[i] * (kernel_size[i] - 1) - 1) // stride[i] + 1 + ) + for i in range(dim) + ] + return ((batch,) if has_batch else ()) + (out_channels, *new_img_shape) + + +def shape_poolnd( + dim, input_shape, kernel_size, stride=None, padding=0, dilation=1, has_batch=False +): + """ + The only difference from infer_shape_convnd is that `stride` default is None + """ + if has_batch: + out_channels = input_shape[1] + else: + out_channels = input_shape[0] + return shape_convnd( + dim, + input_shape, + out_channels, + kernel_size, + stride, + padding, + dilation, + has_batch, + ) + + +def shape_transpose_convnd( + dim, + input_shape, + out_channels, + kernel_size, + stride=1, + padding=0, + output_padding=0, + dilation=1, + has_batch=False, +): + """ + http://pytorch.org/docs/nn.html#convtranspose1d + http://pytorch.org/docs/nn.html#convtranspose2d + http://pytorch.org/docs/nn.html#convtranspose3d + + Args: + dim: supports 1D to 3D + input_shape: + - 1D: [channel, length] + - 2D: [channel, height, width] + - 3D: [channel, depth, height, width] + has_batch: whether the first dim is batch size or not + """ + if has_batch: + assert ( + len(input_shape) == dim + 2 + ), "input shape with batch should be {}-dimensional".format(dim + 2) + else: + assert ( + len(input_shape) == dim + 1 + ), "input shape without batch should be {}-dimensional".format(dim + 1) + kernel_size, stride, padding, output_padding, dilation = _expands( + dim, kernel_size, stride, padding, output_padding, dilation + ) + if has_batch: + batch = input_shape[0] + input_shape = input_shape[1:] + else: + batch = None + _, *img = input_shape + new_img_shape = [ + (img[i] - 1) * stride[i] - 2 * padding[i] + kernel_size[i] + output_padding[i] + for i in range(dim) + ] + return ((batch,) if has_batch else ()) + (out_channels, *new_img_shape) + + +shape_conv1d = partial(shape_convnd, 1) +shape_conv2d = partial(shape_convnd, 2) +shape_conv3d = partial(shape_convnd, 3) + + +shape_transpose_conv1d = partial(shape_transpose_convnd, 1) +shape_transpose_conv2d = partial(shape_transpose_convnd, 2) +shape_transpose_conv3d = partial(shape_transpose_convnd, 3) + + +shape_maxpool1d = partial(shape_poolnd, 1) +shape_maxpool2d = partial(shape_poolnd, 2) +shape_maxpool3d = partial(shape_poolnd, 3) + + +""" +http://pytorch.org/docs/nn.html#avgpool1d +http://pytorch.org/docs/nn.html#avgpool2d +http://pytorch.org/docs/nn.html#avgpool3d +""" +shape_avgpool1d = partial(shape_maxpool1d, dilation=1) +shape_avgpool2d = partial(shape_maxpool2d, dilation=1) +shape_avgpool3d = partial(shape_maxpool3d, dilation=1) diff --git a/groot/vla/common/utils/data_structure/tree_utils.py b/groot/vla/common/utils/data_structure/tree_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9cce53b6c14b1e7cf4ccd93790e964f47b2c84d8 --- /dev/null +++ b/groot/vla/common/utils/data_structure/tree_utils.py @@ -0,0 +1,219 @@ +""" +Utils to handle nested data structures + +Install dm_tree first: +https://tree.readthedocs.io/en/latest/api.html +""" + +import collections +from typing import Any, Iterable, List, Tuple, TypeVar + +import numpy as np + +try: + import tree + +except ImportError: + raise ImportError("Please install dm_tree first: `pip install dm_tree`") + + +def is_sequence(obj): + """ + Returns: + True if the sequence is a collections.Sequence and not a string. + """ + return isinstance(obj, collections.abc.Sequence) and not isinstance(obj, str) + + +def is_mapping(obj): + """ + Returns: + True if the sequence is a collections.Mapping + """ + return isinstance(obj, collections.abc.Mapping) + + +def tree_value_at_path(obj, paths: Tuple): + try: + for p in paths: + obj = obj[p] + return obj + except Exception as e: + raise ValueError(f"{e}\n\n-- Incorrect nested path {paths} for object: {obj}.") + + +def tree_assign_at_path(obj, paths: Tuple, value): + try: + for p in paths[:-1]: + obj = obj[p] + if len(paths) > 0: + obj[paths[-1]] = value + except Exception as e: + raise ValueError(f"{e}\n\n-- Incorrect nested path {paths} for object: {obj}.") + + +def copy_non_leaf(obj): + """ + Deepcopy the nested structure, but does NOT copy the leaf values like Tensors + """ + return tree.map_structure(lambda x: x, obj) + + +# ======================================================================= +# Copyright 2018 DeepMind Technologies Limited. 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. + +# Tensor framework-agnostic utilities for manipulating nested structures. + +ElementType = TypeVar("ElementType") + + +def fast_map_structure(func, *structure): + """Faster map_structure implementation which skips some error checking.""" + flat_structure = (tree.flatten(s) for s in structure) + entries = zip(*flat_structure) + # Arbitrarily choose one of the structures of the original sequence (the last) + # to match the structure for the flattened sequence. + return tree.unflatten_as(structure[-1], [func(*x) for x in entries]) + + +def stack_sequence_fields(sequence: Iterable[ElementType]) -> ElementType: + """Stacks a list of identically nested objects. + + This takes a sequence of identically nested objects and returns a single + nested object whose ith leaf is a stacked numpy array of the corresponding + ith leaf from each element of the sequence. + + For example, if `sequence` is: + + ```python + [{ + 'action': np.array([1.0]), + 'observation': (np.array([0.0, 1.0, 2.0]),), + 'reward': 1.0 + }, { + 'action': np.array([0.5]), + 'observation': (np.array([1.0, 2.0, 3.0]),), + 'reward': 0.0 + }, { + 'action': np.array([0.3]),1 + 'observation': (np.array([2.0, 3.0, 4.0]),), + 'reward': 0.5 + }] + ``` + + Then this function will return: + + ```python + { + 'action': np.array([....]) # array shape = [3 x 1] + 'observation': (np.array([...]),) # array shape = [3 x 3] + 'reward': np.array([...]) # array shape = [3] + } + ``` + + Note that the 'observation' entry in the above example has two levels of + nesting, i.e it is a tuple of arrays. + + Args: + sequence: a list of identically nested objects. + + Returns: + A nested object with numpy. + + Raises: + ValueError: If `sequence` is an empty sequence. + """ + # Handle empty input sequences. + if not sequence: + raise ValueError("Input sequence must not be empty") + + # Default to asarray when arrays don't have the same shape to be compatible + # with old behaviour. + try: + return fast_map_structure(lambda *values: np.stack(values), *sequence) + except ValueError: + return fast_map_structure(lambda *values: np.asarray(values), *sequence) + + +def unstack_sequence_fields(struct: ElementType, batch_size: int) -> List[ElementType]: + """Converts a struct of batched arrays to a list of structs. + + This is effectively the inverse of `stack_sequence_fields`. + + Args: + struct: An (arbitrarily nested) structure of arrays. + batch_size: The length of the leading dimension of each array in the struct. + This is assumed to be static and known. + + Returns: + A list of structs with the same structure as `struct`, where each leaf node + is an unbatched element of the original leaf node. + """ + + return [tree.map_structure(lambda s, i=i: s[i], struct) for i in range(batch_size)] + + +def broadcast_structures(*args: Any) -> Any: + """Returns versions of the arguments that give them the same nested structure. + + Any nested items in *args must have the same structure. + + Any non-nested item will be replaced with a nested version that shares that + structure. The leaves will all be references to the same original non-nested + item. + + If all *args are nested, or all *args are non-nested, this function will + return *args unchanged. + + Example: + ``` + a = ('a', 'b') + b = 'c' + tree_a, tree_b = broadcast_structure(a, b) + tree_a + > ('a', 'b') + tree_b + > ('c', 'c') + ``` + + Args: + *args: A Sequence of nested or non-nested items. + + Returns: + `*args`, except with all items sharing the same nest structure. + """ + if not args: + return + + reference_tree = None + for arg in args: + if tree.is_nested(arg): + reference_tree = arg + break + + if reference_tree is None: + reference_tree = args[0] + + def mirror_structure(value, reference_tree): + if tree.is_nested(value): + # Use check_types=True so that the types of the trees we construct aren't + # dependent on our arbitrary choice of which nested arg to use as the + # reference_tree. + tree.assert_same_structure(value, reference_tree, check_types=True) + return value + else: + return tree.map_structure(lambda _: value, reference_tree) + + return tuple(mirror_structure(arg, reference_tree) for arg in args) diff --git a/groot/vla/common/utils/io/__init__.py b/groot/vla/common/utils/io/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8dc853ff2849d027d12c92d09d2ebaa553cfd97e --- /dev/null +++ b/groot/vla/common/utils/io/__init__.py @@ -0,0 +1,6 @@ +from .config_utils import * # noqa: F403 +from .file_utils import * # noqa: F403 +from .hdf5_utils import * # noqa: F403 +from .json_utils import * # noqa: F403 +from .print_utils import * # noqa: F403 +from .termcolor import * # noqa: F403 diff --git a/groot/vla/common/utils/io/config_utils.py b/groot/vla/common/utils/io/config_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d8d541efa9c55107c37c0cd3bcb32566b0ce2f8b --- /dev/null +++ b/groot/vla/common/utils/io/config_utils.py @@ -0,0 +1,260 @@ +from copy import deepcopy +import importlib.resources +import os +import sys + +import hydra +from omegaconf import DictConfig, OmegaConf +import tree + +from ..misc.functional_utils import call_once, is_mapping, is_sequence, meta_decorator +from .print_utils import to_scientific_str + +_CLASS_REGISTRY = {} # for instantiation + + +def resource_file_path(pkg_name, fname) -> str: + with importlib.resources.path(pkg_name, fname) as p: + return str(p) + + +def print_config(cfg: DictConfig): + print(cfg.pretty(resolve=True)) + + +def is_hydra_initialized(): + return hydra.utils.HydraConfig.initialized() + + +def hydra_config(): + # https://github.com/facebookresearch/hydra/issues/377 + # HydraConfig() is a singleton + if is_hydra_initialized(): + return hydra.utils.HydraConfig().cfg.hydra + else: + return None + + +def hydra_override_arg_list() -> list[str]: + """ + Returns: + list ["lr=0.2", "batch=64", ...] + """ + if is_hydra_initialized(): + return hydra_config().overrides.task + else: + return [] + + +def hydra_override_name(): + if is_hydra_initialized(): + return hydra_config().job.override_dirname + else: + return "" + + +def hydra_original_dir(*subpaths): + return os.path.join(hydra.utils.get_original_cwd(), *subpaths) + + +@call_once(on_second_call="noop") +def register_omegaconf_resolvers(): + import numpy as np + + OmegaConf.register_new_resolver("scientific", lambda v, i=0: to_scientific_str(v, i)) + OmegaConf.register_new_resolver("_optional", lambda v: f"_{v}" if v else "") + OmegaConf.register_new_resolver("optional_", lambda v: f"{v}_" if v else "") + OmegaConf.register_new_resolver("_optional_", lambda v: f"_{v}_" if v else "") + OmegaConf.register_new_resolver("__optional", lambda v: f"__{v}" if v else "") + OmegaConf.register_new_resolver("optional__", lambda v: f"{v}__" if v else "") + OmegaConf.register_new_resolver("__optional__", lambda v: f"__{v}__" if v else "") + OmegaConf.register_new_resolver("iftrue", lambda cond, v_default: cond if cond else v_default) + OmegaConf.register_new_resolver("ifelse", lambda cond, v1, v2="": v1 if cond else v2) + OmegaConf.register_new_resolver( + "ifequal", lambda query, key, v1, v2: v1 if query == key else v2 + ) + OmegaConf.register_new_resolver("intbool", lambda cond: 1 if cond else 0) + OmegaConf.register_new_resolver("mult", lambda *x: np.prod(x).tolist()) + OmegaConf.register_new_resolver("add", lambda *x: sum(x)) + OmegaConf.register_new_resolver("div", lambda x, y: x / y) + OmegaConf.register_new_resolver("intdiv", lambda x, y: x // y) + + # try each key until the key exists. Useful for multiple classes that have different + # names for the same key + def _try_key(cfg, *keys): + for k in keys: + if k in cfg: + return cfg[k] + raise KeyError(f"no key in {keys} is valid") + + OmegaConf.register_new_resolver("trykey", _try_key) + # replace `resnet.gn.ws` -> `resnet_gn_ws`, because omegaconf doesn't support + # keys with dots. Useful for generating run name with dots + OmegaConf.register_new_resolver("underscore_to_dots", lambda s: s.replace("_", ".")) + + def _no_instantiate(cfg): + cfg = deepcopy(cfg) + cfg[_NO_INSTANTIATE] = True + return cfg + + OmegaConf.register_new_resolver("no_instantiate", _no_instantiate) + + +# ======================================================== +# ================== Instantiation tools ================ +# ======================================================== + + +def register_callable(name, class_type): + if isinstance(class_type, str): + class_type, name = name, class_type + assert callable(class_type) + _CLASS_REGISTRY[name] = class_type + + +@meta_decorator +def register_class(cls, alias=None): + """ + Decorator + """ + assert callable(cls) + _CLASS_REGISTRY[cls.__name__] = cls + if alias: + assert is_sequence(alias) + for a in alias: + _CLASS_REGISTRY[str(a)] = cls + return cls + + +def omegaconf_to_dict(cfg, resolve: bool = True, enum_to_str: bool = False): + """ + Convert arbitrary nested omegaconf objects to primitive containers + + WARNING: cannot use tree lib because it gets confused on DictConfig and ListConfig + """ + kw = dict(resolve=resolve, enum_to_str=enum_to_str) + if OmegaConf.is_config(cfg): + return OmegaConf.to_container(cfg, **kw) + elif is_sequence(cfg): + return type(cfg)(omegaconf_to_dict(c, **kw) for c in cfg) + elif is_mapping(cfg): + return {k: omegaconf_to_dict(c, **kw) for k, c in cfg.items()} + else: + return cfg + + +def omegaconf_save(cfg, *paths: str, resolve: bool = True): + """ + Save omegaconf to yaml + """ + from .file_utils import f_join + + OmegaConf.save(cfg, f_join(*paths), resolve=resolve) + + +def get_class(path): + """ + First try to find the class in the registry first, + if it doesn't exist, use importlib to locate it + """ + if path in _CLASS_REGISTRY: + return _CLASS_REGISTRY[path] + else: + assert "." in path, ( + f"Because {path} is not found in class registry, " f"it must be a full module path" + ) + try: + from importlib import import_module + + module_path, _, class_name = path.rpartition(".") + mod = import_module(module_path) + try: + class_type = getattr(mod, class_name) + except AttributeError: + raise ImportError("Class {} is not in module {}".format(class_name, module_path)) + return class_type + except ValueError as e: + print("Error initializing class " + path, file=sys.stderr) + raise e + + +_DELETE_ARG = "__delete__" +_NO_INSTANTIATE = "__no_instantiate__" # return config as-is +_OMEGA_MISSING = "???" + + +def _get_instantiate_params(cfg, kwargs=None): + params = cfg + f_args, f_kwargs = (), {} + for k, value in params.items(): + if k in ["cls", "class"]: + continue + elif k == "*args": + assert is_sequence(value), '"*args" value must be a sequence' + f_args = list(value) + continue + if value == _OMEGA_MISSING: + if kwargs and k in kwargs: + value = kwargs[k] + else: + raise ValueError(f'Missing required keyword arg "{k}" in cfg: {cfg}') + if value == _DELETE_ARG: + continue + else: + f_kwargs[k] = value + return f_args, f_kwargs + + +def _instantiate_single(cfg): + if is_mapping(cfg) and ("cls" in cfg or "class" in cfg): + assert bool("cls" in cfg) != bool("class" in cfg), ( + "to instantiate from config, " + 'one and only one of "cls" or "class" key should be provided' + ) + if _NO_INSTANTIATE in cfg: + no_instantiate = cfg.pop(_NO_INSTANTIATE) + if no_instantiate: + cfg = deepcopy(cfg) + return cfg + else: + return _instantiate_single(cfg) + + cls = cfg.get("class", cfg.get("cls")) + args, kwargs = _get_instantiate_params(cfg) + try: + class_type = get_class(cls) + return class_type(*args, **kwargs) + except Exception as e: + raise RuntimeError(f"Error instantiating {cls}: {e}") + else: + return None + + +def instantiate(_cfg_, **kwargs): + """ + Any dict with "cls" or "class" key is considered instantiable. + + Any key that has the special value "__delete__" + will not be passed to the constructor + + **kwargs only apply to the top level object if it's a dict, otherwise raise error + """ + assert OmegaConf.is_config(_cfg_) or isinstance(_cfg_, (list, tuple)) or is_mapping(_cfg_), ( + '"cfg" must be a dict, list, tuple, or OmegaConf config to be instantiated. ' + f"Current its type is {type(_cfg_)}" + ) + + _cfg_ = omegaconf_to_dict(_cfg_, resolve=True) + + if kwargs: + if is_mapping(_cfg_): + _cfg_ = _cfg_.copy() + _cfg_.update(kwargs) + _cfg_ = {k: v for k, v in _cfg_.items() if v != _DELETE_ARG} + else: + raise RuntimeError( + f"**kwargs specified, but the top-level cfg is not a dict. " + f"It has type {type(_cfg_)}" + ) + + return tree.traverse(_instantiate_single, _cfg_, top_down=False) diff --git a/groot/vla/common/utils/io/file_utils.py b/groot/vla/common/utils/io/file_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..3e53246b39997b519b1f131a440c64d500c00db2 --- /dev/null +++ b/groot/vla/common/utils/io/file_utils.py @@ -0,0 +1,707 @@ +""" +File system utils. +""" + +import glob +import os +import pickle +import shutil +import sys +from typing import Callable, Union + +from ..data_structure.tree_utils import is_sequence + +__all__ = [ + "create_tar", + "dump_pickle", + "dump_text", + "dump_text_lines", + "extract_tar", + "f_add_ext", + "f_append_before_ext", + "f_copy", + "f_copytree", + "f_exists", + "f_expand", + "f_ext", + "f_glob", + "f_has_ext", + "f_join", + "f_listdir", + "f_mkdir", + "f_mkdir_in_path", + "f_move", + "f_not_empty", + "f_remove", + "f_size", + "f_split_path", + "f_time", + "get_dir", + "get_file_lock", + "get_package_root", + "get_parent_dir", + "get_script_dir", + "get_script_file_name", + "get_script_self_path", + "host_id", + "host_name", + "insert_before_ext", + "is_abs_path", + "is_dir", + "is_file", + "is_relative_path", + "last_part_in_path", + "load_pickle", + "load_text", + "load_text_lines", + "md5_checksum", + "move_with_backup", + "next_available_file_name", + "owner_name", + "pickle_dump", + "pickle_load", + "read_text", + "read_text_lines", + "text_dump", + "text_load", + "timestamp_file_name", + "utf_open", + "write_text", + "write_text_lines", +] + +f_ext = os.path.splitext + +f_size = os.path.getsize + +is_file = os.path.isfile + +is_dir = os.path.isdir + +get_dir = os.path.dirname + + +def owner_name(filepath): + """ + Returns: file owner name, unix only + """ + import pwd + + return pwd.getpwuid(os.stat(filepath).st_uid).pw_name + + +def host_name(): + "Get host name, alias with ``socket.gethostname()``" + from socket import gethostname + + return gethostname() + + +def host_id(): + """ + Returns: first part of hostname up to '.' + """ + return host_name().split(".")[0] + + +def utf_open(fname, mode): + """ + Wrapper for codecs.open + """ + import codecs + + return codecs.open(fname, mode=mode, encoding="utf-8") + + +def f_not_empty(*fpaths): + """ + Returns: + True if and only if the file exists and file size > 0 + if fpath is a dir, if and only if dir exists and has at least 1 file + """ + fpath = f_join(*fpaths) + if not os.path.exists(fpath): + return False + + if os.path.isdir(fpath): + return len(os.listdir(fpath)) > 0 + else: + return os.path.getsize(fpath) > 0 + + +def f_expand(fpath): + return os.path.expandvars(os.path.expanduser(fpath)) + + +def f_exists(*fpaths): + return os.path.exists(f_join(*fpaths)) + + +def f_join(*fpaths): + """ + Join file paths and expand special symbols like `~` for home dir + """ + + def pack_varargs(args): + """ + Pack *args or a single list arg as list + + def f(*args): + arg_list = pack_varargs(args) + # arg_list is now packed as a list + """ + assert isinstance(args, tuple), "please input the tuple `args` as in *args" + if len(args) == 1 and is_sequence(args[0]): + return args[0] + else: + return args + + fpaths = pack_varargs(fpaths) + fpath = f_expand(os.path.join(*fpaths)) + if isinstance(fpath, str): + fpath = fpath.strip() + return fpath + + +def f_listdir( + *fpaths, + filter_ext=None, + filter=None, + sort=True, + full_path=False, + nonexist_ok=True, + recursive=False, +): + """ + Args: + full_path: True to return full paths to the dir contents + filter: function that takes in file name and returns True to include + nonexist_ok: True to return [] if the dir is non-existent, False to raise + sort: sort the file names by alphabetical + recursive: True to use os.walk to recursively list files. Note that `filter` + will be applied to the relative path string to the root dir. + e.g. filter will take "a/data1.txt" and "a/b/data3.txt" as input, instead of + just the base file names "data1.txt" and "data3.txt". + if False, will simply call os.listdir() + """ + assert not (filter_ext and filter), "filter_ext and filter are mutually exclusive" + dir_path = f_join(*fpaths) + if not os.path.exists(dir_path) and nonexist_ok: + return [] + if recursive: + files = [ + os.path.join(os.path.relpath(root, dir_path), file) + for root, _, files in os.walk(dir_path) + for file in files + ] + else: + files = os.listdir(dir_path) + if filter is not None: + files = [f for f in files if filter(f)] + elif filter_ext is not None: + files = [f for f in files if f.endswith(filter_ext)] + if sort: + files.sort() + if full_path: + return [os.path.join(dir_path, f) for f in files] + else: + return files + + +def f_mkdir(*fpaths): + """ + Recursively creates all the subdirs + If exist, do nothing. + """ + fpath = f_join(*fpaths) + os.makedirs(fpath, exist_ok=True) + return fpath + + +def f_mkdir_in_path(*fpaths): + """ + fpath is a file, + recursively creates all the parent dirs that lead to the file + If exist, do nothing. + """ + os.makedirs(get_dir(f_join(*fpaths)), exist_ok=True) + + +def last_part_in_path(fpath): + """ + https://stackoverflow.com/questions/3925096/how-to-get-only-the-last-part-of-a-path-in-python + """ + return os.path.basename(os.path.normpath(f_expand(fpath))) + + +def is_abs_path(*fpath): + return os.path.isabs(f_join(*fpath)) + + +def is_relative_path(*fpath): + return not is_abs_path(f_join(*fpath)) + + +def f_time(*fpath): + "File modification time" + return str(os.path.getctime(f_join(*fpath))) + + +def f_append_before_ext(fpath, suffix): + """ + Append a suffix to file name and retain its extension + """ + name, ext = f_ext(fpath) + return name + suffix + ext + + +def f_add_ext(fpath, ext): + """ + Append an extension if not already there + Args: + ext: will add a preceding `.` if doesn't exist + """ + if not ext.startswith("."): + ext = "." + ext + if fpath.endswith(ext): + return fpath + else: + return fpath + ext + + +def f_has_ext(fpath, ext): + "Test if file path is a text file" + _, actual_ext = f_ext(fpath) + return actual_ext == "." + ext.lstrip(".") + + +def f_glob(*fpath): + return glob.glob(f_join(*fpath), recursive=True) + + +def f_remove(*fpath, verbose=False, dry_run=False): + """ + If exist, remove. Supports both dir and file. Supports glob wildcard. + """ + import errno + + assert isinstance(verbose, bool) + fpath = f_join(fpath) + if dry_run: + print("Dry run, delete:", fpath) + return + for f in glob.glob(fpath): + try: + shutil.rmtree(f) + except OSError as e: + if e.errno == errno.ENOTDIR: + try: + os.remove(f) + except Exception as e: # final resort safeguard + pass + if verbose: + print(f'Deleted "{fpath}"') + + +def f_copy(fsrc, fdst, ignore=None, include=None, exists_ok=True, verbose=False): + """ + Supports both dir and file. Supports glob wildcard. + """ + import errno + + fsrc, fdst = f_expand(fsrc), f_expand(fdst) + for f in glob.glob(fsrc): + try: + f_copytree(f, fdst, ignore=ignore, include=include, exist_ok=exists_ok) + except OSError as e: + if e.errno == errno.ENOTDIR: + shutil.copy(f, fdst) + else: + raise + if verbose: + print(f'Copied "{fsrc}" to "{fdst}"') + + +def _f_copytree( + src, + dst, + symlinks=False, + ignore=None, + exist_ok=True, + copy_function=shutil.copy2, + ignore_dangling_symlinks=False, +): + """Copied from python standard lib shutil.copytree + except that we allow exist_ok + Use f_copytree as entry + """ + names = os.listdir(src) + if ignore is not None: + ignored_names = ignore(src, names) + else: + ignored_names = set() + + os.makedirs(dst, exist_ok=exist_ok) + errors = [] + for name in names: + if name in ignored_names: + continue + srcname = os.path.join(src, name) + dstname = os.path.join(dst, name) + try: + if os.path.islink(srcname): + linkto = os.readlink(srcname) + if symlinks: + # We can't just leave it to `copy_function` because legacy + # code with a custom `copy_function` may rely on copytree + # doing the right thing. + os.symlink(linkto, dstname) + shutil.copystat(srcname, dstname, follow_symlinks=not symlinks) + else: + # ignore dangling symlink if the flag is on + if not os.path.exists(linkto) and ignore_dangling_symlinks: + continue + # otherwise let the copy occurs. copy2 will raise an error + if os.path.isdir(srcname): + _f_copytree(srcname, dstname, symlinks, ignore, exist_ok, copy_function) + else: + copy_function(srcname, dstname) + elif os.path.isdir(srcname): + _f_copytree(srcname, dstname, symlinks, ignore, exist_ok, copy_function) + else: + # Will raise a SpecialFileError for unsupported file types + copy_function(srcname, dstname) + # catch the Error from the recursive copytree so that we can + # continue with other files + except shutil.Error as err: + errors.extend(err.args[0]) + except OSError as why: + errors.append((srcname, dstname, str(why))) + try: + shutil.copystat(src, dst) + except OSError as why: + # Copying file access times may fail on Windows + if getattr(why, "winerror", None) is None: + errors.append((src, dst, str(why))) + if errors: + raise shutil.Error(errors) + return dst + + +def _include_patterns(*patterns): + """Factory function that can be used with copytree() ignore parameter. + + Arguments define a sequence of glob-style patterns + that are used to specify what files to NOT ignore. + Creates and returns a function that determines this for each directory + in the file hierarchy rooted at the source directory when used with + shutil.copytree(). + """ + + def _ignore_patterns(path, names): + import fnmatch + + keep = set(name for pattern in patterns for name in fnmatch.filter(names, pattern)) + ignore = set( + name + for name in names + if name not in keep and not os.path.isdir(os.path.join(path, name)) + ) + return ignore + + return _ignore_patterns + + +def f_copytree(fsrc, fdst, symlinks=False, ignore=None, include=None, exist_ok=True): + fsrc, fdst = f_expand(fsrc), f_expand(fdst) + assert (ignore is None) or (include is None), "ignore= and include= are mutually exclusive" + if ignore: + ignore = shutil.ignore_patterns(*ignore) + elif include: + ignore = _include_patterns(*include) + _f_copytree(fsrc, fdst, ignore=ignore, symlinks=symlinks, exist_ok=exist_ok) + + +def f_move(fsrc, fdst): + fsrc, fdst = f_expand(fsrc), f_expand(fdst) + for f in glob.glob(fsrc): + shutil.move(f, fdst) + + +def f_split_path(fpath, normpath=True): + """ + Splits path into a list of its component folders + + Args: + normpath: call os.path.normpath to remove redundant '/' and + up-level references like ".." + """ + if normpath: + fpath = os.path.normpath(fpath) + allparts = [] + while 1: + parts = os.path.split(fpath) + if parts[0] == fpath: # sentinel for absolute paths + allparts.insert(0, parts[0]) + break + elif parts[1] == fpath: # sentinel for relative paths + allparts.insert(0, parts[1]) + break + else: + fpath = parts[0] + allparts.insert(0, parts[1]) + return allparts + + +def get_script_dir(): + """ + Returns: the dir of current script + """ + return os.path.dirname(os.path.realpath(sys.argv[0])) + + +def get_script_file_name(): + """ + Returns: the dir of current script + """ + return os.path.basename(sys.argv[0]) + + +def get_script_self_path(): + """ + Returns: the dir of current script + """ + return os.path.realpath(sys.argv[0]) + + +def get_parent_dir(location, abspath=False): + """ + Args: + location: current directory or file + + Returns: + parent directory absolute or relative path + """ + _path = os.path.abspath if abspath else os.path.relpath + return _path(f_join(location, os.pardir)) + + +def md5_checksum(*fpath): + """ + File md5 signature + """ + import hashlib + + hash_md5 = hashlib.md5() + with open(f_join(*fpath), "rb") as f: + for chunk in iter(lambda: f.read(65536), b""): + hash_md5.update(chunk) + return hash_md5.hexdigest() + + +def create_tar(fsrc, output_tarball, include=None, ignore=None, compress_mode="gz"): + """ + Args: + fsrc: source file or folder + output_tarball: output tar file name + compress_mode: ``gz``, ``bz2``, ``xz`` or ``''`` (empty for uncompressed write) + include: include pattern, will trigger copy to temp directory + ignore: ignore pattern, will trigger copy to temp directory + """ + import tarfile + import tempfile + + fsrc, output_tarball = f_expand(fsrc), f_expand(output_tarball) + assert compress_mode in ["gz", "bz2", "xz", ""] + src_base = os.path.basename(fsrc) + + tempdir = None + if include or ignore: + tempdir = tempfile.mkdtemp() + tempdest = f_join(tempdir, src_base) + f_copy(fsrc, tempdest, include=include, ignore=ignore) + fsrc = tempdest + + with tarfile.open(output_tarball, "w:" + compress_mode) as tar: + tar.add(fsrc, arcname=src_base) + + if tempdir: + f_remove(tempdir) + + +def extract_tar(source_tarball, output_dir=".", members=None): + """ + Args: + source_tarball: extract members from archive + output_dir: default to current working dir + members: must be a subset of the list returned by getmembers() + """ + import tarfile + + source_tarball, output_dir = f_expand(source_tarball), f_expand(output_dir) + with tarfile.open(source_tarball, "r:*") as tar: + tar.extractall(output_dir, members=members) + + +def move_with_backup(*fpath, suffix=".bak"): + """ + Ensures that a path is not occupied. If there is a file, rename it by + adding @suffix. Resursively backs up everything. + + Args: + fpath: file path to clear + suffix: Add to backed up files (default: {'.bak'}) + """ + fpath = str(f_join(*fpath)) + if os.path.exists(fpath): + move_with_backup(fpath + suffix) + shutil.move(fpath, fpath + suffix) + + +def insert_before_ext(name, insert): + """ + log.txt -> log.ep50.txt + """ + name, ext = os.path.splitext(name) + return name + insert + ext + + +def timestamp_file_name(fname): + from datetime import datetime + + timestr = datetime.now().strftime("_%H-%M-%S_%m-%d-%y") + return insert_before_ext(fname, timestr) + + +def next_available_file_name( + *fpath, + suffix_template: Union[str, Callable[[int], str]] = "_v{i+1}", + before_ext: bool = True, +): + """ + Args: + suffix_template: a format string using "i" variable or + lambda int -> str + before_ext: True to insert suffix before the extension + """ + + def fstring(fmt_str, **kwargs): + """ + Simulate python f-string but without `f` + """ + import shlex + + locals().update(kwargs) + return eval("f" + shlex.quote(fmt_str)) + + orig_file_path = f_join(*fpath) + i = 0 + fpath = orig_file_path + while os.path.exists(fpath): + if isinstance(suffix_template, str): + suffix = fstring(suffix_template, i=i) + elif callable(suffix_template): + suffix = suffix_template(i) + assert isinstance(suffix, str) + else: + raise NotImplementedError(f"Unsupported suffix template {suffix_template}") + if before_ext: + fpath = insert_before_ext(orig_file_path, suffix) + else: + fpath = orig_file_path + suffix + i += 1 + return fpath + + +def get_file_lock(*fpath, timeout: int = 15, logging_level="critical"): + """ + NFS-safe filesystem-backed lock. `pip install flufl.lock` + https://flufllock.readthedocs.io/en/stable/apiref.html + + Args: + fpath: should be a path on NFS so that every process can see it + timeout: seconds + """ + import logging + + from flufl.lock import Lock + + logging.getLogger("flufl.lock").setLevel(logging_level.upper()) + return Lock(f_join(*fpath), lifetime=timeout) + + +def load_pickle(*fpaths): + with open(f_join(*fpaths), "rb") as fp: + return pickle.load(fp) + + +def dump_pickle(data, *fpaths): + with open(f_join(*fpaths), "wb") as fp: + pickle.dump(data, fp) + + +def load_text(*fpaths, by_lines=False): + with open(f_join(*fpaths), "r") as fp: + if by_lines: + return fp.readlines() + else: + return fp.read() + + +def load_text_lines(*fpaths): + return load_text(*fpaths, by_lines=True) + + +def dump_text(s, *fpaths): + with open(f_join(*fpaths), "w") as fp: + fp.write(s) + + +def dump_text_lines(lines: list[str], *fpaths, add_newline=True): + with open(f_join(*fpaths), "w") as fp: + for line in lines: + print(line, file=fp, end="\n" if add_newline else "") + + +def get_package_root() -> str: + import importlib.util + import inspect + + # Get the current frame + current_frame = inspect.currentframe() + if current_frame is None: + raise ImportError("Cannot determine the package name from __package__") + + # Get the caller module + caller_module = inspect.getmodule(current_frame.f_back) + if caller_module is None: + raise ImportError("Cannot determine the package name from __package__") + + # Get the package name + package_name = caller_module.__package__ + if not package_name: + raise ImportError("Cannot determine the package name from __package__") + + # Get the top-level package name + top_package_name = package_name.split(".")[0] + + # Find the package specification + spec = importlib.util.find_spec(top_package_name) + + if spec and spec.origin: + # Get the directory containing the package's __init__.py file + package_dir = os.path.dirname(spec.origin) + return package_dir + else: + raise ImportError(f"Cannot find the package {top_package_name}") + + +# aliases to be consistent with other load_* and dump_* +pickle_load = load_pickle +pickle_dump = dump_pickle +text_load = load_text +read_text = load_text +read_text_lines = load_text_lines +write_text = dump_text +write_text_lines = dump_text_lines +text_dump = dump_text diff --git a/groot/vla/common/utils/io/hdf5_utils.py b/groot/vla/common/utils/io/hdf5_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..03775968a73e9c3a39d1371472e10eb8f323323a --- /dev/null +++ b/groot/vla/common/utils/io/hdf5_utils.py @@ -0,0 +1,84 @@ +import h5py +import numpy as np +from pydantic import BaseModel + + +def hdf5_save(data: BaseModel | dict, group: h5py.Group) -> None: + """Recursively save Pydantic model or dict to HDF5 group.""" + if isinstance(data, BaseModel): + # Convert to dict and exclude None values + data_dict = data.model_dump(mode="python", exclude_none=True) + else: + data_dict = data + + for key, value in data_dict.items(): + if isinstance(value, np.ndarray): + group.create_dataset(key, data=value) + elif isinstance(value, (BaseModel, dict)): + subgroup = group.create_group(key) + hdf5_save(value, subgroup) + else: + # For primitive types, convert to numpy array + try: + group.create_dataset(key, data=np.array(value)) + except TypeError: + raise ValueError(f"Unsupported type: {type(value)} for key: {key}") + + +def hdf5_load(group: h5py.Group) -> dict: + """Recursively load HDF5 group to Pydantic model or dict.""" + data_dict = {} + for key, value in group.items(): + if isinstance(value, h5py.Dataset): + data_dict[key] = value[()] + elif isinstance(value, h5py.Group): + data_dict[key] = hdf5_load(value) + return data_dict + + +def hdf5_is_subset(this: h5py.Group, other: h5py.Group, verbose: bool = False) -> bool: + """Check if this HDF5 group is a subset of another HDF5 group.""" + for key, value in this.items(): + if key not in other: + if verbose: + print(f"Key {key} not in other") + return False + elif isinstance(value, h5py.Group): + if not isinstance(other[key], h5py.Group): + if verbose: + print(f"Key {key} is not a group in other") + return False + if not hdf5_is_subset(value, other[key], verbose): + if verbose: + print(f"Key {key} is not a subset of other") + return False + elif isinstance(value, h5py.Dataset): + if not isinstance(other[key], h5py.Dataset): + if verbose: + print(f"Key {key} is not a dataset in other") + return False + if not np.array_equal(value, other[key]): + if verbose: + print(f"Key {key} is not equal in other") + return False + elif isinstance(value, h5py.Datatype): + if not isinstance(other[key], h5py.Datatype): + if verbose: + print(f"Key {key} is not a datatype in other") + return False + if value != other[key]: + if verbose: + print(f"Key {key} is not equal in other") + return False + else: + # try to compare + if value != other[key]: + if verbose: + print(f"Key {key} is not equal in other") + return False + return True + + +def hdf5_is_equal(this: h5py.Group, other: h5py.Group, verbose: bool = False) -> bool: + """Check if this HDF5 group is equal to another HDF5 group.""" + return hdf5_is_subset(this, other, verbose) and hdf5_is_subset(other, this, verbose) diff --git a/groot/vla/common/utils/io/json_utils.py b/groot/vla/common/utils/io/json_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..34eb226264b6ae98f017ea46bbdd0bbdb549604f --- /dev/null +++ b/groot/vla/common/utils/io/json_utils.py @@ -0,0 +1,270 @@ +""" +JSON, YAML, and python config file utilities +""" + +from io import StringIO +import json +import os.path as path + +import yaml + +from ..misc.functional_utils import make_recursive_func +from .file_utils import f_join + +__all__ = [ + "json_load", + "json_loads", + "jsonl_load", + "yaml_load", + "yaml_loads", + "json_dump", + "json_dumps", + "jsonl_dump", + "yaml_dump", + "yaml_dumps", + "json_or_yaml_load", + "json_or_yaml_dump", + "Jsonl", + # ---------------- Aliases ----------------- + "load_json", + "loads_json", + "load_jsonl", + "load_yaml", + "loads_yaml", + "dump_json", + "dumps_json", + "dump_jsonl", + "dump_yaml", + "dumps_yaml", + "load_json_or_yaml", + "dump_json_or_yaml", +] + +from typing import Dict, List + +from typing_extensions import Literal + + +def json_load(*file_path, **kwargs): + file_path = f_join(file_path) + with open(file_path, "r") as fp: + return json.load(fp, **kwargs) + + +def json_loads(string, **kwargs): + return json.loads(string, **kwargs) + + +def jsonl_load(*file_path, **kwargs): + file_path = f_join(file_path) + data = [] + for line in open(file_path): + data.append(json.loads(line, **kwargs)) + return data + + +@make_recursive_func +def any_to_primitive(x): + try: + import torch + except ImportError: + raise ImportError("torch is required for any_to_primitive") + import numpy as np + + if isinstance(x, (np.ndarray, np.number, torch.Tensor)): + return x.tolist() + else: + return x + + +def json_dump(data, *file_path, convert_to_primitive=False, **kwargs): + if convert_to_primitive: + data = any_to_primitive(data) + file_path = f_join(file_path) + with open(file_path, "w") as fp: + json.dump(data, fp, **kwargs) + + +def json_dumps(data, convert_to_primitive=False, **kwargs): + """ + Returns: string + """ + if convert_to_primitive: + data = any_to_primitive(data) + return json.dumps(data, **kwargs) + + +def jsonl_dump(data, *file_path): + from .file_utils import is_sequence + + assert is_sequence(data) + data = any_to_primitive(data) + file_path = f_join(file_path) + with open(file_path, "w") as fp: + for line in data: + print(json.dumps(line), file=fp, flush=True) + + +def yaml_load(*file_path, loader=yaml.safe_load, **kwargs): + file_path = f_join(file_path) + with open(file_path, "r") as fp: + return loader(fp, **kwargs) + + +def yaml_loads(string, *, loader=yaml.safe_load, **kwargs): + return loader(string, **kwargs) + + +def yaml_dump(data, *file_path, dumper=yaml.safe_dump, convert_to_primitive=False, **kwargs): + if convert_to_primitive: + data = any_to_primitive(data) + file_path = f_join(file_path) + indent = kwargs.pop("indent", 2) + default_flow_style = kwargs.pop("default_flow_style", False) + sort_keys = kwargs.pop("sort_keys", False) # preserves original dict order + with open(file_path, "w") as fp: + dumper( + data, + stream=fp, + indent=indent, + default_flow_style=default_flow_style, + sort_keys=sort_keys, + **kwargs, + ) + + +def yaml_dumps(data, *, dumper=yaml.safe_dump, convert_to_primitive=False, **kwargs): + "Returns: string" + if convert_to_primitive: + data = any_to_primitive(data) + stream = StringIO() + indent = kwargs.pop("indent", 2) + default_flow_style = kwargs.pop("default_flow_style", False) + sort_keys = kwargs.pop("sort_keys", False) # preserves original dict order + dumper( + data, + stream, + indent=indent, + default_flow_style=default_flow_style, + sort_keys=sort_keys, + **kwargs, + ) + return stream.getvalue() + + +# ==================== auto-recognize extension ==================== +def json_or_yaml_load(*file_path, **loader_kwargs): + """ + Args: + file_path: JSON or YAML loader depends on the file extension + + Raises: + IOError: if extension is not ".json", ".yml", or ".yaml" + """ + file_path = str(f_join(file_path)) + if file_path.endswith(".json"): + return json_load(file_path, **loader_kwargs) + elif file_path.endswith(".yml") or file_path.endswith(".yaml"): + return yaml_load(file_path, **loader_kwargs) + else: + raise IOError( + f'unknown file extension: "{file_path}", ' + f'loader supports only ".json", ".yml", ".yaml"' + ) + + +def json_or_yaml_dump(data, *file_path, **dumper_kwargs): + """ + Args: + file_path: JSON or YAML loader depends on the file extension + + Raises: + IOError: if extension is not ".json", ".yml", or ".yaml" + """ + file_path = str(f_join(file_path)) + if file_path.endswith(".json"): + return json_dump(data, file_path, **dumper_kwargs) + elif file_path.endswith(".yml") or file_path.endswith(".yaml"): + return yaml_dump(data, file_path, **dumper_kwargs) + else: + raise IOError( + f'unknown file extension: "{file_path}", ' + f'dumper supports only ".json", ".yml", ".yaml"' + ) + + +# ---------------- Aliases ----------------- +# add aliases where verb goes first, json_load -> load_json +load_json = json_load +load_yaml = yaml_load +load_jsonl = jsonl_load +loads_json = json_loads +loads_yaml = yaml_loads +dump_json = json_dump +dump_jsonl = jsonl_dump +dump_yaml = yaml_dump +dumps_json = json_dumps +dumps_yaml = yaml_dumps +load_json_or_yaml = json_or_yaml_load +dump_json_or_yaml = json_or_yaml_dump + +# ==================== Jsonl ==================== + + +class Jsonl: + """ + Both reader and writer, as if everything's in-memory + """ + + def __init__(self, *file_path, mode: Literal["r", "w", "a"] = "a"): + """ + Args: + mode: + - 'r': file must already exists + - 'w': overwrite the file regardless of whether it exists or not + - 'a': create a new file if doesn't exist, or append to an existing file + """ + assert mode in "rwa" + self._file_path = str(f_join(file_path)) + self._mode = mode + if mode == "r": + assert path.exists(self._file_path) + self._fp = None + else: + self._fp = open(self._file_path, mode) + if path.exists(self._file_path) and mode != "w": + self.data = jsonl_load(self._file_path) + else: + self.data = [] + + def append(self, data: Dict): + if self._mode == "r": + raise RuntimeError("Jsonl read mode cannot call append()") + self.data.append(data) + print(json_dumps(data), file=self._fp, flush=True) + + def extend(self, data_list: List[Dict]): + for data in data_list: + self.append(data) + + def close(self): + if self._fp is not None: + self._fp.close() + + def __getitem__(self, idx): + return self.data[idx] + + def __len__(self): + return len(self.data) + + def __iter__(self): + return iter(self.data) + + def __enter__(self): + return self + + def __exit__(self, type, value, traceback): + self.close() + + def __bool__(self): + return bool(self.data) diff --git a/groot/vla/common/utils/io/print_utils.py b/groot/vla/common/utils/io/print_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..161746c9c83179dbb4a61902478e767854efaa36 --- /dev/null +++ b/groot/vla/common/utils/io/print_utils.py @@ -0,0 +1,362 @@ +from datetime import datetime +import io +import logging +import os +import pprint +import shlex +import string +import sys +import textwrap +import time +import traceback +from typing import Callable, Union + +import numpy as np +from typing_extensions import Literal + +from ..misc.functional_utils import meta_decorator +from ..misc.misc_utils import match_patterns + + +def to_readable_count_str(value: int, precision: int = 2) -> str: + assert value >= 0 + labels = [" ", "K", "M", "B", "T"] + num_digits = int(np.floor(np.log10(value)) + 1 if value > 0 else 1) + num_groups = int(np.ceil(num_digits / 3)) + num_groups = min(num_groups, len(labels)) # don't abbreviate beyond trillions + shift = -3 * (num_groups - 1) + value = value * (10**shift) + index = num_groups - 1 + rem = value - int(value) + if precision > 0 and rem > 0.01: + fmt = f"{{:.{precision}f}}" + rem_str = fmt.format(rem).lstrip("0") + else: + rem_str = "" + return f"{int(value):,d}{rem_str} {labels[index]}" + + +def to_scientific_str(value, precision: int = 1, capitalize: bool = False) -> str: + """ + 0.0015 -> "1.5e-3" + """ + if value == 0: + return "0" + return f"{value:.{precision}e}".replace("e-0", "E-" if capitalize else "e-") + + +def print_str(*args, **kwargs): + """ + Same as print() signature but returns a string + """ + sstream = io.StringIO() + kwargs.pop("file", None) + print(*args, **kwargs, file=sstream) + return sstream.getvalue() + + +def fstring(fmt_str, **kwargs): + """ + Simulate python f-string but without `f` + """ + locals().update(kwargs) + return eval("f" + shlex.quote(fmt_str)) + + +def get_format_keys(fmt_str): + keys = [] + for literal, field_name, fmt_spec, conversion in string.Formatter().parse(fmt_str): + if field_name: + keys.append(field_name) + return keys + + +def get_timestamp(milli_precision: int = 3): + fmt = "%y-%m-%d %H:%M:%S" + if milli_precision > 0: + fmt += ".%f" + stamp = datetime.now().strftime(fmt) + if milli_precision > 0: + stamp = stamp[:-milli_precision] + return stamp + + +def pretty_repr_str(obj, **kwargs): + """ + Useful to produce __repr__() + """ + if isinstance(obj, str): + cls_name = obj + else: + cls_name = obj.__class__.__name__ + kw_strs = [k + "=" + pprint.pformat(v, indent=2, compact=True) for k, v in kwargs.items()] + new_line = len(cls_name) + sum(len(kw) for kw in kw_strs) > 84 + if new_line: + kw = ",\n".join(kw_strs) + return f"{cls_name}(\n{textwrap.indent(kw, ' ')}\n)" + else: + kw = ", ".join(kw_strs) + return f"{cls_name}({kw})" + + +def pprint_(*objs, **kwargs): + """ + Use pprint to format the objects + """ + print( + *[pprint.pformat(obj, indent=2) if not isinstance(obj, str) else obj for obj in objs], + **kwargs, + ) + + +def get_exception_info(to_str: bool = False): + """ + Returns: + {'type': ExceptionType, 'value': ExceptionObject, 'trace': } + """ + typ_, value, trace = sys.exc_info() + return { + "type": typ_.__name__ if to_str else typ_, + "value": str(value) if to_str else value, + "trace": "".join(traceback.format_exception(typ_, value, trace)), + } + + +class DebugPrinter: + """ + Debug print, usage: dprint = DebugPrint(enabled=True) + dprint(...) + """ + + def __init__(self, enabled, tensor_summary: Literal["shape", "shape+dtype", "none"] = "shape"): + """ + Args: + tensor_summary: + - shape: only prints shape + - shape+dtype: also prints dtype and device + - none: print full tensor + """ + self.enabled = enabled + assert tensor_summary in ["shape", "shape+dtype", "none"] + self.tensor_summary = tensor_summary + + def __call__(self, *args, **kwargs): + if not self.enabled: + return + args = [self._process_arg(a) for a in args] + pprint_(*args, **kwargs) + + def _process_arg(self, arg): + import numpy as np + import torch + + if torch.is_tensor(arg): + if self.tensor_summary == "shape": + return str(list(arg.size())) + elif self.tensor_summary == "shape+dtype": + return f"{arg.dtype}{list(arg.size())}|{arg.device}" + elif isinstance(arg, np.ndarray): + if self.tensor_summary == "shape": + return str(list(arg.shape)) + elif self.tensor_summary == "shape+dtype": + return f"{arg.dtype}{list(arg.shape)}" + return arg + + +@meta_decorator +def watch(func, seconds: int = 5, max_times: int = 0, keep_returns: bool = False): + """ + Decorator: executes a function repeated with the args and + emulate `watch -n` capability + + See `gpustat` repo: https://github.com/wookayin/gpustat/pull/41/files + + Args: + max_times: watch for `max_times` and then exit. If 0, never exits + keep_returns: if True, will keep the return value from the function + and return as a list at the end + """ + from blessings import Terminal + + def _wrapped(*args, **kwargs): + term = Terminal() + N = 0 + returns = [] + with term.fullscreen(): + while True: + try: + with term.location(0, 0): + ret = func(*args, **kwargs) + print(term.clear_eos, end="") + if keep_returns: + returns.append(ret) + N += 1 + if max_times > 0 and N >= max_times: + break + time.sleep(seconds) + except KeyboardInterrupt: + break + return returns + + return _wrapped + + +class PrintRedirection(object): + """ + Context manager: temporarily redirects stdout and stderr + """ + + def __init__(self, stdout=None, stderr=None): + """ + Args: + stdout: if None, defaults to sys.stdout, unchanged + stderr: if None, defaults to sys.stderr, unchanged + """ + if stdout is None: + stdout = sys.stdout + if stderr is None: + stderr = sys.stderr + self._stdout, self._stderr = stdout, stderr + + def __enter__(self): + self._old_out, self._old_err = sys.stdout, sys.stderr + self._old_out.flush() + self._old_err.flush() + sys.stdout, sys.stderr = self._stdout, self._stderr + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.flush() + # restore the normal stdout and stderr + sys.stdout, sys.stderr = self._old_out, self._old_err + + def flush(self): + "Manually flush the replaced stdout/stderr buffers." + self._stdout.flush() + self._stderr.flush() + + +class PrintToFile(PrintRedirection): + """ + Print to file and save/close the handle at the end. + """ + + def __init__(self, out_file=None, err_file=None): + """ + Args: + out_file: file path + err_file: file path. If the same as out_file, print both stdout + and stderr to one file in order. + """ + self.out_file, self.err_file = out_file, err_file + if out_file: + out_file = os.path.expanduser(out_file) + self.out_file = open(out_file, "w") + if err_file: + err_file = os.path.expanduser(out_file) + if err_file == out_file: # redirect both stdout/err to one file + self.err_file = self.out_file + else: + self.err_file = open(os.path.expanduser(out_file), "w") + super().__init__(stdout=self.out_file, stderr=self.err_file) + + def __exit__(self, *args): + super().__exit__(*args) + if self.out_file: + self.out_file.close() + if self.err_file: + self.err_file.close() + + +def PrintSuppress(no_out=True, no_err=False): + """ + Args: + no_out: stdout writes to sys.devnull + no_err: stderr writes to sys.devnull + """ + out_file = os.devnull if no_out else None + err_file = os.devnull if no_err else None + return PrintToFile(out_file=out_file, err_file=err_file) + + +class PrintString(PrintRedirection): + """ + Redirect stdout and stderr to strings. + """ + + def __init__(self): + self.out_stream = io.StringIO() + self.err_stream = io.StringIO() + super().__init__(stdout=self.out_stream, stderr=self.err_stream) + + def stdout(self): + "Returns: stdout as one string." + return self.out_stream.getvalue() + + def stderr(self): + "Returns: stderr as one string." + return self.err_stream.getvalue() + + def stdout_by_line(self): + "Returns: a list of stdout line by line, ignore trailing blanks" + return self.stdout().rstrip().split("\n") + + def stderr_by_line(self): + "Returns: a list of stderr line by line, ignore trailing blanks" + return self.stderr().rstrip().split("\n") + + +# ==================== Logging filters ==================== +class ExcludeLoggingFilter(logging.Filter): + """ + Usage: logging.getLogger('name').addFilter( + ExcludeLoggingFilter(['info mess*age', 'Warning: *']) + ) + Supports wildcard. + https://relaxdiego.com/2014/07/logging-in-python.html + """ + + def __init__(self, patterns): + super().__init__() + self._patterns = patterns + + def filter(self, record): + if match_patterns(record.msg, include=self._patterns): + return False + else: + return True + + +class ReplaceStringLoggingFilter(logging.Filter): + def __init__(self, patterns, replacer: Callable): + super().__init__() + self._patterns = patterns + assert callable(replacer) + self._replacer = replacer + + def filter(self, record): + if match_patterns(record.msg, include=self._patterns): + record.msg = self._replacer(record.msg) + + +def logging_exclude_pattern( + logger_name, + patterns: Union[str, list[str], Callable, list[Callable], None], +): + """ + Args: + patterns: see groot.vla.common.utils.misc_utils.match_patterns + """ + logging.getLogger(logger_name).addFilter(ExcludeLoggingFilter(patterns)) + + +def logging_replace_string( + logger_name, + patterns: Union[str, list[str], Callable, list[Callable], None], + replacer: Callable, +): + """ + Args: + patterns: see groot.vla.common.utils.misc_utils.match_patterns + """ + logging.getLogger(logger_name).addFilter(ReplaceStringLoggingFilter(patterns, replacer)) diff --git a/groot/vla/common/utils/io/termcolor.py b/groot/vla/common/utils/io/termcolor.py new file mode 100644 index 0000000000000000000000000000000000000000..bf5a13d3aebb39b02cee7d6f709696cafca92197 --- /dev/null +++ b/groot/vla/common/utils/io/termcolor.py @@ -0,0 +1,186 @@ +# coding: utf-8 +# Copyright (c) 2008-2011 Volvox Development 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. +# +# 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. +# +# Original Author: Konstantin Lepa +# Updated by Jim Fan + +"""ANSII Color formatting for output in terminal.""" +import io +import os +from typing import List, Optional, Union + +__ALL__ = ["color_text", "cprint"] + +STYLES = dict( + list( + zip( + ["bold", "dark", "", "underline", "blink", "", "reverse", "concealed"], + list(range(1, 9)), + ) + ) +) +del STYLES[""] + + +HIGHLIGHTS = dict( + list( + zip( + ["grey", "red", "green", "yellow", "blue", "magenta", "cyan", "white"], + list(range(40, 48)), + ) + ) +) + + +COLORS = dict( + list( + zip( + ["grey", "red", "green", "yellow", "blue", "magenta", "cyan", "white"], + list(range(30, 38)), + ) + ) +) + + +def _strip_bg_prefix(color): + "on_red -> red" + if color.startswith("on_"): + return color[len("on_") :] + else: + return color + + +RESET = "\033[0m" + + +def color_text( + text, + color: Optional[str] = None, + bg_color: Optional[str] = None, + styles: Optional[Union[str, List[str]]] = None, +): + """Colorize text. + + Available text colors: + red, green, yellow, blue, magenta, cyan, white. + + Available text highlights: + on_red, on_green, on_yellow, on_blue, on_magenta, on_cyan, on_white. + + Available attributes: + bold, dark, underline, blink, reverse, concealed. + + Example: + colored('Hello, World!', 'red', 'on_grey', ['blue', 'blink']) + colored('Hello, World!', 'green') + """ + if os.getenv("ANSI_COLORS_DISABLED") is None: + fmt_str = "\033[%dm%s" + if color is not None: + text = fmt_str % (COLORS[color], text) + + if bg_color is not None: + bg_color = _strip_bg_prefix(bg_color) + text = fmt_str % (HIGHLIGHTS[bg_color], text) + + if styles is not None: + if isinstance(styles, str): + styles = [styles] + for style in styles: + text = fmt_str % (STYLES[style], text) + + text += RESET + return text + + +def cprint( + *args, + color: Optional[str] = None, + bg_color: Optional[str] = None, + styles: Optional[Union[str, List[str]]] = None, + **kwargs, +): + """Print colorize text. + + It accepts arguments of print function. + """ + sstream = io.StringIO() + print(*args, sep=kwargs.pop("sep", None), end="", file=sstream) + text = sstream.getvalue() + print((color_text(text, color, bg_color, styles)), **kwargs) + + +if __name__ == "__main__": + print("Current terminal type: %s" % os.getenv("TERM")) + print("Test basic colors:") + cprint("Grey color", color="grey") + cprint("Red color", color="red") + cprint("Green color", color="green") + cprint("Yellow color", color="yellow") + cprint("Blue color", color="blue") + cprint("Magenta color", color="magenta") + cprint("Cyan color", color="cyan") + cprint("White color", color="white") + print(("-" * 78)) + + print("Test highlights:") + cprint("On grey color", bg_color="on_grey") + cprint("On red color", bg_color="on_red") + cprint("On green color", bg_color="on_green") + cprint("On yellow color", bg_color="on_yellow") + cprint("On blue color", bg_color="on_blue") + cprint("On magenta color", bg_color="on_magenta") + cprint("On cyan color", bg_color="on_cyan") + cprint("On white color", color="grey", bg_color="on_white") + print("-" * 78) + + print("Test attributes:") + cprint("Bold grey color", color="grey", styles="bold") + cprint("Dark red color", color="red", styles=["dark"]) + cprint("Underline green color", color="green", styles=["underline"]) + cprint("Blink yellow color", color="yellow", styles=["blink"]) + cprint("Reversed blue color", color="blue", styles=["reverse"]) + cprint("Concealed Magenta color", color="magenta", styles=["concealed"]) + cprint( + "Bold underline reverse cyan color", + color="cyan", + styles=["bold", "underline", "reverse"], + ) + cprint( + "Dark blink concealed white color", + color="white", + styles=["dark", "blink", "concealed"], + ) + print(("-" * 78)) + + print("Test mixing:") + cprint( + "Underline red on grey color", + color="red", + bg_color="on_grey", + styles="underline", + ) + cprint( + "Reversed green on red color", + color="green", + bg_color="on_red", + styles="reverse", + ) diff --git a/groot/vla/common/utils/misc/__init__.py b/groot/vla/common/utils/misc/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a3ab939903e7c546f80c32175d84560e747557a1 --- /dev/null +++ b/groot/vla/common/utils/misc/__init__.py @@ -0,0 +1,5 @@ +from .functional_utils import * # noqa: F403 +from .image_utils import * # noqa: F403 +from .misc_utils import * # noqa: F403 +from .torch_utils import * # noqa: F403 +from .video_utils import * # noqa: F403 diff --git a/groot/vla/common/utils/misc/array_tensor_utils.py b/groot/vla/common/utils/misc/array_tensor_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..6585ca8a2693c38390b6236eb7c2a337bf3f2926 --- /dev/null +++ b/groot/vla/common/utils/misc/array_tensor_utils.py @@ -0,0 +1,372 @@ +""" +Functions that work on nested structures of torch.Tensor or numpy array +""" + +from typing import Any, Dict, List, Optional, Union + +import numpy as np +import torch +import tree + +from ..data_structure.tree_utils import ( + copy_non_leaf, + is_sequence, + tree_assign_at_path, + tree_value_at_path, +) +from .functional_utils import make_recursive_func + + +def is_array_tensor(obj): + return isinstance(obj, (np.ndarray, torch.Tensor)) + + +def is_numpy(obj): + return isinstance(obj, np.ndarray) + + +def is_tensor(obj): + return torch.is_tensor(obj) + + +def any_stack(xs: List, *, dim: int = 0): + """ + Works for both torch Tensor and numpy array + """ + + def _any_stack_helper(*xs): + x = xs[0] + if isinstance(x, np.ndarray): + return np.stack(xs, axis=dim) + elif torch.is_tensor(x): + return torch.stack(xs, dim=dim) + elif isinstance(x, float): + # special treatment for float, defaults to float32 + return np.array(xs, dtype=np.float32) + else: + return np.array(xs) + + return tree.map_structure(_any_stack_helper, *xs) + + +def any_concat(xs: List, *, dim: int = 0): + """ + Works for both torch Tensor and numpy array + """ + + def _any_concat_helper(*xs): + x = xs[0] + if isinstance(x, np.ndarray): + return np.concatenate(xs, axis=dim) + elif torch.is_tensor(x): + return torch.cat(xs, dim=dim) + elif isinstance(x, float): + # special treatment for float, defaults to float32 + return np.array(xs, dtype=np.float32) + else: + return np.array(xs) + + return tree.map_structure(_any_concat_helper, *xs) + + +def any_chunk(x, chunks: int, *, dim: int = 0, strict: bool = True) -> List[Any]: + """ + Works for both torch Tensor and numpy array + + Returns: + list of chunked nested structures + """ + assert chunks >= 1 + + x_copies = [copy_non_leaf(x) for _ in range(chunks)] + + def _any_chunk_helper(path, x): + if is_array_tensor(x): + if isinstance(x, np.ndarray): + chunked_values = np.split(x, chunks, axis=dim) + else: + chunked_values = torch.chunk(x, chunks, dim=dim) + + if path: + for xc, chunked in zip(x_copies, chunked_values): + tree_assign_at_path(xc, path, chunked) + else: # top-level, no nested path + for i, chunked in enumerate(chunked_values): + x_copies[i] = chunked + else: + if strict: + raise NotImplementedError(f"Cannot chunk type {type(x)}") + else: + return + + tree.map_structure_with_path(_any_chunk_helper, x) + return x_copies + + +def chunk_seq(arr, chunks: int, check_divide=True): + """ + Args: + check_divide: True to force arr must divide n + """ + k, m = divmod(len(arr), chunks) + if check_divide and m != 0: + raise ValueError(f"Array len {len(arr)} does not divide chunks {chunks}") + return (arr[i * k + min(i, m) : (i + 1) * k + min(i + 1, m)] for i in range(chunks)) + + +@make_recursive_func +def any_zeros_like(x: Union[Dict, np.ndarray, torch.Tensor, int, float, np.number]): + """Returns a zero-filled object of the same (d)type and shape as the input. + + The difference between this and `np.zeros_like()` is that this works well + with `np.number`, `int`, `float`, and `jax.numpy.DeviceArray` objects without + converting them to `np.ndarray`s. + + Args: + x: The object to replace with 0s. + + Returns: + A zero-filed object of the same (d)type and shape as the input. + """ + if isinstance(x, (int, float, np.number)): + return type(x)(0) + elif is_tensor(x): + return torch.zeros_like(x) + elif is_numpy(x): + return np.zeros_like(x) + else: + raise ValueError( + f"Input ({type(x)}) must be either a numpy array, a tensor, an int, or a float." + ) + + +@make_recursive_func +def any_ones_like(x: Union[Dict, np.ndarray, torch.Tensor, int, float, np.number]): + """Returns a one-filled object of the same (d)type and shape as the input. + The difference between this and `np.ones_like()` is that this works well + with `np.number`, `int`, `float`, and `jax.numpy.DeviceArray` objects without + converting them to `np.ndarray`s. + Args: + x: The object to replace with 1s. + Returns: + A one-filed object of the same (d)type and shape as the input. + """ + if isinstance(x, (int, float, np.number)): + return type(x)(1) + elif is_tensor(x): + return torch.ones_like(x) + elif is_numpy(x): + return np.ones_like(x) + else: + raise ValueError( + f"Input ({type(x)}) must be either a numpy array, a tensor, an int, or a float." + ) + + +@make_recursive_func +def any_zero_(x: Union[Dict, np.ndarray, torch.Tensor]): + """ + Apply in-place zero-out to a tensor, i.e. x.zero_() + """ + if is_tensor(x): + x.zero_() + elif is_numpy(x): + x.fill(0) + else: + raise ValueError(f"Input ({type(x)}) must be either a numpy array or a tensor") + + +@make_recursive_func +def any_fill_(x: Union[Dict, np.ndarray, torch.Tensor], value): + """ + Apply in-place zero-out to a tensor, i.e. x.zero_() + """ + if is_tensor(x): + x.fill_(value) + elif is_numpy(x): + x.fill(value) + else: + raise ValueError(f"Input ({type(x)}) must be either a numpy array or a tensor") + + +def get_batch_size(x, strict: bool = False) -> int: + """ + Args: + x: can be any arbitrary nested structure of np array and torch tensor + strict: True to check all batch sizes are the same + """ + + def _get_batch_size(x): + if isinstance(x, np.ndarray): + return x.shape[0] + elif torch.is_tensor(x): + return x.size(0) + else: + return len(x) + + xs = tree.flatten(x) + + if strict: + batch_sizes = [_get_batch_size(x) for x in xs] + assert all( + b == batch_sizes[0] for b in batch_sizes + ), f"batch sizes must all be the same in nested structure: {batch_sizes}" + return batch_sizes[0] + else: + return _get_batch_size(xs[0]) + + +@make_recursive_func +def add_batch_dim(x): + if is_numpy(x): + return np.expand_dims(x, axis=0) + elif is_tensor(x): + return x.unsqueeze(0) + else: + raise NotImplementedError(f"Unsupported data structure: {type(x)}") + + +@make_recursive_func +def remove_batch_dim(x): + if is_numpy(x): + return np.squeeze(x, axis=0) + elif is_tensor(x): + return x.squeeze(0) + else: + raise NotImplementedError(f"Unsupported data structure: {type(x)}") + + +@make_recursive_func +def any_to_primitive(x): + if isinstance(x, (np.ndarray, np.number, torch.Tensor)): + return x.tolist() + else: + return x + + +@make_recursive_func +def any_get_shape(x): + if is_numpy(x): + return tuple(x.shape) + elif is_tensor(x): + return tuple(x.size()) + else: + raise NotImplementedError(f"Unsupported data structure: {type(x)}") + + +@make_recursive_func +def any_mean(x, dim: Optional[int] = None, keepdim: bool = False): + if is_numpy(x): + return np.mean(x, axis=dim, keepdims=keepdim) + elif is_tensor(x): + return torch.mean(x, dim=dim, keepdim=keepdim) + else: + raise NotImplementedError(f"Unsupported data structure: {type(x)}") + + +@make_recursive_func +def any_variance(x, dim: Optional[int] = None, keepdim: bool = False, unbiased: bool = False): + if is_numpy(x): + return np.var(x, axis=dim, keepdims=keepdim, ddof=1 if unbiased else 0) + elif is_tensor(x): + return torch.var(x, dim=dim, keepdim=keepdim, unbiased=unbiased) + else: + raise NotImplementedError(f"Unsupported data structure: {type(x)}") + + +@make_recursive_func +def any_describe_str(x, shape_only=False): + """ + Describe type, shape, device, data type (of np array/tensor) + Very useful for debugging + """ + t = type(x) + tname = type(x).__name__ + if is_numpy(x): + shape = list(x.shape) + if x.size == 1: + if shape_only: + return f"np scalar: {x.item()} {shape}" + else: + return f"np scalar: {x.item()} {shape} {x.dtype}" + else: + if shape_only: + return f"np: {shape}" + else: + return f"np: {shape} {x.dtype}" + elif is_tensor(x): + shape = list(x.size()) + if x.numel() == 1: + if shape_only: + return f"torch scalar: {x.item()} {shape}" + else: + return f"torch scalar: {x.item()} {shape} {x.dtype} {x.device}" + else: + if shape_only: + return f"torch: {shape}" + else: + return f"torch: {shape} {x.dtype} {x.device}" + elif is_sequence(x): + return f"{tname}[{len(x)}]" + elif isinstance(x, str): + return x + elif x is None: + return "None" + elif np.issubdtype(t, np.number) or np.issubdtype(t, np.bool_): + return f"{tname}: {x}" + else: + return f"{tname}" + + +def any_describe(x, msg="", *, shape_only=False): + # from omlet.utils import yaml_dumps + from pprint import pprint + + if isinstance(x, str) and msg != "": + x, msg = msg, x + + if msg: + msg += ": " + print(msg, end="") + pprint(any_describe_str(x, shape_only=shape_only)) + + +@make_recursive_func +def any_slice(x, slice): + """ + Args: + slice: you can use np.s_[...] to return the slice object + """ + if is_array_tensor(x): + return x[slice] + else: + return x + + +def any_assign(x, assign_value, slice): + """ + Recursive version of x[slice] = assign_value + If structures of x and assign_value do not match, we will respect `assign_value` + E.g. x = {'a': ..., 'b': ...}, assign_value = {'a': ...}, then 'b' will not change + + Use np.s_[...] to get advanced slicing + """ + + def _any_assign_helper(path, v): + y = tree_value_at_path(x, path) + y[slice] = v + + tree.map_structure_with_path(_any_assign_helper, assign_value) + + +@make_recursive_func +def any_transpose_first_two_axes(x): + """ + util to convert between (L, B, ...) and (B, L, ...) + """ + if is_numpy(x): + return np.swapaxes(x, 0, 1) + elif is_tensor(x): + return torch.swapaxes(x, 0, 1) + else: + raise ValueError(f"Input ({type(x)}) must be either a numpy array or a tensor.") diff --git a/groot/vla/common/utils/misc/functional_utils.py b/groot/vla/common/utils/misc/functional_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9a5a509feac792751ccfd7956736f932bea33679 --- /dev/null +++ b/groot/vla/common/utils/misc/functional_utils.py @@ -0,0 +1,635 @@ +""" +Inspect, meta, etc. +""" + +from __future__ import annotations + +import functools +import inspect +import pprint +import sys +import types +from typing import Any, Dict, Literal +import warnings + +from ..data_structure.tree_utils import is_mapping, is_sequence + + +def state_dict_class(keys: list[str]): + """ + Just like pytorch nn.Module + Add the following methods to the class: + state_dict() -> dict of attribute keys + load_state_dict(sdict) restore states + """ + + def _wrap_class(cls): + assert inspect.isclass(cls) + + def state_dict(self): + return {k: getattr(self, k) for k in keys} + + def load_state_dict(self, states: Dict[str, Any]): + if not set(keys).issubset(set(states.keys())): + raise ValueError(f"states does not have all the required keys: {keys}") + for k in keys: + setattr(self, k, states[k]) + + @property + def state_keys(self): + return keys + + cls.state_dict = state_dict + cls.load_state_dict = load_state_dict + cls.state_keys = state_keys + return cls + + return _wrap_class + + +def implements_method(object, method: str): + """ + Returns: + True if object implements a method + """ + return hasattr(object, method) and callable(getattr(object, method)) + + +def assert_implements_method(object, method: str | list[str]): + if isinstance(method, str): + method = [method] + for m in method: + assert implements_method(object, m), ( + f"object {object.__class__} does not " f"implement method {m}()" + ) + + +def meta_decorator(decor): + """ + a decorator, allowing the wrapped decorator to be used as: + @decorator(*args, **kwargs) + def callable() + -- or -- + @decorator # without parenthesis, args and kwargs will use default + def callable() + + Args: + decor: a decorator whose first argument is a callable (function or class + to be decorated), and the rest of the arguments can be omitted as default. + decor(f, ... the other arguments must have default values) + + Warning: + decor can NOT be a function that receives a single, callable argument. + See stackoverflow: http://goo.gl/UEYbDB + """ + import functools + + def single_callable(args, kwargs): + return len(args) == 1 and len(kwargs) == 0 and callable(args[0]) + + @functools.wraps(decor) + def new_decor(*args, **kwargs): + if single_callable(args, kwargs): + # this is the double-decorated f. + # It should not run on a single callable. + return decor(args[0]) + else: + # decorator arguments + return lambda real_f: decor(real_f, *args, **kwargs) + + return new_decor + + +@meta_decorator +def make_recursive_func(fn, *, with_path=False): + """ + Decorator that turns a function that works on a single array/tensor to working on + arbitrary nested structures. + """ + import functools + + import tree + + @functools.wraps(fn) + def _wrapper(tensor_struct, *args, **kwargs): + if with_path: + return tree.map_structure_with_path( + lambda paths, x: fn(paths, x, *args, **kwargs), tensor_struct + ) + else: + return tree.map_structure(lambda x: fn(x, *args, **kwargs), tensor_struct) + + return _wrapper + + +@meta_decorator +def deprecated(func, msg="", action="warning", type=""): + """ + Function/class decorator: designate deprecation. + + Args: + msg: string message. + action: string mode + - 'warning': (default) prints `msg` to stderr + - 'noop': do nothing, just for source code annotation purposes + - 'raise': raise DeprecatedError(`msg`) + """ + action = action.lower() + type = type.lower() + ALL_ACTIONS = ["warn", "warning", "noop", "raise"] + if action not in ALL_ACTIONS: + raise ValueError(f"Unknown action {action}. Choose from {ALL_ACTIONS}.") + ALL_TYPES = { + "": DeprecationWarning, + "pending": PendingDeprecationWarning, + "future": FutureWarning, + } + if type not in ALL_TYPES: + raise ValueError(f"Unknown type {type}. Choose from {ALL_TYPES.keys()}.") + if not msg: + msg = "This is a deprecated feature." + + WarningExceptionCls = ALL_TYPES[type] + + # only does the deprecation when being called + @functools.wraps(func) + def _deprecated(*args, **kwargs): + if action in ["warning", "warn"]: + warnings.warn(msg, WarningExceptionCls) + elif action == "raise": + raise WarningExceptionCls(msg) + return func(*args, **kwargs) + + return _deprecated + + +@meta_decorator +def call_once(func, on_second_call: Literal["noop", "raise", "warn"] = "noop"): + """ + Decorator to ensure that a function is only called once. + + Args: + on_second_call (str): what happens when the function is called a second time. + """ + assert on_second_call in [ + "noop", + "raise", + "warn", + ], "mode must be one of 'noop', 'raise', 'warn'" + + @functools.wraps(func) + def wrapper(*args, **kwargs): + if wrapper._called: + if on_second_call == "raise": + raise RuntimeError(f"{func.__name__} has already been called. Can only call once.") + elif on_second_call == "warn": + warnings.warn(f"{func.__name__} has already been called. Should only call once.") + else: + wrapper._called = True + return func(*args, **kwargs) + + wrapper._called = False + return wrapper + + +class NoopObject: + """ + Object that does nothing when called any method + """ + + def __init__(self, *args, **kwargs): + self.init_args = args + self.init_kwargs = kwargs + + def __getattr__(self, name): + def _func(*args, **kwargs): + pass + + return _func + + +class NoopContext: + """ + Placeholder context manager that does nothing. + We could have written simply as: + + @contextmanager + def noop_context(*args, **kwargs): + yield + + but the returned context manager cannot be called twice, i.e. + my_noop = NoopContext() + with my_noop: + do1() + with my_noop: # trigger generator error + do2() + """ + + def __init__(self, *args, **kwargs): + self.args = args + self.kwargs = kwargs + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + pass + + +def make_registry_metaclass(class_name): + """ + Usage: + + TrainerRegistry = make_registry_metaclass('TrainerRegistry') + + class BaseTrainer(metaclass=TrainerRegistry): + pass + + class MyTrainer(BaseTrainer): + pass + + TrainerRegistry['MyTrainer'] -> MyTrainer class # syntax enabled by metaclass + TrainerRegistry.get_class('MyTrainer') # same as above + TrainerRegistry.registry -> full dict of {name: trainer_class} + + Templated definition: + class TrainerRegistry(type): + registry = {} + + def __new__(cls, name, bases, attr): + new_cls = super().__new__(cls, name, bases, attr) + TrainerRegistry.registry[name] = new_cls + return new_cls + + def get_class(cls, name): + if name not in cls.registry: + raise KeyError( + f"Trainer class {name} not found in registry: " + f"{pprint.pformat(cls.registry)}" + ) + return cls.registry[name]""" + + def new__(cls, name, bases, attr): + """ + Change the attr dict to dynamically add methods and attributes + """ + new_cls = type.__new__(cls, name, bases, attr) + cls.registry[name] = new_cls + return new_cls + + def get_class(cls, name): + if name not in cls.registry: + existing_cls = list(cls.registry.keys()) + raise KeyError(f"{class_name} class '{name}' not found in registry: {existing_cls}") + return cls.registry[name] + + def instantiate(cls_, cls, **kwargs): + Cls = cls_.get_class(cls) + return Cls(**kwargs) + + class _BracketOperator(type): + def __getitem__(cls, name): + return get_class(cls, name) + + return types.new_class( + class_name, + bases=(type,), + kwds={"metaclass": _BracketOperator}, + exec_body=lambda ns: ns.update( + { + "registry": {}, + "__new__": new__, + "get_class": classmethod(get_class), + "instantiate": classmethod(instantiate), + } + ), + ) + + +class ClassRegistry: + """ + May be a preferred way over make_registry_metaclass if your code does not support + metaclass well, e.g. pickle or Ray + + Use in conjunction with `__init_subclass__` hook in your base class + + class BaseClass: + registry = ClassRegistry() + + def __init_subclass__(cls, **kwargs): + cls.registry.add(cls) + super().__init_subclass__(**kwargs) + + print(registry) + """ + + def __init__(self, base_class_name: str = None): + self.registry = {} + self._base_class_name = base_class_name + + def add(self, cls): + self.registry[cls.__name__] = cls + + def get(self, name): + if name not in self.registry: + existing_cls = list(self.registry.keys()) + base_name = self._base_class_name + " " if self._base_class_name else "" + raise KeyError(f"{base_name} subclass '{name}' not found in registry: {existing_cls}") + return self.registry[name] + + def __str__(self): + return pprint.pformat(self.registry) + + def __getitem__(self, name): + return self.get(name) + + def instantiate(self, cls, **kwargs): + return self.get(cls)(**kwargs) + + +# ======================================================== +# =================== Inspect utils ==================== +# ======================================================== + + +def func_parameters(func): + return inspect.signature(func).parameters + + +def func_has_arg(func, arg_name): + return arg_name in func_parameters(func) + + +def pack_varargs(args): + """ + Pack *args or a single list arg as list + + def f(*args): + arg_list = pack_varargs(args) + # arg_list is now packed as a list + """ + assert isinstance(args, tuple), "please input the tuple `args` as in *args" + if len(args) == 1 and is_sequence(args[0]): + return args[0] + else: + return args + + +def enable_list_arg(func): + """ + Function decorator. + If a function only accepts varargs (*args), + make it support a single list arg as well + """ + + @functools.wraps(func) + def wrapper(*args, **kwargs): + args = pack_varargs(args) + return func(*args, **kwargs) + + return wrapper + + +def enable_varargs(func): + """ + Function decorator. + If a function only accepts a list arg, make it support varargs as well + """ + + @functools.wraps(func) + def wrapper(*args, **kwargs): + args = pack_varargs(args) + return func(args, **kwargs) + + return wrapper + + +def pack_kwargs(args, kwargs): + """ + Pack **kwargs or a single dict arg as dict + + def f(*args, **kwargs): + kwdict = pack_kwargs(args, kwargs) + # kwdict is now packed as a dict + """ + if len(args) == 1 and is_mapping(args[0]): + assert not kwargs, "cannot have both **kwargs and a dict arg" + return args[0] # single-dict + else: + assert not args, "cannot have positional args if **kwargs exist" + return kwargs + + +def merge_kwargs(args, kwargs) -> Dict: + """ + Merge all dicts in `args` and keywords in kwargs. + + E.g. merge_kwargs({"a.b": 1, "a.c": 2}, foo=6, bar=8) + -> {"a.b": 1, "a.c": 2, "foo": 6, "bar": 8} + """ + kw_all = {} + for arg in args: + assert is_mapping(arg), f"{arg} is not a dict." + kw_all.update(arg) + kw_all.update(kwargs) + return kw_all + + +def enable_dict_arg(func): + """ + Function decorator. + If a function only accepts varargs (*args), + make it support a single list arg as well + """ + + @functools.wraps(func) + def wrapper(*args, **kwargs): + kwargs = pack_kwargs(args, kwargs) + return func(**kwargs) + + return wrapper + + +def enable_kwargs(func): + """ + Function decorator. + If a function only accepts a dict arg, make it support kwargs as well + """ + + @functools.wraps(func) + def wrapper(*args, **kwargs): + kwargs = pack_kwargs(args, kwargs) + return func(kwargs) + + return wrapper + + +def has_keys(D, keys: list): + assert is_mapping(D) + return all(key in D for key in keys) + + +def assert_has_keys(D, keys: list): + assert is_mapping(D), "Input is not a dict" + for key in keys: + if key not in D: + raise KeyError(f'Required key "{key}" is missing in dict {D}') + return True + + +def method_decorator(decorator): + """ + Decorator of decorator: transform a decorator that only works on normal + functions to a decorator that works on class methods + From Django form: https://goo.gl/XLjxKK + """ + + @functools.wraps(decorator) + def wrapped_decorator(method): + @functools.wraps(method) + def wrapper(self, *args, **kwargs): + def bound_func(*args2, **kwargs2): + return method(self, *args2, **kwargs2) + + return decorator(bound_func)(*args, **kwargs) + + return wrapper + + return wrapped_decorator + + +def accepts_varargs(func): + """ + If a function accepts *args + """ + params = inspect.signature(func).parameters + return any(param.kind == inspect.Parameter.VAR_POSITIONAL for param in params.values()) + + +def accepts_kwargs(func): + """ + If a function accepts **kwargs + """ + params = inspect.signature(func).parameters + return any(param.kind == inspect.Parameter.VAR_KEYWORD for param in params.values()) + + +def is_signature_compatible(func, *args, **kwargs): + sig = inspect.signature(func) + try: + sig.bind(*args, **kwargs) + return True + except TypeError: + return False + + +def make_list(x): + """ + Turns a singleton object to a list. If already a list, no change. + """ + if is_sequence(x): + return x + else: + return [x] + + +def make_tuple(elem, repeats): + """ + E.g. expand a singleton x into (x, x, x) + useful for things like image_size or kernal, which can be a single int/float + or a tuple of fixed size + """ + if is_sequence(elem): + assert len(elem) == repeats, f"length of input must be {repeats}: {elem}" + return elem + else: + return (elem,) * repeats + + +def accumulate(iterable, fn=lambda x, y: x + y): + """ + Return running totals + # _accumulate([1,2,3,4,5]) --> 1 3 6 10 15 + # _accumulate([1,2,3,4,5], operator.mul) --> 1 2 6 24 120 + """ + it = iter(iterable) + try: + total = next(it) + except StopIteration: + return + yield total + for element in it: + total = fn(total, element) + yield total + + +class DecoratorContextManager: + """ + Allow a context manager to be used as a decorator + From torch.auto_grad.grad_mode + """ + + def __call__(self, func): + if inspect.isgeneratorfunction(func): + return self._wrap_generator(func) + + @functools.wraps(func) + def decorate_context(*args, **kwargs): + with self.__class__(): + return func(*args, **kwargs) + + return decorate_context + + def _wrap_generator(self, func): + """Wrap each generator invocation with the context manager""" + + @functools.wraps(func) + def generator_context(*args, **kwargs): + gen = func(*args, **kwargs) + + # Generators are suspended and unsuspended at `yield`, hence we + # make sure the grad mode is properly set every time the execution + # flow returns into the wrapped generator and restored when it + # returns through our `yield` to our caller (see PR #49017). + cls = type(self) + try: + # Issuing `None` to a generator fires it up + with cls(): + response = gen.send(None) + + while True: + try: + # Forward the response to our caller and get its next request + request = yield response + + except GeneratorExit: + # Inform the still active generator about its imminent closure + with cls(): + gen.close() + raise + + except BaseException: + # Propagate the exception thrown at us by the caller + with cls(): + response = gen.throw(*sys.exc_info()) + + else: + # Pass the last request to the generator and get its response + with cls(): + response = gen.send(request) + + # We let the exceptions raised above by the generator's `.throw` or + # `.send` methods bubble up to our caller, except for StopIteration + except StopIteration as e: + # The generator informed us that it is done: take whatever its + # returned value (if any) was and indicate that we're done too + # by returning it (see docs for python's return-statement). + return e.value + + return generator_context + + def __enter__(self) -> None: + raise NotImplementedError + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + raise NotImplementedError diff --git a/groot/vla/common/utils/misc/image_utils.py b/groot/vla/common/utils/misc/image_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..46ae18cf8b9b89e577da99df8dd0919d07318c71 --- /dev/null +++ b/groot/vla/common/utils/misc/image_utils.py @@ -0,0 +1,225 @@ +""" +Visualizations +""" + +from __future__ import annotations + +import os +import time +from typing import Literal +import warnings + +import cv2 +import imageio +import matplotlib.pyplot as plt +import numpy as np +import torch + +from .array_tensor_utils import any_describe +from .misc_utils import global_once +from .torch_utils import torch_normalize + + +def to_image(img, channel_order="auto"): + """ + Returns: + numpy image of shape [H, W, C] + in "auto" mode, we assume C == 3 + """ + assert channel_order in ["hwc", "chw", "auto"] + if torch.is_tensor(img): + img = img.cpu().numpy() + assert isinstance(img, np.ndarray) + if img.ndim == 4: + assert img.shape[0] == 1 + img = img[0] + assert img.ndim == 3 + if channel_order == "auto": + # use C==3 to detect order + if img.shape[0] == 3: + channel_order = "chw" + else: + assert img.shape[-1] == 3, "image should either have [3,H,W] or [H,W,3]" + channel_order = "hwc" + img = img.astype(np.uint8) + if channel_order == "chw": + return np.transpose(img, (1, 2, 0)) + else: + return img + + +def imshow(img): + plt.imshow(to_image(img)) + + +def imsave(img, path): + imageio.imsave(os.path.expanduser(path), to_image(img)) + + +def imread(path, channel_order="chw", format="torch"): + assert channel_order in ["hwc", "chw"] + assert format in ["numpy", "torch"] + img = imageio.imread(path) + if channel_order == "chw": + img = np.transpose(img, (2, 0, 1)) # hwc -> chw + if format == "torch": + return torch.from_numpy(img) + else: + return img + + +class Cv2Display: + def __init__( + self, + window_name="display", + image_size=None, + channel_order="auto", + bgr2rgb=True, + step_sleep=0, + enabled=True, + ): + """ + Use cv2.imshow() to pop a window, requires virtual desktop GUI + + Args: + channel_order: auto, hwc, or chw + image_size: None to use the original image size, otherwise resize + step_sleep: sleep for a few seconds + """ + self._window_name = window_name + if isinstance(image_size, int): + image_size = (image_size, image_size) + else: + assert image_size is None or len(image_size) == 2 + self._image_size = image_size + assert channel_order in ["auto", "chw", "hwc"] + self._channel_order = channel_order + self._bgr2rgb = bgr2rgb + self._step_sleep = step_sleep + self._enabled = enabled + + def _resize(self, img): + if self._image_size is None: + return img + H, W = img.shape[:2] + Ht, Wt = self._image_size # target + return cv2.resize( + img, + self._image_size, + interpolation=cv2.INTER_AREA if Ht < H else cv2.INTER_LINEAR, + ) + + def _reorder(self, img): + if self._channel_order == "chw": + return np.transpose(img, (1, 2, 0)) + elif self._channel_order == "hwc": + return img + else: + if img.shape[0] in [1, 3]: # chw + return np.transpose(img, (1, 2, 0)) + else: + return img + + def __call__(self, img): + if not self._enabled: + return + import torch + + # prevent segfault in IsaacGym + display_var = os.environ.get("DISPLAY", None) + if not display_var: + os.environ["DISPLAY"] = ":0.0" + + if torch.is_tensor(img): + img = img.detach().cpu().numpy() + + img = self._resize(self._reorder(img)) + if self._bgr2rgb: + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + time.sleep(self._step_sleep) + cv2.imshow(self._window_name, img) + cv2.waitKey(1) + + if display_var is not None: + os.environ["DISPLAY"] = display_var + + def close(self): + if not self._enabled: + return + cv2.destroyWindow(self._window_name) + + +# ---------------- Image tensor handling ----------------- +def sanity_check_image_tensor( + img: torch.Tensor, on_error: Literal["raise", "warn", "ignore"] = "raise" +): + """ + Check if the input image tensor is all integers, which is wrong for any NN input. + This is a common case if the user forgets to normalize the image first + """ + assert on_error in [ + "raise", + "warn", + "ignore", + ], 'on_error must be "raise", "warn", or "ignore"' + if not img.dtype.is_floating_point: + msg = f"Image tensor is not floating point format, but {img.dtype}!" + if on_error == "raise": + raise ValueError(msg) + elif on_error == "warn": + warnings.warn(msg) + else: + return False + # check if all values in the image are close to an integer + if (img - torch.round(img)).abs().max() < 1e-5: + msg = ( + "Input image is all close to integers, " + "are you sure you have normalized it before passing it to a NN?" + ) + if on_error == "raise": + raise ValueError(msg) + elif on_error == "warn": + warnings.warn(msg) + else: + return False + return True + + +@torch.no_grad() +def basic_image_tensor_preprocess( + img, + mean: tuple[float, float, float] = (0.5, 0.5, 0.5), + std: tuple[float, float, float] = (0.5, 0.5, 0.5), + shape: tuple[int, int] | None = None, +): + """ + Check for resize, and divide by 255 + """ + import kornia + + assert torch.is_tensor(img) + assert img.dim() >= 4, any_describe(img) + original_shape = list(img.size()) + img = img.float() + img = img.flatten(0, img.dim() - 4) + assert img.dim() == 4 + + input_size = img.size()[-2:] + if global_once("groot.vla.common.utils.image_utils.basic_image_preprocess:input_size"): + assert img.max() > 2, "img should be between [0, 255] before normalize" + + if shape and input_size != shape: + if global_once("groot.vla.common.utils.image_utils.basic_image_preprocess:transform"): + warnings.warn( + f'{"Down" if shape < input_size else "Up"}sampling image' + f" from original resolution {input_size}x{input_size}" + f" to {shape}x{shape}" + ) + img = kornia.geometry.transform.resize(img, shape).clamp(0.0, 255.0) + + B, C, H, W = img.size() + assert C % 3 == 0, "channel must divide 3" + img = img.view(B * C // 3, 3, H, W) + img = torch_normalize(img / 255.0, mean=mean, std=std) + original_shape[-2:] = H, W + return img.view(original_shape) diff --git a/groot/vla/common/utils/misc/misc_utils.py b/groot/vla/common/utils/misc/misc_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ed486fec52a63fe1c97039c813ceaf03f5db5315 --- /dev/null +++ b/groot/vla/common/utils/misc/misc_utils.py @@ -0,0 +1,261 @@ +import codecs +from collections import Counter +import fnmatch +import hashlib +import os +import pickle +from typing import Any, Callable, Dict, List, Optional, Union + +from typing_extensions import Literal + + +def set_os_envs(envs: Optional[Dict[str, Any]] = None): + """ + Special value __delete__ or None indicates that the ENV_VAR should be removed + """ + if envs is None: + envs = {} + DEL = {None, "__delete__"} + for k, v in envs.items(): + if v in DEL: + os.environ.pop(k, None) + os.environ.update({k: str(v) for k, v in envs.items() if v not in DEL}) + + +def argmax(L): + return max(zip(L, range(len(L))))[1] + + +def _match_patterns_helper(element, patterns): + for p in patterns: + if callable(p) and p(element): + return True + if fnmatch.fnmatch(element, p): + return True + return False + + +def match_patterns( + item: str, + include: Union[str, List[str], Callable, List[Callable], None] = None, + exclude: Union[str, List[str], Callable, List[Callable], None] = None, + *, + precedence: Literal["include", "exclude"] = "exclude", +): + """ + Args: + include: None to disable `include` filter and delegate to exclude + precedence: "include" or "exclude" + """ + assert precedence in ["include", "exclude"] + if exclude is None: + exclude = [] + if isinstance(exclude, (str, Callable)): + exclude = [exclude] + if isinstance(include, (str, Callable)): + include = [include] + if include is None: + # exclude is the sole veto vote + return not _match_patterns_helper(item, exclude) + + if precedence == "include": + return _match_patterns_helper(item, include) + else: + if _match_patterns_helper(item, exclude): + return False + else: + return _match_patterns_helper(item, include) + + +def filter_patterns( + items: List[str], + include: Union[str, List[str], Callable, List[Callable], None] = None, + exclude: Union[str, List[str], Callable, List[Callable], None] = None, + *, + precedence: Literal["include", "exclude"] = "exclude", + ordering: Literal["original", "include"] = "original", +): + """ + Args: + ordering: affects the order of items in the returned list. Does not affect the + content of the returned list. + - "original": keep the ordering of items in the input list + - "include": order items by the order of include patterns + """ + assert ordering in ["original", "include"] + if include is None or isinstance(include, str) or ordering == "original": + return [ + x + for x in items + if match_patterns(x, include=include, exclude=exclude, precedence=precedence) + ] + else: + items = items.copy() + ret = [] + for inc in include: + for i, item in enumerate(items): + if item is None: + continue + if match_patterns(item, include=inc, exclude=exclude, precedence=precedence): + ret.append(item) + items[i] = None + return ret + + +def getitem_nested(cfg, key: str): + """ + Recursively get key, if key has '.' in it + """ + keys = key.split(".") + for k in keys: + assert k in cfg, f'{k} in key "{key}" does not exist in config' + cfg = cfg[k] + return cfg + + +def setitem_nested(cfg, key: str, value): + """ + Recursively get key, if key has '.' in it + """ + keys = key.split(".") + for k in keys[:-1]: + assert k in cfg, f'{k} in key "{key}" does not exist in config' + cfg = cfg[k] + cfg[keys[-1]] = value + + +def getattr_nested(obj, key: str): + """ + Recursively get attribute + """ + keys = key.split(".") + for k in keys: + assert hasattr(obj, k), f'{k} in attribute "{key}" does not exist' + obj = getattr(obj, k) + return obj + + +def setattr_nested(obj, key: str, value): + """ + Recursively set attribute + """ + keys = key.split(".") + for k in keys[:-1]: + assert hasattr(obj, k), f'{k} in attribute "{key}" does not exist' + obj = getattr(obj, k) + setattr(obj, keys[-1], value) + + +class PeriodicEvent: + """ + triggers every period + """ + + def __init__(self, period: int, initial_value=0): + self._period = period + assert self._period >= 1 + self._last_threshold = initial_value + self._last_value = initial_value + self._trigger_counts = 0 + + def __call__(self, new_value=None, increment=None): + assert bool(new_value is None) != bool(increment is None), ( + "you must specify one and only one of new_value or increment, " "but not both" + ) + d = self._period + if new_value is None: + new_value = self._last_value + increment + assert new_value >= self._last_value, ( + f"value must be monotonically increasing. " + f"Current value {new_value} < last value {self._last_value}" + ) + self._last_value = new_value + if new_value - self._last_threshold >= d: + self._last_threshold += (new_value - self._last_threshold) // d * d + self._trigger_counts += 1 + return True + else: + return False + + @property + def trigger_counts(self): + return self._trigger_counts + + @property + def current_value(self): + return self._last_value + + +class Once: + def __init__(self): + self._triggered = False + + def __call__(self): + if not self._triggered: + self._triggered = True + return True + else: + return False + + def __bool__(self): + raise RuntimeError("`Once` objects should be used by calling ()") + + +_GLOBAL_ONCE_SET = set() +_GLOBAL_NTIMES_COUNTER = Counter() + + +def global_once(name): + """ + Try this to automate the name: + https://gist.github.com/techtonik/2151727#gistcomment-2333747 + """ + if name in _GLOBAL_ONCE_SET: + return False + else: + _GLOBAL_ONCE_SET.add(name) + return True + + +def global_n_times(name, n: int): + """ + Triggers N times + """ + assert n >= 1 + if _GLOBAL_NTIMES_COUNTER[name] < n: + _GLOBAL_NTIMES_COUNTER[name] += 1 + return True + else: + return False + + +class Every: + def __init__(self, n: int, on_first: bool = False): + assert n > 0 + self._i = 0 if on_first else 1 + self._n = n + + def __call__(self): + return self._i % self._n == 0 + + def __bool__(self): + raise RuntimeError("`Every` objects should be used by calling ()") + + +def encode_base64(obj) -> str: + return codecs.encode(pickle.dumps(obj), "base64").decode() + + +def decode_base64(s: str): + return pickle.loads(codecs.decode(s.encode(), "base64")) + + +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 diff --git a/groot/vla/common/utils/misc/torch_utils.py b/groot/vla/common/utils/misc/torch_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e735532c4289262d632adfee93ca68581f248533 --- /dev/null +++ b/groot/vla/common/utils/misc/torch_utils.py @@ -0,0 +1,748 @@ +from __future__ import annotations + +from copy import deepcopy +import os +import random +import time +from typing import List, Optional, Tuple, Union + +import numpy as np +import torch +import torch.nn as nn +import tree +from typing_extensions import Literal + +from ..data_structure.tree_utils import tree_value_at_path +from ..io.file_utils import f_join +from ..io.print_utils import to_readable_count_str +from .functional_utils import assert_implements_method, implements_method + + +def weight_init(m): + """Custom weight init for Conv2D and Linear layers.""" + if isinstance(m, nn.Linear): + nn.init.orthogonal_(m.weight.data) + if hasattr(m.bias, "data"): + m.bias.data.fill_(0.0) + elif isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d): + gain = nn.init.calculate_gain("relu") + nn.init.orthogonal_(m.weight.data, gain) + if hasattr(m.bias, "data"): + m.bias.data.fill_(0.0) + + +def get_seed( + seed: Union[int, str, None], + handle_invalid_seed: Literal["none", "system", "raise"] = "none", +) -> Optional[int]: + """ + Args: + seed: + "system": use scrambled int based on system time + None or int < 0: invalid seed values, see `handle_invalid_seed` + int >= 0: returns seed + handle_invalid_seed: None or int < 0 + - "none": returns None + - "system": returns scrambled int based on system time + - "raise": raise Exception + """ + handle_invalid_seed = handle_invalid_seed.lower() + assert handle_invalid_seed in ["none", "system", "raise"] + if isinstance(seed, str): + assert seed in ["system"] + invalid = False + else: + assert seed is None or isinstance(seed, int) + invalid = seed is None or seed < 0 + + if seed == "system" or invalid and handle_invalid_seed == "system": + # https://stackoverflow.com/questions/27276135/python-random-system-time-seed + t = int(time.time() * 100000) + return ( + ((t & 0xFF000000) >> 24) + + ((t & 0x00FF0000) >> 8) + + ((t & 0x0000FF00) << 8) + + ((t & 0x000000FF) << 24) + ) + elif invalid: + if handle_invalid_seed == "none": + return None + elif handle_invalid_seed == "raise": + raise ValueError( + f"Invalid random seed: {seed}, " f'must be a non-negative integer or "system"' + ) + else: + raise NotImplementedError + else: + return seed + + +def set_deterministic(flag: bool = True): + if not flag: + return + + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + os.environ["HOROVOD_FUSION_THRESHOLD"] = "0" + import torch.backends.cudnn as cudnn + + cudnn.deterministic = True + cudnn.benchmark = False + if hasattr(torch, "use_deterministic_algorithms"): + # only available in PyTorch >= 1.9 + torch.use_deterministic_algorithms(True) + elif hasattr(torch, "set_deterministic"): + # only available in PyTorch >= 1.7 + torch.set_deterministic(True) + + +def set_seed_everywhere( + seed: Optional[Union[int, str]], + deterministic=False, + set_tensorflow=False, + handle_invalid_seed: Literal["none", "system", "raise"] = "none", +) -> Optional[int]: + """ + References: + - https://github.com/NVIDIA/framework-determinism/blob/master/pytorch.md + - https://pytorch.org/docs/stable/notes/randomness.html + - CUBLAS env var: + https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility + + Args: + seed: see `get_seed()` + handle_invalid_seed: see `get_seed()` + """ + set_deterministic(deterministic) + + seed = get_seed(seed, handle_invalid_seed=handle_invalid_seed) + if seed is None: + return None + + os.environ["PYTHONHASHSEED"] = str(seed) + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + if set_tensorflow: + try: + import tensorflow as tf + + tf.random.set_seed(seed) + except ImportError: + pass + return seed + + +class eval_mode(object): + def __init__(self, *models): + self.models = models + + def __enter__(self): + self.prev_states = [] + for model in self.models: + self.prev_states.append(model.training) + model.train(False) + + def __exit__(self, *args): + for model, state in zip(self.models, self.prev_states): + model.train(state) + return False + + +def get_device(x, strict: bool = False) -> int: + """ + Args: + x: can be any arbitrary nested structure of np array and torch tensor + strict: True to check all batch sizes are the same + """ + xs = tree.flatten(x) + + def _get_device(x): + if torch.is_tensor(x): + return x.device + elif isinstance(x, nn.Module): + return get_module_device(x) + else: + return None + + if strict: + devices = [_get_device(x) for x in xs] + assert all( + b == devices[0] for b in devices + ), f"devices must all be the same in nested structure: {devices}" + return devices[0] + else: + return _get_device(xs[0]) + + +def load_torch(*fpath: str, map_location="cpu") -> dict: + """ + Default maps to "cpu" + """ + fpath = str(f_join(fpath)) + try: + return torch.load(fpath, map_location=map_location) + except RuntimeError as e: + raise RuntimeError(f"{e}\n\n --- Error loading {fpath}") + + +def save_torch(D, *fpath): + """ + Supports both (D, fpath) and (fpath, D) arg order, as long as one of them is a str + """ + if isinstance(D, str): + assert not isinstance(fpath, str), "Either torch_save(D, fpath) " "or torch_save(fpath, D)" + fpath, D = D, fpath + torch.save(D, str(f_join(fpath))) + + +# Aliases for consistency with load_pickle, load_text, load_json/yaml, etc. +torch_load = load_torch +torch_save = save_torch +dump_torch = save_torch + + +def torch_compute_stats(x, precision: int = 2): + x = x.to(dtype=torch.float32) + return ( + f"mean|std: {torch.mean(x):.{precision}f} +/- {torch.std(x):.{precision}f}, " + f"median: {torch.median(x):.{precision}f}, " + f"max: {torch.max(x):.{precision}f}, min: {torch.min(x):.{precision}f}" + ) + + +def tensor_hash(x: torch.Tensor, mode: str = "mean"): + if isinstance(x, np.ndarray): + x = torch.from_numpy(x) + x = x.float().abs() + if mode == "sum": + x = x.sum() + elif mode == "mean": + x = x.mean() + else: + raise NotImplementedError + return float(x) + + +def torch_flatten_indices(indices: torch.Tensor, shape: Tuple[int]): + """ + Convert M dim indices to 1D indices with the given shape + + Args: + indices: BxM, batch_size x M-dimensional + """ + offsets = np.array(shape) # e.g. [3, 4, 5, 6] + offsets = np.append(offsets[1:], 1) # [4, 5, 6, 1] + offsets = np.cumprod(offsets[::-1])[::-1] # [4*5*6, 5*6, 6, 1] + offsets = torch.tensor(offsets.copy(), dtype=torch.long) + assert offsets.size() == (len(shape),) + return (indices * offsets.to(device=indices.device)).sum(dim=1) + + +def torch_multi_index_select(x: torch.Tensor, indices: torch.Tensor): + """ + Args: + x: N dim + indices: [B x M], M <= N, will select the first M-D from N-D + + Returns: + (N - M + 1) dim + """ + assert indices.ndim == 2 + B, idx_dim = indices.size() + x_shape = x.size() + assert len(x_shape) >= idx_dim + remainder_dim = len(x_shape) - idx_dim + if remainder_dim == 0: + x = torch.flatten(x) + else: + x = x.view(-1, *x_shape[-remainder_dim:]) # flatten the first M dims + # convert indices to a 1D flattened array + indices = torch_flatten_indices(indices, x_shape[:idx_dim]) + selected = x[indices] + return selected + + +# ========== module operations ========= +def set_requires_grad(model, requires_grad): + if torch.is_tensor(model): + model.requires_grad = requires_grad + else: + for param in model.parameters(): + param.requires_grad = requires_grad + + +def freeze_params(model): + set_requires_grad(model, False) + if not torch.is_tensor(model): + model.eval() + + +def unfreeze_params(model): + set_requires_grad(model, True) + if not torch.is_tensor(model): + model.train() + + +def clip_grad_value(model, max_value): + with torch.no_grad(): + nn.utils.clip_grad_value_(model.parameters(), max_value) + + +def clip_grad_norm(model, max_norm, norm_type=2): + """ + Returns: + Total norm of the parameters (viewed as a single vector). + """ + with torch.no_grad(): + return nn.utils.clip_grad_norm_(model.parameters(), max_norm=max_norm, norm_type=norm_type) + + +def implements_state_dict(object, requires_load_method: bool = False): + cond = implements_method(object, "state_dict") + if requires_load_method: + return cond and implements_method(object, "load_state_dict") + else: + return cond + + +def unwrap_ddp_model(model): + if hasattr(model, "module") and len(list(model.children())) == 1: + model = model.module + return model + + +class DDPMethodWrapper(nn.Module): + """ + Wraps another module's method as forward(), because DDP only works on forward() + This module can be wrapped with DDP and directly called. + It will not save any extra parameters + """ + + def __init__(self, net: nn.Module, method_name: str): + super().__init__() + self.net = net + assert_implements_method(net, method_name) + self._method_name = method_name + + def forward(self, *args, **kwargs): + return getattr(self.net, self._method_name)(*args, **kwargs) + + def state_dict(self): + return {} + + +def to_state_dict(objects, to_cpu: bool = False, copy: bool = False, unwrap_ddp: bool = False): + """ + Anything that has state_dict() method, e.g. nn.Module, Optimizer, LRScheduler, etc. + + Args: + to_cpu: True to copy to CPU. The original tensors will still be on GPU. + copy: takes effect if and only if to_cpu is False + """ + + def _transfer(x): + if torch.is_tensor(x): + x = x.detach() + if to_cpu: + return x.cpu() + elif copy: + return x.clone() + return x + + def _to_state_dict(m): + if implements_state_dict(m): + if isinstance(m, nn.Module) and unwrap_ddp: + m = unwrap_ddp_model(m) + return tree.map_structure(_transfer, m.state_dict()) + else: + return _transfer(m) + + return tree.map_structure(_to_state_dict, objects) + + +def load_state_dict(objects, states, strip_prefix=None, strict=False): + """ + Args: + strict: objects and states must match exactly + strip_prefix: only match the keys that have the prefix, and strip it + """ + + def _load(paths, obj): + if not implements_method(obj, "load_state_dict"): + raise ValueError(f"Object {type(obj)} does not support load_state_dict() method") + try: + state = tree_value_at_path(states, paths) + except ValueError: # paths do not exist in `states` structure + if strict: + raise + else: + return + if strip_prefix: + assert isinstance(strip_prefix, str) + state = { + k[len(strip_prefix) :]: v for k, v in state.items() if k.startswith(strip_prefix) + } + if isinstance(obj, nn.Module): + return obj.load_state_dict(state, strict=strict) + else: + return obj.load_state_dict(state) + + return tree.map_structure_with_path(_load, objects) + + +def count_parameters(model): + return sum(x.numel() for x in model.parameters()) + + +def readable_count_parameters(model, precision: int = 2): + return to_readable_count_str(count_parameters(model), precision=precision) + + +def get_module_device(model): + """ + Returns: + first model parameter's device + """ + return next(model.parameters()).device + + +def maybe_transfer_module(model, device): + """ + Transfer a module to another device if and only if they are on different devices. + Assumes that the module's first parameter determines the module device, i.e. + no model parallelism. + + Returns: + True if module is transferred to a different device, False otherwise + """ + if device is None: + return False + device = torch.device(device) + if get_module_device(model) != device: + model.to(device=device) + return True + else: + return False + + +def clone_model(model): + with torch.no_grad(): + new_model = deepcopy(model).to(get_module_device(model)) + # new_model.load_state_dict(model.state_dict()) + return new_model + + +def update_soft_params(net, target_net, tau): + for param, target_param in zip(net.parameters(), target_net.parameters()): + target_param.data.copy_(tau * param.data + (1 - tau) * target_param.data) + + +def tie_weights(src, trg): + # TODO deprecate this + assert type(src) is type(trg) + trg.weight = src.weight + trg.bias = src.bias + + +def torch_normalize(tensor: torch.Tensor, mean, std, inplace=False): + """ + Adapted from https://pytorch.org/docs/stable/_modules/torchvision/transforms/functional.html#normalize + + Normalize a tensor image with mean and standard deviation. + + .. note:: + This transform acts out of place by default, i.e., it does not mutates the input tensor. + + See :class:`~torchvision.transforms.Normalize` for more details. + + Args: + tensor (Tensor): Tensor image of size (C, H, W) to be normalized. + mean (sequence): Sequence of means for each channel. + std (sequence): Sequence of standard deviations for each channel. + inplace(bool,optional): Bool to make this operation inplace. + + Returns: + Tensor: Normalized Tensor image. + """ + if not torch.is_tensor(tensor): + raise TypeError("tensor should be a torch tensor. Got {}.".format(type(tensor))) + + if not inplace: + tensor = tensor.clone() + + dtype = tensor.dtype + mean = torch.as_tensor(mean, dtype=dtype, device=tensor.device) + std = torch.as_tensor(std, dtype=dtype, device=tensor.device) + if (std == 0).any(): + raise ValueError( + f"std evaluated to zero after conversion to {dtype}, leading to division by zero." + ) + if mean.ndim == 1: + mean = mean[:, None, None] + if std.ndim == 1: + std = std[:, None, None] + tensor.sub_(mean).div_(std) + return tensor + + +def contains_rnn(net: nn.Module) -> bool: + for m in net.modules(): + if isinstance(m, nn.RNNBase): + return True + return False + + +def multi_one_hot(x, num_classes: List[int], to_float=True): + """ + Concatenates multiple one-hot matrices, useful for embedding MultiDiscrete action space + + Args: + x: torch.long, [*N, D] + num_classes: list len == D, match the last dim of x + + Returns: + [*N, sum(num_classes)] + """ + from torch.nn.functional import one_hot + + assert x.dtype == torch.long + assert x.dim() >= 2, x.size() + assert len(num_classes) == x.size(-1), f"{len(num_classes)} != {x.size(1)}" + result = torch.cat( + [one_hot(t, c) for t, c in zip(torch.unbind(x, dim=-1), num_classes)], dim=-1 + ) + if to_float: + return result.float() + else: + return result + + +def _random_derangement(n): + while True: + v = [i for i in range(n)] + for j in range(n - 1, -1, -1): + p = random.randint(0, j) + if v[p] == j: + break + else: + v[j], v[p] = v[p], v[j] + else: + if v[0] != 0: + return tuple(v) + + +def random_derangement(n, format: Literal["list", "numpy", "torch"] = "torch"): + """ + Early refusal algorithm, described at + https://stackoverflow.com/questions/25200220/generate-a-random-derangement-of-a-list + Derangement is permuation without fixed point, useful for constructing negative + pairs in contrastive learning. + """ + assert format in ["list", "numpy", "torch"] + D = _random_derangement(n) + if format == "list": + return D + elif format == "numpy": + return np.array(D, dtype=np.long) + elif format == "torch": + return torch.tensor(D, dtype=torch.long) + else: + raise NotImplementedError(f"Unknown format {format}") + + +def classify_accuracy( + output, + target, + topk: Union[int, List[int], Tuple[int]] = 1, + mask=None, + reduction="mean", + scale_100=False, +): + """ + Computes the accuracy over the k top predictions for the specified values of k. + Accuracy is a float between 0.0 and 1.0 + + Args: + topk: if int, return a single acc. If tuple, return a tuple of accs + mask: shape [batch_size,], binary mask of whether to include this sample or not + """ + if isinstance(topk, int): + topk = [topk] + is_int = True + else: + is_int = False + + batch_size = target.size(0) + assert output.size(0) == batch_size + if mask is not None: + assert mask.dim() == 1 + assert mask.size(0) == batch_size + + assert reduction in ["sum", "mean", "none"] + if reduction != "mean": + assert not scale_100, f"reduce={reduction} does not support scale_100=True" + + with torch.no_grad(): + maxk = max(topk) + + _, pred = output.topk(maxk, 1, True, True) + pred = pred.t() + correct = pred.eq(target.view(1, -1).expand_as(pred)) + if mask is not None: + correct = mask * correct + + mult = 100.0 if scale_100 else 1.0 + res = [] + for k in topk: + correct_k = correct[:k].int().sum(dim=0) + if reduction == "mean": + if mask is not None: + # fmt: off + res.append( + float(correct_k.float().sum().mul_(mult / mask.sum().item()).item()) + ) + # fmt: on + else: + res.append(float(correct_k.float().sum().mul_(mult / batch_size).item())) + elif reduction == "sum": + res.append(int(correct_k.sum().item())) + elif reduction == "none": + res.append(correct_k) + else: + raise NotImplementedError(f"Unknown reduce={reduction}") + + if is_int: + assert len(res) == 1, "INTERNAL" + return res[0] + else: + return res + + +def sequential_split_dataset(dataset: torch.utils.data.Dataset, split_portions: list[float]): + """ + Split a dataset into multiple datasets, each with a different portion of the + original dataset. Uses torch.utils.data.Subset. + """ + from .functional_utils import accumulate + + assert len(split_portions) > 0, "split_portions must be a non-empty list" + assert all(0.0 <= p <= 1.0 for p in split_portions), f"{split_portions=}" + assert abs(sum(split_portions) - 1.0) < 1e-6, f"{sum(split_portions)=} != 1.0" + L = len(dataset) + assert L > 0, "dataset must be non-empty" + # split the list with proportions + lengths = [int(p * L) for p in split_portions] + # make sure the last split fills the full dataset + lengths[-1] += L - sum(lengths) + indices = list(range(L)) + + return [ + torch.utils.data.Subset(dataset, indices[offset - length : offset]) + for offset, length in zip(accumulate(lengths), lengths) + ] + + +class RunningMeanStd: + def __init__(self): + """ + Calulates the running mean and std of a data stream + https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm + """ + self._mean = None + self._var = None + self._count = 0 + + @property + def mean(self): + return self._mean + + @property + def var(self): + return self._var + + @property + def std(self): + if isinstance(self._var, np.ndarray): + return np.sqrt(self._var) + else: + return self._var.sqrt() + + @property + def count(self): + return self._count + + def update(self, values: np.ndarray | torch.Tensor) -> None: + from .array_tensor_utils import any_mean, any_variance, get_batch_size + + batch_mean = any_mean(values, dim=0) + # our running var calculation currently only supports unbiased=False + batch_var = any_variance(values, dim=0, unbiased=False) + batch_count = get_batch_size(values) + self.update_from_moments(batch_mean, batch_var, batch_count) + + def update_from_moments( + self, + batch_mean: np.ndarray | torch.Tensor, + batch_var: np.ndarray | torch.Tensor, + batch_count: int, + ) -> None: + from .array_tensor_utils import any_get_shape + + is_tensor = torch.is_tensor(batch_mean) + _zeros = batch_mean.new_zeros if is_tensor else np.zeros + if self._mean is None: + self._mean = _zeros(any_get_shape(batch_mean)) + if self._var is None: + self._var = _zeros(any_get_shape(batch_var)) + 1.0 + + delta = batch_mean - self._mean + tot_count = self._count + batch_count + assert tot_count > 0, "count must be > 0" + + new_mean = self._mean + delta * batch_count / tot_count + m_a = self._var * self._count + m_b = batch_var * batch_count + m_2 = m_a + m_b + delta * delta * self._count * batch_count / tot_count + new_var = m_2 / tot_count + + self._mean = new_mean + self._var = new_var + self._count = tot_count + + +class AverageMeter: + """Computes and stores the average and current value""" + + def __init__(self, name="", fmt="f"): + self._name = name + self._fmt = fmt + self.reset() + + def reset(self): + self._sum = 0.0 + self._count = 0.0 + + @torch.no_grad() + def update(self, value, n=1): + if torch.is_tensor(value): + value = value.detach() + self._sum += value * n + self._count += n + + @torch.no_grad() + def compute(self): + return float(self._sum / self._count) + + def __float__(self): + return self.compute() + + def __str__(self): + if self._fmt: + s = f"{float(self):{self._fmt}}" + else: + s = str(float(self)) + if self._name: + return f"{self._name}: {s}" + return s diff --git a/groot/vla/common/utils/misc/video_utils.py b/groot/vla/common/utils/misc/video_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..481f169cef3adf518df273e1175603c31cfcfdd0 --- /dev/null +++ b/groot/vla/common/utils/misc/video_utils.py @@ -0,0 +1,487 @@ +import json +import subprocess + +import av +import cv2 +import numpy as np +import torchvision + +# Import decord with graceful fallback +try: + import decord + + DECORD_AVAILABLE = True +except ImportError: + DECORD_AVAILABLE = False + +try: + import torchcodec + + TORCHCODEC_AVAILABLE = True +except (ImportError, RuntimeError): + TORCHCODEC_AVAILABLE = False + + +def _get_video_info_ffmpeg(video_path: str) -> dict: + """Get video metadata using ffprobe.""" + cmd = [ + "ffprobe", + "-v", + "error", + "-select_streams", + "v:0", + "-show_entries", + "stream=nb_frames,duration,r_frame_rate", + "-of", + "json", + video_path, + ] + + try: + output = subprocess.check_output(cmd, stderr=subprocess.STDOUT).decode("utf-8") + probe_data = json.loads(output) + stream = probe_data["streams"][0] + + # Parse frame rate (comes as fraction like "15/1") + if "/" in stream["r_frame_rate"]: + num, den = map(int, stream["r_frame_rate"].split("/")) + fps = num / den + else: + fps = float(stream["r_frame_rate"]) + + # Get frame count and duration + nb_frames = int(stream.get("nb_frames", 0)) + duration = float(stream.get("duration", 0)) + + # If nb_frames is not available, estimate from duration and fps + if nb_frames == 0 and duration > 0: + nb_frames = int(duration * fps) + + return { + "nb_frames": nb_frames, + "fps": fps, + "duration": duration, + } + except (subprocess.CalledProcessError, json.JSONDecodeError, KeyError) as e: + raise ValueError(f"Failed to get video info for {video_path}: {e}") + + +def _extract_frames_ffmpeg(video_path: str, frame_indices: list[int]) -> np.ndarray: + """Extract specific frames using ffmpeg.""" + frames = [] + + for idx in frame_indices: + # Use ffmpeg to extract a specific frame + cmd = [ + "ffmpeg", + "-i", + video_path, + "-vf", + f"select=eq(n\\,{idx})", + "-vframes", + "1", + "-f", + "image2pipe", + "-pix_fmt", + "rgb24", + "-vcodec", + "rawvideo", + "-", + ] + + try: + output = subprocess.check_output(cmd, stderr=subprocess.DEVNULL) + + # Check if output is empty (frame doesn't exist) + if len(output) == 0: + raise subprocess.CalledProcessError(1, cmd) + + # Get frame dimensions by probing first + if len(frames) == 0: + info_cmd = [ + "ffprobe", + "-v", + "error", + "-select_streams", + "v:0", + "-show_entries", + "stream=width,height", + "-of", + "json", + video_path, + ] + info_output = subprocess.check_output(info_cmd).decode("utf-8") + info_data = json.loads(info_output) + width = info_data["streams"][0]["width"] + height = info_data["streams"][0]["height"] + + # Decode raw RGB data + frame_data = np.frombuffer(output, dtype=np.uint8) + frame = frame_data.reshape((height, width, 3)) + frames.append(frame) + + except subprocess.CalledProcessError: + # Frame might not exist, create a black frame + if len(frames) > 0: + frames.append(np.zeros_like(frames[0])) + else: + # Default fallback frame + frames.append(np.zeros((480, 640, 3), dtype=np.uint8)) + + return np.array(frames) + + +def _extract_frames_at_timestamps_ffmpeg(video_path: str, timestamps: list[float]) -> np.ndarray: + """Extract frames at specific timestamps using ffmpeg.""" + frames = [] + + for timestamp in timestamps: + cmd = [ + "ffmpeg", + "-ss", + str(timestamp), + "-i", + video_path, + "-vframes", + "1", + "-f", + "image2pipe", + "-pix_fmt", + "rgb24", + "-vcodec", + "rawvideo", + "-", + ] + + try: + output = subprocess.check_output(cmd, stderr=subprocess.DEVNULL) + + # Check if output is empty (timestamp doesn't exist) + if len(output) == 0: + raise subprocess.CalledProcessError(1, cmd) + + # Get frame dimensions + if len(frames) == 0: + info_cmd = [ + "ffprobe", + "-v", + "error", + "-select_streams", + "v:0", + "-show_entries", + "stream=width,height", + "-of", + "json", + video_path, + ] + info_output = subprocess.check_output(info_cmd).decode("utf-8") + info_data = json.loads(info_output) + width = info_data["streams"][0]["width"] + height = info_data["streams"][0]["height"] + + # Decode raw RGB data + frame_data = np.frombuffer(output, dtype=np.uint8) + frame = frame_data.reshape((height, width, 3)) + frames.append(frame) + + except subprocess.CalledProcessError: + # Timestamp might be out of bounds, use last frame or black frame + if len(frames) > 0: + frames.append(frames[-1]) + else: + frames.append(np.zeros((480, 640, 3), dtype=np.uint8)) + + return np.array(frames) + + +def _extract_all_frames_ffmpeg(video_path: str) -> tuple[np.ndarray, np.ndarray]: + """Extract all frames and their timestamps using ffmpeg.""" + # Get video info + info = _get_video_info_ffmpeg(video_path) + fps = info["fps"] + + # Extract all frames + cmd = [ + "ffmpeg", + "-i", + video_path, + "-f", + "image2pipe", + "-pix_fmt", + "rgb24", + "-vcodec", + "rawvideo", + "-", + ] + + try: + output = subprocess.check_output(cmd, stderr=subprocess.DEVNULL) + + # Get frame dimensions + info_cmd = [ + "ffprobe", + "-v", + "error", + "-select_streams", + "v:0", + "-show_entries", + "stream=width,height", + "-of", + "json", + video_path, + ] + info_output = subprocess.check_output(info_cmd).decode("utf-8") + info_data = json.loads(info_output) + width = info_data["streams"][0]["width"] + height = info_data["streams"][0]["height"] + + # Decode all frames + frame_data = np.frombuffer(output, dtype=np.uint8) + total_pixels = len(frame_data) // 3 + actual_frames = total_pixels // (width * height) + + frames = frame_data[: actual_frames * width * height * 3].reshape( + (actual_frames, height, width, 3) + ) + + # Generate timestamps + timestamps = np.arange(actual_frames) / fps + + return frames, timestamps + + except subprocess.CalledProcessError as e: + raise ValueError(f"Failed to extract frames from {video_path}: {e}") + + +def get_frames_by_indices( + video_path: str, + indices: list[int] | np.ndarray, + video_backend: str = "ffmpeg", + video_backend_kwargs: dict = {}, +) -> np.ndarray: + if video_backend == "decord": + if not DECORD_AVAILABLE: + raise ImportError("decord is not available. Install it with: pip install decord") + 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 == "ffmpeg": + return _extract_frames_ffmpeg(video_path, list(indices)) + 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 + else: + raise NotImplementedError + + +def get_frames_by_timestamps( + video_path: str, + timestamps: list[float] | np.ndarray, + video_backend: str = "ffmpeg", + video_backend_kwargs: dict = {}, + fps: None | float = None, +) -> 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 "ffmpeg". + fps (float, optional): FPS of the video. Defaults to 30. + Returns: + np.ndarray: Frames at the specified timestamps. + """ + if video_backend == "decord": + if not DECORD_AVAILABLE: + raise ImportError("decord is not available. Install it with: pip install decord") + 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 + ) + + # https://docs.pytorch.org/torchcodec/stable/generated/torchcodec.decoders.VideoStreamMetadata.html#torchcodec.decoders.VideoStreamMetadata + # Temporary fix: use 30 fps as the fps of the video (agibot) + # TODO: get fps as parameter + if fps is None: + fps = decoder.metadata.average_fps + interval = 1 / fps + timestamps = np.array(timestamps).astype(np.float64) + + if np.all(timestamps == 0): + timestamps = np.arange(len(timestamps)) / fps + + # Get video duration range from first and last frames + # This is a robust way to get valid timestamp range without depending on specific metadata attributes + first_frame = decoder.get_frames_at(indices=[0]) + last_frame = decoder.get_frames_at(indices=[len(decoder) - 1]) + min_pts = float(first_frame.pts_seconds[0]) + max_pts = float(last_frame.pts_seconds[0]) + + # Clamp timestamps to valid range to avoid RuntimeError + timestamps = np.clip(timestamps, min_pts, max_pts) + + # Correct float precision issues in timestamps + # E.g. for 5fps video: [1.0, 1.20000005, 1.39999998] -> [1.0, 1.2, 1.4] + # Without this, the torchcodec will read the delayed frame (e.g. 1.39999998 -> 1.2) + # Round to nearest frame interval to prevent torchcodec from reading wrong frames + # Allow max 1% error from expected interval + if fps is None: + closest_timestamps = np.round(timestamps / interval) * interval + # Re-clamp after rounding to ensure still in valid range + closest_timestamps = np.clip(closest_timestamps, min_pts, max_pts) + timestamp_errors = np.abs(closest_timestamps - timestamps) / interval + invalid_mask = timestamp_errors >= 0.01 + if np.any(invalid_mask): + invalid_indices = np.where(invalid_mask)[0] + invalid_timestamps = timestamps[invalid_indices] + raise ValueError( + f"Try to read invalid timestamps {invalid_timestamps} from video {video_path} (FPS: {fps})" + ) + + timestamps = closest_timestamps + + return decoder.get_frames_played_at(seconds=timestamps).data.numpy() + elif video_backend == "ffmpeg": + return _extract_frames_at_timestamps_ffmpeg(video_path, list(timestamps)) + 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 == "torchvision_av": + # set backend + torchvision.set_video_backend("pyav") + + # set a video stream reader + reader = torchvision.io.VideoReader(video_path, "video") + + # set the first and last requested timestamps + # Note: previous timestamps are usually loaded, since we need to access the previous key frame + first_ts = timestamps[0] + last_ts = timestamps[-1] + + # access closest key frame of the first requested frame + # Note: closest key frame timestamp is usally smaller than `first_ts` (e.g. key frame can be the first frame of the video) + # for details on what `seek` is doing see: https://pyav.basswood-io.com/docs/stable/api/container.html?highlight=inputcontainer#av.container.InputContainer.seek + reader.seek(first_ts, keyframes_only=True) + + # Decode frames sequentially, storing the ones we need in a dictionary + # to map timestamps to frame data. This allows for easy re-ordering later. + found_frames_map = {} + tolerance = 0.001 # 1ms tolerance for timestamp matching + + for frame in reader: + current_ts = frame["pts"] + + # Use tolerance-based matching instead of exact match + for ts in timestamps: + if ts not in found_frames_map and abs(current_ts - ts) < tolerance: + found_frames_map[ts] = frame["data"] + break + + if current_ts >= last_ts + tolerance or len(found_frames_map) == len(timestamps): + break + + reader.container.close() + reader = None + + # Debug: print timestamp matching results + print(f"[video_utils] Requested {len(timestamps)} timestamps: {timestamps[:4]}{'...' if len(timestamps) > 4 else ''}") + print(f"[video_utils] Found {len(found_frames_map)} frames with tolerance={tolerance}s") + if len(found_frames_map) < len(timestamps): + missing = [ts for ts in timestamps if ts not in found_frames_map] + print(f"[video_utils] WARNING: Missing timestamps: {missing[:4]}{'...' if len(missing) > 4 else ''}") + + frames = np.array(list(found_frames_map.values())) + return frames.transpose(0, 2, 3, 1) + + else: + raise NotImplementedError + + +def get_all_frames( + video_path: str, + video_backend: str = "ffmpeg", + video_backend_kwargs: dict = {}, +) -> tuple[np.ndarray, np.ndarray]: + """Get all frames from a video. + + Returns: + tuple[np.ndarray, np.ndarray]: Frames and timestamps. + """ + if video_backend == "decord": + if not DECORD_AVAILABLE: + raise ImportError("decord is not available. Install it with: pip install decord") + vr = decord.VideoReader(video_path, **video_backend_kwargs) + frames = vr.get_batch(range(len(vr))).asnumpy() + return frames, vr.get_frame_timestamp(range(len(vr)))[:, 0] + 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 == "ffmpeg": + return _extract_all_frames_ffmpeg(video_path) + elif video_backend == "pyav": + container = av.open(video_path) + stream = container.streams.video[0] + assert stream.time_base is not None + frames = [] + timestamps = [] + for frame in container.decode(video=0): + frames.append(frame.to_ndarray(format="rgb24")) + timestamps.append(frame.pts * stream.time_base) + container.close() + return np.stack(frames), np.array(timestamps) + + else: + raise NotImplementedError diff --git a/groot/vla/configs/conf.yaml b/groot/vla/configs/conf.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0334a8dce35473d6fd9dcc4d7b122fa882b41f5c --- /dev/null +++ b/groot/vla/configs/conf.yaml @@ -0,0 +1,151 @@ +defaults: + - _self_ # all below configs will override this conf.yaml + - model: dreamzero/vla + - data: dreamzero/droid_horizon_relative + - override hydra/hydra_logging: disabled # disable hydra logging + - override hydra/job_logging: disabled # disable hydra job logging + + +# === Model Arguments === +model: ??? + +# === Data Arguments === +train_dataset: ??? + +# ======== Trainer ======== +trainer: + _target_: groot.vla.experiment.VLATrainer + _partial_: true + _recursive_: false + callbacks: + model: ??? # model + train_dataset: ??? # train_dataset + compute_dtype: ??? # dtype_from_string(model.config.model_dtype) + benchmark_time: false # whether or not to benchmark time for training + # Legacy per-step profiling (profiles every N steps) + enable_profiling: false # (legacy) enable per-step profiling in training_step + profiling_steps: 5 # (legacy) profile every N steps + # ProfCallback: window-based profiling + enable_prof_callback: false # enable ProfCallback for window-based profiling + profile_start_step: 50 # session step to start profiling + profile_warmup_steps: 1 # warmup steps for profiler + profile_active_steps: 3 # active profiling steps + profile_record_shapes: false # record tensor shapes (adds overhead) + profile_with_stack: false # record Python stack traces + profile_memory: false # record memory allocation + +# === Training Arguments === + +wandb_project: ??? # needs to be specified by user +output_dir: ??? # need to be specified by user +load_from_yaml: # need to be specified by user, will override the current config +gear_credentials: null +upload_checkpoints: false +upload_every: 1000 +upload_last_n_checkpoints: 5 +remove_unused_columns: false +bf16: false +tf32: false +global_batch_size: null +raise_error_if_global_batch_size_not_set: false +per_device_train_batch_size: 256 +per_device_eval_batch_size: 64 +gradient_accumulation_steps: 1 +dataloader_num_workers: 10 +dataloader_pin_memory: true +dataloader_persistent_workers: true +optim: adamw_torch +learning_rate: 1e-4 +adam_beta1: 0.95 +adam_beta2: 0.999 +adam_epsilon: 1e-8 +weight_decay: 1e-6 +lr_scheduler_type: cosine +warmup_ratio: 0.05 +logging_steps: 10.0 +num_train_epochs: 1000 +max_steps: -1 +save_strategy: steps +save_steps: 500 +eval_strategy: "no" # there has to be a double quote; otherwise a bare `no` will be interpreted as False +save_total_limit: 8 +report_to: wandb +seed: 42 +do_eval: false +gradient_checkpointing: false +ddp_find_unused_parameters: false +ddp_bucket_cap_mb: 100 +ray_num_workers: ??? +eval_bf16: true +torch_compile_mode: null + +pretrained_model_path: null +only_tune_projectors: false + +save_llm: false +save_lora_only: false +save_value_model: false +save_q_model: false + +download_cache: false + +training_args: + _target_: transformers.TrainingArguments + output_dir: ${output_dir} + run_name: ??? # training_args.output_dir.split("/")[-1] + remove_unused_columns: ${remove_unused_columns} + deepspeed: "" + gradient_checkpointing: ${gradient_checkpointing} + bf16: ${bf16} + tf32: ${tf32} + per_device_train_batch_size: ${per_device_train_batch_size} + per_device_eval_batch_size: ${per_device_eval_batch_size} + gradient_accumulation_steps: ${gradient_accumulation_steps} + dataloader_num_workers: ${dataloader_num_workers} + dataloader_pin_memory: ${dataloader_pin_memory} + dataloader_persistent_workers: ${dataloader_persistent_workers} + optim: ${optim} + adam_beta1: ${adam_beta1} + adam_beta2: ${adam_beta2} + adam_epsilon: ${adam_epsilon} + learning_rate: ${learning_rate} + weight_decay: ${weight_decay} + warmup_ratio: ${warmup_ratio} + lr_scheduler_type: ${lr_scheduler_type} + logging_steps: ${logging_steps} + num_train_epochs: ${num_train_epochs} + max_steps: ${max_steps} + save_strategy: ${save_strategy} + save_steps: ${save_steps} + save_total_limit: ${save_total_limit} + report_to: ${report_to} + seed: ${seed} + do_eval: ${do_eval} + ddp_find_unused_parameters: ${ddp_find_unused_parameters} + ddp_bucket_cap_mb: ${ddp_bucket_cap_mb} + torch_compile_mode: ${torch_compile_mode} + +# === T-Rex wandb video reconstruction (optional) === +enable_wandb_video_reconstruction: false +wandb_video_reconstruction_steps: 100 +wandb_video_reconstruction_episode: 0 +wandb_video_reconstruction_num_chunks: 4 +wandb_video_reconstruction_fps: 5 +wandb_video_use_dataset_prompt: true +wandb_video_prompt: "perform the task" +wandb_video_overlay_tracks: false +wandb_video_save_local: true +wandb_video_save_tracks: false +wandb_video_track_trail_steps: 8 +wandb_video_reconstruction_inference_steps: 1 +wandb_video_start_chunk_index: 0 +wandb_video_compare_tracks_on_gt: true + +# === Profiling Arguments === +profile_dir: null + +# === Disable Hydra Config === +hydra: + output_subdir: null + run: + dir: . diff --git a/groot/vla/configs/data/dreamzero/agibot_relative.yaml b/groot/vla/configs/data/dreamzero/agibot_relative.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1972af4e9ab846f94bc9304a238b731e4c2deb14 --- /dev/null +++ b/groot/vla/configs/data/dreamzero/agibot_relative.yaml @@ -0,0 +1,54 @@ +# @package _global_ + +defaults: + - dreamzero/base_48_wan_fine_aug_relative + - _self_ # this file will override the base + +max_state_dim: 64 +use_global_metadata: false +relative_action: true +relative_action_per_horizon: false +relative_action_keys: + - left_arm_joint_position + - right_arm_joint_position + - left_effector_position + - right_effector_position + - head_position + - waist_position +max_chunk_size: 5 +# Use 10% of data in shards before moving to next shard +dataset_shard_sampling_rate: 0.1 +mixture_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotMixtureDataset.from_mixture_spec +single_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotSubLangSingleActionChunkDatasetDROID + +# Set your AGIbot dataset path here or override via CLI: +# agibot_data_root=/path/to/your/agibot_dataset +agibot_data_root: ??? + +train_dataset: + _target_: ${mixture_dataset_cls} + _convert_: object + mixture_spec: + - dataset_path: + agibot: + - ${agibot_data_root} + dataset_weight: 1.0 + distribute_weights: true + + dataset_class: ${single_dataset_cls} + all_modality_configs: ${modality_configs} + all_transforms: ${transforms} + metadata_versions: ${metadata_versions} + fps: ${fps} + dataset_kwargs: + video_backend: decord + use_global_metadata: ${use_global_metadata} + max_chunk_size: ${max_chunk_size} + relative_action: ${relative_action} + relative_action_keys: ${relative_action_keys} + relative_action_per_horizon: ${relative_action_per_horizon} + mixture_kwargs: + training: true + balance_dataset_weights: false + seed: 42 + shard_sampling_rate: ${dataset_shard_sampling_rate} diff --git a/groot/vla/configs/data/dreamzero/base_48_wan_fine_aug_relative.yaml b/groot/vla/configs/data/dreamzero/base_48_wan_fine_aug_relative.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6ab84e2511f5b2324434303014846632c2c592d8 --- /dev/null +++ b/groot/vla/configs/data/dreamzero/base_48_wan_fine_aug_relative.yaml @@ -0,0 +1,464 @@ +# @package _global_ + +# Assume `model_specific_transform` is defined in the model config file + +################################################################################ +# Normalization Statistics +# By default, we compute the normalization statistics for the datasets actually +# used in the mixture. If you want to use the global metadata, set this to true. +################################################################################ + +use_global_metadata: false + +################################################################################ +# Dimension Information +################################################################################ + +num_frames: 49 +action_horizon: 48 +state_horizon: 1 + +# image_resolution_width: 832 +# image_resolution_height: 480 + +image_resolution_width: 480 +image_resolution_height: 256 + +image_resolution_width_single_frame: 256 +image_resolution_height_single_frame: 256 + +################################################################################ +# Anchored Video Transforms +################################################################################ +totensor_cfg: &totensor_cfg + _target_: groot.vla.data.transform.VideoToTensor + apply_to: ??? + +crop_cfg: &crop_cfg + _target_: groot.vla.data.transform.VideoCrop + apply_to: ??? + scale: 0.95 + mode: random + + +resize_cfg: &resize_cfg + _target_: groot.vla.data.transform.VideoResize + apply_to: ??? + height: ${image_resolution_height} + width: ${image_resolution_width} + interpolation: linear + +resize_cfg_single_frame: &resize_cfg_single_frame + _target_: groot.vla.data.transform.VideoResize + apply_to: ??? + height: ${image_resolution_height_single_frame} + width: ${image_resolution_width_single_frame} + interpolation: linear + +color_jitter_cfg: &color_jitter_cfg + _target_: groot.vla.data.transform.VideoColorJitter + apply_to: ??? + brightness: 0.3 + contrast: 0.4 + saturation: 0.5 + hue: 0.08 + +random_grayscale_cfg: &random_grayscale_cfg + _target_: groot.vla.data.transform.VideoRandomGrayscale + apply_to: ??? + p: 0.1 + +random_posterize_cfg: &random_posterize_cfg + _target_: groot.vla.data.transform.VideoRandomPosterize + apply_to: ??? + bits: 4 + p: 0.1 + +normalize_cfg: &normalize_cfg + _target_: groot.vla.data.transform.VideoNormalize + apply_to: ??? + mean: [0.5, 0.5, 0.5] + std: [0.5, 0.5, 0.5] + + +to_numpy_cfg: &to_numpy_cfg + _target_: groot.vla.data.transform.VideoToNumpy + apply_to: ??? + + +################################################################################ +# oxe_droid (OXE Droid) +################################################################################ + +# Modality Configs +modality_config_oxe_droid: + video: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] + eval_delta_indices: [0] + modality_keys: + - video.exterior_image_1_left + - video.exterior_image_2_left + - video.wrist_image_left + state: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0] + modality_keys: + - state.joint_position + - state.gripper_position + action: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23] + modality_keys: + - action.joint_position + - action.gripper_position + language: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0] + modality_keys: + - annotation.language.language_instruction + - annotation.language.language_instruction_2 + - annotation.language.language_instruction_3 + lapa_action: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0] + modality_keys: + - lapa_action + +# Transforms +transform_oxe_droid: + _target_: groot.vla.data.transform.ComposedModalityTransform + transforms: + # Video transforms + - <<: *totensor_cfg + apply_to: ${modality_config_oxe_droid.video.modality_keys} + - <<: *crop_cfg + apply_to: ${modality_config_oxe_droid.video.modality_keys} + - <<: *resize_cfg + apply_to: ${modality_config_oxe_droid.video.modality_keys} + - <<: *color_jitter_cfg + apply_to: ${modality_config_oxe_droid.video.modality_keys} + - <<: *to_numpy_cfg + apply_to: ${modality_config_oxe_droid.video.modality_keys} + + # State transforms + - _target_: groot.vla.data.transform.StateActionToTensor + apply_to: ${modality_config_oxe_droid.state.modality_keys} + - _target_: groot.vla.data.transform.StateActionTransform + apply_to: ${modality_config_oxe_droid.state.modality_keys} + normalization_modes: + state.joint_position: q99 + state.gripper_position: q99 + + # Action transforms + - _target_: groot.vla.data.transform.StateActionToTensor + apply_to: ${modality_config_oxe_droid.action.modality_keys} + - _target_: groot.vla.data.transform.StateActionTransform + apply_to: ${modality_config_oxe_droid.action.modality_keys} + normalization_modes: + action.joint_position: q99 + action.gripper_position: q99 + + # ConcatTransform + - _target_: groot.vla.data.transform.ConcatTransform + video_concat_order: ${modality_config_oxe_droid.video.modality_keys} + state_concat_order: ${modality_config_oxe_droid.state.modality_keys} + action_concat_order: ${modality_config_oxe_droid.action.modality_keys} + + # Model-specific transform + - ${model_specific_transform} + +################################################################################ +# agibot (AGIbot: state 32, action 22, 3 views) +################################################################################ + +modality_config_agibot: + video: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] + eval_delta_indices: [-3, -2, -1, 0] + modality_keys: + - video.top_head + - video.hand_left + - video.hand_right + state: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0] + modality_keys: + - state.left_arm_joint_position + - state.right_arm_joint_position + - state.left_effector_position + - state.right_effector_position + - state.head_position + - state.waist_position + action: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23] + modality_keys: + - action.left_arm_joint_position + - action.right_arm_joint_position + - action.left_effector_position + - action.right_effector_position + - action.head_position + - action.waist_position + - action.robot_velocity + language: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0] + modality_keys: + - annotation.language.action_text + +transform_agibot: + _target_: groot.vla.data.transform.ComposedModalityTransform + transforms: + # Video transforms + - <<: *totensor_cfg + apply_to: ${modality_config_agibot.video.modality_keys} + - <<: *crop_cfg + apply_to: ${modality_config_agibot.video.modality_keys} + - <<: *resize_cfg + apply_to: ${modality_config_agibot.video.modality_keys} + - <<: *color_jitter_cfg + apply_to: ${modality_config_agibot.video.modality_keys} + - <<: *to_numpy_cfg + apply_to: ${modality_config_agibot.video.modality_keys} + + # State transforms + - _target_: groot.vla.data.transform.StateActionToTensor + apply_to: ${modality_config_agibot.state.modality_keys} + - _target_: groot.vla.data.transform.StateActionTransform + apply_to: ${modality_config_agibot.state.modality_keys} + normalization_modes: + state.left_arm_joint_position: q99 + state.right_arm_joint_position: q99 + state.left_effector_position: q99 + state.right_effector_position: q99 + state.head_position: q99 + state.waist_position: q99 + + # Action transforms + - _target_: groot.vla.data.transform.StateActionToTensor + apply_to: ${modality_config_agibot.action.modality_keys} + - _target_: groot.vla.data.transform.StateActionTransform + apply_to: ${modality_config_agibot.action.modality_keys} + normalization_modes: + action.left_arm_joint_position: q99 + action.right_arm_joint_position: q99 + action.left_effector_position: q99 + action.right_effector_position: q99 + action.head_position: q99 + action.waist_position: q99 + action.robot_velocity: q99 + + # ConcatTransform + - _target_: groot.vla.data.transform.ConcatTransform + video_concat_order: ${modality_config_agibot.video.modality_keys} + state_concat_order: ${modality_config_agibot.state.modality_keys} + action_concat_order: ${modality_config_agibot.action.modality_keys} + + # Model-specific transform + - ${model_specific_transform} + +################################################################################ +# yam (YAM: joint+gripper only from Dataset/YAM_play_data/meta/modality.json; +# state 14 dims, action 14 dims, 3 views) +################################################################################ + +modality_config_yam: + video: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] + eval_delta_indices: [0] + modality_keys: + - video.top_camera-images-rgb + - video.left_camera-images-rgb + - video.right_camera-images-rgb + state: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0] + modality_keys: + - state.left_joint_pos + - state.left_gripper_pos + - state.right_joint_pos + - state.right_gripper_pos + action: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23] + modality_keys: + - action.left_joint_pos + - action.left_gripper_pos + - action.right_joint_pos + - action.right_gripper_pos + language: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0] + modality_keys: + - annotation.task + +transform_yam: + _target_: groot.vla.data.transform.ComposedModalityTransform + transforms: + # Video transforms + - <<: *totensor_cfg + apply_to: ${modality_config_yam.video.modality_keys} + - <<: *crop_cfg + apply_to: ${modality_config_yam.video.modality_keys} + - <<: *resize_cfg + apply_to: ${modality_config_yam.video.modality_keys} + - <<: *color_jitter_cfg + apply_to: ${modality_config_yam.video.modality_keys} + - <<: *to_numpy_cfg + apply_to: ${modality_config_yam.video.modality_keys} + + # State transforms + - _target_: groot.vla.data.transform.StateActionToTensor + apply_to: ${modality_config_yam.state.modality_keys} + - _target_: groot.vla.data.transform.StateActionTransform + apply_to: ${modality_config_yam.state.modality_keys} + normalization_modes: + state.left_joint_pos: q99 + state.left_gripper_pos: q99 + state.right_joint_pos: q99 + state.right_gripper_pos: q99 + + # Action transforms + - _target_: groot.vla.data.transform.StateActionToTensor + apply_to: ${modality_config_yam.action.modality_keys} + - _target_: groot.vla.data.transform.StateActionTransform + apply_to: ${modality_config_yam.action.modality_keys} + normalization_modes: + action.left_joint_pos: q99 + action.left_gripper_pos: q99 + action.right_joint_pos: q99 + action.right_gripper_pos: q99 + + # ConcatTransform + - _target_: groot.vla.data.transform.ConcatTransform + video_concat_order: ${modality_config_yam.video.modality_keys} + state_concat_order: ${modality_config_yam.state.modality_keys} + action_concat_order: ${modality_config_yam.action.modality_keys} + + # Model-specific transform + - ${model_specific_transform} + + +################################################################################ +# trex (T-Rex: Dexmate Vega-1 dual-arm + 2x Sharpa Wave hands; +# state 58, action 58 (7 arm + 22 hand per side), 3 views @ 30fps) +################################################################################ + +modality_config_trex: + video: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] + eval_delta_indices: [0] + modality_keys: + - video.head_left + - video.left_wrist + - video.right_wrist + state: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0] + modality_keys: + - state.left_arm + - state.left_hand + - state.right_arm + - state.right_hand + action: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23] + modality_keys: + - action.left_arm + - action.left_hand + - action.right_arm + - action.right_hand + language: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0] + modality_keys: + - annotation.task + +transform_trex: + _target_: groot.vla.data.transform.ComposedModalityTransform + transforms: + # Video transforms + - <<: *totensor_cfg + apply_to: ${modality_config_trex.video.modality_keys} + - <<: *crop_cfg + apply_to: ${modality_config_trex.video.modality_keys} + - <<: *resize_cfg + apply_to: ${modality_config_trex.video.modality_keys} + - <<: *color_jitter_cfg + apply_to: ${modality_config_trex.video.modality_keys} + - <<: *to_numpy_cfg + apply_to: ${modality_config_trex.video.modality_keys} + + # State transforms + - _target_: groot.vla.data.transform.StateActionToTensor + apply_to: ${modality_config_trex.state.modality_keys} + - _target_: groot.vla.data.transform.StateActionTransform + apply_to: ${modality_config_trex.state.modality_keys} + normalization_modes: + state.left_arm: q99 + state.left_hand: q99 + state.right_arm: q99 + state.right_hand: q99 + + # Action transforms + - _target_: groot.vla.data.transform.StateActionToTensor + apply_to: ${modality_config_trex.action.modality_keys} + - _target_: groot.vla.data.transform.StateActionTransform + apply_to: ${modality_config_trex.action.modality_keys} + normalization_modes: + action.left_arm: q99 + action.left_hand: q99 + action.right_arm: q99 + action.right_hand: q99 + + # ConcatTransform + - _target_: groot.vla.data.transform.ConcatTransform + video_concat_order: ${modality_config_trex.video.modality_keys} + state_concat_order: ${modality_config_trex.state.modality_keys} + action_concat_order: ${modality_config_trex.action.modality_keys} + + # Model-specific transform + - ${model_specific_transform} + + +################################################################################ +# Modality Configs +################################################################################ + +modality_configs: + oxe_droid: ${modality_config_oxe_droid} + agibot: ${modality_config_agibot} + yam: ${modality_config_yam} + trex: ${modality_config_trex} + +################################################################################ +# Transforms +################################################################################ + +transforms: + oxe_droid: ${transform_oxe_droid} + agibot: ${transform_agibot} + yam: ${transform_yam} + trex: ${transform_trex} + +################################################################################ +# Metadata Versions +################################################################################ + +metadata_versions: + oxe_droid: '0221' + agibot: '0221' + yam: '0221' + trex: '0221' + +################################################################################ +# FPS (per embodiment, null means use dataset default) +################################################################################ + +fps: + yam: 30 + trex: 30 diff --git a/groot/vla/configs/data/dreamzero/droid_relative.yaml b/groot/vla/configs/data/dreamzero/droid_relative.yaml new file mode 100644 index 0000000000000000000000000000000000000000..197ff702b465dddd0ef54cd572e59921bae04302 --- /dev/null +++ b/groot/vla/configs/data/dreamzero/droid_relative.yaml @@ -0,0 +1,49 @@ +# @package _global_ + +defaults: + - dreamzero/base_48_wan_fine_aug_relative + - _self_ # this file will override the base + +max_state_dim: 64 +use_global_metadata: false +relative_action: true +relative_action_per_horizon: false +relative_action_keys: + - joint_position +max_chunk_size: 5 +# Use 10% of data in shards before moving to next shard +dataset_shard_sampling_rate: 0.1 +mixture_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotMixtureDataset.from_mixture_spec +single_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotSubLangSingleActionChunkDatasetDROID + +# Set your DROID dataset path here or override via CLI: +# droid_data_root=/path/to/your/droid_dataset +droid_data_root: ??? + +train_dataset: + _target_: ${mixture_dataset_cls} + _convert_: object + mixture_spec: + - dataset_path: + oxe_droid: + - ${droid_data_root} + dataset_weight: 1.0 + distribute_weights: true + + dataset_class: ${single_dataset_cls} + all_modality_configs: ${modality_configs} + all_transforms: ${transforms} + metadata_versions: ${metadata_versions} + fps: ${fps} + dataset_kwargs: + video_backend: decord + use_global_metadata: ${use_global_metadata} + max_chunk_size: ${max_chunk_size} + relative_action: ${relative_action} + relative_action_keys: ${relative_action_keys} + relative_action_per_horizon: ${relative_action_per_horizon} + mixture_kwargs: + training: true + balance_dataset_weights: false + seed: 42 + shard_sampling_rate: ${dataset_shard_sampling_rate} diff --git a/groot/vla/configs/data/dreamzero/droid_relative_wan22.yaml b/groot/vla/configs/data/dreamzero/droid_relative_wan22.yaml new file mode 100644 index 0000000000000000000000000000000000000000..42537e2b7b66be9b55c4b936eaa931b5fb3f17d1 --- /dev/null +++ b/groot/vla/configs/data/dreamzero/droid_relative_wan22.yaml @@ -0,0 +1,51 @@ +# @package _global_ +# DROID data config for Wan 5B (Wan2.2): 320x160 so latent is 20x10 (even H,W) with WanVideoVAE38 (16x), frame_seqlen=50. +# Extends base and droid settings directly to avoid Hydra nesting (dreamzero/droid_relative would double-resolve defaults). + +defaults: + - dreamzero/base_48_wan_fine_aug_relative + - _self_ + +# Wan 5B: 160x320 (HxW) -> latent 10x20, (10//2)*(20//2)=50. Use H,W divisible by 32 so latent is even (no crop in loss). +image_resolution_width: 320 +image_resolution_height: 160 + +max_state_dim: 64 +use_global_metadata: false +relative_action: true +relative_action_per_horizon: false +relative_action_keys: + - joint_position +max_chunk_size: 5 +dataset_shard_sampling_rate: 0.1 +mixture_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotMixtureDataset.from_mixture_spec +single_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotSubLangSingleActionChunkDatasetDROID +droid_data_root: ??? + +train_dataset: + _target_: ${mixture_dataset_cls} + _convert_: object + mixture_spec: + - dataset_path: + oxe_droid: + - ${droid_data_root} + dataset_weight: 1.0 + distribute_weights: true + + dataset_class: ${single_dataset_cls} + all_modality_configs: ${modality_configs} + all_transforms: ${transforms} + metadata_versions: ${metadata_versions} + fps: ${fps} + dataset_kwargs: + video_backend: decord + use_global_metadata: ${use_global_metadata} + max_chunk_size: ${max_chunk_size} + relative_action: ${relative_action} + relative_action_keys: ${relative_action_keys} + relative_action_per_horizon: ${relative_action_per_horizon} + mixture_kwargs: + training: true + balance_dataset_weights: false + seed: 42 + shard_sampling_rate: ${dataset_shard_sampling_rate} diff --git a/groot/vla/configs/data/dreamzero/trex_relative_wan22.yaml b/groot/vla/configs/data/dreamzero/trex_relative_wan22.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c13ed9f33fce338dc3f4da47c322d7315292ce55 --- /dev/null +++ b/groot/vla/configs/data/dreamzero/trex_relative_wan22.yaml @@ -0,0 +1,54 @@ +# @package _global_ +# T-Rex data config for Wan2.2-TI2V-5B: 320x160 so latent is 20x10 (even H,W) with WanVideoVAE38 (16x), frame_seqlen=50. +# T-Rex: Dexmate Vega-1 dual-arm + 2x Sharpa Wave hands, state/action 58-dim, 3 RGB views @ 30fps. +# Dataset: LeRobot v2 layout converted from v3 by scripts/data/convert_trex_v3_to_v2.py + convert_lerobot_to_gear.py. + +defaults: + - dreamzero/base_48_wan_fine_aug_relative + - _self_ + +# Wan 5B: 160x320 (HxW) -> latent 10x20, (10//2)*(20//2)=50. Use H,W divisible by 32 so latent is even (no crop in loss). +image_resolution_width: 320 +image_resolution_height: 160 + +max_state_dim: 64 +use_global_metadata: false +relative_action: true +relative_action_per_horizon: false +# Relative actions for the arms only; dexterous hands (22 dof each) stay absolute. +relative_action_keys: + - left_arm + - right_arm +max_chunk_size: 5 +dataset_shard_sampling_rate: 0.1 +mixture_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotMixtureDataset.from_mixture_spec +single_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotSubLangSingleActionChunkDatasetDROID +trex_data_root: ??? + +train_dataset: + _target_: ${mixture_dataset_cls} + _convert_: object + mixture_spec: + - dataset_path: + trex: + - ${trex_data_root} + dataset_weight: 1.0 + distribute_weights: true + + dataset_class: ${single_dataset_cls} + all_modality_configs: ${modality_configs} + all_transforms: ${transforms} + metadata_versions: ${metadata_versions} + fps: ${fps} + dataset_kwargs: + video_backend: decord + use_global_metadata: ${use_global_metadata} + max_chunk_size: ${max_chunk_size} + relative_action: ${relative_action} + relative_action_keys: ${relative_action_keys} + relative_action_per_horizon: ${relative_action_per_horizon} + mixture_kwargs: + training: true + balance_dataset_weights: false + seed: 42 + shard_sampling_rate: ${dataset_shard_sampling_rate} diff --git a/groot/vla/configs/data/dreamzero/trex_track_force_wan22.yaml b/groot/vla/configs/data/dreamzero/trex_track_force_wan22.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0e4853b8145c69f4b7c8f1b28791819e1ba8a460 --- /dev/null +++ b/groot/vla/configs/data/dreamzero/trex_track_force_wan22.yaml @@ -0,0 +1,139 @@ +# @package _global_ +# T-Rex 20 Hz EEF62 + 250-point tracks + 5 Hz raw tactile history. + +defaults: + - /data/dreamzero/base_48_wan_fine_aug_relative + - _self_ + +image_resolution_width: 320 +image_resolution_height: 160 +max_state_dim: 64 +max_action_dim: 64 +action_horizon: 16 +state_horizon: 1 +num_frames: 33 +relative_action: true +relative_action_per_horizon: false +relative_action_keys: [eef62] +use_global_metadata: false +max_chunk_size: 4 +dataset_shard_sampling_rate: 0.1 +# Optional overfit controls (None = use every valid causal anchor). +max_training_anchors: null +pin_anchor_rank: null +mixture_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotMixtureDataset.from_mixture_spec +single_dataset_cls: groot.vla.model.trex_track_force.dataset.TrexTrackForceShardedDataset +trex_data_root: ??? + +# The parquet stores absolute EEF targets. The dedicated loader converts each +# 16-step chunk into T-Rex delta-base actions relative to its start state. +modality_config_trex: + video: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0] + eval_delta_indices: [0] + modality_keys: + - video.head_left + - video.left_wrist + - video.right_wrist + state: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0] + modality_keys: + - state.eef62 + action: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0] + modality_keys: + - action.eef62 + language: + _target_: groot.vla.data.dataset.ModalityConfig + delta_indices: [0] + modality_keys: + - annotation.task + +transform_trex: + _target_: groot.vla.data.transform.ComposedModalityTransform + transforms: + - _target_: groot.vla.data.transform.VideoToTensor + apply_to: ${modality_config_trex.video.modality_keys} + - _target_: groot.vla.data.transform.VideoResize + apply_to: ${modality_config_trex.video.modality_keys} + height: ${image_resolution_height} + width: ${image_resolution_width} + interpolation: linear + - _target_: groot.vla.data.transform.VideoToNumpy + apply_to: ${modality_config_trex.video.modality_keys} + - _target_: groot.vla.data.transform.StateActionToTensor + apply_to: ${modality_config_trex.state.modality_keys} + - _target_: groot.vla.data.transform.StateActionTransform + apply_to: ${modality_config_trex.state.modality_keys} + normalization_modes: + state.eef62: q99 + - _target_: groot.vla.data.transform.StateActionToTensor + apply_to: ${modality_config_trex.action.modality_keys} + - _target_: groot.vla.data.transform.StateActionTransform + apply_to: ${modality_config_trex.action.modality_keys} + normalization_modes: + action.eef62: q99 + - _target_: groot.vla.data.transform.ConcatTransform + video_concat_order: ${modality_config_trex.video.modality_keys} + state_concat_order: ${modality_config_trex.state.modality_keys} + action_concat_order: ${modality_config_trex.action.modality_keys} + - ${model_specific_transform} + +# Explicit model/data hand-off contract. +track_force_columns: + track_xy: observation.track_xy + track_visibility: observation.track_visibility + tactile_force_history: observation.tactile_force + track_points: 250 + track_history_frames: 16 + track_future_steps: 16 + force_shape: [10, 6] + force_history_frames: 16 + force_code_tokens: 10 + force_codebook_size: 64 + action_rate_hz: 20 + tactile_rate_hz: 5 + force_stride: 4 + force_offsets: [0, 4, 8, 12] + autoregressive_blocks: 4 + video_frames_per_block: 8 + +fps: + # This is the source mp4 rate; logical control samples use timestamp-based + # 20 Hz action/track and 5 Hz force grids inside the dedicated loader. + trex: 30 + +train_dataset: + _target_: ${mixture_dataset_cls} + _convert_: object + mixture_spec: + - dataset_path: + trex: + - ${trex_data_root} + dataset_weight: 1.0 + distribute_weights: true + dataset_class: ${single_dataset_cls} + all_modality_configs: ${modality_configs} + all_transforms: ${transforms} + metadata_versions: ${metadata_versions} + fps: ${fps} + dataset_kwargs: + video_backend: decord + use_global_metadata: ${use_global_metadata} + max_chunk_size: ${max_chunk_size} + relative_action: ${relative_action} + relative_action_keys: ${relative_action_keys} + relative_action_per_horizon: ${relative_action_per_horizon} + action_rate_hz: 20 + tactile_rate_hz: 5 + video_rate_hz: 10 + max_training_anchors: ${max_training_anchors} + pin_anchor_rank: ${pin_anchor_rank} + mixture_kwargs: + training: true + balance_dataset_weights: false + seed: 42 + shard_sampling_rate: ${dataset_shard_sampling_rate} diff --git a/groot/vla/configs/data/dreamzero/yam_relative.yaml b/groot/vla/configs/data/dreamzero/yam_relative.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2670479f2c2369f1440f5c8333adf2310391bbd2 --- /dev/null +++ b/groot/vla/configs/data/dreamzero/yam_relative.yaml @@ -0,0 +1,52 @@ +# @package _global_ + +defaults: + - dreamzero/base_48_wan_fine_aug_relative + - _self_ # this file will override the base + +max_state_dim: 64 +use_global_metadata: false +relative_action: true +relative_action_per_horizon: false +relative_action_keys: + - left_joint_pos + - left_gripper_pos + - right_joint_pos + - right_gripper_pos +max_chunk_size: 5 +# Use 10% of data in shards before moving to next shard +dataset_shard_sampling_rate: 0.1 +mixture_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotMixtureDataset.from_mixture_spec +single_dataset_cls: groot.vla.data.dataset.lerobot_sharded.ShardedLeRobotSubLangSingleActionChunkDatasetDROID + +# Set your YAM dataset path here or override via CLI: +# yam_data_root=/path/to/your/yam_dataset +yam_data_root: ??? + +train_dataset: + _target_: ${mixture_dataset_cls} + _convert_: object + mixture_spec: + - dataset_path: + yam: + - ${yam_data_root} + dataset_weight: 1.0 + distribute_weights: true + + dataset_class: ${single_dataset_cls} + all_modality_configs: ${modality_configs} + all_transforms: ${transforms} + metadata_versions: ${metadata_versions} + fps: ${fps} + dataset_kwargs: + video_backend: decord + use_global_metadata: ${use_global_metadata} + max_chunk_size: ${max_chunk_size} + relative_action: ${relative_action} + relative_action_keys: ${relative_action_keys} + relative_action_per_horizon: ${relative_action_per_horizon} + mixture_kwargs: + training: true + balance_dataset_weights: false + seed: 42 + shard_sampling_rate: ${dataset_shard_sampling_rate} diff --git a/groot/vla/configs/deepspeed/zero2.json b/groot/vla/configs/deepspeed/zero2.json new file mode 100644 index 0000000000000000000000000000000000000000..b1a5a0b41b082ff36c8a0cec197bb61c7d78c2b1 --- /dev/null +++ b/groot/vla/configs/deepspeed/zero2.json @@ -0,0 +1,28 @@ +{ + "checkpoint": { + "load_universal": false + }, + "train_batch_size": "auto", + "train_micro_batch_size_per_gpu": "auto", + "gradient_accumulation_steps": "auto", + "gradient_clipping": "auto", + "zero_allow_untested_optimizer": true, + "fp16": { + "enabled": "auto", + "loss_scale": 0, + "loss_scale_window": 1000, + "initial_scale_power": 16, + "hysteresis": 2, + "min_loss_scale": 1 + }, + "bf16": { + "enabled": "auto" + }, + "zero_optimization": { + "stage": 2, + "overlap_comm": false, + "contiguous_gradients": true, + "sub_group_size": 1e9, + "reduce_bucket_size": 1e8 + } +} diff --git a/groot/vla/configs/deepspeed/zero2_offload.json b/groot/vla/configs/deepspeed/zero2_offload.json new file mode 100644 index 0000000000000000000000000000000000000000..29c7f2f22577fa6ef48a9958285a4bbf5dc4026d --- /dev/null +++ b/groot/vla/configs/deepspeed/zero2_offload.json @@ -0,0 +1,32 @@ +{ + "checkpoint": { + "load_universal": false + }, + "train_batch_size": "auto", + "train_micro_batch_size_per_gpu": "auto", + "gradient_accumulation_steps": "auto", + "gradient_clipping": "auto", + "zero_allow_untested_optimizer": true, + "fp16": { + "enabled": "auto", + "loss_scale": 0, + "loss_scale_window": 1000, + "initial_scale_power": 16, + "hysteresis": 2, + "min_loss_scale": 1 + }, + "bf16": { + "enabled": "auto" + }, + "zero_optimization": { + "stage": 2, + "offload_optimizer": { + "device": "cpu", + "pin_memory": true + }, + "overlap_comm": false, + "contiguous_gradients": true, + "sub_group_size": 1e9, + "reduce_bucket_size": 1e8 + } +} diff --git a/groot/vla/configs/deepspeed/zero3.json b/groot/vla/configs/deepspeed/zero3.json new file mode 100644 index 0000000000000000000000000000000000000000..cfd068a9f91485787b0895e029bc3f0b1577adf1 --- /dev/null +++ b/groot/vla/configs/deepspeed/zero3.json @@ -0,0 +1,33 @@ +{ + "checkpoint": { + "load_universal": false + }, + "train_batch_size": "auto", + "train_micro_batch_size_per_gpu": "auto", + "gradient_accumulation_steps": "auto", + "gradient_clipping": "auto", + "zero_allow_untested_optimizer": true, + "fp16": { + "enabled": "auto", + "loss_scale": 0, + "loss_scale_window": 1000, + "initial_scale_power": 16, + "hysteresis": 2, + "min_loss_scale": 1 + }, + "bf16": { + "enabled": "auto" + }, + "zero_optimization": { + "stage": 3, + "overlap_comm": true, + "contiguous_gradients": true, + "sub_group_size": 1e9, + "reduce_bucket_size": 5e8, + "stage3_prefetch_bucket_size": 5e8, + "stage3_param_persistence_threshold": 1e6, + "stage3_max_live_parameters": 1e9, + "stage3_max_reuse_distance": 1e9, + "stage3_gather_16bit_weights_on_model_save": true + } +} diff --git a/groot/vla/configs/model/dreamzero/action_head/wan_flow_matching_action_tf.yaml b/groot/vla/configs/model/dreamzero/action_head/wan_flow_matching_action_tf.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8cccdc6319ac822f6f93ed9af9746034e74215ea --- /dev/null +++ b/groot/vla/configs/model/dreamzero/action_head/wan_flow_matching_action_tf.yaml @@ -0,0 +1,108 @@ +# @package _global_ +add_pos_embed: true +hidden_size: 64 +attn_dropout: 0.2 +repa_layer: 8 +repa_coeff: 1.0 +load_pretrained_det_decode_layer_path: null +expand_batch: null +dit_version: null +text_encoder_pretrained_path: null +image_encoder_pretrained_path: null +vae_pretrained_path: null +train_architecture: "lora" +num_frame_per_block: 1 +num_action_per_block: 32 +num_state_per_block: 1 +frame_seqlen: 880 +action_head_cfg: + _target_: groot.vla.model.dreamzero.action_head.wan_flow_matching_action_tf.WANPolicyHead + _convert_: object # to convert to a dict so that it can save the config to json + config: + _target_: groot.vla.model.dreamzero.action_head.wan_flow_matching_action_tf.WANPolicyHeadConfig + _recursive_: false # avoid instantiating noise_scheduler_cfg and obs_encoder_cfg + tiled: false + tile_size_height: 34 + tile_size_width: 34 + tile_stride_height: 18 + tile_stride_width: 16 + lora_rank: 4 + lora_alpha: 4 + num_frames: ${num_frames} + num_frame_per_block: ${num_frame_per_block} + lora_target_modules: "q,k,v,o,ffn.0,ffn.2" + init_lora_weights: "kaiming" + train_architecture: ${train_architecture} + use_gradient_checkpointing: true + add_pos_embed: ${add_pos_embed} + model_dtype: float32 + max_state_dim: ${max_state_dim} + max_action_dim: ${max_action_dim} + action_loss_embodiment_ids: [26, 17, 32] + hidden_size: ${hidden_size} + input_embedding_dim: 1536 + backbone_embedding_dim: ${backbone_hidden_size} + repa_layer: ${repa_layer} + repa_coeff: ${repa_coeff} + load_pretrained_det_decode_layer_path: ${load_pretrained_det_decode_layer_path} + freeze_decode_layer: false + expand_batch: ${expand_batch} + use_vlln: true + vl_self_attention_cfg: + _target_: groot.vla.model.n1_5.modules.cross_attention_dit.SelfAttentionTransformer + positional_embeddings: null + num_layers: 4 + num_attention_heads: 24 + attention_head_dim: 64 + dropout: ${attn_dropout} + final_dropout: true + diffusion_model_cfg: + _target_: groot.vla.model.dreamzero.modules.wan_video_dit_action_casual_chunk.CausalWanModel + _convert_: object + diffusion_model_pretrained_path: ${dit_version} + model_type: "i2v" + frame_seqlen: ${frame_seqlen} + dim: 5120 + in_dim: 36 + ffn_dim: 13824 + out_dim: 16 + freq_dim: 256 + eps: 1e-6 + num_heads: 40 + num_layers: 40 + max_chunk_size: ${max_chunk_size} + num_frame_per_block: ${num_frame_per_block} + num_action_per_block: ${num_action_per_block} + num_state_per_block: ${num_state_per_block} + text_encoder_cfg: + _target_: groot.vla.model.dreamzero.modules.wan_video_text_encoder.WanTextEncoder + _convert_: object + text_encoder_pretrained_path: ${text_encoder_pretrained_path} + image_encoder_cfg: + _target_: groot.vla.model.dreamzero.modules.wan_video_image_encoder.WanImageEncoder + _convert_: object + image_encoder_pretrained_path: ${image_encoder_pretrained_path} + vae_cfg: + _target_: groot.vla.model.dreamzero.modules.wan_video_vae.WanVideoVAE + _convert_: object + vae_pretrained_path: ${vae_pretrained_path} + + action_dim: ${max_action_dim} + action_horizon: ${action_horizon} + num_inference_timesteps: 4 # not used during training + noise_beta_alpha: 1.5 + noise_beta_beta: 1.0 + noise_s: 0.999 + num_timestep_buckets: 1000 + + # ========== DECOUPLED NOISE SAMPLING CONFIG ========== + # When true: video uses Beta(alpha,beta) biased high, action uses independent uniform + # When false (default): original behavior - video and action share same timestep + decouple_video_action_noise: false + # Beta params (only used when decouple_video_action_noise=true) + # Beta(3,1): mean=0.75, Beta(5,1): mean=0.83 (higher alpha = more bias to high noise) + video_noise_beta_alpha: 3.0 + video_noise_beta_beta: 1.0 + + tune_projector: true + tune_diffusion_model: true diff --git a/groot/vla/configs/model/dreamzero/action_head/wan_flow_matching_action_tf_wan22.yaml b/groot/vla/configs/model/dreamzero/action_head/wan_flow_matching_action_tf_wan22.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1fa6c8e8de9ee0e88f9ea5fdf066f69f46f04cda --- /dev/null +++ b/groot/vla/configs/model/dreamzero/action_head/wan_flow_matching_action_tf_wan22.yaml @@ -0,0 +1,38 @@ +# @package _global_ +# Wan2.2-TI2V-5B backbone config for DreamZero DROID training +# Uses CausalWanModel with Wan2.2 architecture (dim=3072, 48-channel VAE) +# Override base config with: model/dreamzero/action_head=wan_flow_matching_action_tf_wan22 + +defaults: + - wan_flow_matching_action_tf + +# Wan2.2-TI2V-5B specific overrides +# frame_seqlen must match DiT patch_embedding output per frame: (H//2)*(W//2) for stride (1,2,2). +# WanVideoVAE38 has 16x spatial downscale. For even latent H,W (no crop in dynamics loss) use video +# resolution with H,W divisible by 32: 160x320 -> latent 10x20 -> patch 5x10=50. Else 176x320 -> 11x20 (odd H). +# Same as 5B pretrained: latent only (in_dim=48), first-frame via CLIP; no [x;y] concat. +# Action/state invariant: need (noisy_frames-1)//num_frame_per_block >= 1, so min 3 frames per half. +frame_seqlen: 50 + +action_head_cfg: + config: + # 160x320 gives latent 10x20 (even); avoids dynamics-loss crop. Match data=dreamzero/droid_relative_wan22. + target_video_height: 160 + target_video_width: 320 + diffusion_model_cfg: + model_type: "ti2v" + dim: 3072 + in_dim: 48 + concat_first_frame_latent: false + ffn_dim: 14336 + out_dim: 48 + freq_dim: 256 + eps: 1e-6 + num_heads: 24 + num_layers: 30 + vae_cfg: + _target_: groot.vla.model.dreamzero.modules.wan_video_vae.WanVideoVAE38 + _convert_: object + z_dim: 48 + dim: 160 + vae_pretrained_path: ${vae_pretrained_path} diff --git a/groot/vla/configs/model/dreamzero/backbone/identity.yaml b/groot/vla/configs/model/dreamzero/backbone/identity.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2e4cf9ced4117a2dc622308691e39551ef0ff2c6 --- /dev/null +++ b/groot/vla/configs/model/dreamzero/backbone/identity.yaml @@ -0,0 +1,9 @@ +# @package _global_ + +backbone_hidden_size: 0 +backbone_cfg: + _target_: groot.vla.model.dreamzero.backbone.identity.IdentityBackbone + +action_head_cfg: + config: + backbone_features_projector_cfg: null # no projection because no backbone features diff --git a/groot/vla/configs/model/dreamzero/transform/base.yaml b/groot/vla/configs/model/dreamzero/transform/base.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5e3d1916116663a252fe7643315cd9e364c1feb0 --- /dev/null +++ b/groot/vla/configs/model/dreamzero/transform/base.yaml @@ -0,0 +1,43 @@ +# @package _global_ + +embodiment_tag_to_projector_index: + real_gr1_arms_only: 0 + real_gr1_arms_only_annotated: 1 + real_gr1_arms_waist: 2 + real_gr1_arms_waist_annotated: 3 + dexmg_gr1_arms_only_inspire: 4 + dexmg_gr1_arms_only_fourier: 5 + dexmg_gr1_arms_waist_fourier: 6 + robocasa_single_arm: 7 + onex_eve_gripper: 8 + robocasa_gr1_arms_only_inspire_hands: 9 + robocasa_gr1_arms_only_fourier_hands: 10 + robocasa_gr1_fixed_lower_body_inspire_hands: 11 + robocasa_gr1_fixed_lower_body_fourier_hands: 12 + robocasa_panda_omron: 13 + # s: 14 + robocasa_bimanual_panda_parallel_gripper: 15 + robocasa_bimanual_panda_inspire_hand: 16 + oxe_droid: 17 + oxe_fractal: 18 + oxe_language_table: 19 + oxe_bridge: 20 + real_panda_single_arm: 21 + hot3d_hands_only: 23 + gr1_unified: 24 + robocasa_gr1_arms_waist_fourier_hands: 25 + agibot: 26 + lapa: 27 + oxe_mutex: 28 + oxe_roboset: 29 + oxe_plex: 30 + dream: 31 + yam: 32 + xdof: 22 + gr1_unified_segmentation: 14 + language_table_sim: 7 + gr1_isaac: 0 + sim_behavior_r1_pro: 31 + mecka_hands: 27 + real_r1_pro_sharpa: 28 + trex: 28 diff --git a/groot/vla/configs/model/dreamzero/transform/dreamzero_cotrain.yaml b/groot/vla/configs/model/dreamzero/transform/dreamzero_cotrain.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b39a62f300410c921c54056f54754b1e3ec20a8c --- /dev/null +++ b/groot/vla/configs/model/dreamzero/transform/dreamzero_cotrain.yaml @@ -0,0 +1,44 @@ +# @package _global_ + +defaults: + - base + - _self_ # this file will override model/transform/base + + +max_length: 512 +num_views: 1 +tokenizer_path: "google/umt5-xxl" +################################################################################ +# Data Collator +################################################################################ + +data_collator: + _target_: groot.vla.model.dreamzero.transform.dreamzero_cotrain.DefaultDataCollator + tokenizer_path: ${tokenizer_path} + max_length : ${max_length} + num_views: ${num_views} + embodiment_tag_mapping: ${embodiment_tag_to_projector_index} + +################################################################################ +# Model-Specific Transform +# Needs to change this model-specific transform for different models +################################################################################ + +num_visual_tokens_per_frame: 16 +max_state_dim: 44 +max_action_dim: 32 +language_dropout_prob: 0.0 + + +model_specific_transform: + _target_: groot.vla.model.dreamzero.transform.dreamzero_cotrain.DreamTransform + default_instruction: "Perform the default behavior." + language_dropout_prob: ${language_dropout_prob} + always_use_default_instruction: false + max_state_dim: ${max_state_dim} + max_action_dim: ${max_action_dim} + max_length: ${max_length} + state_horizon: ${state_horizon} + action_horizon: ${action_horizon} + embodiment_tag_mapping: ${embodiment_tag_to_projector_index} + tokenizer_path: ${tokenizer_path} diff --git a/groot/vla/configs/model/dreamzero/vla.yaml b/groot/vla/configs/model/dreamzero/vla.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4ab9fac6999409208720e71754ba1a9a0207b39f --- /dev/null +++ b/groot/vla/configs/model/dreamzero/vla.yaml @@ -0,0 +1,20 @@ +# @package _global_ + +defaults: + - _self_ # all below configs will override this file + - /model/dreamzero/backbone: identity + - /model/dreamzero/action_head: wan_flow_matching_action_tf + - /model/dreamzero/transform: dreamzero_cotrain + +model: + _target_: groot.vla.model.dreamzero.base_vla.VLA + _convert_: object + config: + _target_: groot.vla.model.dreamzero.base_vla.VLAConfig + _recursive_: false # avoid instantiating backbone_cfg and action_head_cfg + model_dtype: float32 + hidden_size: ${backbone_hidden_size} # this is defined in the backbone's yaml + action_horizon: ${action_head_cfg.config.action_horizon} # this is defined in the action_head's yaml + action_dim: ${action_head_cfg.config.action_dim} # this is defined in the action_head's yaml + backbone_cfg: ${backbone_cfg} # this is defined in the backbone's yaml + action_head_cfg: ${action_head_cfg} # this is defined in the action_head's yaml diff --git a/groot/vla/configs/model/trex_track_force/action_head/wan22.yaml b/groot/vla/configs/model/trex_track_force/action_head/wan22.yaml new file mode 100644 index 0000000000000000000000000000000000000000..692ce4fac17b6e3cba6fbb3f70789aaf0db20ab5 --- /dev/null +++ b/groot/vla/configs/model/trex_track_force/action_head/wan22.yaml @@ -0,0 +1,109 @@ +# @package _global_ + +# Shared checkpoint paths. Keep component loading opt-in so config composition +# and CPU unit tests never download weights. +dit_version: null +text_encoder_pretrained_path: null +image_encoder_pretrained_path: null +vae_pretrained_path: null +load_pretrained_components: false +allow_checkpoint_download: false +train_architecture: lora + +frame_seqlen: 50 +num_frame_per_block: 2 +num_action_per_block: 16 +num_state_per_block: 1 +max_chunk_size: 4 + +action_head_cfg: + _target_: groot.vla.model.trex_track_force.policy.WANTrackForcePolicyHead + _convert_: object + config: + _target_: groot.vla.model.trex_track_force.policy.WANTrackForcePolicyHeadConfig + _recursive_: false + action_dim: 64 + physical_action_dim: 62 + action_horizon: 16 + track_points: 250 + track_horizon: 16 + track_loss_weight: 1.0 + track_position_loss_weight: 0.0 + track_smoothness_loss_weight: 0.0 + track_motion_loss_boost: 0.0 + zero_invisible_track_xy: true + force_loss_weight: 1.0 + vq_loss_weight: 1.0 + vq_commitment_weight: 0.25 + tactile_dropout: 0.1 + dynamics_loss_weight: 1.0 + require_force_inputs: true + target_video_height: 160 + target_video_width: 320 + train_architecture: ${train_architecture} + lora_rank: 4 + lora_alpha: 4 + lora_target_modules: "q,k,v,o,k_img,v_img,ffn.0,ffn.2" + defer_lora_injection: false + load_pretrained_components: ${load_pretrained_components} + allow_checkpoint_download: ${allow_checkpoint_download} + + diffusion_model_cfg: + _target_: groot.vla.model.trex_track_force.model.CausalWanTrackForceModel + _convert_: object + diffusion_model_pretrained_path: ${dit_version} + model_type: ti2v + patch_size: [1, 2, 2] + frame_seqlen: ${frame_seqlen} + in_dim: 48 + out_dim: 48 + dim: 3072 + ffn_dim: 14336 + freq_dim: 256 + num_heads: 24 + num_layers: 30 + max_chunk_size: ${max_chunk_size} + num_frame_per_block: ${num_frame_per_block} + action_dim: 64 + action_horizon: ${num_action_per_block} + max_state_dim: 64 + num_state_per_block: ${num_state_per_block} + track_points: 250 + track_horizon: 16 + max_track_history: 16 + text_dim: 4096 + clip_dim: 1280 + gradient_checkpointing: true + + force_model_cfg: + _target_: groot.vla.model.trex_track_force.force.ForceOnlyTransformer + _convert_: object + dim: 768 + num_heads: 12 + num_layers: 6 + ffn_dim: 3072 + action_horizon: 16 + action_dim: 64 + force_sensor_count: 10 + force_sensor_dim: 6 + history_frames: 16 + vq_codebook_size: 64 + vq_code_dim: 256 + coarse_memory_dim: 3072 + force_stride: 4 + + # Reuse the existing Wan2.2 VAE38, UMT5, and Wan CLIP implementations. + text_encoder_cfg: + _target_: groot.vla.model.dreamzero.modules.wan_video_text_encoder.WanTextEncoder + _convert_: object + text_encoder_pretrained_path: ${text_encoder_pretrained_path} + image_encoder_cfg: + _target_: groot.vla.model.dreamzero.modules.wan_video_image_encoder.WanImageEncoder + _convert_: object + image_encoder_pretrained_path: ${image_encoder_pretrained_path} + vae_cfg: + _target_: groot.vla.model.dreamzero.modules.wan_video_vae.WanVideoVAE38 + _convert_: object + z_dim: 48 + dim: 160 + vae_pretrained_path: ${vae_pretrained_path} diff --git a/groot/vla/configs/model/trex_track_force/transform/track_force.yaml b/groot/vla/configs/model/trex_track_force/transform/track_force.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ea0954a42013f9ccd437b838d1e3420a6662261a --- /dev/null +++ b/groot/vla/configs/model/trex_track_force/transform/track_force.yaml @@ -0,0 +1,37 @@ +# @package _global_ + +defaults: + - /model/dreamzero/transform: base + - _self_ + +max_length: 512 +num_views: 3 +tokenizer_path: google/umt5-xxl +max_state_dim: 64 +max_action_dim: 64 +language_dropout_prob: 0.0 + +data_collator: + _target_: groot.vla.model.trex_track_force.transform.TrackForceDataCollator + tokenizer_path: ${tokenizer_path} + max_length: ${max_length} + num_views: ${num_views} + embodiment_tag_mapping: ${embodiment_tag_to_projector_index} + +model_specific_transform: + _target_: groot.vla.model.trex_track_force.transform.TrexTrackForceTransform + default_instruction: "Perform the requested bimanual manipulation." + language_dropout_prob: ${language_dropout_prob} + always_use_default_instruction: false + max_state_dim: ${max_state_dim} + max_action_dim: ${max_action_dim} + max_length: ${max_length} + state_horizon: ${state_horizon} + action_horizon: ${action_horizon} + embodiment_tag_mapping: ${embodiment_tag_to_projector_index} + tokenizer_path: ${tokenizer_path} + track_history_frames: 16 + track_horizon: 16 + track_points: 250 + require_track_inputs: true + require_force_inputs: true diff --git a/groot/vla/configs/model/trex_track_force/vla.yaml b/groot/vla/configs/model/trex_track_force/vla.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6293eaa1d896e5836f726507822e52caa73dd6c3 --- /dev/null +++ b/groot/vla/configs/model/trex_track_force/vla.yaml @@ -0,0 +1,22 @@ +# @package _global_ + +defaults: + - _self_ + - /model/dreamzero/backbone: identity + - /model/trex_track_force/action_head: wan22 + - /model/trex_track_force/transform: track_force + +# Selected independently with: model=trex_track_force/vla +model: + _target_: groot.vla.model.trex_track_force.vla.TrexTrackForceVLA + _convert_: object + config: + _target_: groot.vla.model.trex_track_force.vla.TrexTrackForceVLAConfig + _recursive_: false + model_dtype: float32 + compute_dtype: float32 + hidden_size: ${backbone_hidden_size} + action_horizon: ${action_head_cfg.config.action_horizon} + action_dim: ${action_head_cfg.config.action_dim} + backbone_cfg: ${backbone_cfg} + action_head_cfg: ${action_head_cfg} diff --git a/groot/vla/data/__init__.py b/groot/vla/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/data/conversion/__init__.py b/groot/vla/data/conversion/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/data/conversion/gr1/__init__.py b/groot/vla/data/conversion/gr1/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/data/conversion/gr1/constants.py b/groot/vla/data/conversion/gr1/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..55dbd88fe0f14222ef91b2d23ebce40838d477b1 --- /dev/null +++ b/groot/vla/data/conversion/gr1/constants.py @@ -0,0 +1,27 @@ +from groot.vla.data.schema import EmbodimentTag + +RAW_HDF5_FILENAME = "state_action.hdf5" +RAW_METADATA_FILENAME = "metadata.json" +RAW_VIDEO_FILENAME = "ego_view.mp4" +RAW_ANNOTATION_FILENAME = "annotation.json" + +TRAINABLE_HDF5_FILENAME = "dataset.hdf5" +TRAINABLE_METADATA_FILENAME = "metadata.json" + +RAW_DATA_CONTROL_FREQUENCY = 20 + +INITIAL_ACTIONS_FILENAME = "initial_actions.npz" + +EMBODIMENT_TAG_TO_ANNOTATED_VERSION = { + EmbodimentTag.REAL_GR1_ARMS_ONLY: EmbodimentTag.REAL_GR1_ARMS_ONLY_ANNOTATED, + EmbodimentTag.REAL_GR1_ARMS_WAIST: EmbodimentTag.REAL_GR1_ARMS_WAIST_ANNOTATED, + # Special for 5DC-S + EmbodimentTag.ROBOCASA_GR1_ARMS_ONLY_FOURIER_HANDS: EmbodimentTag.REAL_GR1_ARMS_ONLY_ANNOTATED, + EmbodimentTag.ROBOCASA_GR1_ARMS_WAIST_FOURIER_HANDS: EmbodimentTag.REAL_GR1_ARMS_WAIST_ANNOTATED, +} + +# For chopped data +EPISODE_LENGTH_FILENAME = "episode_length.json" + +# Processed +PROCESSED_VIDEO_FILENAME = "ego_view_pad_res224_freq20.mp4" diff --git a/groot/vla/data/conversion/gr1/get_initial_actions.py b/groot/vla/data/conversion/gr1/get_initial_actions.py new file mode 100644 index 0000000000000000000000000000000000000000..885aed3e4c4ddbdf94ba56e4da2e17bd75102c9d --- /dev/null +++ b/groot/vla/data/conversion/gr1/get_initial_actions.py @@ -0,0 +1,128 @@ +from pathlib import Path + +import h5py +import numpy as np +import pandas as pd + +import groot.vla.common.utils as U +from groot.vla.data.conversion.gr1.constants import ( + INITIAL_ACTIONS_FILENAME, + TRAINABLE_HDF5_FILENAME, +) +from groot.vla.data.dataset.macro import ( + LE_ROBOT_EPISODE_FILENAME, + LE_ROBOT_INFO_FILENAME, + LE_ROBOT_METADATA_DIR, + LE_ROBOT_MODALITY_FILENAME, +) + + +def get_initial_actions(data_dir: str | Path): + hdf5_file = h5py.File(Path(data_dir) / TRAINABLE_HDF5_FILENAME, "r") + initial_actions = [] + + """ + initial_actions: dict[str, dict[str, np.ndarray]] + 0: (the dataset dimension) + trajectory_name: + action_key: + action: np.ndarray + """ + initial_actions = {} + for demo_name in hdf5_file["data"].keys(): + demo_group = hdf5_file["data"][demo_name] + initial_actions[demo_name] = {} + action_keys = list(demo_group["action"].keys()) + for action_key in action_keys: + initial_actions[demo_name][action_key] = demo_group["action"][action_key][0] + return [initial_actions] + + +def get_initial_actions_from_lerobot(data_dir: str | Path): + data_dir = Path(data_dir) + + # 1. Get modality for slicing action + meta_modality_path = data_dir / LE_ROBOT_METADATA_DIR / LE_ROBOT_MODALITY_FILENAME + meta_modality = U.load_json(meta_modality_path) + action_keys = meta_modality["action"].keys() + + # 2. Get episode paths + # 2.1. Get data_path_pattern + meta_info_path = data_dir / LE_ROBOT_METADATA_DIR / LE_ROBOT_INFO_FILENAME + meta_info = U.load_json(meta_info_path) + data_path_pattern = meta_info["data_path"] + chunk_size = meta_info["chunks_size"] + + # 2.2. Get episode info + episode_metadata_path = data_dir / LE_ROBOT_METADATA_DIR / LE_ROBOT_EPISODE_FILENAME + episode_metadata = U.load_jsonl(episode_metadata_path) + + initial_actions = {} + for episode_info in episode_metadata: + episode_index = episode_info["episode_index"] + episode_chunk = episode_index // chunk_size + episode_path = data_dir / data_path_pattern.format( + episode_chunk=episode_chunk, episode_index=episode_index + ) + if not episode_path.exists(): + raise ValueError(f"Episode path {episode_path} does not exist") + + episode_data = pd.read_parquet(episode_path) + + initial_action_concat = episode_data["action"].iloc[0] + trajectory_id = episode_info["episode_index"] + initial_actions[trajectory_id] = {} + for action_key in action_keys: + start = meta_modality["action"][action_key]["start"] + end = meta_modality["action"][action_key]["end"] + initial_actions[trajectory_id][action_key] = initial_action_concat[start:end] + return [initial_actions] + + +def save_initial_actions( + initial_actions: dict[str, dict[str, np.ndarray]], initial_actions_path: str | Path +): + np.savez(str(initial_actions_path), initial_actions) + + +def load_initial_actions(initial_actions_path: str | Path): + """ + initial_actions: list[dict[str, dict[str, np.ndarray]]] + 0: (the first dataset) + trajectory_name: + action_key: + action: np.ndarray + 1: (the second dataset) + ... + """ + initial_actions_npz = np.load(str(initial_actions_path), allow_pickle=True) + initial_actions = [] + initial_actions_array = initial_actions_npz[ + "arr_0" + ] # This is the default key when np.savez saves a list + for dataset_initial_actions in initial_actions_array: + initial_actions_for_this_dataset = {} + for trajectory_name, action_dict in dataset_initial_actions.items(): + initial_actions_for_this_dataset[trajectory_name] = action_dict + initial_actions.append(initial_actions_for_this_dataset) + return initial_actions + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="Generate initial_actions.npz for a LeRobot dataset") + parser.add_argument("data_dir", type=str, help="Path to LeRobot dataset directory") + args = parser.parse_args() + + initial_actions = get_initial_actions_from_lerobot(args.data_dir) + save_initial_actions( + initial_actions, + Path(args.data_dir) / LE_ROBOT_METADATA_DIR / INITIAL_ACTIONS_FILENAME, + ) + + # Verify + loaded_initial_actions = load_initial_actions( + Path(args.data_dir) / LE_ROBOT_METADATA_DIR / INITIAL_ACTIONS_FILENAME + ) + print(f"Saved initial actions for {len(loaded_initial_actions[0])} trajectories") diff --git a/groot/vla/data/dataset/__init__.py b/groot/vla/data/dataset/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..db17fc385679fa2b4e40e5c5c7631e43d3ae3638 --- /dev/null +++ b/groot/vla/data/dataset/__init__.py @@ -0,0 +1,5 @@ +from .lerobot import ModalityConfig + +__all__ = [ + "ModalityConfig", +] diff --git a/groot/vla/data/dataset/lerobot.py b/groot/vla/data/dataset/lerobot.py new file mode 100644 index 0000000000000000000000000000000000000000..f8efc6ed0f424f32cd4273be064680c33059106c --- /dev/null +++ b/groot/vla/data/dataset/lerobot.py @@ -0,0 +1,2499 @@ +from collections import defaultdict +import copy +import glob +import hashlib +import importlib +import json +from pathlib import Path +import time +from typing import Sequence, TypeVar + +import numpy as np +import pandas as pd +from pydantic import BaseModel, Field, ValidationError, field_validator +from torch.utils.data import Dataset +from tqdm import tqdm +import yaml +import torch + +from groot.vla.common.utils import get_all_frames, get_frames_by_timestamps +from groot.vla.data.conversion.gr1.get_initial_actions import load_initial_actions +from groot.vla.data.schema import ( + DatasetMetadata, + DatasetStatisticalValues, + EmbodimentTag, + LeRobotModalityMetadata, + LeRobotStateActionMetadata, +) +from groot.vla.data.transform import ComposedModalityTransform + +T_LeRobotMixtureDataset = TypeVar("T_LeRobotMixtureDataset", bound="LeRobotMixtureDataset") + +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.json" +LE_ROBOT_DATA_FILENAME = "data/*/*.parquet" +LE_ROBOT_TASK_EMBEDDINGS_FILENAME = "meta/task_embeddings.pt" +LE_ROBOT_DETAILED_GLOBAL_INSTRUCTION_FILENAME = "meta/episodes_detail_global_instruction.jsonl" +INITIAL_ACTIONS_FILENAME = "meta/initial_actions.npz" +METADATA_DIR = Path(importlib.import_module("groot.vla.data").__file__).parent / "metadata" # type: ignore +STEP_FILTER_FILENAME = "meta/step_filter.jsonl" +LEROBOT_RELATIVE_STATS_FILE_NAME = "meta/relative_stats_dreamzero.json" +LEROBOT_RELATIVE_HORIZON_STATS_FILE_NAME = "meta/relative_horizon_stats_dreamzero.json" + +# Special language keys that load from metadata files instead of parquet columns +METADATA_LANG_KEYS = ["detailed_global_instruction_medium", "detailed_global_instruction_concise"] + + +def calculate_dataset_statistics( + parquet_paths: list[Path], features: list[str] | None = None +) -> dict[str, DatasetStatisticalValues]: + """Calculate the dataset statistics of all columns for a list of parquet files. + + Args: + parquet_paths (list[Path]): List of paths to parquet files to process. + features (list[str] | None): List of feature names to compute statistics for. + If None, computes statistics for all columns in the data. + + Returns: + dict[str, DatasetStatisticalValues]: Dictionary mapping feature names to their + statistical values (mean, std, min, max, q01, q99). + """ + # Dataset statistics + all_low_dim_data_list = [] + # Collect all the data + 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) + all_low_dim_data = pd.concat(all_low_dim_data_list, axis=0) + # Compute dataset statistics + dataset_statistics = {} + if features is None: + features = list(all_low_dim_data.columns) + for le_modality in features: + print(f"Computing statistics for {le_modality}...") + np_data = np.vstack( + [np.asarray(x, dtype=np.float32) for x in all_low_dim_data[le_modality]] + ) + dataset_statistics[le_modality] = DatasetStatisticalValues( + 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 + + +class ModalityConfig(BaseModel): + """Configuration for a modality defining how data should be sampled and loaded. + + This class specifies which indices to sample relative to a base index and which + keys to load for a particular modality (e.g., video, state, action). + """ + + 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.""" + eval_delta_indices: list[int] | None = None + """Delta indices to sample relative to the current index for evaluation. If None, uses the same indices as delta_indices.""" + modality_keys: list[str] + """The keys to load for the modality in the dataset.""" + + def model_post_init(self, *args, **kwargs): + """Initialize eval_delta_indices to delta_indices if not provided.""" + super().model_post_init(*args, **kwargs) + if self.eval_delta_indices is None: + self.eval_delta_indices = self.delta_indices + + +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, + use_global_metadata: bool = True, + metadata_version: str | None = None, + video_backend: str = "ffmpeg", + video_backend_kwargs: dict | None = None, + transforms: ComposedModalityTransform | None = None, + discard_bad_trajectories: bool = True, + fps: float = None, + max_chunk_size: int = None, + relative_action: bool = False, + relative_action_keys: list[str] | None = None, + relative_action_per_horizon: bool = False, + ): + """ + 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. + use_global_metadata (bool): Whether to use global metadata for normalization. + metadata_version (str): The version of the metadata, if `use_global_metadata` is True. + 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" + relative_action (bool): Whether to use relative action stats for normalization. If True, will load or calculate + relative action stats from relative_stats_dreamzero.json. If the file doesn't exist, stats will be calculated. + relative_action_keys (list[str] | None): List of action keys to apply relative action to (e.g., ['joint_position']). + If None and relative_action is True, applies to all action keys except those containing 'gripper'. + relative_action_per_horizon (bool): Whether to use per-horizon relative action stats. If True, will load or calculate + separate stats for each action horizon index from relative_horizon_stats_dreamzero.json. + """ + # first check if the path directory exists + if not Path(dataset_path).exists(): + raise FileNotFoundError(f"Dataset path {dataset_path} does not exist") + + self.modality_configs = modality_configs + self.use_global_metadata = use_global_metadata + self.metadata_version = metadata_version + self.video_backend = video_backend + self.video_backend_kwargs = video_backend_kwargs if video_backend_kwargs is not None else {} + self.fps = fps + self.max_chunk_size = max_chunk_size + self.transforms = ( + transforms if transforms is not None else ComposedModalityTransform(transforms=[]) + ) + self.discard_bad_trajectories = discard_bad_trajectories + self.relative_action = relative_action + self.relative_action_per_horizon = relative_action_per_horizon + # Determine which action keys should use relative action + if relative_action_keys is not None: + self.relative_action_keys = relative_action_keys + else: + # Default: apply to all action keys except those containing 'gripper' + self.relative_action_keys = None # Will be set after modality_configs is available + self._relative_action_keys_input = relative_action_keys # Store original input + self._dataset_path = Path(dataset_path) + self._dataset_name = self._dataset_path.name + self.tag = EmbodimentTag(embodiment_tag) + # For dream and lapa, we use the global metadata since the lapa_actions and dream_actions are already normalized + if self.tag == EmbodimentTag.DREAM or self.tag == EmbodimentTag.LAPA: + self.use_global_metadata = True + self._lerobot_modality_meta = self._get_lerobot_modality_meta() + self._lerobot_info_meta = self._get_lerobot_info_meta() + # Notice: We also include discarded trajectories in stats for larger state coverage, for questions please ask @Fengyuan Hu @Yuqi Xie + self._lerobot_stats_meta = self._get_lerobot_stats_meta() + + # Initialize trajectory info and chunk size early (needed for relative stats calculation) + self._trajectory_ids, self._trajectory_lengths = self._get_trajectories() + self._data_path_pattern = self._get_data_path_pattern() + self._chunk_size = self._get_chunk_size() + + # Set default relative_action_keys if not provided + if self.relative_action and self._relative_action_keys_input is None: + # Default: apply to all action keys except those containing 'gripper' + action_keys = self.modality_configs.get("action", ModalityConfig(delta_indices=[0], modality_keys=[])).modality_keys + self.relative_action_keys = [ + k.replace("action.", "") for k in action_keys + if "gripper" not in k.lower() + ] + print(f"Relative action will be applied to keys: {self.relative_action_keys}") + # Load relative action stats if relative_action is enabled + print("relative_action", self.relative_action) + self._lerobot_relative_stats_meta = self._get_lerobot_relative_stats_meta() if self.relative_action else {} + # Load per-horizon relative action stats if relative_action_per_horizon is enabled + print("relative_action_per_horizon", self.relative_action_per_horizon) + self._lerobot_relative_horizon_stats_meta = self._get_lerobot_relative_horizon_stats_meta() if self.relative_action_per_horizon else {} + self._metadata = self._get_metadata() + self._step_filter = self._get_step_filter() + self._all_steps = self._get_all_steps() + self._modality_keys = self._get_modality_keys() + self._delta_indices = self._get_delta_indices() + self._max_delta_index = self._get_max_delta_index() + self._dataset_name = self._dataset_path.name + + # NOTE(YL): method to predict the task progress + if "action.task_progress" in self._modality_keys["action"]: + from groot.vla.data.schema import StateActionMetadata + + print("we will add task progress to the action modality") + self._modality_keys["action"].append("action.task_progress") + self._metadata.modalities.action["task_progress"] = StateActionMetadata( + absolute=True, rotation_type=None, shape=(1,), continuous=True + ) + # assume the task progress is uniformly distributed between 0 and 1 + self._metadata.statistics.action["task_progress"] = DatasetStatisticalValues( + max=[1.0], min=[0.0], mean=[0.5], std=[0.2887], q01=[0.01], q99=[0.99] + ) + + self.set_transforms_metadata(self.metadata) + self.set_epoch(0) + + print(f"Initialized dataset {self.dataset_name} with {embodiment_tag}") + + # LeRobot-specific config (some already initialized above for relative stats) + self._video_path_pattern = self._get_video_path_pattern() + self._tasks = self._get_tasks() + self._detailed_global_instructions = self._get_detailed_global_instructions() + self.curr_traj_data = None + self.curr_traj_id = None + + # 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 + + def _get_max_delta_index(self) -> int: + """Calculate the maximum delta index across all modalities. + + Returns: + int: The maximum delta index value. + """ + max_delta_index = 0 + for delta_index in self.delta_indices.values(): + max_delta_index = max(max_delta_index, delta_index.max()) + return max_delta_index + + @property + def max_delta_index(self) -> int: + """The maximum delta index across all modalities.""" + return self._max_delta_index + + @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 lerobot_stats_meta(self) -> dict[str, DatasetStatisticalValues]: + """The metadata for the LeRobot dataset.""" + return self._lerobot_stats_meta + + @property + def lerobot_relative_stats_meta(self) -> dict[str, DatasetStatisticalValues]: + """The relative action stats metadata for the LeRobot dataset.""" + return self._lerobot_relative_stats_meta + + @property + def lerobot_relative_horizon_stats_meta(self) -> dict[str, dict[str, list]]: + """The per-horizon relative action stats metadata for the LeRobot dataset. + + Format: {action_key: {stat_name: [[h0_vals], [h1_vals], ...]}} + """ + return self._lerobot_relative_horizon_stats_meta + + @property + def step_filter(self) -> dict[int, np.ndarray]: + """The step filter for the dataset.""" + return self._step_filter + + @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_lerobot_modality_meta(self) -> LeRobotModalityMetadata: + """Get the metadata for the LeRobot dataset.""" + if self.use_global_metadata: + assert ( + self.metadata_version is not None + ), "metadata_version must be provided if use_global_metadata is True" + modality_meta_path = ( + METADATA_DIR + / self.tag.value + / self.metadata_version + / Path(LE_ROBOT_MODALITY_FILENAME).name + ) + assert ( + modality_meta_path.exists() + ), f"Please provide a {Path(LE_ROBOT_MODALITY_FILENAME).name} file in {METADATA_DIR / self.tag.value / self.metadata_version}" + with open(modality_meta_path, "r") as f: + modality_meta = LeRobotModalityMetadata.model_validate(json.load(f)) + return modality_meta + else: + 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_lerobot_stats_meta(self) -> dict[str, DatasetStatisticalValues]: + """Get the metadata for the LeRobot dataset.""" + if self.use_global_metadata: + assert ( + self.metadata_version is not None + ), "metadata_version must be provided if use_global_metadata is True" + stats_path = ( + METADATA_DIR + / self.tag.value + / self.metadata_version + / Path(LE_ROBOT_STATS_FILENAME).name + ) + else: + stats_path = self.dataset_path / LE_ROBOT_STATS_FILENAME + try: + with open(stats_path, "r") as f: + stats: dict = json.load(f) + for name in ["num_trajectories", "total_trajectory_length"]: + stats.pop(name, None) + for name, stat in stats.items(): + stats[name] = DatasetStatisticalValues.model_validate(stat) + return stats + except (FileNotFoundError, ValidationError) as e: + if self.use_global_metadata: + raise ValueError( + f"{e}: Please provide a {Path(LE_ROBOT_STATS_FILENAME).name} file in {stats_path}" + " and ensure the metadata format is correct." + ) + print(f"Failed to load dataset statistics: {e}") + print(f"Calculating dataset statistics for {self.dataset_name}") + # Get all parquet files in the dataset paths + parquet_files = list((self.dataset_path).glob(LE_ROBOT_DATA_FILENAME)) + lowdim_features = [] + le_features = self.lerobot_info_meta["features"] + for feature in le_features: + if "float" in le_features[feature]["dtype"]: + lowdim_features.append(feature) + + stats = calculate_dataset_statistics(parquet_files, lowdim_features) + stats_serialized = {k: v.model_dump(mode="json") for k, v in stats.items()} + with open(stats_path, "w") as f: + json.dump(stats_serialized, f, indent=4) + return stats + + def _get_lerobot_relative_stats_meta(self) -> dict[str, DatasetStatisticalValues]: + """Get the relative action stats metadata for the LeRobot dataset. + + Returns: + dict[str, DatasetStatisticalValues]: Dictionary mapping action keys to their relative stats. + """ + # Determine the path for relative stats file + if self.use_global_metadata: + assert ( + self.metadata_version is not None + ), "metadata_version must be provided if use_global_metadata is True" + stats_path = ( + METADATA_DIR + / self.tag.value + / self.metadata_version + / Path(LEROBOT_RELATIVE_STATS_FILE_NAME).name + ) + assert ( + stats_path.exists() + ), f"Please provide a {Path(LEROBOT_RELATIVE_STATS_FILE_NAME).name} file in {METADATA_DIR / self.tag.value / self.metadata_version}" + else: + stats_path = self.dataset_path / LEROBOT_RELATIVE_STATS_FILE_NAME + + # Try to load existing relative stats + if stats_path.exists(): + print(f"Loading relative action stats from {stats_path}") + with open(stats_path, "r") as f: + stats: dict = json.load(f) + for name, stat in stats.items(): + stats[name] = DatasetStatisticalValues.model_validate(stat) + return stats + + # Calculate relative stats if file doesn't exist + print(f"Relative stats file not found at {stats_path}") + print(f"Calculating relative action stats for {self.dataset_name}") + + # Get action keys from modality configs, filtered by relative_action_keys + all_action_keys = self.modality_configs.get("action", ModalityConfig(delta_indices=[0], modality_keys=[])).modality_keys + if not all_action_keys: + print("No action keys found in modality configs, skipping relative stats calculation") + return {} + + # Filter to only the keys that should use relative action + action_keys_to_process = [] + for key in all_action_keys: + subkey = key.replace("action.", "") + if self.relative_action_keys is None or subkey in self.relative_action_keys: + action_keys_to_process.append(subkey) + + if not action_keys_to_process: + print("No action keys to process for relative stats") + return {} + + print(f"Will calculate relative stats for: {action_keys_to_process}") + + stats = {} + for action_key in action_keys_to_process: + print(f"Calculating relative stats for action key: {action_key}") + try: + relative_stats = self._calculate_relative_stats_for_key(action_key) + stats[action_key] = relative_stats + except Exception as e: + print(f"Failed to calculate relative stats for {action_key}: {e}") + continue + + if stats: + # Save the calculated stats + stats_serialized = {k: v.model_dump(mode="json") for k, v in stats.items()} + # Only save to dataset path (not global metadata path) + save_path = self.dataset_path / LEROBOT_RELATIVE_STATS_FILE_NAME + print(f"Saving relative action stats to {save_path}") + with open(save_path, "w") as f: + json.dump(stats_serialized, f, indent=4) + + return stats + + def _get_lerobot_relative_horizon_stats_meta(self) -> dict[str, dict[str, list]]: + """Get the per-horizon relative action stats metadata for the LeRobot dataset. + + Similar to _get_lerobot_relative_stats_meta but calculates separate stats for each + action horizon index. Will load from file if exists, otherwise calculate and save. + + Returns: + dict[str, dict[str, list]]: Nested dictionary where: + - Outer key is the action key (e.g., 'joint_position') + - Inner key is the stat name (e.g., 'max', 'min', 'mean', 'std', 'q01', 'q99') + - Value is a list of stat values per horizon index + """ + # Determine the path for per-horizon relative stats file + if self.use_global_metadata: + assert ( + self.metadata_version is not None + ), "metadata_version must be provided if use_global_metadata is True" + stats_path = ( + METADATA_DIR + / self.tag.value + / self.metadata_version + / Path(LEROBOT_RELATIVE_HORIZON_STATS_FILE_NAME).name + ) + assert ( + stats_path.exists() + ), f"Please provide a {Path(LEROBOT_RELATIVE_HORIZON_STATS_FILE_NAME).name} file in {METADATA_DIR / self.tag.value / self.metadata_version}" + else: + stats_path = self.dataset_path / LEROBOT_RELATIVE_HORIZON_STATS_FILE_NAME + + # Try to load existing per-horizon relative stats + if stats_path.exists(): + print(f"Loading per-horizon relative action stats from {stats_path}") + with open(stats_path, "r") as f: + stats: dict = json.load(f) + return stats + + # Calculate per-horizon relative stats if file doesn't exist + print(f"Per-horizon relative stats file not found at {stats_path}") + print(f"Calculating per-horizon relative action stats for {self.dataset_name}") + + # Get action keys from modality configs, filtered by relative_action_keys + all_action_keys = self.modality_configs.get("action", ModalityConfig(delta_indices=[0], modality_keys=[])).modality_keys + if not all_action_keys: + print("No action keys found in modality configs, skipping per-horizon relative stats calculation") + return {} + + # Filter to only the keys that should use relative action + action_keys_to_process = [] + for key in all_action_keys: + subkey = key.replace("action.", "") + if self.relative_action_keys is None or subkey in self.relative_action_keys: + action_keys_to_process.append(subkey) + + if not action_keys_to_process: + print("No action keys to process for per-horizon relative stats") + return {} + + print(f"Will calculate per-horizon relative stats for: {action_keys_to_process}") + + stats = {} + for action_key in action_keys_to_process: + print(f"Calculating per-horizon relative stats for action key: {action_key}") + try: + relative_stats = self._calculate_relative_stats_for_key_per_horizon(action_key) + stats[action_key] = relative_stats + except Exception as e: + print(f"Failed to calculate per-horizon relative stats for {action_key}: {e}") + continue + + if stats: + # Only save to dataset path (not global metadata path) + save_path = self.dataset_path / LEROBOT_RELATIVE_HORIZON_STATS_FILE_NAME + print(f"Saving per-horizon relative action stats to {save_path}") + with open(save_path, "w") as f: + json.dump(stats, f, indent=4) + + return stats + + def _calculate_relative_stats_for_key(self, action_key: str) -> DatasetStatisticalValues: + """Calculate relative action statistics for a specific action key. + + Args: + action_key: The action key to calculate stats for (e.g., 'joint_position') + + Returns: + DatasetStatisticalValues: The calculated statistics for the relative action. + """ + # Get state and action metadata from lerobot modality config + state_key = action_key # Assume state key matches action key + + # Get the modality metadata to find original column names and indices + state_meta = self.lerobot_modality_meta.state.get(state_key) + action_meta = self.lerobot_modality_meta.action.get(action_key) + + if state_meta is None: + raise ValueError(f"State key '{state_key}' not found in modality metadata") + if action_meta is None: + raise ValueError(f"Action key '{action_key}' not found in modality metadata") + + # Get the original column names (e.g., 'observation.state', 'action') + state_original_key = state_meta.original_key + action_original_key = action_meta.original_key + + # Get the indices to slice from the concatenated vectors + state_start, state_end = state_meta.start, state_meta.end + action_start, action_end = action_meta.start, action_meta.end + + state_delta_indices = self.modality_configs.get("state", ModalityConfig(delta_indices=[0], modality_keys=[])).delta_indices + action_delta_indices = self.modality_configs["action"].delta_indices + + print(f"Calculating relative stats for {action_key}:") + print(f" State: column='{state_original_key}', indices=[{state_start}:{state_end}]") + print(f" Action: column='{action_original_key}', indices=[{action_start}:{action_end}]") + + # # Calculate relative actions for all trajectories + all_relative_actions = [] + + # for traj_id in tqdm(self.trajectory_ids, desc=f"Calculating relative stats for {action_key}"): + max_trajs_for_stats = 10000 + traj_ids_to_process = self.trajectory_ids + if len(traj_ids_to_process) > max_trajs_for_stats: + # Randomly sample 500 trajectories + rng = np.random.default_rng(seed=42) + sampled_indices = rng.choice(len(traj_ids_to_process), size=max_trajs_for_stats, replace=False) + traj_ids_to_process = traj_ids_to_process[sampled_indices] + print(f"Sampling {max_trajs_for_stats} trajectories out of {len(self.trajectory_ids)} for stats calculation") + + # Calculate relative actions for sampled trajectories + all_relative_actions = [] + + for traj_id in tqdm(traj_ids_to_process, desc=f"Calculating relative stats for {action_key}"): + try: + # Load trajectory data + traj_data = self._load_trajectory_data(traj_id) + if traj_data is None: + continue + + # Check if columns exist + if state_original_key not in traj_data.columns or action_original_key not in traj_data.columns: + print(f"Missing columns: state='{state_original_key}' exists={state_original_key in traj_data.columns}, " + f"action='{action_original_key}' exists={action_original_key in traj_data.columns}") + continue + + # Load full state and action arrays, then slice to get the specific component + full_state_data = np.stack(traj_data[state_original_key].values) + full_action_data = np.stack(traj_data[action_original_key].values) + + # Slice to get just the component we care about (e.g., joint_position) + state_data = full_state_data[:, state_start:state_end] + action_data = full_action_data[:, action_start:action_end] + + # Calculate usable length based on action delta indices + usable_length = len(traj_data) - max(action_delta_indices) + + for i in range(usable_length): + # Get reference state (last state before action chunk) + ref_state_idx = state_delta_indices[-1] + i + if ref_state_idx >= len(state_data): + continue + ref_state = state_data[ref_state_idx] + + # Get action chunk + action_indices = [idx + i for idx in action_delta_indices] + if max(action_indices) >= len(action_data): + continue + actions = action_data[action_indices] + + # print("actions shape", actions.shape, "ref_state shape", ref_state.shape) + + # Calculate relative actions (action - reference state) + relative_actions = actions - ref_state + all_relative_actions.extend(relative_actions) + + except Exception as e: + print(f"Error processing trajectory {traj_id}: {e}") + continue + + if not all_relative_actions: + raise ValueError(f"No relative actions calculated for {action_key}") + + all_relative_actions = np.array(all_relative_actions) + print(f"Collected {len(all_relative_actions)} relative action samples for {action_key}") + + return DatasetStatisticalValues( + max=np.max(all_relative_actions, axis=0).tolist(), + min=np.min(all_relative_actions, axis=0).tolist(), + mean=np.mean(all_relative_actions, axis=0).tolist(), + std=np.std(all_relative_actions, axis=0).tolist(), + q01=np.quantile(all_relative_actions, 0.01, axis=0).tolist(), + q99=np.quantile(all_relative_actions, 0.99, axis=0).tolist(), + ) + + def _calculate_relative_stats_for_key_per_horizon( + self, action_key: str + ) -> dict[str, list]: + """Calculate relative action statistics for each delta index (horizon step) separately. + + Unlike `_calculate_relative_stats_for_key` which pools all horizon steps together, + this method calculates separate statistics for each action horizon index. + + Args: + action_key: The action key to calculate stats for (e.g., 'joint_position') + + Returns: + dict[str, list]: Dictionary where keys are stat names (max, min, mean, std, q01, q99) + and values are lists of stat values per horizon index. + Format: {"max": [[h0_vals], [h1_vals], ...], "min": [...], ...} + """ + # Get state and action metadata from lerobot modality config + state_key = action_key # Assume state key matches action key + + # Get the modality metadata to find original column names and indices + state_meta = self.lerobot_modality_meta.state.get(state_key) + action_meta = self.lerobot_modality_meta.action.get(action_key) + + if state_meta is None: + raise ValueError(f"State key '{state_key}' not found in modality metadata") + if action_meta is None: + raise ValueError(f"Action key '{action_key}' not found in modality metadata") + + # Get the original column names (e.g., 'observation.state', 'action') + state_original_key = state_meta.original_key + action_original_key = action_meta.original_key + + # Get the indices to slice from the concatenated vectors + state_start, state_end = state_meta.start, state_meta.end + action_start, action_end = action_meta.start, action_meta.end + + state_delta_indices = self.modality_configs.get("state", ModalityConfig(delta_indices=[0], modality_keys=[])).delta_indices + action_delta_indices = self.modality_configs["action"].delta_indices + + print(f"Calculating per-horizon relative stats for {action_key}:") + print(f" State: column='{state_original_key}', indices=[{state_start}:{state_end}]") + print(f" Action: column='{action_original_key}', indices=[{action_start}:{action_end}]") + print(f" Action delta indices: {action_delta_indices}") + + # Initialize separate lists for each horizon index + all_relative_actions_per_horizon: dict[int, list] = { + delta_idx: [] for delta_idx in action_delta_indices + } + + max_trajs_for_stats = 10000 + traj_ids_to_process = self.trajectory_ids + if len(traj_ids_to_process) > max_trajs_for_stats: + # Randomly sample trajectories + rng = np.random.default_rng(seed=42) + sampled_indices = rng.choice(len(traj_ids_to_process), size=max_trajs_for_stats, replace=False) + traj_ids_to_process = traj_ids_to_process[sampled_indices] + print(f"Sampling {max_trajs_for_stats} trajectories out of {len(self.trajectory_ids)} for stats calculation") + + for traj_id in tqdm(traj_ids_to_process, desc=f"Calculating per-horizon relative stats for {action_key}"): + try: + # Load trajectory data + traj_data = self._load_trajectory_data(traj_id) + if traj_data is None: + continue + + # Check if columns exist + if state_original_key not in traj_data.columns or action_original_key not in traj_data.columns: + continue + + # Load full state and action arrays, then slice to get the specific component + full_state_data = np.stack(traj_data[state_original_key].values) + full_action_data = np.stack(traj_data[action_original_key].values) + + # Slice to get just the component we care about (e.g., joint_position) + state_data = full_state_data[:, state_start:state_end] + action_data = full_action_data[:, action_start:action_end] + + # Calculate usable length based on action delta indices + usable_length = len(traj_data) - max(action_delta_indices) + + for i in range(usable_length): + # Get reference state (last state before action chunk) + ref_state_idx = state_delta_indices[-1] + i + if ref_state_idx >= len(state_data): + continue + ref_state = state_data[ref_state_idx] + + # Get action for each horizon index separately + for delta_idx in action_delta_indices: + action_idx = delta_idx + i + if action_idx >= len(action_data): + continue + action = action_data[action_idx] + + # Calculate relative action (action - reference state) + relative_action = action - ref_state + all_relative_actions_per_horizon[delta_idx].append(relative_action) + + except Exception as e: + print(f"Error processing trajectory {traj_id}: {e}") + continue + + # Calculate stats for each horizon index and organize by stat name + stat_names = ["max", "min", "mean", "std", "q01", "q99"] + stats_by_name: dict[str, list] = {name: [] for name in stat_names} + + for delta_idx in action_delta_indices: + relative_actions = all_relative_actions_per_horizon[delta_idx] + if not relative_actions: + print(f"Warning: No relative actions calculated for {action_key} at horizon index {delta_idx}") + # Add empty/placeholder values + for name in stat_names: + stats_by_name[name].append([]) + continue + + relative_actions_array = np.array(relative_actions) + print(f"Collected {len(relative_actions_array)} relative action samples for {action_key} at horizon {delta_idx}") + + stats_by_name["max"].append(np.max(relative_actions_array, axis=0).tolist()) + stats_by_name["min"].append(np.min(relative_actions_array, axis=0).tolist()) + stats_by_name["mean"].append(np.mean(relative_actions_array, axis=0).tolist()) + stats_by_name["std"].append(np.std(relative_actions_array, axis=0).tolist()) + stats_by_name["q01"].append(np.quantile(relative_actions_array, 0.01, axis=0).tolist()) + stats_by_name["q99"].append(np.quantile(relative_actions_array, 0.99, axis=0).tolist()) + + return stats_by_name + + def get_relative_stats_per_horizon( + self, + action_keys: list[str] | None = None, + save_to_file: bool = True, + ) -> dict[str, dict[str, list]]: + """Get relative action stats calculated separately for each horizon index. + + This is useful when you want different normalization for different action horizon + steps, e.g., near-future actions vs far-future actions might have different distributions. + + Args: + action_keys: List of action keys to calculate stats for. If None, uses + relative_action_keys (all action keys except gripper by default). + save_to_file: Whether to save the calculated stats to a file. + + Returns: + dict[str, dict[str, list]]: Nested dictionary where: + - Outer key is the action key (e.g., 'joint_position') + - Inner key is the stat name (e.g., 'max', 'min', 'mean', 'std', 'q01', 'q99') + - Value is a list of stat values per horizon index + + Example output format: + { + "joint_position": { + "max": [[h0_vals], [h1_vals], ...], + "min": [[h0_vals], [h1_vals], ...], + ... + } + } + """ + # Determine which action keys to process + if action_keys is None: + all_action_keys = self.modality_configs.get( + "action", ModalityConfig(delta_indices=[0], modality_keys=[]) + ).modality_keys + if not all_action_keys: + print("No action keys found in modality configs") + return {} + # Default: apply to all action keys except those containing 'gripper' + action_keys = [ + k.replace("action.", "") for k in all_action_keys + if "gripper" not in k.lower() + ] + + if not action_keys: + print("No action keys to process for per-horizon relative stats") + return {} + + print(f"Calculating per-horizon relative stats for: {action_keys}") + + all_stats: dict[str, dict[str, list]] = {} + + for action_key in action_keys: + print(f"Processing action key: {action_key}") + try: + stats_per_horizon = self._calculate_relative_stats_for_key_per_horizon(action_key) + all_stats[action_key] = stats_per_horizon + except Exception as e: + print(f"Failed to calculate per-horizon relative stats for {action_key}: {e}") + continue + + if save_to_file and all_stats: + # Save to the designated file + save_path = self.dataset_path / LEROBOT_RELATIVE_HORIZON_STATS_FILE_NAME + print(f"Saving per-horizon relative action stats to {save_path}") + + with open(save_path, "w") as f: + json.dump(all_stats, f, indent=4) + + return all_stats + + def load_relative_stats_per_horizon( + self, + stats_path: Path | str | None = None, + ) -> dict[str, dict[str, list]]: + """Load pre-computed per-horizon relative stats from a file. + + Args: + stats_path: Path to the stats file. If None, uses the default path in the dataset. + + Returns: + dict[str, dict[str, list]]: Nested dictionary of stats per horizon. + """ + if stats_path is None: + stats_path = self.dataset_path / LEROBOT_RELATIVE_HORIZON_STATS_FILE_NAME + else: + stats_path = Path(stats_path) + + if not stats_path.exists(): + print(f"Per-horizon relative stats file not found at {stats_path}") + return {} + + print(f"Loading per-horizon relative stats from {stats_path}") + with open(stats_path, "r") as f: + all_stats = json.load(f) + + return all_stats + + def _load_trajectory_data(self, traj_id: int) -> pd.DataFrame | None: + """Load trajectory data from parquet file. + + Args: + traj_id: The trajectory ID to load. + + Returns: + pd.DataFrame or None if loading fails. + """ + try: + chunk_index = traj_id // self.chunk_size + parquet_path = self.dataset_path / f"data/chunk-{chunk_index:03d}/episode_{traj_id:06d}.parquet" + if not parquet_path.exists(): + # Try alternative pattern + parquet_files = list(self.dataset_path.glob(f"data/*/episode_{traj_id:06d}.parquet")) + if parquet_files: + parquet_path = parquet_files[0] + else: + return None + return pd.read_parquet(parquet_path) + except Exception: + return None + + def _get_step_filter(self) -> dict[int, np.ndarray]: + """Get the step filter for the dataset.""" + step_filter_path = self.dataset_path / STEP_FILTER_FILENAME + step_filter = {} + if step_filter_path.exists(): + with open(step_filter_path, "r") as f: + for line in f: + episode_step_filter = json.loads(line) + trajectory_id = episode_step_filter["episode_index"] + all_indices = np.arange(self.trajectory_lengths[trajectory_id].item()) + indices_to_filter = np.array(episode_step_filter["step_indices"]) + step_filter[trajectory_id] = np.setdiff1d(all_indices, indices_to_filter) + else: + for trajectory_id in self.trajectory_ids: + step_filter[trajectory_id] = np.arange( + self.trajectory_lengths[trajectory_id].item() + ) + return step_filter + + def _get_metadata(self) -> DatasetMetadata: + """Get the metadata for the dataset. + + Returns: + dict: The metadata for the dataset. + """ + + # 1. Modality metadata + # 1.1. State and action modalities + simplified_modality_meta: dict[str, dict] = {} + for modality in ["state", "action"]: + simplified_modality_meta[modality] = {} + le_state_action_meta: dict[str, LeRobotStateActionMetadata] = getattr( + self.lerobot_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 self.lerobot_modality_meta.video: + original_key = self.lerobot_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): + # channels = le_video_meta["shape"][le_video_meta["names"].index("channels")] + channels = le_video_meta["info"]["video.channels"] + fps = le_video_meta["info"]["video.fps"] + simplified_modality_meta["video"][new_key] = { + "resolution": [width, height], + "channels": channels, + "fps": fps, + } + + # 2. Dataset statistics + dataset_statistics = {} + le_statistics = {k: v.model_dump() for k, v in self.lerobot_stats_meta.items()} + # Prepare relative stats if available + relative_stats = {} + if self.relative_action and hasattr(self, '_lerobot_relative_stats_meta'): + relative_stats = {k: v.model_dump() for k, v in self._lerobot_relative_stats_meta.items()} + + # Prepare per-horizon relative stats if available + per_horizon_stats = {} + if self.relative_action_per_horizon and hasattr(self, '_lerobot_relative_horizon_stats_meta'): + per_horizon_stats = self._lerobot_relative_horizon_stats_meta + + 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 = self.lerobot_modality_meta.get_key_meta( + f"{our_modality}.{subkey}" + ) + assert isinstance(state_action_meta, LeRobotStateActionMetadata) + + # Check if we should use per-horizon relative stats for this action key + should_use_per_horizon = ( + our_modality == "action" + and self.relative_action_per_horizon + and subkey in per_horizon_stats + and (self.relative_action_keys is None or subkey in self.relative_action_keys) + ) + + # Use relative stats for action modality if relative_action is enabled and stats are available + # Also check if this subkey is in the list of keys that should use relative action + should_use_relative = ( + our_modality == "action" + and self.relative_action + and subkey in relative_stats + and (self.relative_action_keys is None or subkey in self.relative_action_keys) + ) + + if should_use_per_horizon: + # Use per-horizon relative action stats (format: {stat_name: [[h0_vals], [h1_vals], ...]}) + for stat_name in per_horizon_stats[subkey]: + dataset_statistics[our_modality][subkey][stat_name] = per_horizon_stats[subkey][stat_name] + print(f"Using per-horizon relative stats for {subkey}") + elif should_use_relative: + # Use relative action stats directly + for stat_name in relative_stats[subkey]: + dataset_statistics[our_modality][subkey][stat_name] = relative_stats[subkey][stat_name] + print(f"Using relative stats for {subkey}: {dataset_statistics[our_modality][subkey]}") + else: + # Use original absolute stats + 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=self.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 + episode_path = self.dataset_path / LE_ROBOT_EPISODE_FILENAME + with open(episode_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) + + 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[int, int]]: A list of (trajectory_id, base_index) tuples. + + Example: + self.trajectory_ids: [0, 1, 2] + self.step_filter: { + 0: [0, 1, 2], + 1: [0, 1], + 2: [0, 2, 3] + } + return: [ + (0, 0), (0, 1), (0, 2), + (1, 0), (1, 1), + (2, 0), (2, 2), (2, 3) + ] + """ + all_steps: list[tuple[int, int]] = [] + # All steps is used in single dataset, so we need to discard bad trajectories + # Mixture dataset directly use trajectory_ids, so we handle it by changing the sampling weights + discarded_episode_indices = [] + if self.discard_bad_trajectories: + discarded_episode_indices = self._lerobot_info_meta.get("discarded_episode_indices", []) + + for trajectory_id in self.trajectory_ids: + if trajectory_id in discarded_episode_indices: + continue + for base_index in self.step_filter[trajectory_id]: + all_steps.append((trajectory_id, base_index)) + return all_steps + + def _get_modality_keys(self) -> dict: + """Get the modality keys for the dataset. + + Returns: + dict: Dictionary mapping modality names to their keys. + """ + 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 _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.""" + 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") + + def _get_task_embeddings(self) -> dict: + """Get the task embeddings for the dataset.""" + task_embeddings_path = self.dataset_path / LE_ROBOT_TASK_EMBEDDINGS_FILENAME + return torch.load(task_embeddings_path) + + def _get_detailed_global_instructions(self) -> dict[int, dict]: + """Get the detailed global instructions for the dataset. + + Loads from episodes_detail_global_instruction.jsonl if it exists. + + Returns: + dict[int, dict]: Mapping from episode_index to detailed instruction dict. + """ + detailed_instruction_path = self.dataset_path / LE_ROBOT_DETAILED_GLOBAL_INSTRUCTION_FILENAME + if not detailed_instruction_path.exists(): + return {} + with open(detailed_instruction_path, "r") as f: + instructions_list = [json.loads(line) for line in f] + return {entry["episode_index"]: entry for entry in instructions_list} + + 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, modality_config in self.modality_configs.items(): + if modality in ["lapa_action", "dream_actions", "rl_info", "task_embedding"]: + continue + for key in modality_config.modality_keys: + + if key == "action.task_progress": + continue + # Skip metadata-based language keys (they don't need modality metadata) + if modality == "language" and key.startswith("annotation."): + lang_subkey = key.replace("annotation.", "") + if lang_subkey in METADATA_LANG_KEYS: + continue + # 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) + # Also set per-horizon statistics if available + if self.relative_action_per_horizon and hasattr(self, '_lerobot_relative_horizon_stats_meta'): + if hasattr(self.transforms, 'set_per_horizon_statistics'): + self.transforms.set_per_horizon_statistics(self._lerobot_relative_horizon_stats_meta) + + 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] + indices = { + key: delta_indices + base_index for key, delta_indices in self.delta_indices.items() + } + return self.transforms(self.get_step_data(trajectory_id, indices)) + + def get_step_data(self, trajectory_id: int, indices: dict[str, np.ndarray]) -> 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. + indices (dict[str, np.ndarray]): The indices for each modality. + + 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 + self.curr_traj_data = self.get_trajectory_data(trajectory_id) + for modality in self.modality_keys: + # Get the data corresponding to each key in the modality + for key in self.modality_keys[modality]: + # Only load the data if the key is in the indices + if key in indices: + data[key] = self.get_data_by_modality( + trajectory_id, modality, key, indices[key] + ) + return data + + def get_parquet_path(self, trajectory_id: int) -> Path: + """Get the parquet path for a trajectory.""" + chunk_index = self.get_episode_chunk(trajectory_id) + return self.dataset_path / self.data_path_pattern.format( + episode_chunk=chunk_index, episode_index=trajectory_id + ) + + def get_trajectory_data(self, trajectory_id: int) -> pd.DataFrame: + """Get the data for a trajectory.""" + if self.curr_traj_id == trajectory_id and self.curr_traj_data is not None: + return self.curr_traj_data + else: + parquet_path = self.get_parquet_path(trajectory_id) + assert parquet_path.exists(), f"Parquet file not found at {parquet_path}" + return pd.read_parquet(parquet_path) + + 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 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 trajectory. + padding_strategy (str): The padding strategy, either "first_last" or "zero". + "first_last" uses first/last step data for padding, "zero" uses zero padding. + + Returns: + np.ndarray: The retrieved and padded data. + """ + # 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: + """Get the video file path for a specific trajectory and video key. + + Args: + trajectory_id (int): The ID of the trajectory. + key (str): The video key (without 'video.' prefix). + + Returns: + Path: Path to the video file. + """ + 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 + video_filename = self.video_path_pattern.format( + episode_chunk=chunk_index, episode_index=trajectory_id, video_key=original_key + ) + return self.dataset_path / video_filename + + def get_video( + self, + trajectory_id: int, + key: str, + step_indices: np.ndarray, + ) -> 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) + """ + # 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.", "") + 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] + + # try: + return get_frames_by_timestamps( + video_path.as_posix(), + video_timestamp, + video_backend=self.video_backend, + video_backend_kwargs=self.video_backend_kwargs, + ) + # except: + # self.video_backend = "torchvision_av" + # return get_frames_by_timestamps( + # video_path.as_posix(), + # video_timestamp, + # video_backend=self.video_backend, + # video_backend_kwargs=self.video_backend_kwargs, + # ) + + + def get_state_or_action( + self, + trajectory_id: int, + modality: str, + key: str, + step_indices: np.ndarray, + ) -> 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 trajectory index + trajectory_index = self.get_trajectory_index(trajectory_id) + # Get the maximum length of the trajectory + max_length = self.trajectory_lengths[trajectory_index] + + # Note [YL]: this handles action.task_progress if specified + if key == "action.task_progress": + # Get frame_index array and apply proper bounds checking and padding + frame_index_array = self.curr_traj_data["frame_index"].to_numpy() + # Use retrieve_data_and_pad to handle out-of-bounds indices + frame_index = self.retrieve_data_and_pad( + array=frame_index_array, + step_indices=step_indices, + max_length=max_length, + padding_strategy="first_last", # Use first/last for task progress + ) + # get the task progress by using "frame index / trajectory length" + progress = frame_index / max_length + progress = progress.reshape(-1, 1) + return progress + + assert key.startswith(modality + "."), f"{key} must start with {modality + '.'}, got {key}" + # Get the sub-key, e.g. state.joint_angles -> joint_angles + subkey = 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[subkey].original_key + if le_key is None: + le_key = subkey + # 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 + if data_array.ndim == 1: + assert ( + data_array.shape[0] == max_length + ), f"Expected 1D array with length {max_length}, got {data_array.shape} array" + data_array = data_array.reshape(-1, 1) + assert data_array.ndim == 2, f"Expected 2D array, got {data_array.shape} array" + le_indices = np.arange( + le_state_or_action_cfg[subkey].start, + le_state_or_action_cfg[subkey].end, + ) + data_array = data_array[:, le_indices] + # Get the state or action configuration + state_or_action_cfg = getattr(self.metadata.modalities, modality)[subkey] + + # 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", + ) + + def get_lapa_action( + self, + trajectory_id: int, + key: str, + step_indices: np.ndarray, + ) -> np.ndarray | None: + """Get LAPA action data for a trajectory by step indices. + + Args: + trajectory_id (int): The ID of the trajectory. + key (str): The key of the LAPA action data. + step_indices (np.ndarray): The step indices to retrieve data for. + + Returns: + np.ndarray | None: The LAPA action data, or None if the key is not found. + """ + # 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] + # Check key in the current trajectory data + assert self.curr_traj_data is not None, f"No data found for {trajectory_id=}" + if ( + key not in self.curr_traj_data.columns + ): # this ensures that we can still load data w/o lapa actions. will store values that are None. + return None + # assert key in self.curr_traj_data.columns, f"{key} not found in {trajectory_id=}" + # Get the data array, shape: (T, D) + data_array: np.ndarray = np.stack(self.curr_traj_data[key]) # type: ignore + assert data_array.ndim == 2, f"Expected 2D array, got {data_array.shape} array" + # Pad the data + return self.retrieve_data_and_pad( + array=data_array, + step_indices=step_indices, + max_length=max_length, + padding_strategy="first_last", + ) + + def get_dream_actions( + self, + trajectory_id: int, + key: str, + step_indices: np.ndarray, + ) -> np.ndarray | None: + """Get DREAM action data for a trajectory by step indices. + + Args: + trajectory_id (int): The ID of the trajectory. + key (str): The key of the DREAM action data. + step_indices (np.ndarray): The step indices to retrieve data for. + + Returns: + np.ndarray | None: The DREAM action data, or None if the key is not found. + """ + # 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] + # Check key in the current trajectory data + assert self.curr_traj_data is not None, f"No data found for {trajectory_id=}" + if ( + key not in self.curr_traj_data.columns + ): # this ensures that we can still load data w/o lapa actions. will store values that are None. + return None + # assert key in self.curr_traj_data.columns, f"{key} not found in {trajectory_id=}" + # Get the data array, shape: (T, D) + data_array: np.ndarray = np.stack(self.curr_traj_data[key]) # type: ignore + assert data_array.ndim == 2, f"Expected 2D array, got {data_array.shape} array" + # Pad the data + return self.retrieve_data_and_pad( + array=data_array, + step_indices=step_indices, + max_length=max_length, + padding_strategy="first_last", + ) + + def get_language( + self, + trajectory_id: int, + key: str, + step_indices: np.ndarray, + ) -> list[str]: + """Get the language annotation data for a trajectory by step indices. + + Args: + trajectory_id (int): The ID of the trajectory. + key (str): The key of the annotation. + step_indices (np.ndarray): The step indices to retrieve data for. + + 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 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 + assert key.startswith( + "annotation." + ), f"Language key must start with 'annotation.', got {key}" + subkey = key.replace("annotation.", "") + # print("subkey", subkey) + + # Check if this is a metadata-based language key (detailed_global_instruction_medium/concise) + if subkey in METADATA_LANG_KEYS: + # print("return metadata language") + return self._get_language_from_metadata(trajectory_id, subkey, len(step_indices)) + + # Otherwise, load from parquet columns (original behavior) + 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 + if pd.api.types.is_numeric_dtype(self.curr_traj_data[original_key]): + # Stored as list of integers + task_indices: list[int] = self.curr_traj_data[original_key].iloc[step_indices].tolist() + return self.tasks.loc[task_indices]["task"].tolist() + else: + # Stored as list of strings + return self.curr_traj_data[original_key].iloc[step_indices].astype(str).tolist() + + def _get_language_from_metadata( + self, + trajectory_id: int, + lang_key: str, + nframes: int, + ) -> list[str]: + """Get language instruction from metadata files for special language keys. + + Supports: + - detailed_global_instruction_medium: Longer, detailed description + - detailed_global_instruction_concise: Short summary + + Args: + trajectory_id (int): The ID of the trajectory (episode_index). + lang_key (str): The language key (e.g., "detailed_global_instruction_medium"). + nframes (int): Number of frames to return the instruction for. + + Returns: + list[str]: The instruction repeated for each frame (empty string if not found). + """ + if trajectory_id in self._detailed_global_instructions: + instruction = self._detailed_global_instructions[trajectory_id].get(lang_key, "") + # print("instruction", instruction) + else: + instruction = "" + return [instruction] * nframes + + def get_rl_info( + self, + trajectory_id: int, + key: str, + step_indices: np.ndarray, + ) -> np.ndarray: + """Get the reward data for a trajectory by step indices. + + If the step indices are out of range, pad with first/last step data. + + Args: + trajectory_id (int): The ID of the trajectory. + key (str): The key of the reward data. + step_indices (np.ndarray): The step indices to retrieve data for. + + Returns: + np.ndarray: The reward data for the trajectory and step indices. + """ + # 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] + data_array: np.ndarray = np.stack(self.curr_traj_data[key]) # type: ignore + + if key == "rl_info.next.reward": + padding_strategy = "zero" + else: + padding_strategy = "first_last" + + # Pad the data + return self.retrieve_data_and_pad( + array=data_array, + step_indices=step_indices, + max_length=max_length, + padding_strategy=padding_strategy, + ) + + def get_data_by_modality( + self, + trajectory_id: int, + modality: str, + key: str, + step_indices: np.ndarray, + ) -> np.ndarray | list[str] | None: + """Get the data corresponding to the modality for a trajectory by step indices. + + This method dispatches to the appropriate specialized method based on the modality. + For the language modality, empty strings are returned if no matching data is found. + + Args: + trajectory_id (int): The ID of the trajectory. + modality (str): The modality of the data (video, state, action, language, etc.). + key (str): The key of the data. + step_indices (np.ndarray): The step indices of the trajectory. + + Returns: + np.ndarray | list[str] | None: The data for the specified modality. + """ + if modality == "video": + return self.get_video(trajectory_id, key, step_indices) + elif modality == "state" or modality == "action": + return self.get_state_or_action(trajectory_id, modality, key, step_indices) + elif modality == "language": + return self.get_language(trajectory_id, key, step_indices) + elif modality == "lapa_action": + return self.get_lapa_action(trajectory_id, key, step_indices) + elif modality == "dream_actions": + return self.get_dream_actions(trajectory_id, key, step_indices) + elif modality == "rl_info": + return self.get_rl_info(trajectory_id, key, step_indices) + else: + raise ValueError(f"Invalid modality: {modality}") + + def get_initial_actions(self): + """Load initial actions from the dataset if available. + + Returns: + list: List containing initial actions if the file exists, empty list otherwise. + """ + initial_actions_path = self.dataset_path / INITIAL_ACTIONS_FILENAME + if initial_actions_path.exists(): + initial_actions = load_initial_actions(initial_actions_path) + return initial_actions # a single-element list of dict[str, dict[str, np.ndarray]] + else: + return [] + + + +class CachedLeRobotSingleDataset(LeRobotSingleDataset): + """A cached version of LeRobotSingleDataset that preloads all video frames into memory. + + This class caches video frames for each trajectory and key to improve access speed + when video frames need to be accessed multiple times. Recommended for small datasets + or when memory usage is not a concern. + """ + + def __init__(self, *args, **kwargs): + """Initialize the cached dataset and preload all video frames. + + Args: + *args: Arguments passed to parent LeRobotSingleDataset. + **kwargs: Keyword arguments passed to parent LeRobotSingleDataset. + """ + # 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 = [] + 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, + ) + 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" + 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, step_indices: np.ndarray) -> np.ndarray: + """Get video frames from the cached data. + + Args: + trajectory_id (int): The ID of the trajectory. + key (str): The video key (with 'video.' prefix). + step_indices (np.ndarray): The step indices to retrieve frames for. + + Returns: + np.ndarray: The video frames with shape (T, H, W, C). + """ + # 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 safe_hash(input_tuple): + """Generate a safe hash from an input tuple. + + Creates a deterministic hash using SHA256 and returns the lower 128 bits. + This is used for deterministic random seed generation. + + Args: + input_tuple: The tuple to hash. + + Returns: + int: A 128-bit hash value. + """ + # 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): + """Specification element for a dataset mixture defining paths and weights. + + This class validates dataset paths by embodiment tag and handles weight distribution + across multiple dataset paths if requested. + """ + + dataset_path: dict[str, 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.", + ) + + @field_validator("dataset_path", mode="after") + def validate_dataset_path_keys(cls, v: dict[str, list[Path] | Path]) -> dict[str, list[Path]]: + """Validate dataset paths and expand glob patterns. + + Args: + v (dict[str, list[Path] | Path]): Dictionary mapping embodiment tags to paths. + + Returns: + dict[str, list[Path]]: Validated and expanded paths. + + Raises: + ValueError: If an invalid embodiment tag is provided. + """ + all_globbed_paths: dict[str, list[Path]] = {} + for embodiment_tag, paths in v.items(): + try: + _ = EmbodimentTag(embodiment_tag) + except ValueError: + raise ValueError(f"Invalid embodiment tag: {embodiment_tag}") + if isinstance(paths, Path): + paths = [paths] + globbed_paths = [] + for path in paths: + globbed_paths.extend(glob.glob(str(path))) + all_globbed_paths[embodiment_tag] = globbed_paths + return all_globbed_paths + + +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]], + training: bool, + balance_dataset_weights: bool = True, + balance_trajectory_weights: bool = True, + seed: int = 42, + allow_padding_at_end: bool = False, + metadata_config: dict = { + "percentile_mixing_method": "min_max", + }, + ): + """ + Initialize the mixture dataset. + + Args: + data_mixture (list[tuple[LeRobotSingleDataset, float]]): Datasets and their corresponding weights. + training (bool): If True, __getitem__ will return different samples every epoch; if False, __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. + allow_padding_at_end (bool): If True, allow padding at the end of the dataset. + """ + datasets: list[LeRobotSingleDataset] = [] + dataset_sampling_weights: list[float] = [] + for dataset, weight in data_mixture: + datasets.append(dataset) + dataset_sampling_weights.append(weight) + self.datasets = datasets + self.balance_dataset_weights = balance_dataset_weights + self.balance_trajectory_weights = balance_trajectory_weights + self.seed = seed + self.training = training + self.allow_padding_at_end = allow_padding_at_end + + # Set properties for sampling + + # 1. Dataset lengths + self._dataset_lengths = np.array([len(dataset) for dataset in self.datasets]) + + # 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 + self._dataset_sampling_weights /= self._dataset_sampling_weights.sum() + + # 3. Trajectory sampling weights + self._trajectory_sampling_weights: list[np.ndarray] = [] + for dataset in self.datasets: + trajectory_sampling_weights = np.ones(len(dataset.trajectory_ids)) + if self.balance_trajectory_weights: + trajectory_sampling_weights *= np.array( + [ + len(dataset.step_filter[trajectory_id]) + for trajectory_id in dataset.trajectory_ids + ] + ) + + if dataset.discard_bad_trajectories: + bad_trajectory_indices = dataset.lerobot_info_meta.get( + "discarded_episode_indices", [] + ) + trajectory_sampling_weights[bad_trajectory_indices] = 0.0 + + if trajectory_sampling_weights.sum() == 0: + raise ValueError(f"No valid trajectories found for dataset {dataset}") + + trajectory_sampling_weights /= trajectory_sampling_weights.sum() + self._trajectory_sampling_weights.append(trajectory_sampling_weights) + + # 4. Primary dataset indices + self._primary_dataset_indices = np.array(dataset_sampling_weights) == 1.0 + + # Set the epoch and sample the first epoch + self.set_epoch(0) + + # Create a merged metadata for the mixture dataset (we don't need this in the future as eval will directly use `get_metadata`) + self.update_metadata(metadata_config) + + # Set the transforms to training or evaluation mode + if self.training: + for dataset in self.datasets: + dataset.transforms.train() + else: + for dataset in self.datasets: + dataset.transforms.eval() + + @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: + """Return a string representation of the mixture dataset with weights.""" + 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 yaml.dump({"Mixture dataset": dataset_descriptions}) + + @classmethod + def from_mixture_spec( + cls: type[T_LeRobotMixtureDataset], + mixture_spec: Sequence[MixtureSpecElement | dict], + dataset_class: type[LeRobotSingleDataset] | str, + all_modality_configs: dict[str, dict[str, ModalityConfig]], + all_transforms: dict[str, ComposedModalityTransform], + metadata_versions: dict[str, str], + fps: float = None, + dataset_kwargs: dict | None = None, + mixture_kwargs: dict | None = None, + ) -> T_LeRobotMixtureDataset: + """Initialize the mixture dataset from a specification. + + Args: + mixture_spec (Sequence[MixtureSpecElement | dict]): The specification for the mixture dataset. + dataset_class (type[LeRobotSingleDataset] | str): The dataset class or its string path. + all_modality_configs (dict[str, dict[str, ModalityConfig]]): The modality configs for each embodiment. + all_transforms (dict[str, ComposedModalityTransform]): The transforms for each embodiment. + metadata_versions (dict[str, str]): The metadata versions for each embodiment. + dataset_kwargs (dict | None): Additional keyword arguments for the dataset classes. + mixture_kwargs (dict | None): Additional keyword arguments for the mixture dataset. + + Returns: + LeRobotMixtureDataset: The initialized mixture dataset. + """ + if isinstance(dataset_class, str): + module_name, class_name = dataset_class.rsplit(".", 1) + module = importlib.import_module(module_name) + dataset_class = getattr(module, class_name) + assert not isinstance(dataset_class, str), f"{dataset_class} is a string" + assert issubclass( + dataset_class, LeRobotSingleDataset + ), f"{dataset_class} is not a subclass of LeRobotSingleDataset" + data_mixture = [] + + for dataset_spec in tqdm( + mixture_spec, + total=len(mixture_spec), + desc="Initializing datasets", + ): + start_time = time.time() + if isinstance(dataset_spec, dict): + dataset_spec = MixtureSpecElement.model_validate(dataset_spec) + datasets = [] + for embodiment_tag, paths in dataset_spec.dataset_path.items(): + if isinstance(paths, Path): + paths = [paths] + for dataset_path in paths: + if '.sh' in dataset_path or '.json' in dataset_path: + continue + assert ( + embodiment_tag in all_modality_configs + ), f"{embodiment_tag} not in modality_configs: {all_modality_configs.keys()}" + assert ( + embodiment_tag in all_transforms + ), f"{embodiment_tag} not in transforms: {all_transforms.keys()}" + dataset = dataset_class( + dataset_path=dataset_path, + embodiment_tag=EmbodimentTag(embodiment_tag), + modality_configs=copy.copy(all_modality_configs[embodiment_tag]), + transforms=copy.copy(all_transforms[embodiment_tag]), + metadata_version=metadata_versions[embodiment_tag], + fps=fps[embodiment_tag] if embodiment_tag in fps else None, + **(dataset_kwargs if dataset_kwargs is not None else {}), + ) + datasets.append(dataset) + dataset_lengths = np.array([len(dataset) for dataset in datasets]) + dataset_relative_lengths = dataset_lengths / dataset_lengths.sum() + for dataset, relative_length in zip(datasets, dataset_relative_lengths): + if dataset_spec.distribute_weights: + weight = relative_length * dataset_spec.dataset_weight + else: + weight = dataset_spec.dataset_weight + data_mixture.append((dataset, weight)) + + print( + f"Time taken to initialize {len(datasets)} datasets: {time.time() - start_time:.2f} seconds" + ) + + return cls( + data_mixture=data_mixture, + **(mixture_kwargs if mixture_kwargs is not None else {}), + ) + + 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 sample_step(self, index: int) -> tuple[LeRobotSingleDataset, int, int]: + """Sample a single step from the mixture dataset. + + Args: + index (int): The index to sample (used for deterministic sampling). + + Returns: + tuple[LeRobotSingleDataset, int, int]: A tuple of (dataset, trajectory_id, step_index). + """ + # return self.sampled_steps[index] + + # Set seed + if self.training: + seed = 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] + + if self.allow_padding_at_end: + # Sample trajectory + trajectory_index = rng.choice( + len(dataset.trajectory_ids), p=self.trajectory_sampling_weights[dataset_index] + ) + trajectory_id = dataset.trajectory_ids[trajectory_index] + + allowed_length = dataset.trajectory_lengths[trajectory_index] + else: + # Avoid padding at the end of the trajectory + max_delta_index = dataset.max_delta_index + trajectory_length = 0 + trajectory_id = None + while trajectory_length < max_delta_index + 1: + # Sample trajectory + trajectory_index = rng.choice( + len(dataset.trajectory_ids), + p=self.trajectory_sampling_weights[dataset_index], + ) + trajectory_id = dataset.trajectory_ids[trajectory_index] + trajectory_length = dataset.trajectory_lengths[trajectory_index] + assert trajectory_id is not None + + # Sample step + assert ( + trajectory_length >= max_delta_index + 1 + ), f"{trajectory_length=}, {max_delta_index=}" + allowed_length = trajectory_length - max_delta_index + # Get the allowed indices from the step filter + allowed_indices = dataset.step_filter[trajectory_id] + # Remove indices that are too large + allowed_indices = allowed_indices[allowed_indices <= allowed_length] + step_index = rng.choice(allowed_indices) + return dataset, trajectory_id, step_index + else: + length_cumsum = np.cumsum(self.dataset_lengths) + dataset_index = np.searchsorted(length_cumsum, index) + dataset = self.datasets[dataset_index] + assert ( + len(dataset._lerobot_info_meta.get("discarded_episode_indices", [])) == 0 + ), f"Find discarded episode indices in evaluation dataset {dataset.dataset_path}" + trajectory_id, step_index = dataset.all_steps[index - length_cumsum[dataset_index]] + return dataset, trajectory_id, step_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. + """ + dataset, trajectory_id, step_index = self.sample_step(index) + indices = { + key: delta_indices + step_index for key, delta_indices in dataset.delta_indices.items() + } + return dataset.transforms(dataset.get_step_data(trajectory_id, indices)) + + 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. + """ + if self.training: + return int((self.dataset_lengths * self.dataset_sampling_weights).sum()) + else: + return int(self.dataset_lengths.sum()) + + @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: + # Check if stats are per-horizon (2D) by examining the first task's mean + first_mean = np.array(per_task_stats[0][modality]["mean"]) + is_per_horizon = first_mean.ndim == 2 # Shape (horizon_len, action_dim) + + if is_per_horizon: + # Handle per-horizon stats (2D arrays) + stats_shape = first_mean.shape # (horizon_len, action_dim) + + # Initialize accumulators for means and variances + weighted_means = np.zeros(stats_shape) + weighted_squares = np.zeros(stats_shape) + + # 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(np.array(stats["min"])) + max_list.append(np.array(stats["max"])) + q01_list.append(np.array(stats["q01"])) + q99_list.append(np.array(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(np.maximum(overall_variance, 0)).tolist() + + # Compute overall min and max per dimension + # Stack along new axis: (num_tasks, horizon_len, action_dim) + overall_min = np.min(np.stack(min_list, axis=0), axis=0).tolist() + overall_max = np.max(np.stack(max_list, axis=0), axis=0).tolist() + + # Compute overall q01 and q99 per dimension + q01_array = np.stack(q01_list, axis=0) # (num_tasks, horizon_len, action_dim) + q99_array = np.stack(q99_list, axis=0) + if percentile_mixing_method == "weighted_average": + # Weighted average along task axis + weighted_q01 = np.average(q01_array, axis=0, weights=normalized_weights).tolist() + weighted_q99 = np.average(q99_array, axis=0, weights=normalized_weights).tolist() + 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}") + else: + # Handle regular stats (1D arrays) + num_dims = len(first_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(np.maximum(overall_variance, 0)).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() + 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) -> 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. + """ + + self.merged_metadata: dict[str, DatasetMetadata] = {} + # Group metadata by tag + all_metadatas: dict[str, list[DatasetMetadata]] = {} + for dataset in self.datasets: + if dataset.tag.value not in all_metadatas: + all_metadatas[dataset.tag.value] = [] + all_metadatas[dataset.tag.value].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.value]) + + def get_initial_actions(self): + initial_actions = [] + for dataset in self.datasets: + if hasattr(dataset, "get_initial_actions"): + initial_actions.extend(dataset.get_initial_actions()) + return initial_actions diff --git a/groot/vla/data/dataset/lerobot_sharded.py b/groot/vla/data/dataset/lerobot_sharded.py new file mode 100644 index 0000000000000000000000000000000000000000..409fa3786449666e84ad54e12412c041c26e4066 --- /dev/null +++ b/groot/vla/data/dataset/lerobot_sharded.py @@ -0,0 +1,1542 @@ +from concurrent.futures import Future, ThreadPoolExecutor +import json +from pathlib import Path +import time + +import numpy as np +import pandas as pd +import torch +import torch.distributed as dist +from torch.utils.data import IterableDataset, get_worker_info +import yaml + +from groot.vla.common.utils import get_frames_by_timestamps + +from .lerobot import LE_ROBOT_EPISODE_FILENAME, LeRobotMixtureDataset, LeRobotSingleDataset + + +class ShardedLeRobotSingleDataset(LeRobotSingleDataset): + """ + A single dataset with shards. + """ + + def __init__( + self, + *args, + num_steps_per_shard: int = int(1e4), + **kwargs, + ): + self.args = args + self.kwargs = kwargs + super().__init__(*args, **kwargs) + self.num_steps_per_shard = num_steps_per_shard + self.all_video_paths = self.get_all_video_paths() + self.all_parquet_paths = self.get_all_parquet_paths() + self.sharded_trajectories, self.shard_lengths = self.generate_shards() + self.frames_to_load = self.get_all_frames_to_load() + + # Set shard caching properties + self.shard_start_indices: dict[int, int] | None = None + self.cached_shard: dict[str, np.ndarray] | None = None + self.cached_df: pd.DataFrame | None = None + self.frame_indices_map: dict[int, dict[str, np.ndarray]] | None = None + self._executor = ThreadPoolExecutor(max_workers=1) + self._cache_job: Future | None = None + + @property + def num_shards(self) -> int: + """The number of shards.""" + return len(self.sharded_trajectories) + + def get_all_video_paths(self) -> dict[int, dict[str, Path]]: + """Get the video paths for all trajectories and all views. + + Returns: + dict[int, dict[str, Path]]: The video paths for all trajectories. + """ + video_paths = {} + for trajectory_id in self.trajectory_ids: + if isinstance(trajectory_id, np.integer): + trajectory_id = trajectory_id.item() + assert isinstance( + trajectory_id, int + ), f"trajectory_id must be an integer, got {type(trajectory_id)}" + video_paths[trajectory_id] = {} + for key in self.modality_keys["video"]: + assert key.startswith("video."), f"Video key must start with 'video.', got {key}" + video_paths[trajectory_id][key] = self.get_video_path( + trajectory_id, key.replace("video.", "") + ) + return video_paths + + def get_all_parquet_paths(self) -> dict[int, Path]: + """Get the parquet paths for all trajectories. + + Returns: + dict[int, Path]: The parquet paths for all trajectories. + """ + return { + trajectory_id: self.get_parquet_path(trajectory_id) + for trajectory_id in self.trajectory_ids + } + + def generate_shards(self) -> tuple[list[list[int]], np.ndarray]: + """Generate shards of trajectories. We recommend num_steps_per_shard >> average trajectory length. + + Args: + num_steps_per_shard (int): The number of steps per shard. + + Returns: + list[list[str]]: The shards of trajectories. + """ + sharded_trajectories = [[]] + curr_num_steps = 0 + curr_shard_index = 0 + discarded_episode_indices = [] + trajectory_ids = self.trajectory_ids + if self.discard_bad_trajectories: + discarded_episode_indices = self._lerobot_info_meta.get("discarded_episode_indices", []) + trajectory_ids = [ + trajectory_id + for trajectory_id in trajectory_ids + if trajectory_id not in discarded_episode_indices + ] + + assert ( + len(trajectory_ids) > 0 + ), f"No valid trajectories found for dataset {self.dataset_path}" + total_steps = np.sum( + [len(self.step_filter[trajectory_id]) for trajectory_id in trajectory_ids] + ).astype(int) + num_shards = np.ceil(total_steps / self.num_steps_per_shard).astype(int) + cutoffs = np.linspace(0, total_steps, num_shards + 1)[1:] # Exclude the first cutoff (0) + shard_lengths = [] + last_num_steps = 0 + for trajectory_id in trajectory_ids: + sharded_trajectories[-1].append(trajectory_id) + curr_num_steps += len(self.step_filter[trajectory_id]) + if curr_num_steps > cutoffs[curr_shard_index]: + sharded_trajectories.append([]) + curr_shard_index += 1 + shard_lengths.append(curr_num_steps - last_num_steps) + last_num_steps = curr_num_steps + shard_lengths.append(curr_num_steps - last_num_steps) + assert ( + curr_num_steps == total_steps + ), "Total steps not equal to the sum of trajectory lengths" + assert ( + len(shard_lengths) == num_shards + ), "Number of shards not equal to the number of cutoffs" + assert ( + len(sharded_trajectories) == num_shards + ), "Number of shards not equal to the number of cutoffs" + print(f"Generated {len(sharded_trajectories)} shards for dataset {self.dataset_path}") + return sharded_trajectories, np.array(shard_lengths) + + def get_all_frames_to_load(self): + """Generate a map of video frame indices to trajectory indices.""" + all_frames_to_load = {} + for trajectory_id in self.trajectory_ids: + all_frames_to_load[trajectory_id] = {} + for key in self.modality_keys["video"]: + assert key.startswith("video."), f"Video key must start with 'video.', got {key}" + filtered_indices = self.step_filter[trajectory_id] + if len(filtered_indices) > 0: + frames_to_load = np.unique( + np.concatenate( + [i + np.array(self.delta_indices[key]) for i in filtered_indices] + ) + ) + # Cap within the length of the trajectory and >= 0 + frames_to_load = frames_to_load[ + (frames_to_load < self.trajectory_lengths[trajectory_id]) + & (frames_to_load >= 0) + ] + else: + frames_to_load = np.array([]) + all_frames_to_load[trajectory_id][key] = frames_to_load + return all_frames_to_load + + @staticmethod + def get_shard( + trajectory_ids: list[int] | np.ndarray, + modality_keys: dict, + video_paths: dict[int, dict[str, Path]], + parquet_paths: dict[int, Path], + frames_to_load: dict[int, dict[str, np.ndarray]], + video_backend: str = "pyav", + video_backend_kwargs: dict | None = None, + ) -> tuple[ + dict[str, np.ndarray], dict[int, int], pd.DataFrame, dict[int, dict[str, np.ndarray]] + ]: + print("Caching shard") + start_time = time.time() + assert "video" in modality_keys, "No video modality found. No need to use caching." + cached_frames = {} + trajectory_start_indices = {} + frame_indices_map = {} + curr_step_index = 0 + cached_df = None + curr_frame_index = {key: 0 for key in modality_keys["video"]} + for trajectory_id in trajectory_ids: + trajectory_start_indices[trajectory_id] = curr_step_index + parquet_path = parquet_paths[trajectory_id] + parquet_df = pd.read_parquet(parquet_path) + # Check timestamps are in sync + parquet_timestamps = parquet_df["timestamp"].to_numpy() + trajectory_length = len(parquet_timestamps) + if isinstance(trajectory_id, np.integer): + trajectory_id = trajectory_id.item() + assert isinstance( + trajectory_id, int + ), f"trajectory_id must be an integer, got {type(trajectory_id)}" + frame_indices_map[trajectory_id] = {} + for key in modality_keys["video"]: + # Only load the frames that are needed + this_frames_to_load = frames_to_load[trajectory_id][key] + if len(this_frames_to_load) == 0: + continue + load_timestamps = parquet_timestamps[this_frames_to_load] + assert key.startswith("video."), f"Video key must start with 'video.', got {key}" + # Store a mapping that frame_indices_map[trajectory_id][key][frame_index] = index_in_concat_video_frames + frame_indices_map[trajectory_id][key] = ( + np.ones(len(parquet_timestamps), dtype=np.int32) * -1 + ) + frame_indices_map[trajectory_id][key][this_frames_to_load] = np.arange( + curr_frame_index[key], + curr_frame_index[key] + len(this_frames_to_load), + dtype=np.int32, + ) + curr_frame_index[key] += len(this_frames_to_load) + if key not in cached_frames: + cached_frames[key] = [] + frames = get_frames_by_timestamps( + video_paths[trajectory_id][key].as_posix(), + timestamps=load_timestamps, + video_backend=video_backend, + video_backend_kwargs=video_backend_kwargs or {}, + ) + cached_frames[key].append(frames) + if cached_df is None: + cached_df = parquet_df + else: + cached_df = pd.concat([cached_df, parquet_df]) + curr_step_index += trajectory_length + + # Concatenate the frames + for key in cached_frames: + cached_frames[key] = np.concatenate(cached_frames[key], axis=0) + end_time = time.time() + print(f"Cached shard in {end_time - start_time:.2f} seconds") + assert cached_df is not None, "Cached dataframe is None" + # Add global "index" column if missing (some dataset formats omit it) + if "index" not in cached_df.columns: + cached_df = cached_df.reset_index(drop=True) + cached_df["index"] = cached_df.index + return cached_frames, trajectory_start_indices, cached_df, frame_indices_map + + def start_cache_shard(self, shard_index: int) -> None: + """Start caching a shard in a background thread.""" + self._cache_job = self._executor.submit( + self.get_shard, + self.sharded_trajectories[shard_index], + self.modality_keys, + self.all_video_paths, + self.all_parquet_paths, + self.frames_to_load, + self.video_backend, + self.video_backend_kwargs, + ) + + def finish_cache_shard(self): + """Get the cached shard.""" + assert self._cache_job is not None + self.cached_shard, self.shard_start_indices, self.cached_df, self.frame_indices_map = ( + self._cache_job.result() + ) + self._cache_job = None # Clear the future to allow memory to be freed + + def delete_cached_shard(self): + """Delete the cached shard.""" + del self.cached_shard + del self.shard_start_indices + del self.cached_df + + def get_trajectories_in_shard(self) -> list[int]: + """Get the trajectories in a shard.""" + assert self.shard_start_indices is not None + return list(self.shard_start_indices.keys()) + + def get_video(self, trajectory_id: int, key: str, step_indices: np.ndarray) -> np.ndarray: + """Get the video frames from cached shards for a trajectory by a base index. + + Args: + 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 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) + # Calculate the absolute indices + assert ( + self.shard_start_indices is not None + and self.cached_shard is not None + and trajectory_id in self.shard_start_indices + and self.frame_indices_map is not None + and trajectory_id in self.frame_indices_map + and key in self.frame_indices_map[trajectory_id] + ), "Shard not cached. Please call `cache_next_shard` and `use_next_shard` first." + indices_in_shard = self.frame_indices_map[trajectory_id][key][step_indices] + assert np.all( + indices_in_shard != -1 + ), f"Indices in shard are not loaded for {trajectory_id=}, {key=}, {step_indices=}" + return self.cached_shard[key][indices_in_shard] + + def get_trajectory_data(self, trajectory_id: int) -> pd.DataFrame: + """Get the trajectory data.""" + assert self.cached_df is not None, "Cached dataframe is None" + traj_data = self.cached_df.loc[self.cached_df["episode_index"] == trajectory_id] + trajectory_index = self.get_trajectory_index(trajectory_id) + trajectory_length = self.trajectory_lengths[trajectory_index] + assert ( + len(traj_data) == trajectory_length + ), f"Trajectory length mismatch: {len(traj_data)} != {trajectory_length} {self.args} {self.kwargs}" + indices = traj_data["index"].to_numpy() + if len(indices) > 0: + start_index = indices[0] + expected_indices = np.arange(start_index, start_index + len(indices)) + assert np.array_equal( + indices, expected_indices + ), f"[{self}] Index sequence mismatch in trajectory data, {trajectory_id=}" + return traj_data + + +class ShardedLeRobotSubLangSingleActionChunkDatasetDROID(LeRobotSingleDataset): + """ + A single dataset with shards. + """ + + def __init__( + self, + *args, + num_steps_per_shard: int = int(1e4), + **kwargs, + ): + self.args = args + self.kwargs = kwargs + super().__init__(*args, **kwargs) + self.num_steps_per_shard = num_steps_per_shard + self.all_video_paths = self.get_all_video_paths() + self.all_parquet_paths = self.get_all_parquet_paths() + self.sharded_trajectories, self.shard_lengths = self.generate_shards() + + # Set shard caching properties + self.shard_start_indices: dict[int, int] | None = None + self.cached_shard: dict[str, np.ndarray] | None = None + self.cached_df: pd.DataFrame | None = None + self._executor = ThreadPoolExecutor(max_workers=1) + self._cache_job: Future | None = None + # self._traj_cache: dict[int, pd.DataFrame] = {} + # # Precompute language key once to avoid repeated scans + # self.language_key: str | None = next( + # ( + # modality_key + # for modality_name in self.modality_keys + # for modality_key in self.modality_keys[modality_name] + # if modality_key.startswith("annotation.") + # ), + # None, + # ) + # # Track current sample's chunk count for alignment across modalities + # self._current_num_chunks: dict[int, int] = {} + + @property + def num_shards(self) -> int: + """The number of shards.""" + return len(self.sharded_trajectories) + + def get_all_video_paths(self) -> dict[int, dict[str, Path]]: + """Get the video paths for all trajectories and all views. + + Returns: + dict[int, dict[str, Path]]: The video paths for all trajectories. + """ + video_paths = {} + for trajectory_id in self.trajectory_ids: + if isinstance(trajectory_id, np.integer): + trajectory_id = trajectory_id.item() + assert isinstance( + trajectory_id, int + ), f"trajectory_id must be an integer, got {type(trajectory_id)}" + video_paths[trajectory_id] = {} + for key in self.modality_keys["video"]: + assert key.startswith("video."), f"Video key must start with 'video.', got {key}" + video_paths[trajectory_id][key] = self.get_video_path( + trajectory_id, key.replace("video.", "") + ) + return video_paths + + def get_all_parquet_paths(self) -> dict[int, Path]: + """Get the parquet paths for all trajectories. + + Returns: + dict[int, Path]: The parquet paths for all trajectories. + """ + return { + trajectory_id: self.get_parquet_path(trajectory_id) + for trajectory_id in self.trajectory_ids + } + + def generate_shards(self) -> tuple[list[list[int]], np.ndarray]: + """Generate shards of trajectories. We recommend num_steps_per_shard >> average trajectory length. + + Args: + num_steps_per_shard (int): The number of steps per shard. + + Returns: + list[list[str]]: The shards of trajectories. + """ + sharded_trajectories = [[]] + curr_num_steps = 0 + curr_shard_index = 0 + discarded_episode_indices = [] + trajectory_ids = self.trajectory_ids + if self.discard_bad_trajectories: + discarded_episode_indices = self._lerobot_info_meta.get("discarded_episode_indices", []) + trajectory_ids = [ + trajectory_id + for trajectory_id in trajectory_ids + if trajectory_id not in discarded_episode_indices + ] + + assert len(trajectory_ids) > 0, "No valid trajectories found for dataset" + total_steps = np.sum( + [len(self.step_filter[trajectory_id]) for trajectory_id in trajectory_ids] + ).astype(int) + num_shards = np.ceil(total_steps / self.num_steps_per_shard).astype(int) + cutoffs = np.linspace(0, total_steps, num_shards + 1)[1:] # Exclude the first cutoff (0) + shard_lengths = [] + last_num_steps = 0 + for trajectory_id in trajectory_ids: + sharded_trajectories[-1].append(trajectory_id) + curr_num_steps += len(self.step_filter[trajectory_id]) + if curr_num_steps > cutoffs[curr_shard_index]: + sharded_trajectories.append([]) + curr_shard_index += 1 + shard_lengths.append(curr_num_steps - last_num_steps) + last_num_steps = curr_num_steps + shard_lengths.append(curr_num_steps - last_num_steps) + assert ( + curr_num_steps == total_steps + ), "Total steps not equal to the sum of trajectory lengths" + assert ( + len(shard_lengths) == num_shards + ), "Number of shards not equal to the number of cutoffs" + assert ( + len(sharded_trajectories) == num_shards + ), "Number of shards not equal to the number of cutoffs" + print(f"Generated {len(sharded_trajectories)} shards for dataset {self.dataset_path}") + return sharded_trajectories, np.array(shard_lengths) + + @staticmethod + def get_shard( + trajectory_ids: list[int] | np.ndarray, + modality_keys: dict, + video_paths: dict[int, dict[str, Path]], + parquet_paths: dict[int, Path], + video_backend: str = "pyav", + video_backend_kwargs: dict | None = None, + fps: float = None, + ) -> tuple[dict[str, np.ndarray], dict[int, int], pd.DataFrame]: + # Optional logging to avoid stdout overhead during tight loops + # (controlled by instance-level verbose flag) + # Using a staticmethod, we cannot read self.verbose; defer to caller to control prints + print("Caching shard") + start_time = time.time() + assert "video" in modality_keys, "No video modality found. No need to use caching." + cached_frames = {} + trajectory_start_indices = {} + curr_step_index = 0 + cached_df = None + for trajectory_id in trajectory_ids: + trajectory_start_indices[trajectory_id] = curr_step_index + parquet_path = parquet_paths[trajectory_id] + parquet_df = pd.read_parquet(parquet_path) + # Check timestamps are in sync + parquet_timestamps = parquet_df["timestamp"].to_numpy() + trajectory_length = len(parquet_timestamps) + if isinstance(trajectory_id, np.integer): + trajectory_id = trajectory_id.item() + assert isinstance( + trajectory_id, int + ), f"trajectory_id must be an integer, got {type(trajectory_id)}" + for key in modality_keys["video"]: + assert key.startswith("video."), f"Video key must start with 'video.', got {key}" + if key not in cached_frames: + cached_frames[key] = [] + frames = get_frames_by_timestamps( + video_paths[trajectory_id][key].as_posix(), + timestamps=parquet_timestamps, + video_backend=video_backend, + video_backend_kwargs=video_backend_kwargs, + fps=fps, + ) + cached_frames[key].append(frames) + if cached_df is None: + cached_df = parquet_df + else: + cached_df = pd.concat([cached_df, parquet_df]) + curr_step_index += trajectory_length + + # Concatenate the frames + for key in cached_frames: + cached_frames[key] = np.concatenate(cached_frames[key], axis=0) + end_time = time.time() + print(f"Cached shard in {end_time - start_time:.2f} seconds") + assert cached_df is not None, "Cached dataframe is None" + # Add global "index" column if missing (some dataset formats omit it) + if "index" not in cached_df.columns: + cached_df = cached_df.reset_index(drop=True) + cached_df["index"] = cached_df.index + return cached_frames, trajectory_start_indices, cached_df + + def start_cache_shard(self, shard_index: int) -> None: + """Start caching a shard in a background thread.""" + self._cache_job = self._executor.submit( + self.get_shard, + self.sharded_trajectories[shard_index], + self.modality_keys, + self.all_video_paths, + self.all_parquet_paths, + self.video_backend, + self.video_backend_kwargs, + self.fps, + ) + + def finish_cache_shard(self): + """Get the cached shard.""" + assert self._cache_job is not None + self.cached_shard, self.shard_start_indices, self.cached_df = self._cache_job.result() + self._cache_job = None # Clear the future to allow memory to be freed + + def delete_cached_shard(self): + """Delete the cached shard.""" + del self.cached_shard + del self.shard_start_indices + del self.cached_df + # self._traj_cache.clear() + + def get_trajectories_in_shard(self) -> list[int]: + """Get the trajectories in a shard.""" + assert self.shard_start_indices is not None + return list(self.shard_start_indices.keys()) + + def get_step_data(self, trajectory_id: int, indices: dict[str, np.ndarray]) -> 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. + indices (dict[str, np.ndarray]): The indices for each modality. + + 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 + self.curr_traj_data = self.get_trajectory_data(trajectory_id) + for modality in self.modality_keys: + # Get the data corresponding to each key in the modality + for key in self.modality_keys[modality]: + # Only load the data if the key is in the indices + if key in indices: + data[key] = self.get_data_by_modality( + trajectory_id, modality, key, indices[key] + ) + # Skip this sample if state or action data is empty + if data[key] is not None and hasattr(data[key], '__len__') and len(data[key]) == 0: + return None + return data + + def get_video(self, trajectory_id: int, key: str, step_indices: np.ndarray) -> np.ndarray: + """Get the video frames from cached shards for a trajectory by uniformly sampling from language-consistent ranges. + + Args: + trajectory_id (int): The ID of the trajectory. + key (str): The key of the video. + step_indices (np.ndarray): The step indices to retrieve data for. + + Returns: + np.ndarray: The video frames for the trajectory and frame indices. Shape: (T, H, W, C) + """ + # Get the trajectory index + trajectory_index = self.get_trajectory_index(trajectory_id) + trajectory_length = self.trajectory_lengths[trajectory_index] + + # Get trajectory data to access language annotations (reuse if already loaded) + # traj_data = ( + # self.curr_traj_data + # if getattr(self, "curr_traj_data", None) is not None + # else self.get_trajectory_data(trajectory_id) + # ) + traj_data = self.get_trajectory_data(trajectory_id) + # print("trajectory id", trajectory_id, step_indices, trajectory_index) + + # Get language annotations for all steps in the trajectory + # language_key = self.language_key + for modality in self.modality_keys: + for modality_key in self.modality_keys[modality]: + if modality_key.startswith("annotation."): + subkey = modality_key.replace("annotation.", "") + annotation_meta = self.lerobot_modality_meta.annotation + subkey_meta = annotation_meta[subkey] + language_key = subkey_meta.original_key + break + assert language_key is not None, "Language key not found" + if language_key in traj_data.columns: + language_annotations = traj_data[language_key].values + else: + # Fallback to original behavior if language annotations are not available + step_indices = np.maximum(step_indices, 0) + step_indices = np.minimum(step_indices, trajectory_length - 1) + assert ( + self.shard_start_indices is not None + and self.cached_shard is not None + and trajectory_id in self.shard_start_indices + ), "Shard not cached. Please call `cache_next_shard` and `use_next_shard` first." + indices_in_shard = self.shard_start_indices[trajectory_id] + step_indices + return self.cached_shard[key][indices_in_shard] + + # Find language-consistent ranges and uniformly sample from them + sampled_indices = self._uniform_sample_from_language_ranges( + step_indices, language_annotations, trajectory_length + ) + + # Ensure the sampled indices are within the valid range + sampled_indices = np.maximum(sampled_indices, 0) + sampled_indices = np.minimum(sampled_indices, trajectory_length - 1) + # print("sampled indices", sampled_indices) + + # Calculate the absolute indices + assert ( + self.shard_start_indices is not None + and self.cached_shard is not None + and trajectory_id in self.shard_start_indices + ), "Shard not cached. Please call `cache_next_shard` and `use_next_shard` first." + indices_in_shard = self.shard_start_indices[trajectory_id] + sampled_indices + return self.cached_shard[key][indices_in_shard] + + def get_data_by_modality( + self, + trajectory_id: int, + modality: str, + key: str, + step_indices: np.ndarray, + ) -> np.ndarray | list[str] | None: + """Get the data corresponding to the modality for a trajectory by step indices. + + This method dispatches to the appropriate specialized method based on the modality. + For the language modality, empty strings are returned if no matching data is found. + + Args: + trajectory_id (int): The ID of the trajectory. + modality (str): The modality of the data (video, state, action, language, etc.). + key (str): The key of the data. + step_indices (np.ndarray): The step indices of the trajectory. + + Returns: + np.ndarray | list[str] | None: The data for the specified modality. + """ + if modality == "video": + return self.get_video(trajectory_id, key, step_indices) + elif modality == "state": + return self.get_state(trajectory_id, modality, key, step_indices) + elif modality == "action": + return self.get_action(trajectory_id, modality, key, step_indices) + elif modality == "language": + return self.get_language(trajectory_id, key, step_indices) + elif modality == "lapa_action": + return self.get_lapa_action(trajectory_id, key, step_indices) + elif modality == "dream_actions": + return self.get_dream_actions(trajectory_id, key, step_indices) + elif modality == "rl_info": + return self.get_rl_info(trajectory_id, key, step_indices) + else: + raise ValueError(f"Invalid modality: {modality}") + + + def get_state( + self, + trajectory_id: int, + modality: str, + key: str, + step_indices: np.ndarray, + ) -> np.ndarray: + """Get the state 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 trajectory index + trajectory_index = self.get_trajectory_index(trajectory_id) + # Get the maximum length of the trajectory + max_length = self.trajectory_lengths[trajectory_index] + + # Note [YL]: this handles action.task_progress if specified + if key == "action.task_progress": + # Get frame_index array and apply proper bounds checking and padding + frame_index_array = self.curr_traj_data["frame_index"].to_numpy() + # Use retrieve_data_and_pad to handle out-of-bounds indices + frame_index = self.retrieve_data_and_pad( + array=frame_index_array, + step_indices=step_indices, + max_length=max_length, + padding_strategy="first_last", # Use first/last for task progress + ) + # get the task progress by using "frame index / trajectory length" + progress = frame_index / max_length + progress = progress.reshape(-1, 1) + return progress + + assert key.startswith(modality + "."), f"{key} must start with {modality + '.'}, got {key}" + # Get the sub-key, e.g. state.joint_angles -> joint_angles + subkey = 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[subkey].original_key + if le_key is None: + le_key = subkey + # 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 + if data_array.ndim == 1: + assert ( + data_array.shape[0] == max_length + ), f"Expected 1D array with length {max_length}, got {data_array.shape} array" + data_array = data_array.reshape(-1, 1) + assert data_array.ndim == 2, f"Expected 2D array, got {data_array.shape} array" + le_indices = np.arange( + le_state_or_action_cfg[subkey].start, + le_state_or_action_cfg[subkey].end, + ) + data_array = data_array[:, le_indices] + # Get the state or action configuration + state_or_action_cfg = getattr(self.metadata.modalities, modality)[subkey] + + # Build sampled indices for state aligned with language and video sampling + # For state, select only the anchor index per 30-frame chunk (stride 30): + # [..., first_idx-30, first_idx, first_idx+30, ...] + # Stop on language change at the step anchor, bounds, or when reaching 16 anchors (to match 16 chunks). + trajectory_index = self.get_trajectory_index(trajectory_id) + trajectory_length = self.trajectory_lengths[trajectory_index] + # traj_data = ( + # self.curr_traj_data + # if getattr(self, "curr_traj_data", None) is not None + # else self.get_trajectory_data(trajectory_id) + # ) + # language_key = self.language_key + traj_data = self.get_trajectory_data(trajectory_id) + language_key = None + for modality_name in self.modality_keys: + for modality_key in self.modality_keys[modality_name]: + if modality_key.startswith("annotation."): + subkey = modality_key.replace("annotation.", "") + annotation_meta = self.lerobot_modality_meta.annotation + subkey_meta = annotation_meta[subkey] + language_key = subkey_meta.original_key + break + if language_key is not None and language_key in traj_data.columns and len(step_indices) > 0: + language_annotations = traj_data[language_key].values + first_idx = max(0, min(int(step_indices[0]), trajectory_length - 1)) + target_language = language_annotations[first_idx] + + # Get the number of chunks from video sampling to ensure alignment + target_num_chunks = None + # if first_idx in self._current_num_chunks: + if hasattr(self, '_current_num_chunks') and first_idx in self._current_num_chunks: + target_num_chunks = self._current_num_chunks[first_idx] + # print(f"State: Using target_num_chunks from video: {target_num_chunks}") + + max_frames = self.max_chunk_size # 16 anchors to align with 16 chunks as video/action + sampled_list: list[int] = [] + + def add_anchor(anchor_index: int) -> None: + nonlocal sampled_list + if len(sampled_list) >= max_frames: + return + # If we have a target number of chunks, stop when we reach it + if target_num_chunks is not None and len(sampled_list) >= target_num_chunks: + return + # Require full 32-length window to exist for alignment with action/video + if 0 <= anchor_index and anchor_index + 24 < trajectory_length: + sampled_list.append(int(anchor_index)) + + # Always include first_idx anchor + add_anchor(first_idx) + + # Expand outward in 32-frame steps + step = 1 + back_done = False + fwd_done = False + while len(sampled_list) < max_frames and (not back_done or not fwd_done): + # Stop if we've reached the target number of chunks + if target_num_chunks is not None and len(sampled_list) >= target_num_chunks: + break + + if not back_done: + back_anchor = first_idx - 24 * step + if back_anchor < 0: + back_done = True + elif language_annotations[back_anchor] != target_language: + back_done = True + else: + add_anchor(back_anchor) + if len(sampled_list) >= max_frames: + break + if not fwd_done: + fwd_anchor = first_idx + 24 * step + if fwd_anchor >= trajectory_length: + fwd_done = True + elif language_annotations[fwd_anchor] != target_language: + fwd_done = True + else: + add_anchor(fwd_anchor) + step += 1 + + if len(sampled_list) > 0: + sampled_indices = np.array(sorted(set(sampled_list)), dtype=int) + else: + sampled_indices = np.array([], dtype=int) + else: + # Fallback: use provided indices with bounds + sampled_indices = np.maximum(step_indices, 0) + sampled_indices = np.minimum(sampled_indices, trajectory_length - 1) + + # print("sampled indices for state", sampled_indices) + + # Pad the data using the computed sampled indices + return self.retrieve_data_and_pad( + array=data_array, + step_indices=sampled_indices, + max_length=max_length, + padding_strategy="first_last" if state_or_action_cfg.absolute else "zero", + ) + + def get_action( + self, + trajectory_id: int, + modality: str, + key: str, + step_indices: np.ndarray, + ) -> np.ndarray: + """Get the 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 trajectory index + trajectory_index = self.get_trajectory_index(trajectory_id) + # Get the maximum length of the trajectory + max_length = self.trajectory_lengths[trajectory_index] + + # Note [YL]: this handles action.task_progress if specified + if key == "action.task_progress": + # Get frame_index array and apply proper bounds checking and padding + frame_index_array = self.curr_traj_data["frame_index"].to_numpy() + # Use retrieve_data_and_pad to handle out-of-bounds indices + frame_index = self.retrieve_data_and_pad( + array=frame_index_array, + step_indices=step_indices, + max_length=max_length, + padding_strategy="first_last", # Use first/last for task progress + ) + # get the task progress by using "frame index / trajectory length" + progress = frame_index / max_length + progress = progress.reshape(-1, 1) + return progress + + assert key.startswith(modality + "."), f"{key} must start with {modality + '.'}, got {key}" + # Get the sub-key, e.g. state.joint_angles -> joint_angles + subkey = 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[subkey].original_key + if le_key is None: + le_key = subkey + # 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 + if data_array.ndim == 1: + assert ( + data_array.shape[0] == max_length + ), f"Expected 1D array with length {max_length}, got {data_array.shape} array" + data_array = data_array.reshape(-1, 1) + assert data_array.ndim == 2, f"Expected 2D array, got {data_array.shape} array" + le_indices = np.arange( + le_state_or_action_cfg[subkey].start, + le_state_or_action_cfg[subkey].end, + ) + data_array = data_array[:, le_indices] + # Get the state or action configuration + state_or_action_cfg = getattr(self.metadata.modalities, modality)[subkey] + + + # Build sampled indices for action aligned with language and video sampling + # Action runs at 30fps, so for each ±30-frame step around first_idx, + # collect a 30-length chunk with stride 1: [anchor ... anchor+29]. + # Stop on language change at the step anchor, bounds, or when reaching 480 frames (16 chunks * 30). + trajectory_index = self.get_trajectory_index(trajectory_id) + trajectory_length = self.trajectory_lengths[trajectory_index] + # traj_data = ( + # self.curr_traj_data + # if getattr(self, "curr_traj_data", None) is not None + # else self.get_trajectory_data(trajectory_id) + # ) + # language_key = self.language_key + traj_data = self.get_trajectory_data(trajectory_id) + language_key = None + for modality_name in self.modality_keys: + for modality_key in self.modality_keys[modality_name]: + if modality_key.startswith("annotation."): + subkey = modality_key.replace("annotation.", "") + annotation_meta = self.lerobot_modality_meta.annotation + subkey_meta = annotation_meta[subkey] + language_key = subkey_meta.original_key + break + if language_key is not None and language_key in traj_data.columns and len(step_indices) > 0: + language_annotations = traj_data[language_key].values + first_idx = max(0, min(int(step_indices[0]), trajectory_length - 1)) + target_language = language_annotations[first_idx] + + # Get the number of chunks from video sampling to ensure alignment + target_num_chunks = None + # if first_idx in self._current_num_chunks: + if hasattr(self, '_current_num_chunks') and first_idx in self._current_num_chunks: + target_num_chunks = self._current_num_chunks[first_idx] + # print(f"Using target_num_chunks from video: {target_num_chunks}") + + max_frames = 24 * self.max_chunk_size + per_step_offsets = list(range(24)) # 0..23 + sampled_list: list[int] = [] + + def add_step_set(anchor_index: int) -> None: + nonlocal sampled_list + # Ensure the whole 32-length chunk fits within bounds + if anchor_index < 0 or anchor_index + 24 >= trajectory_length: + return + # Ensure we don't overrun the max_frames cap with a partial chunk + if len(sampled_list) + 24 > max_frames: + return + # If we have a target number of chunks, stop when we reach it + if target_num_chunks is not None and len(sampled_list) // 24 >= target_num_chunks: + return + for offset in per_step_offsets: + idx = anchor_index + offset + sampled_list.append(int(idx)) + + # Always include first_idx chunk + add_step_set(first_idx) + + # Expand outward in 32-frame steps + step = 1 + back_done = False + fwd_done = False + while len(sampled_list) < max_frames and (not back_done or not fwd_done): + # Stop if we've reached the target number of chunks + if target_num_chunks is not None and len(sampled_list) // 24 >= target_num_chunks: + break + + if not back_done: + back_anchor = first_idx - 24 * step + if back_anchor < 0: + back_done = True + elif language_annotations[back_anchor] != target_language: + back_done = True + else: + add_step_set(back_anchor) + if len(sampled_list) >= max_frames: + break + if not fwd_done: + fwd_anchor = first_idx + 24 * step + if fwd_anchor >= trajectory_length: + fwd_done = True + elif language_annotations[fwd_anchor] != target_language: + fwd_done = True + else: + add_step_set(fwd_anchor) + step += 1 + + if len(sampled_list) > 0: + unique_sorted = np.array(sorted(set(sampled_list)), dtype=int) + # Enforce divisibility by 30 and the 480 cap + capped_size = min(unique_sorted.size, max_frames) + divisible_size = (capped_size // 24) * 24 + sampled_indices = unique_sorted[:divisible_size] + else: + sampled_indices = np.array([], dtype=int) + else: + # Fallback: use provided indices with bounds + sampled_indices = np.maximum(step_indices, 0) + sampled_indices = np.minimum(sampled_indices, trajectory_length - 1) + + # print("sampled indices for action", first_idx, sampled_indices, trajectory_length) + + # Pad the data using the computed sampled indices + action_data = self.retrieve_data_and_pad( + array=data_array, + step_indices=sampled_indices, + max_length=max_length, + padding_strategy="first_last" if state_or_action_cfg.absolute else "zero", + ) + # print("action data before convert", key) + # Calculate relative action on the fly if relative_action is enabled + # Only apply to keys that are in relative_action_keys + subkey = key.replace("action.", "") + should_convert_to_relative = ( + (self.relative_action or self.relative_action_per_horizon) + and len(sampled_indices) > 0 + and (self.relative_action_keys is None or subkey in self.relative_action_keys) + ) + if should_convert_to_relative: + # print("action data before convert", action_data[0], action_data[-1], key) + action_data = self._convert_to_relative_action( + action_data=action_data, + action_key=key, + sampled_indices=sampled_indices, + trajectory_id=trajectory_id, + chunk_size=24, + ) + # print("action data after convert", action_data[0], action_data[-1], key) + + return action_data + + def _convert_to_relative_action( + self, + action_data: np.ndarray, + action_key: str, + sampled_indices: np.ndarray, + trajectory_id: int, + chunk_size: int = 24, + ) -> np.ndarray: + """Convert absolute action to relative action by subtracting reference state. + + Args: + action_data: Absolute action data, shape (T, D) + action_key: The action key (e.g., 'action.left_arm_joints') + sampled_indices: The sampled indices for the action + trajectory_id: The trajectory ID + chunk_size: Size of each action chunk (default 24) + + Returns: + np.ndarray: Relative action data, shape (T, D) + """ + # Get corresponding state key (assume state key matches action key) + state_key = action_key.replace("action.", "state.") + subkey = action_key.replace("action.", "") + + # Get state data from trajectory + traj_data = self.get_trajectory_data(trajectory_id) + le_state_cfg = getattr(self.lerobot_modality_meta, "state", None) + + if le_state_cfg is None or subkey not in le_state_cfg: + # If no corresponding state key, return original action data + return action_data + + le_state_key = le_state_cfg[subkey].original_key + if le_state_key is None: + le_state_key = subkey + + if le_state_key not in traj_data.columns: + # If state column doesn't exist, return original action data + return action_data + + # Get state data array + state_array: np.ndarray = np.stack(traj_data[le_state_key]) + if state_array.ndim == 1: + state_array = state_array.reshape(-1, 1) + + # Apply same indices as action + le_indices = np.arange( + le_state_cfg[subkey].start, + le_state_cfg[subkey].end, + ) + state_array = state_array[:, le_indices] + + # Calculate relative action for each chunk + relative_action_data = action_data.copy() + num_chunks = len(sampled_indices) // chunk_size + + for chunk_idx in range(num_chunks): + chunk_start = chunk_idx * chunk_size + chunk_end = chunk_start + chunk_size + + # Get anchor index (first index of the chunk) + anchor_idx = sampled_indices[chunk_start] + + # Get reference state at anchor index + if anchor_idx < len(state_array): + reference_state = state_array[anchor_idx] + + # Subtract reference state from all actions in this chunk + relative_action_data[chunk_start:chunk_end] = ( + action_data[chunk_start:chunk_end] - reference_state + ) + + return relative_action_data + + def _uniform_sample_from_language_ranges( + self, + step_indices: np.ndarray, + language_annotations: np.ndarray, + trajectory_length: int + ) -> np.ndarray: + """Uniformly sample from language-consistent ranges based on the first index's language. + + Args: + step_indices (np.ndarray): Original step indices to sample. + language_annotations (np.ndarray): Language annotations for each step in the trajectory. + trajectory_length (int): Total length of the trajectory. + + Returns: + np.ndarray: New indices sampled uniformly from the language-consistent range of the first index. + """ + if len(step_indices) == 0: + return np.array([]) + + # Use only the first index to determine the target language + first_idx = max(0, min(step_indices[0], trajectory_length - 1)) + target_language = language_annotations[first_idx] + + # Build sampled indices by moving in ±32-frame steps from first_idx + # and adding 4 frames at 8-frame strides for each step, while: + # - staying within trajectory bounds, + # - keeping language consistent with target_language at the anchor step, + # - and limiting the total collected frames to 81. + max_frames = 8 * self.max_chunk_size + 1 + per_step_offsets = [0, 3, 6, 9, 12, 15, 18, 21] + sampled_list: list[int] = [] + + def add_step_set(anchor_index: int) -> None: + # Only add a complete 4-frame set if it fully fits and capacity allows + # Require full 32-frame window to exist for alignment with action/state + nonlocal sampled_list + if anchor_index < 0 or anchor_index + 23 >= trajectory_length: + return + if len(sampled_list) + len(per_step_offsets) > max_frames: + return + for offset in per_step_offsets: + idx = anchor_index + offset + sampled_list.append(int(idx)) + + # Always include the set at the first_idx + add_step_set(first_idx) + + # Expand outward in both directions in 32-frame steps + step = 1 + back_done = False + fwd_done = False + while len(sampled_list) < max_frames and (not back_done or not fwd_done): + # Backward step + if not back_done: + back_anchor = first_idx - 24 * step + if back_anchor < 0: + back_done = True + elif language_annotations[back_anchor] != target_language: + back_done = True + else: + add_step_set(back_anchor) + # Forward step + if len(sampled_list) >= max_frames: + break + if not fwd_done: + fwd_anchor = first_idx + 24 * step + if fwd_anchor >= trajectory_length: + fwd_done = True + elif language_annotations[fwd_anchor] != target_language: + fwd_done = True + else: + add_step_set(fwd_anchor) + step += 1 + + # De-duplicate and sort ascending for stable ordering + if len(sampled_list) == 0: + return np.array([]) + unique_sorted = np.array(sorted(set(sampled_list)), dtype=int) + # Ensure we return at most 81 frames + if unique_sorted.size > max_frames: + unique_sorted = unique_sorted[:max_frames] + + # Convert to 4n+1 format by adding one more frame at the end with 8-frame stride + if unique_sorted.size > 0: + # Get the last index and add one more frame with 8-frame stride + last_idx = unique_sorted[-1] + additional_idx = last_idx + 3 + + # Only add if it doesn't exceed trajectory bounds and max_frames + if additional_idx < trajectory_length and unique_sorted.size < max_frames: + unique_sorted = np.append(unique_sorted, additional_idx) + else: + # Trim to 8n+1 format. Require at least 9 frames so (noisy_frames-1)//num_frame_per_block >= 1 + # for action/state model invariant (CausalWanModel); otherwise return empty so sample is skipped. + if unique_sorted.size <= 8: + return np.array([]) + unique_sorted = unique_sorted[:-7] + + # ensure that unique_sorted has 4n+1 frames + assert unique_sorted.size % 8 == 1, f"unique_sorted size {unique_sorted.size} is not 4n+1" + + # Store the number of chunks for alignment with action/state + num_video_chunks = (unique_sorted.size - 1) // 8 + if not hasattr(self, '_current_num_chunks'): + self._current_num_chunks = {} + # Use first_idx as a key to track the current sample's chunk count + self._current_num_chunks[first_idx] = num_video_chunks + + # print("unique_sorted size", unique_sorted.size, "num_video_chunks", num_video_chunks) + return unique_sorted + + + def get_trajectory_data(self, trajectory_id: int) -> pd.DataFrame: + """Get the trajectory data.""" + assert self.cached_df is not None, "Cached dataframe is None" + + # Quick verification + if self.cached_df.empty: + raise ValueError("cached_df is completely empty!") + + # # Fast path: return cached slice if available + # if trajectory_id in self._traj_cache: + # return self._traj_cache[trajectory_id] + + available_episodes = self.cached_df["episode_index"].unique() + if trajectory_id not in available_episodes: + raise ValueError( + f"trajectory_id {trajectory_id} not found in cached_df. " + f"Available episodes: {sorted(available_episodes)}" + ) + + traj_data = self.cached_df.loc[self.cached_df["episode_index"] == trajectory_id] + trajectory_index = self.get_trajectory_index(trajectory_id) + trajectory_length = self.trajectory_lengths[trajectory_index] + assert ( + len(traj_data) == trajectory_length + ), f"Trajectory length mismatch: {len(traj_data)} != {trajectory_length} {self.args} {self.kwargs}" + indices = traj_data["index"].to_numpy() + if len(indices) > 0: + start_index = indices[0] + expected_indices = np.arange(start_index, start_index + len(indices)) + assert np.array_equal( + indices, expected_indices + ), f"[{self}] Index sequence mismatch in trajectory data, {trajectory_id=}" + # Store in cache to avoid repeated filtering on subsequent calls within a batch + # self._traj_cache[trajectory_id] = traj_data + return traj_data + + + +class ShardedLeRobotMixtureDataset(LeRobotMixtureDataset, IterableDataset): + """ + 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: list[tuple[LeRobotSingleDataset, float]], + training: bool, + balance_dataset_weights: bool = True, + balance_trajectory_weights: bool = True, + seed: int = 42, + shard_sampling_rate: float = 0.5, + num_shards_to_sample: int = 2**20, + allow_padding_at_end: bool = False, + ): + """ + Initialize the mixture dataset. + + Args: + data_mixture (list[tuple[ShardedLeRobotSingleDataset, float]]): Datasets and their corresponding weights. + mode (str): If "train", __iter__ will yield different samples every epoch; if "val" or "test", __iter__ will yield 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. + shard_sampling_rate (float): How much data per shard to sample, in a 0-1 scale. + num_shards_to_sample (int): The number of shards to sample. + """ + super().__init__( + data_mixture=data_mixture, + training=training, + balance_dataset_weights=balance_dataset_weights, + balance_trajectory_weights=balance_trajectory_weights, + seed=seed, + allow_padding_at_end=allow_padding_at_end, + ) + # Add type hint + self.datasets: list[ShardedLeRobotSingleDataset] = self.datasets + # Set properties + self.shard_sampling_rate = shard_sampling_rate + self.num_shards_to_sample = num_shards_to_sample + + # Calculate shard sampling weights + all_shard_sampling_weights = [] + all_shards = [] + for dataset_id, (dataset, weight) in enumerate( + zip(self.datasets, self._dataset_sampling_weights) + ): + shard_sampling_weights = dataset.shard_lengths / dataset.shard_lengths.sum() + all_shard_sampling_weights.append(shard_sampling_weights * weight) + all_shards.extend( + [(dataset_id, shard_idx) for shard_idx in range(shard_sampling_weights.shape[0])] + ) + all_shard_sampling_weights = np.concatenate(all_shard_sampling_weights) + all_shard_sampling_weights /= all_shard_sampling_weights.sum() + self._shard_sampling_weights = all_shard_sampling_weights + self._all_shards = all_shards + + # Generate shards sample schedule for all ranks and workers + self._shards_sample_schedule = self.generate_shards_sample_schedule() + + # Check shard sampling rate + assert 0 <= shard_sampling_rate <= 1, "Shard sampling rate must be between 0 and 1" + + # Set properties for distributed training + if dist.is_initialized(): + self.rank = dist.get_rank() + self.world_size = dist.get_world_size() + else: + self.rank = 0 + self.world_size = 1 + self.worker_id = None + self.num_workers = None + + @property + def dataset_sampling_weights(self) -> np.ndarray: + """The dataset sampling weights.""" + return self._dataset_sampling_weights + + @property + def shard_sampling_weights(self) -> list[np.ndarray]: + """The weights of each shard.""" + return self._shard_sampling_weights + + @property + def all_shards(self) -> list[tuple[int, int]]: + """The shards to sample.""" + return self._all_shards + + @property + def shards_sample_schedule(self) -> list[tuple[int, int]]: + """The shards sample schedule. + + Returns: + list[tuple[int, int]]: The shards to sample, in (dataset_index, shard_index). + """ + assert self._shards_sample_schedule is not None, "Shards sample schedule not set." + return self._shards_sample_schedule + + @property + def trajectory_sampling_weights(self): + """The trajectory sampling weights.""" + raise ValueError("ShardedRobotMixtureDataset does not support trajectory sampling weights.") + + @property + def primary_dataset_indices(self): + """The primary dataset indices.""" + raise ValueError("ShardedRobotMixtureDataset does not support primary dataset indices.") + + def reset_seed(self, seed: int): + self.seed = seed + self._shards_sample_schedule = self.generate_shards_sample_schedule() + + def generate_shards_sample_schedule(self): + if self.training: + rng = np.random.default_rng(self.seed) + sampled_shard_ids = rng.choice( + len(self.all_shards), size=self.num_shards_to_sample, p=self.shard_sampling_weights + ) + shards_sample_schedule = [self.all_shards[i] for i in sampled_shard_ids] + rng.shuffle(shards_sample_schedule) + else: + shards_sample_schedule = [ + self.all_shards[i % len(self.all_shards)] for i in range(self.num_shards_to_sample) + ] + return shards_sample_schedule + + def filter_shards_sample_schedule(self): + """Filter the shards sample schedule for each worker. + + Returns: + list[tuple[int, int]]: The shards to sample, in (dataset_index, shard_index). + """ + # Filter shards for each worker + filtered_schedule = [] + worker_info = get_worker_info() + # If we have multiple workers, further split shards among them + if worker_info is not None: + worker_id = worker_info.id + num_workers = worker_info.num_workers + else: + worker_id = 0 + num_workers = 1 + + if self.worker_id is None: + assert self.num_workers is None + self.worker_id = worker_id + self.num_workers = num_workers + else: + assert ( + self.worker_id == worker_id and self.num_workers == num_workers + ), "Worker ID or number of workers has been changed since it was set. This is not allowed." + + for i, shard in enumerate(self.shards_sample_schedule): + if i % (self.world_size * num_workers) == self.rank * num_workers + worker_id: + filtered_schedule.append(shard) + # print(f"Filtered shards for rank {self.rank}, worker {worker_id}: {filtered_schedule}") + return filtered_schedule + + def __str__(self) -> str: + dataset_descriptions = [] + for dataset, weight in zip(self.datasets, self.dataset_sampling_weights): + shard_lengths = dataset.shard_lengths + assert len(shard_lengths.shape) == 1, "Shard lengths must be a 1D array" + num_shards = shard_lengths.shape[0] + max_shard_length = int(shard_lengths.max()) + min_shard_length = int(shard_lengths.min()) + dataset_description = { + "Dataset": str(dataset), + "Sampling weight": float(weight), + "Num shards": num_shards, + "Max shard length": max_shard_length, + "Min shard length": min_shard_length, + } + dataset_descriptions.append(dataset_description) + return yaml.dump( + { + "Mixture dataset": dataset_descriptions, + "Rank": self.rank, + "World size": self.world_size, + } + ) + + def __iter__(self): + """Iterate over the dataset.""" + + # Not supported: balance_trajectory_weights=False + if not self.balance_trajectory_weights: + raise NotImplementedError( + "balance_trajectory_weights=False is not supported. Please use balance_dataset_weights=True instead." + ) + + self._shards_sample_schedule = self.filter_shards_sample_schedule() + self.curr_shard_index = -1 + self.cache_next_shard() + rng = np.random.default_rng(self.seed) + for i, (dataset_index, shard_index) in enumerate(self.shards_sample_schedule): + self.curr_shard_index += 1 + assert ( + i == self.curr_shard_index + ), f"Shard index mismatch: {i} != {self.curr_shard_index}" + dataset = self.datasets[dataset_index] + wait_start = time.time() + dataset.finish_cache_shard() + wait_end = time.time() + print( + f"Rank {self.rank}, Worker {self.worker_id}: Wait for shard {shard_index} in dataset {dataset_index} in {wait_end - wait_start:.2f} seconds" + ) + # Start caching the next shard immediately + self.cache_next_shard() + all_steps: list[tuple[int, int]] = [] + for trajectory_id in dataset.get_trajectories_in_shard(): + trajectory_index = dataset.get_trajectory_index(trajectory_id) + if self.allow_padding_at_end: + allowed_length = dataset.trajectory_lengths[trajectory_index] + else: + max_delta_index = dataset.max_delta_index + trajectory_length = dataset.trajectory_lengths[trajectory_index] + allowed_length = trajectory_length - max_delta_index + # Get the allowed indices from the step filter + allowed_indices = dataset.step_filter[trajectory_id] + # Remove indices that are too large + allowed_indices = allowed_indices[allowed_indices <= allowed_length] + for i in allowed_indices: + all_steps.append((trajectory_id, i)) + if self.training: + rng.shuffle(all_steps) + sampled_steps = all_steps[: int(dataset.num_steps_per_shard * self.shard_sampling_rate)] + for trajectory_id, step_index in sampled_steps: + # print( + # f"Loading step data from rank {self.rank}, worker {self.worker_id}: {dataset_index} {trajectory_id}, {step_index}" + # ) + indices = { + key: delta_indices + step_index + for key, delta_indices in dataset.delta_indices.items() + } + step_data = dataset.get_step_data(trajectory_id, indices) + # Skip samples where state or action would be empty + if step_data is not None: + yield dataset.transforms(step_data) + + # Delete the cached shard and shard start indices to free up memory + dataset.delete_cached_shard() + + def cache_next_shard(self): + """Cache the next shard in a background thread.""" + next_dataset_idx, next_shard_idx = self.shards_sample_schedule[self.curr_shard_index + 1] + self.datasets[next_dataset_idx].start_cache_shard(next_shard_idx) + + def __getitem__(self, index: int) -> dict: + raise NotImplementedError( + "__getitem__ is not supported for CachedRobotMixtureDataset. Please use __iter__ instead." + ) + + def __len__(self) -> int: + """The length of the dataset.""" + total_length = 0 + for dataset_idx, _ in self.shards_sample_schedule: + dataset = self.datasets[dataset_idx] + total_length += int(dataset.num_steps_per_shard * self.shard_sampling_rate) + return total_length diff --git a/groot/vla/data/dataset/macro.py b/groot/vla/data/dataset/macro.py new file mode 100644 index 0000000000000000000000000000000000000000..94077443e7660765e4a56a60ba88280ebcd7e830 --- /dev/null +++ b/groot/vla/data/dataset/macro.py @@ -0,0 +1,16 @@ +# Deprecated +METADATA_FILENAME = "metadata.json" +DATA_FILENAME = "dataset.hdf5" +FULL_SET_NAME = "FullSet" + +# LeRobot +LE_ROBOT_METADATA_FILENAME = "metadata.json" +LE_ROBOT_MODALITY_FILENAME = "modality.json" +LE_ROBOT_FEATURES_FILENAME = "features.json" +LE_ROBOT_STATISTICS_FILENAME = "stats.json" +LE_ROBOT_EMBODIMENT_FILENAME = "embodiment.json" +# LeRobot extras +LE_ROBOT_EPISODE_FILENAME = "episodes.jsonl" +LE_ROBOT_TASKS_FILENAME = "tasks.jsonl" +LE_ROBOT_INFO_FILENAME = "info.json" +LE_ROBOT_METADATA_DIR = "meta" diff --git a/groot/vla/data/dataset/metadata.py b/groot/vla/data/dataset/metadata.py new file mode 100644 index 0000000000000000000000000000000000000000..bc7cefef8796b7dcfea8580f0d20259a0db932e0 --- /dev/null +++ b/groot/vla/data/dataset/metadata.py @@ -0,0 +1,261 @@ +# import importlib +# from pathlib import Path + +# import numpy as np +# import pandas as pd +# from tqdm import tqdm + +# import groot.vla.common.utils as U +# from groot.vla.data.schema import ( +# EmbodimentTag, +# LeRobotModalityMetadata, +# LeRobotStateActionMetadata, +# DatasetMetadata, +# ) + +# from .macro import ( +# FULL_SET_NAME, +# LE_ROBOT_EMBODIMENT_FILENAME, +# LE_ROBOT_FEATURES_FILENAME, +# LE_ROBOT_METADATA_FILENAME, +# LE_ROBOT_MODALITY_FILENAME, +# LE_ROBOT_STATISTICS_FILENAME, +# ) +# from .registry import EMBODIMENT_TAGS_TO_DATASET_PATHS + +# METADATA_DIR = Path(importlib.import_module("groot.vla.data").__file__).parent / "metadata" # type: ignore + + +# def calculate_dataset_statistics( +# parquet_paths: list[Path], features: list[str] | None = None +# ) -> 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 +# 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) +# all_low_dim_data = pd.concat(all_low_dim_data_list, axis=0) +# # Compute dataset statistics +# num_steps = len(all_low_dim_data.index) +# dataset_statistics: dict = { +# "num_trajectories": len(all_low_dim_data_list), +# "total_trajectory_length": num_steps, +# } +# if features is None: +# features = list(all_low_dim_data.columns) +# for le_modality in features: +# print(f"Computing statistics for {le_modality}...") +# np_data = np.vstack( +# [np.asarray(x, dtype=np.float32) for x in all_low_dim_data[le_modality]] # type: ignore +# ) +# 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 get_metadata( +# embodiment_tag: EmbodimentTag, +# metadata_version: str, +# regenerate_stats: bool = False, +# regenerate_metadata: bool = False, +# ) -> DatasetMetadata: +# """Get the metadata corresponding to the given embodiment tag and metadata version.""" +# metadata_dir = METADATA_DIR / embodiment_tag.value / metadata_version +# metadata_path = metadata_dir / LE_ROBOT_METADATA_FILENAME +# if metadata_path.exists() and not regenerate_metadata: +# metadata = DatasetMetadata.model_validate_json(metadata_path.read_text()) +# return metadata + + +# def get_metadata( +# embodiment_tag: EmbodimentTag, +# metadata_version: str, +# regenerate_stats: bool = False, +# regenerate_metadata: bool = False, +# ) -> TrainableDatasetMetadata_V1_2: +# """Get the metadata corresponding to the given embodiment tag and metadata version. + +# Args: +# embodiment_tag: The embodiment tag to load the metadata for. +# metadata_version: The version of the metadata to load. +# generate_metadata: Whether to generate the metadata if it does not exist. +# """ +# metadata_dir = METADATA_DIR / embodiment_tag.value / metadata_version +# metadata_path = metadata_dir / LE_ROBOT_METADATA_FILENAME +# if metadata_path.exists() and not regenerate_metadata: +# metadata = TrainableDatasetMetadata_V1_2.model_validate_json(metadata_path.read_text()) +# return metadata + +# assert ( +# embodiment_tag in EMBODIMENT_TAGS_TO_DATASET_PATHS +# ), f"Embodiment tag {embodiment_tag} not found in dataset registry. Available tags: {EMBODIMENT_TAGS_TO_DATASET_PATHS.keys()}" + +# dataset_paths = EMBODIMENT_TAGS_TO_DATASET_PATHS[embodiment_tag] +# # Load supporting metadata +# le_modality_meta_path = metadata_dir / LE_ROBOT_MODALITY_FILENAME +# le_features_path = metadata_dir / LE_ROBOT_FEATURES_FILENAME +# embodiment_meta_path = metadata_dir / LE_ROBOT_EMBODIMENT_FILENAME +# le_modality_meta = LeRobotModalityMetadata.model_validate_json( +# le_modality_meta_path.read_text() +# ) +# le_features = U.load_json(le_features_path) +# embodiment_meta = U.load_json(embodiment_meta_path) +# # Load stats +# if regenerate_stats: +# le_statistics = None +# else: +# le_statistics_path = metadata_dir / LE_ROBOT_STATISTICS_FILENAME +# le_statistics = U.load_json(le_statistics_path) +# # Generate metadata +# metadata, le_statistics = generate_metadata( +# embodiment_tag=embodiment_tag, +# dataset_paths=dataset_paths, +# le_modality_meta=le_modality_meta, +# le_features=le_features, +# embodiment_meta=embodiment_meta, +# le_statistics=le_statistics, +# ) + +# # Save metadata +# print(f"Generated metadata at {metadata_path}") +# metadata_path.write_text(metadata.model_dump_json(indent=4)) +# # Save stats +# if regenerate_stats: +# le_statistics_path = metadata_dir / LE_ROBOT_STATISTICS_FILENAME +# U.dump_json(le_statistics, le_statistics_path, indent=4) + +# return metadata + + +# def generate_metadata( +# embodiment_tag: EmbodimentTag, +# dataset_paths: list[Path], +# le_modality_meta: LeRobotModalityMetadata, +# le_features: dict, +# embodiment_meta: dict, +# le_statistics: dict | None = None, +# ): +# dataset_name = f"{embodiment_tag.value}:{FULL_SET_NAME}" + +# # Generate our custom modality metadata +# our_modality_meta: dict[str, dict] = {} +# for modality in ["state", "action"]: +# our_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 +# our_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, +# } + +# # Add video modalities +# our_modality_meta["video"] = {} +# for new_key in le_modality_meta.video: +# original_key = le_modality_meta.video[new_key].original_key +# le_video_meta = le_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")] +# channels = le_video_meta["shape"][le_video_meta["names"].index("channel")] +# if "info" in le_video_meta: +# fps = le_video_meta["info"]["video.fps"] +# elif "video_info" in le_video_meta: +# fps = le_video_meta["video_info"]["video.fps"] +# else: +# raise ValueError( +# f"Video modality {new_key} does not contain video_info or info: {le_video_meta.keys()}" +# ) +# our_modality_meta["video"][new_key] = { +# "resolution": [width, height], +# "channels": channels, +# "fps": fps, +# } + +# # Add annotation metadata +# our_modality_meta["annotation"] = {} +# if le_modality_meta.annotation is not None: +# for annotation_key in le_modality_meta.annotation: +# key_split = annotation_key.split(".") +# annotation_source = key_split[0] +# annotation_type = ".".join(key_split[1:]) +# if annotation_source not in our_modality_meta["annotation"]: +# our_modality_meta["annotation"][annotation_source] = [] +# our_modality_meta["annotation"][annotation_source].append(annotation_type) + +# lowdim_features = [] +# for feature in le_features: +# if "float" in le_features[feature]["dtype"]: +# lowdim_features.append(feature) + +# # Dataset statistics +# if le_statistics is None: +# print(f"Calculating dataset statistics for {dataset_name}") +# # Get all parquet files in the dataset paths +# parquet_files = [] +# for dataset_path in dataset_paths: +# parquet_files.extend(list(dataset_path.glob("data/*/*.parquet"))) +# le_statistics = calculate_dataset_statistics(parquet_files, lowdim_features) +# for le_modality in le_statistics: +# if not isinstance(le_statistics[le_modality], dict): +# continue +# for stat in le_statistics[le_modality]: +# le_statistics[le_modality][stat] = np.asarray(le_statistics[le_modality][stat]) + +# # Split statistics keys to our format +# dataset_statistics = { +# "num_trajectories": le_statistics["num_trajectories"], +# "total_trajectory_length": le_statistics["total_trajectory_length"], +# } +# for our_modality in ["state", "action"]: +# dataset_statistics[our_modality] = {} +# for subkey in our_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 in le_statistics[le_modality]: +# indices = np.arange( +# state_action_meta.start, +# state_action_meta.end, +# ) +# dataset_statistics[our_modality][subkey][stat] = le_statistics[le_modality][stat][ +# indices +# ].tolist() + +# # Full dataset metadata +# metadata = TrainableDatasetMetadata_V1_2( +# dataset_name=dataset_name, +# dataset_statistics=dataset_statistics, # type: ignore +# modalities=our_modality_meta, # type: ignore +# embodiment=embodiment_meta, # type: ignore +# ) + +# # Convert stats from numpy to list +# for le_modality in le_statistics: +# if not isinstance(le_statistics[le_modality], dict): +# continue +# for stat in le_statistics[le_modality]: +# le_statistics[le_modality][stat] = le_statistics[le_modality][stat].tolist() + +# return metadata, le_statistics diff --git a/groot/vla/data/dataset/registry.py b/groot/vla/data/dataset/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..47bf672174db5bae46088bdef27fa408628ecdcc --- /dev/null +++ b/groot/vla/data/dataset/registry.py @@ -0,0 +1,14 @@ +from pathlib import Path + +from groot.vla.data.schema import EmbodimentTag + +# Dataset registry: maps embodiment tags to local dataset paths. +# Update these paths to point to your local dataset directories. +# See docs/DATASET_TO_GEAR_AND_TRAIN.md for dataset conversion instructions. +EMBODIMENT_TAGS_TO_DATASET_PATHS: dict[EmbodimentTag, list[Path]] = {} + +DATASET_PATHS_TO_EMBODIMENT_TAGS = { + path: dataset_tag + for dataset_tag, dataset_paths in EMBODIMENT_TAGS_TO_DATASET_PATHS.items() + for path in dataset_paths +} diff --git a/groot/vla/data/schema/__init__.py b/groot/vla/data/schema/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..786e3cb513a829954ce7709786f37a5393e722a2 --- /dev/null +++ b/groot/vla/data/schema/__init__.py @@ -0,0 +1,15 @@ +from .embodiment_tags import EmbodimentTag +from .lerobot import ( + DatasetMetadata, + DatasetModalities, + DatasetStatisticalValues, + DatasetStatistics, + LeRobotActionMetadata, + LeRobotModalityField, + LeRobotModalityMetadata, + LeRobotStateActionMetadata, + LeRobotStateMetadata, + RotationType, + StateActionMetadata, + VideoMetadata, +) diff --git a/groot/vla/data/schema/embodiment_tags.py b/groot/vla/data/schema/embodiment_tags.py new file mode 100644 index 0000000000000000000000000000000000000000..1b7e94683dfacc24554a243808d7dbb12c0804ed --- /dev/null +++ b/groot/vla/data/schema/embodiment_tags.py @@ -0,0 +1,353 @@ +from enum import Enum + + +class EmbodimentTag(Enum): + REAL_GR1_ARMS_ONLY = "real_gr1_arms_only" + """ + The real GR1 robot embodiment with arms only. + """ + + REAL_GR1_ARMS_ONLY_ANNOTATED = "real_gr1_arms_only_annotated" + """ + The real GR1 robot embodiment with arms only with annotations. + """ + + REAL_GR1_ARMS_WAIST = "real_gr1_arms_waist" + """ + The real GR1 robot embodiment with arms and waist. + """ + + REAL_GR1_ARMS_WAIST_ANNOTATED = "real_gr1_arms_waist_annotated" + """ + The real GR1 robot embodiment with arms and waist with annotations. + """ + + DEXMG_GR1_ARMS_ONLY_INSPIRE = "dexmg_gr1_arms_only_inspire" + """ + The DEXMG GR1 dataset with arms only and inspire hand. + """ + + DEXMG_GR1_ARMS_ONLY_FOURIER = "dexmg_gr1_arms_only_fourier" + """ + The DEXMG GR1 dataset with arms only and Fourier hand. + """ + + DEXMG_GR1_ARMS_WAIST_FOURIER = "dexmg_gr1_arms_waist_fourier" + """ + The DEXMG GR1 dataset with arms and waist and Fourier hand. + """ + + ROBOCASA_SINGLE_ARM = "robocasa_single_arm" + """ + The RoboCasa dataset with single arm. + """ + + ONE_X_EVE_GRIPPER = "onex_eve_gripper" + """ + The OneX Eve Robot with gripper. + """ + + ROBOCASA_GR1_ARMS_ONLY_INSPIRE_HANDS = "robocasa_gr1_arms_only_inspire_hands" + """ + The RoboCasa GR1 dataset with arms only and inspire hands. + """ + + ROBOCASA_GR1_ARMS_ONLY_FOURIER_HANDS = "robocasa_gr1_arms_only_fourier_hands" + """ + The RoboCasa GR1 dataset with arms only and Fourier hands. + """ + + ROBOCASA_GR1_FIXED_LOWER_BODY_INSPIRE_HANDS = "robocasa_gr1_fixed_lower_body_inspire_hands" + """ + The RoboCasa GR1 dataset with fixed lower body and inspire hands. + """ + + ROBOCASA_GR1_FIXED_LOWER_BODY_FOURIER_HANDS = "robocasa_gr1_fixed_lower_body_fourier_hands" + """ + The RoboCasa GR1 dataset with fixed lower body and Fourier hands. + """ + + ROBOCASA_GR1_ARMS_WAIST_FOURIER_HANDS = "robocasa_gr1_arms_waist_fourier_hands" + """ + The RoboCasa GR1 dataset with arms and waist and Fourier hands. + """ + + ROBOCASA_PANDA_OMRON = "robocasa_panda_omron" + """ + The RoboCasa dataset with panda omron. + """ + + ROBOCASA_SINGLE_ARM_PANDA_OMRON = "robocasa_single_arm_panda_omron" + """ + The RoboCasa dataset with single arm panda omron. + """ + + ROBOCASA_BIMANUAL_PANDA_PARALLEL_GRIPPER = "robocasa_bimanual_panda_parallel_gripper" + """ + The dexmg bimanual panda dataset with parallel grippers. + """ + + ROBOCASA_BIMANUAL_PANDA_INSPIRE_HAND = "robocasa_bimanual_panda_inspire_hand" + """ + The DEXMG bimanual panda dataset with inspire hands. + """ + + OXE_DROID = "oxe_droid" + """ + The Open X-Embodiment droid dataset. + """ + + OXE_FRACTAL = "oxe_fractal" + """ + The Open X-Embodiment fractal (RT-1) dataset. + """ + + OXE_LANGUAGE_TABLE = "oxe_language_table" + """ + The Open X-Embodiment language table dataset. + """ + + OXE_BRIDGE = "oxe_bridge" + """ + The Open X-Embodiment bridge dataset. + """ + + OXE_MUTEX = "oxe_mutex" + """ + The Open X-Embodiment mutex dataset. + """ + + OXE_ROBOSET = "oxe_roboset" + """ + The Open X-Embodiment Roboset dataset. + """ + + OXE_PLEX = "oxe_plex" + """ + The Open X-Embodiment Plex RoboSuite dataset. + """ + + REAL_PANDA_SINGLE_ARM = "real_panda_single_arm" + """ + The real single arm panda robot. + """ + + HOT3D_HANDS_ONLY = "hot3d_hands_only" + """ + The HOT3D dataset with hands only. + """ + + GR1_UNIFIED = "gr1_unified" + """ + The GR1 unified dataset. + """ + + GR1_ISAAC = "gr1_isaac" + """ + The GR1 Isaac dataset (Shiwei). + """ + + LAPA = "lapa" + """ + The datasets with LAPA actions. + """ + AGIBOT = "agibot" + + YAM = "yam" + + DREAM = "dream" + """ + The datasets with DREAM actions. + """ + + UNKNOWN = "unknown" + + GR1_UNIFIED_SEGMENTATION = "gr1_unified_segmentation" + """ + The GR1 unified dataset with segmentation. + """ + + LANGUAGE_TABLE_SIM = "language_table_sim" + """ + Simulated Language Table. + """ + + SIMPLER_ENV_GOOGLE = "simpler_env_google" + """ + SimplerEnv Google. + """ + + SIMPLER_ENV_WIDOWX = "simpler_env_widowx" + """ + SimplerEnv Widowx. + """ + + LIBERO_SIM = "libero_sim" + """ + The Libero Sim dataset. + """ + + DROID_SIM = "droid_sim" + """ + The Droid dataset in sim. + """ + + UNITREE_G1_UPPER_BODY = "unitree_g1_upper_body" + """ + The Unitree G1 dataset. + """ + + UNITREE_G1_FULL_BODY = "unitree_g1_full_body" + """ + The Unitree G1 dataset with full body. + """ + + UNITREE_G1_FULL_BODY_IN_SIM = "unitree_g1_full_body_in_sim" + """ + The Unitree G1 dataset in sim, data is collected with 50Hz. + """ + + UNITREE_G1_FULL_BODY_WITH_HEIGHT = "unitree_g1_full_body_with_height" + """ + The Unitree G1 dataset with full body and height command. + """ + + UNITREE_G1_FULL_BODY_WITH_HEIGHT_AND_EEF = "unitree_g1_full_body_with_height_and_eef" + """ + The Unitree G1 dataset with full body and height command and eef command. + """ + + UNITREE_G1_FULL_BODY_WITH_HEIGHT_EEF_NAV_CMD = "unitree_g1_full_body_with_height_eef_nav_cmd" + """ + The Unitree G1 dataset with full body and height command and eef command and navigate command. + """ + + UNITREE_G1_FULL_BODY_WITH_HEIGHT_NAV_CMD = "unitree_g1_full_body_with_height_nav_cmd" + """ + The Unitree G1 dataset with full body and height command and navigate command. + """ + + UNITREE_G1_FULL_BODY_WITH_WAIST_HEIGHT_NAV_CMD = ( + "unitree_g1_full_body_with_waist_height_nav_cmd" + ) + """ + The Unitree G1 dataset with full body and waist and height command and navigate command. + """ + + UNITREE_G1_FULL_BODY_WITH_WAIST_HEIGHT_NAV_CMD_AND_TASK_PROG = ( + "unitree_g1_full_body_with_waist_height_nav_cmd_and_task_prog" + ) + """ + The Unitree G1 dataset with full body and waist and height command and navigate command and task progress. + """ + + UNITREE_G1_FULL_BODY_WITH_HEIGHT_NAV_CMD_IN_SIM = ( + "unitree_g1_full_body_with_height_nav_cmd_in_sim" + ) + """ + The Unitree G1 dataset with full body and height command and navigate command in sim, data is collected with 50Hz. + """ + + G1_FIX_LOWER_RIGHT_HAND = "g1_fix_lower_right_hand" + """ + The G1 robot with fixed lower body and right hand. + """ + + GR1_UNIFIED_OFFLINE_RL = "gr1_unified_offline_rl" + """ + The GR1 robot with reward information for oflfine RL. + """ + + UNITREE_G1_FULL_BODY_WITH_HEIGHT_NAV_CMD_OAK_STEREO = ( + "unitree_g1_full_body_with_height_nav_cmd_oak_stereo" + ) + """ + The Unitree G1 dataset with full body and height command and navigate command and oak stereo. + """ + + XDOF = "xdof" + """ + The XDOF robot. + """ + + XDOF_H16 = "xdof_h16" + """ + The XDOF robot with action horizon 16. + """ + + XDOF_OSS_DATA = "xdof_oss_data" + """ + The XDOF data with conversions via jimmywu's processing pipeline. + """ + + SIM_BEHAVIOR_R1_PRO = "sim_behavior_r1_pro" + """ + The sim BEHAVIOR Galaxea R1 Pro robot with grippers. + """ + + # ------------- Deprecated G1 Embodiments ------------- + # below are deprecated G1 embodiments. Why deprecated? + # 1. The neck wasn't locked. So the the camera angles differe slightly between episodes and robots + # 2. Video modality was called `rs_view`. It is now called `ego_view` + # TODO: In future if we want to unify these embodiment tags, we can do so by renaming + # video.rs_view -> video.ego_view. After that we can remove the below deprecated embodiment tags. + DEPRECATED_UNITREE_G1_UPPER_BODY = "deprecated_unitree_g1_upper_body" + """ + The Unitree G1 dataset. + """ + + DEPRECATED_UNITREE_G1_FULL_BODY = "deprecated_unitree_g1_full_body" + """ + The Unitree G1 dataset with full body. + """ + + DEPRECATED_UNITREE_G1_FULL_BODY_WITH_HEIGHT = "deprecated_unitree_g1_full_body_with_height" + """ + The Unitree G1 dataset with full body and height command. + """ + + DEPRECATED_UNITREE_G1_FULL_BODY_WITH_HEIGHT_AND_EEF = ( + "deprecated_unitree_g1_full_body_with_height_and_eef" + ) + """ + The Unitree G1 dataset with full body and height command and eef command. + """ + + DEPRECATED_UNITREE_G1_FULL_BODY_WITH_HEIGHT_EEF_NAV_CMD = ( + "deprecated_unitree_g1_full_body_with_height_eef_nav_cmd" + ) + """ + The Unitree G1 dataset with full body and height command and eef command and navigate command. + """ + + DEPRECATED_UNITREE_G1_FULL_BODY_WITH_HEIGHT_NAV_CMD = ( + "deprecated_unitree_g1_full_body_with_height_nav_cmd" + ) + """ + The Unitree G1 dataset with full body and height command and navigate command. + """ + + GR1_UNIFIED_512 = "gr1_unified_512" + """ + The GR1 unified dataset with 512*320 resolution without cropping. + """ + + R1_PRO = "r1_pro" + """ + Real Galaxea R1 Pro. + """ + + R1_PRO_SINGLE_VIEW = "r1_pro_single-view" + """ + Real Galaxea R1 Pro with single camera view. + """ + + MECKA_HANDS = "mecka_hands" + """ + The Mecka robot with hands. + """ + + TREX = "trex" + """ + The T-Rex dataset: Dexmate Vega-1 dual-arm with two Sharpa Wave dexterous hands. + """ \ No newline at end of file diff --git a/groot/vla/data/schema/lerobot.py b/groot/vla/data/schema/lerobot.py new file mode 100644 index 0000000000000000000000000000000000000000..4139f5697e79a00edf5e845bdd9720677fbed553 --- /dev/null +++ b/groot/vla/data/schema/lerobot.py @@ -0,0 +1,271 @@ +# 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 + +import numpy as np +from pydantic import BaseModel, Field, field_serializer, field_validator, model_validator + +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: str = Field( + ..., + description="The original key of the modality in the LeRobot dataset.", + ) + + +class LeRobotStateMetadata(LeRobotStateActionMetadata): + """Metadata for a LeRobot state modality.""" + + original_key: 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: 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.", + ) + + @model_validator(mode="after") + def check_original_keys(self): + # Generate default original keys for state + for key in self.state: + if self.state[key].original_key is None: + self.state[key].original_key = "observation.state" + + # Generate default original keys for action + for key in self.action: + if self.action[key].original_key is None: + self.action[key].original_key = "action" + + # Generate default original keys for video + for key in self.video: + if self.video[key].original_key is None: + self.video[key].original_key = "observation.images." + key + + # Generate default original keys for annotation + if self.annotation is not None: + for key in self.annotation: + if self.annotation[key].original_key is None: + self.annotation[key].original_key = "annotation." + key + + return self + + 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}") + + def get_original_key(self, key: str) -> str: + """Get the original key for a key in the LeRobot modality metadata. + + Args: + key (str): The key to get the original key for. + + Returns: + str: The original key for the key. + + Example: + lerobot_modality_meta = LeRobotModalityMetadata.model_validate(U.load_json(modality_meta_path)) + lerobot_modality_meta.get_original_key("state.joint_shoulder_y") + lerobot_modality_meta.get_original_key("video.main_camera") + lerobot_modality_meta.get_original_key("annotation.human.action.task_description") + """ + original_key = self.get_key_meta(key).original_key + assert original_key is not None, f"Key: {key}, original key is not set" + return original_key + + +# Dataset schema (parsed from LeRobot schema and simplified) + + +class DatasetStatisticalValues(BaseModel): + max: np.ndarray = Field(..., description="Maximum values") + min: np.ndarray = Field(..., description="Minimum values") + mean: np.ndarray = Field(..., description="Mean values") + std: np.ndarray = Field(..., description="Standard deviation") + q01: np.ndarray = Field(..., description="1st percentile values") + q99: np.ndarray = Field(..., description="99th percentile values") + + model_config = {"arbitrary_types_allowed": True} + + @field_serializer("*", when_used="json") + def serialize_ndarray(self, v: np.ndarray) -> list: + return v.tolist() # type: ignore + + @field_validator("*", mode="before") + def validate_ndarray(cls, v) -> np.ndarray: + return np.array(v) + + +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") diff --git a/groot/vla/data/transform/__init__.py b/groot/vla/data/transform/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..bfa7eeeac139423679a12220cbde4db6161f0b38 --- /dev/null +++ b/groot/vla/data/transform/__init__.py @@ -0,0 +1,29 @@ +from .base import ( + ComposedModalityTransform, + InvertibleModalityTransform, + ModalityTransform, +) +from .concat import ConcatTransform +from .language import LanguageRemovePrefix, LanguageTransform +from .state_action import ( + PerHorizonActionTransform, + StateActionDropout, + StateActionPerturbation, + StateActionSinCosTransform, + StateActionToTensor, + StateActionTransform, +) +from .video import ( + VideoColorJitter, + VideoCrop, + VideoGrayscale, + VideoHorizontalFlip, + VideoRandomGrayscale, + VideoRandomPosterize, + VideoRandomRotation, + VideoResize, + VideoToNumpy, + VideoToTensor, + VideoTransform, + VideoNormalize, +) diff --git a/groot/vla/data/transform/base.py b/groot/vla/data/transform/base.py new file mode 100644 index 0000000000000000000000000000000000000000..228860049653dd3d1b0cb57150d8e561b7d4c935 --- /dev/null +++ b/groot/vla/data/transform/base.py @@ -0,0 +1,156 @@ +from abc import ABC, abstractmethod +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field, PrivateAttr + +from groot.vla.data.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.""" + pass + + 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.""" + pass + + +class IdentityModalityTransform(ModalityTransform): + """Identity transform.""" + + apply_to: list[str] = Field( + default_factory=list, description="Will be ignored for identity transforms." + ) + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + return data + + def unapply(self, data: dict[str, Any]) -> dict[str, Any]: + return 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) + # this is used to pass the list of transforms to concat transform + # concat transform needs needs to know what transforms were applied + # because it needs to compute the correct dimension of features + # post transform (during unapply). + # this attribute can also be used by other transforms to know what + # transforms were applied before it in the pipeline. + if hasattr(transform, "set_transform_pipeline"): + getattr(transform, "set_transform_pipeline")(self.transforms) + + def set_per_horizon_statistics(self, per_horizon_stats: dict[str, dict[str, list]]): + """Set per-horizon statistics for transforms that support it (e.g., PerHorizonActionTransform). + + Args: + per_horizon_stats: Dict from dataset.lerobot_relative_horizon_stats_meta + Format: {action_key: {stat_name: [[h0_vals], [h1_vals], ...]}} + """ + for transform in self.transforms: + if hasattr(transform, "set_per_horizon_statistics"): + transform.set_per_horizon_statistics(per_horizon_stats) + + 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() + self.training = True + + def eval(self): + for transform in self.transforms: + transform.eval() + self.training = False diff --git a/groot/vla/data/transform/concat.py b/groot/vla/data/transform/concat.py new file mode 100644 index 0000000000000000000000000000000000000000..e5474baffd56679fa2098e46abe4300fcff341a7 --- /dev/null +++ b/groot/vla/data/transform/concat.py @@ -0,0 +1,259 @@ +from typing import Any, Dict, List, Optional + +import numpy as np +from pydantic import Field, PrivateAttr +import torch + +from groot.vla.data.schema import DatasetMetadata, StateActionMetadata +from groot.vla.data.transform.base import InvertibleModalityTransform, ModalityTransform + + +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.", + ) + + action_dims_post_transform: dict[str, int] = Field( + default_factory=dict, + description="The new dimensions of the action keys after transform is applied.", + ) + state_dims_post_transform: dict[str, int] = Field( + default_factory=dict, + description="The new dimensions of the state keys after transform is applied.", + ) + # Store the transform pipeline to examine for dimension changes + _transform_pipeline: List[ModalityTransform] = PrivateAttr(default_factory=list) + + 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 set_transform_pipeline(self, transforms: List[ModalityTransform]): + """Set the transform pipeline so this transform can examine it for dimension changes.""" + self._transform_pipeline = transforms + + def _get_target_rotations_from_pipeline(self) -> Dict[str, str]: + """Extract target_rotations from StateActionTransform instances in the pipeline.""" + target_rotations = {} + for transform in self._transform_pipeline: + if hasattr(transform, "target_rotations"): + transform_target_rotations = getattr(transform, "target_rotations", {}) + if transform_target_rotations: + target_rotations.update(transform_target_rotations) + return target_rotations + + def apply(self, data: Dict[str, Any]) -> Dict[str, Any]: + 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.extend( + [3, 4, 6] + ) # 3 -> axis_angle, 4 -> quaternion, 6 -> rotation_6d + 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.extend( + [3, 4, 6] + ) # 3 -> axis_angle, 4 -> quaternion, 6 -> rotation_6d + assert ( + data[key].shape[-1] in target_shapes + ), f"Action dim mismatch for {key=}, {data[key].shape[-1]=}, {target_shapes=}" + # 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.get_state_action_dims_post_transform(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.get_state_action_dims_post_transform(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 get_state_action_dims_post_transform(self, key: str) -> int: + """ + This function is used to get the dims of the state/action keys after transform is applied. + It is different from the `get_state_action_dims` function, because this function accounts for + the case where we apply transforms and the # of dims is change eg. after applying axis_angle transform on + quaternion, the dims change from 4D to 3D. + """ + modality_config = self.get_modality_metadata(key) + shape = modality_config.shape + assert len(shape) == 1, f"{shape=}" + + if self.is_rotation_key(key): + target_rotations = self._get_target_rotations_from_pipeline() + if key in target_rotations: + target_rotation = target_rotations[key] + if target_rotation == "axis_angle": + return 3 + elif target_rotation == "quaternion": + return 4 + elif target_rotation == "rotation_6d": + return 6 + elif target_rotation == "euler_angles": + return 3 + else: + raise ValueError(f"Unknown target rotation type: {target_rotation}") + else: + # No target rotation specified, return original dimension + return shape[0] + else: + 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/groot/vla/data/transform/language.py b/groot/vla/data/transform/language.py new file mode 100644 index 0000000000000000000000000000000000000000..6fbf4d5c6679f6ffacadc0220c3a94418fd23591 --- /dev/null +++ b/groot/vla/data/transform/language.py @@ -0,0 +1,86 @@ +from typing import Any, Union + +import numpy as np +from pydantic import Field, field_validator +from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast + +from groot.vla.data.transform.base import InvertibleModalityTransform, ModalityTransform + +T_Tokenizer = Union[PreTrainedTokenizer, PreTrainedTokenizerFast] + + +class LanguageTransform(InvertibleModalityTransform): + """Transform for language modalities. + + Attributes: + apply_to (list[str]): The keys in the modality to load and transform. + tokenizer (T_Tokenizer): The tokenizer to use. Can be either PreTrainedTokenizer, PreTrainedTokenizerFast, or path to the HuggingFace tokenizer + """ + + apply_to: list[str] = Field(..., description="The keys in the modality to load and transform.") + tokenizer: T_Tokenizer = Field(..., description="The tokenizer to use.") + + @field_validator("tokenizer") + def validate_tokenizer(cls, v: T_Tokenizer | str) -> T_Tokenizer: + if isinstance(v, str): + return AutoTokenizer.from_pretrained(v) + return v + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + """Tokenize the data, pad the tokens to the longest sequence, and concatenate the tokens. + + Args: + data (dict[str, Any]): The complete data dictionary. + + Returns: + dict[str, Any]: The processed data dictionary with the keys in `apply_to` replaced with the tokenized data. + """ + for key in self.apply_to: + data[key] = self.tokenizer( + data[key], return_tensors="pt", padding=True, truncation=True + ).input_ids + return data + + def unapply(self, data: dict[str, Any]) -> dict[str, Any]: + """Untokenize the data. + + Args: + data (dict[str, Any]): The processed data dictionary with the keys in `apply_to` replaced with the tokenized data. + + Returns: + dict[str, Any]: The untokenized data. + """ + for key in self.apply_to: + data[key] = self.tokenizer.decode(data[key], skip_special_tokens=True) + return data + + +class LanguageRemovePrefix(ModalityTransform): + apply_to: list[str] = Field( + ..., description="The keys in the modality to remove the prefix from." + ) + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + """Remove the prefix from the language. + Expects: + - data[key] is a list of strings, shape (T,) + - OR data[key] is a list of lists of strings, shape (B, T) + + Args: + data (dict[str, Any]): The processed data dictionary with the keys in `apply_to` replaced with the tokenized data. + + Returns: + dict[str, Any]: The data with the prefix removed for key in `apply_to`. + """ + for key in self.apply_to: + value = data[key] + # Handle both batched (list of lists) and non-batched (list) language data. + if isinstance(value[0], np.ndarray): + # Batched case: list of lists of strings, shape (B, T) + data[key] = np.array( + [[lang.split(": ")[-1] for lang in sublist] for sublist in value] + ) + else: + # Non-batched case: list of strings, shape (T,) + data[key] = np.array([lang.split(": ")[-1] for lang in value]) + return data diff --git a/groot/vla/data/transform/state_action.py b/groot/vla/data/transform/state_action.py new file mode 100644 index 0000000000000000000000000000000000000000..2be1cc810f1dca0e2c9625d7775bc7b62b0c4816 --- /dev/null +++ b/groot/vla/data/transform/state_action.py @@ -0,0 +1,893 @@ +import functools +import random +from typing import Any, ClassVar + +import numpy as np +from pydantic import Field, PrivateAttr, field_validator, model_validator +import torch + +from groot.vla.data.schema import DatasetMetadata, RotationType, StateActionMetadata +from groot.vla.data.transform.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 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": + import pytorch3d.transforms as pt # noqa: F401 + + 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", "scale"] + + def __init__(self, mode: str, statistics: dict): + self.mode = mode + self.statistics = statistics + 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, -1, 1) + + 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 + + if not (torch.all(normalized <= 1 + 1e-5) and torch.all(normalized >= -1 - 1e-5)): + print( + f"WARNING: normalization out of range: {x.shape=}, max: {torch.max(normalized)}, min: {torch.min(normalized)}" + ) + + # Clip the normalized values to be between -1 and 1 + normalized = torch.clamp(normalized, -1, 1) + + 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 > 0.5).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(dtype=x.dtype, device=x.device) + q99 = self.statistics["q99"].to(dtype=x.dtype, device=x.device) + return (x + 1) / 2 * (q99 - q01) + q01 + elif self.mode == "mean_std": + mean = self.statistics["mean"].to(dtype=x.dtype, device=x.device) + std = self.statistics["std"].to(dtype=x.dtype, device=x.device) + return x * std + mean + elif self.mode == "min_max": + min = self.statistics["min"].to(dtype=x.dtype, device=x.device) + max = self.statistics["max"].to(dtype=x.dtype, device=x.device) + return (x + 1) / 2 * (max - min) + min + elif self.mode == "binary": + return (x > 0.5).to(dtype=x.dtype, device=x.device) + else: + raise ValueError(f"Invalid normalization mode: {self.mode}") + + +class PerHorizonNormalizer: + """Normalizer that applies different statistics for each horizon index. + + Unlike the standard Normalizer which uses single stats for all timesteps, + this normalizer uses per-horizon stats where each action horizon step has + its own normalization parameters. + + Args: + mode: Normalization mode ('q99', 'mean_std', 'min_max', 'binary', 'scale') + statistics: Dict with stat names as keys and 2D tensors of shape (horizon_len, action_dim) as values. + Format: {"q01": [[h0_vals], [h1_vals], ...], "q99": [[h0_vals], [h1_vals], ...], ...} + """ + valid_modes = ["q99", "mean_std", "min_max", "binary", "scale"] + + def __init__(self, mode: str, statistics: dict): + self.mode = mode + self.statistics = {} + for key, value in statistics.items(): + # Convert to tensor: shape (horizon_len, action_dim) + self.statistics[key] = torch.tensor(value) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Normalize tensor with per-horizon statistics. + + Args: + x: Input tensor of shape (horizon_len, action_dim) or (num_chunks * horizon_len, action_dim) + or (..., horizon_len, action_dim) + + Returns: + Normalized tensor of same shape + """ + assert isinstance( + x, torch.Tensor + ), f"Unexpected input type: {type(x)}. Expected type: {torch.Tensor}" + + # Get the stats horizon length + stats_horizon_len = self.statistics["q01"].shape[0] + action_dim = x.shape[-1] + original_shape = x.shape + + # Handle case where x is flattened (num_chunks * horizon_len, action_dim) + # We need to reshape to (..., horizon_len, action_dim) for proper broadcasting + reshaped = False + if len(x.shape) == 2: + total_len = x.shape[0] + if total_len > stats_horizon_len and total_len % stats_horizon_len == 0: + # Flattened batch case: reshape to (num_chunks, horizon_len, action_dim) + num_chunks = total_len // stats_horizon_len + x = x.view(num_chunks, stats_horizon_len, action_dim) + reshaped = True + elif total_len == stats_horizon_len: + # Single chunk, add batch dimension for consistent processing + x = x.unsqueeze(0) + reshaped = True + + # Now x should have shape (..., horizon_len, action_dim) + horizon_len = x.shape[-2] + + if self.mode == "q99": + q01 = self.statistics["q01"].to(dtype=x.dtype, device=x.device)[:horizon_len] # (horizon_len, action_dim) + q99 = self.statistics["q99"].to(dtype=x.dtype, device=x.device)[:horizon_len] + + # Compute normalization with safe division + range_vals = q99 - q01 + range_vals = torch.where(range_vals == 0, torch.ones_like(range_vals), range_vals) + normalized = 2 * (x - q01) / range_vals - 1 + + # For zero-range values, keep original + mask = (q01 == q99).unsqueeze(0).expand_as(x) if len(x.shape) > 2 else (q01 == q99) + normalized = torch.where(mask, x, normalized) + normalized = torch.clamp(normalized, -1, 1) + + elif self.mode == "mean_std": + mean = self.statistics["mean"].to(dtype=x.dtype, device=x.device)[:horizon_len] + std = self.statistics["std"].to(dtype=x.dtype, device=x.device)[:horizon_len] + + std_safe = torch.where(std == 0, torch.ones_like(std), std) + normalized = (x - mean) / std_safe + + mask = (std == 0).unsqueeze(0).expand_as(x) if len(x.shape) > 2 else (std == 0) + normalized = torch.where(mask, x, normalized) + + elif self.mode == "min_max": + min_val = self.statistics["min"].to(dtype=x.dtype, device=x.device)[:horizon_len] + max_val = self.statistics["max"].to(dtype=x.dtype, device=x.device)[:horizon_len] + + range_vals = max_val - min_val + range_vals = torch.where(range_vals == 0, torch.ones_like(range_vals), range_vals) + normalized = 2 * (x - min_val) / range_vals - 1 + + mask = (min_val == max_val).unsqueeze(0).expand_as(x) if len(x.shape) > 2 else (min_val == max_val) + normalized = torch.where(mask, torch.zeros_like(x), normalized) + normalized = torch.clamp(normalized, -1, 1) + + elif self.mode == "scale": + min_val = self.statistics["min"].to(dtype=x.dtype, device=x.device)[:horizon_len] + max_val = self.statistics["max"].to(dtype=x.dtype, device=x.device)[:horizon_len] + abs_max = torch.max(torch.abs(min_val), torch.abs(max_val)) + + abs_max_safe = torch.where(abs_max == 0, torch.ones_like(abs_max), abs_max) + normalized = x / abs_max_safe + + mask = (abs_max == 0).unsqueeze(0).expand_as(x) if len(x.shape) > 2 else (abs_max == 0) + normalized = torch.where(mask, torch.zeros_like(x), normalized) + + elif self.mode == "binary": + normalized = (x > 0.5).to(x.dtype) + else: + raise ValueError(f"Invalid normalization mode: {self.mode}") + + # Reshape back to original shape + if reshaped: + normalized = normalized.view(original_shape) + + return normalized + + def inverse(self, x: torch.Tensor) -> torch.Tensor: + """Denormalize tensor with per-horizon statistics. + + Args: + x: Normalized tensor of shape (horizon_len, action_dim) or (num_chunks * horizon_len, action_dim) + or (..., horizon_len, action_dim) + + Returns: + Denormalized tensor of same shape + """ + assert isinstance( + x, torch.Tensor + ), f"Unexpected input type: {type(x)}. Expected type: {torch.Tensor}" + + # Get the stats horizon length + stats_horizon_len = self.statistics["q01"].shape[0] + action_dim = x.shape[-1] + original_shape = x.shape + + # Handle case where x is flattened (num_chunks * horizon_len, action_dim) + reshaped = False + if len(x.shape) == 2: + total_len = x.shape[0] + if total_len > stats_horizon_len and total_len % stats_horizon_len == 0: + num_chunks = total_len // stats_horizon_len + x = x.view(num_chunks, stats_horizon_len, action_dim) + reshaped = True + elif total_len == stats_horizon_len: + x = x.unsqueeze(0) + reshaped = True + + horizon_len = x.shape[-2] + + if self.mode == "q99": + q01 = self.statistics["q01"].to(dtype=x.dtype, device=x.device)[:horizon_len] + q99 = self.statistics["q99"].to(dtype=x.dtype, device=x.device)[:horizon_len] + denormalized = (x + 1) / 2 * (q99 - q01) + q01 + elif self.mode == "mean_std": + mean = self.statistics["mean"].to(dtype=x.dtype, device=x.device)[:horizon_len] + std = self.statistics["std"].to(dtype=x.dtype, device=x.device)[:horizon_len] + denormalized = x * std + mean + elif self.mode == "min_max": + min_val = self.statistics["min"].to(dtype=x.dtype, device=x.device)[:horizon_len] + max_val = self.statistics["max"].to(dtype=x.dtype, device=x.device)[:horizon_len] + denormalized = (x + 1) / 2 * (max_val - min_val) + min_val + elif self.mode == "binary": + denormalized = (x > 0.5).to(dtype=x.dtype, device=x.device) + else: + raise ValueError(f"Invalid normalization mode: {self.mode}") + + # Reshape back to original shape + if reshaped: + denormalized = denormalized.view(original_shape) + + return denormalized + + +class PerHorizonActionTransform(InvertibleModalityTransform): + """Transform for action normalization with per-horizon statistics. + + This transform applies different normalization parameters for each action + horizon index, allowing for more precise normalization when the distribution + of relative actions varies across the prediction horizon. + + Args: + apply_to: List of action keys to transform (e.g., ['action.joint_position']) + normalization_modes: Dict mapping action keys to normalization modes + per_horizon_statistics: Dict mapping action keys to per-horizon stats. + Format: {action_key: {stat_name: [[h0_vals], [h1_vals], ...]}} + """ + + apply_to: list[str] = Field(..., description="The action keys to transform.") + normalization_modes: dict[str, str] = Field( + default_factory=dict, description="The normalization modes for each action key." + ) + per_horizon_statistics: dict[str, dict[str, list]] = Field( + default_factory=dict, description="Per-horizon statistics for each action key." + ) + + _normalizers: dict[str, PerHorizonNormalizer] = PrivateAttr(default_factory=dict) + _input_dtypes: dict[str, np.dtype | torch.dtype] = PrivateAttr(default_factory=dict) + + def model_dump(self, *args, **kwargs): + if kwargs.get("mode", "python") == "json": + include = {"apply_to", "normalization_modes"} + else: + include = kwargs.pop("include", None) + return super().model_dump(*args, include=include, **kwargs) + + def set_per_horizon_statistics(self, per_horizon_stats: dict[str, dict[str, list]]): + """Set per-horizon statistics and initialize normalizers. + + Args: + per_horizon_stats: Dict from dataset.lerobot_relative_horizon_stats_meta + Format: {action_key: {stat_name: [[h0_vals], [h1_vals], ...]}} + """ + self.per_horizon_statistics = per_horizon_stats + self._initialize_normalizers() + + def _initialize_normalizers(self): + """Initialize per-horizon normalizers from statistics.""" + for key in self.apply_to: + if key not in self.normalization_modes: + continue + + # Extract action subkey (e.g., 'action.joint_position' -> 'joint_position') + subkey = key.replace("action.", "") + + if subkey not in self.per_horizon_statistics: + print(f"Warning: No per-horizon statistics found for {subkey}, skipping") + continue + + statistics = self.per_horizon_statistics[subkey] + self._normalizers[key] = PerHorizonNormalizer( + mode=self.normalization_modes[key], + statistics=statistics + ) + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + """Apply per-horizon normalization to action data. + + Args: + data: Dict with action keys mapping to tensors of shape (horizon_len, action_dim) + + Returns: + Dict with normalized action tensors + """ + for key in self.apply_to: + if key not in data: + 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 + + action = data[key] + if key in self._normalizers: + action = self._normalizers[key].forward(action) + data[key] = action + return data + + def unapply(self, data: dict[str, Any]) -> dict[str, Any]: + """Apply inverse per-horizon normalization (denormalization). + + Args: + data: Dict with normalized action tensors + + Returns: + Dict with denormalized action tensors + """ + for key in self.apply_to: + if key not in data: + continue + + action = data[key] + assert isinstance( + action, torch.Tensor + ), f"Unexpected action type: {type(action)}. Expected type: {torch.Tensor}" + + if key in self._normalizers: + action = self._normalizers[key].inverse(action) + + if key in self._input_dtypes: + original_dtype = self._input_dtypes[key] + if isinstance(original_dtype, np.dtype): + action = action.numpy().astype(original_dtype) + elif isinstance(original_dtype, torch.dtype): + action = action.to(original_dtype) + + data[key] = action + return data + + +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." + ) + 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(".") + 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 in {modality}" + 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(".") + 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(".") + # 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 + ) + + 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/groot/vla/data/transform/video.py b/groot/vla/data/transform/video.py new file mode 100644 index 0000000000000000000000000000000000000000..e5a76df21d1d9b14d25f68de4a63da06cc98869a --- /dev/null +++ b/groot/vla/data/transform/video.py @@ -0,0 +1,1026 @@ +from typing import Any, Callable, ClassVar, Optional, Literal + +import albumentations as A +import cv2 +from einops import rearrange +import functools +import numpy as np +from pydantic import Field, PrivateAttr, field_validator +import torch +import torchvision.transforms.v2 as T + +from groot.vla.data.schema import DatasetMetadata +from groot.vla.data.transform.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 VideoRandomErasing(VideoTransform): + """Adds random rectangles overlaying the video. + + This discourages overfitting to the background. + """ + + probability: float = Field(default=0.2, description="Probability of applying the transform") + scale: tuple[float, float] = Field(default=(0.02, 0.33), description="Scale of the rectangle") + ratio: tuple[float, float] = Field( + default=(0.3, 3.3), description="Aspect ratio of the rectangle" + ) + value: Literal["random"] | tuple[float, float, float] = Field( + default="random", description="Color to fill the erased region with" + ) + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable | None: + """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 transform that adds random rectangles. If mode is "eval", return a no-op. + """ + if mode == "eval": + return None + if self.backend == "torchvision": + return T.RandomErasing( + p=self.probability, scale=self.scale, ratio=self.ratio, value=self.value + ) + elif self.backend == "albumentations": + return A.Erasing( + p=self.probability, scale=self.scale, ratio=self.ratio, value=self.value + ) + else: + raise ValueError(f"Backend {self.backend} not supported") + + +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): + + output_on_cuda: bool = Field( + default=False, + description="Output the tensor on CUDA if True.", + ) + + 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 functools.partial( + self.__class__.to_tensor, + output_on_cuda=self.output_on_cuda, + ) + 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, output_on_cuda: bool) -> torch.Tensor: + """Convert numpy array to tensor efficiently. + + Args: + frames: numpy array of shape [T, H, W, C] in uint8 format + output_on_cuda: whether to output the tensor on CUDA if True. + Returns: + tensor of shape [T, C, H, W] in range [0, 1] + """ + frames = torch.from_numpy(frames) + if output_on_cuda: + frames = frames.cuda() + frames = frames.to(torch.float32) / 255.0 + return frames.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 + """ + frames = (frames.permute(0, 2, 3, 1) * 255).to(torch.uint8) + return frames.cpu().numpy() + + +class VideoMergeTimeBatch(ModalityTransform): + """ + Merge the batch and time dimensions of the video. + """ + + apply_to: list[str] = Field(..., description="The keys of the modalities to merge") + + def apply(self, data: dict) -> dict: + warnings.warn( + "VideoMergeTimeBatch is deprecated. Use ComposedModalityTransform instead.", + DeprecationWarning, + ) + for key in self.apply_to: + data[key] = rearrange(data[key], "b t ... -> (b t) ...") + return data + + +class VideoSplitTimeBatch(ModalityTransform): + """ + Split the batch and time dimensions of the video. + """ + + apply_to: list[str] = Field(..., description="The keys of the modalities to split") + time_dim: int = Field(..., description="The dimension of the time dimension") + + def apply(self, data: dict) -> dict: + warnings.warn( + "VideoSplitTimeBatch is deprecated. Use ComposedModalityTransform instead.", + DeprecationWarning, + ) + for key in self.apply_to: + data[key] = rearrange(data[key], "(b t) ... -> b t ...", t=self.time_dim) + return data + + +class VideoFocusRect(ModalityTransform): + """ + Given a rectangle area in the video, apply focus effects on the target + rectangle, by applying blur and noise to the surrounding region. + + Mainly useful for EgoView + """ + + # Region coordinates in normalized space [0,1] + xtl: float = Field(2 / 12, description="Top-left x coordinate (normalized)", ge=0.0, le=1.0) + ytl: float = Field(3 / 8, description="Top-left y coordinate (normalized)", ge=0.0, le=1.0) + xbr: float = Field(10 / 12, description="Bottom-left x coordinate (normalized)", ge=0.0, le=1.0) + ybr: float = Field(1.0, description="Bottom-left y coordinate (normalized)", ge=0.0, le=1.0) + + # Content region parameters (in pixel coordinates, None means auto-detect) + content_y_min: Optional[int] = Field( + None, description="Top coordinate of content region (pixels)" + ) + content_y_max: Optional[int] = Field( + None, description="Bottom coordinate of content region (pixels)" + ) + content_x_min: Optional[int] = Field( + None, description="Left coordinate of content region (pixels)" + ) + content_x_max: Optional[int] = Field( + None, description="Right coordinate of content region (pixels)" + ) + + # Jitter amount for coordinates (in normalized space) + jitter: float = Field(0.05, description="Amount of random jitter to apply to coordinates") + + # Effect parameters + blur_kernel: int = Field(95, description="Gaussian blur kernel size") + noise_std: float = Field(0.3, description="Standard deviation of Gaussian noise") + blend_size: float = Field(0.1, description="Size of blending region as fraction of image size") + + # Effect probabilities during training + p_blur: float = Field(0.2, description="Probability of applying blur") + p_noise: float = Field(0.2, description="Probability of applying noise") + + def detect_padding(self, image: np.ndarray) -> tuple[slice, slice]: + """ + Detect padding in the image by finding non-black regions. + Returns slices for the content region (y_slice, x_slice). + """ + H, W = image.shape[:2] + + # If all content region parameters are provided, use them + if all( + param is not None + for param in [ + self.content_y_min, + self.content_y_max, + self.content_x_min, + self.content_x_max, + ] + ): + y_min = max(0, self.content_y_min) + y_max = min(H, self.content_y_max) + x_min = max(0, self.content_x_min) + x_max = min(W, self.content_x_max) + return slice(y_min, y_max), slice(x_min, x_max) + + if image.ndim == 3: + gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + else: + gray = image + + # Find non-black regions + rows = np.any(gray > 0.01, axis=1) + cols = np.any(gray > 0.01, axis=0) + + # find first and last non-black pixel indices + rmin, rmax = np.where(rows)[0][[0, -1]] + cmin, cmax = np.where(cols)[0][[0, -1]] + + return slice(rmin, rmax + 1), slice(cmin, cmax + 1) + + def apply(self, data: dict[str, Any]) -> dict[str, Any]: + if not self.training: + # Do nothing in eval mode + return data + + for key in self.apply_to: + video = data[key] + + # Handle numpy array case + assert isinstance( + video, np.ndarray + ), f"Expected numpy array or torch tensor for {key}, got {type(video)}" + assert video.ndim in { + 3, + 4, + }, f"Expected [H, W, C] or [T, H, W, C] array for {key}, got shape {video.shape}" + transformed = self._transform_video(video) + + data[key] = transformed + + return data + + def _transform_video(self, video: np.ndarray) -> np.ndarray: + """ + Apply the focus rectangle transformation to a video. + + Args: + video (np.ndarray): Video tensor of shape [T, H, W, C] + + Returns: + np.ndarray: Transformed video + """ + # Handle both single frame and video inputs + is_single_frame = video.ndim == 3 + if is_single_frame: + video = video[np.newaxis] + + T, H, W, C = video.shape + + assert ( + self.p_blur + self.p_noise <= 1.0 + ), "Sum of blur and noise probabilities must be <= 1.0" + r = np.random.random() + apply_blur = r < self.p_blur + apply_noise = self.p_blur <= r < self.p_blur + self.p_noise + alpha = random.uniform(0.0, 1.0) # Noise blending factor + + # Apply jitter once to rectangle + xtl = self.xtl + np.random.uniform(-self.jitter, self.jitter) + ytl = self.ytl + np.random.uniform(-self.jitter, self.jitter) + xbr = self.xbr + np.random.uniform(-self.jitter, self.jitter) + ybr = self.ybr + np.random.uniform(-self.jitter, self.jitter) + xtl, ytl, xbr, ybr = [np.clip(x, 0.0, 1.0) for x in [xtl, ytl, xbr, ybr]] + + # Detect padding from first frame (assume consistent across frames) + y_slice, x_slice = self.detect_padding(video[0]) + content_h = y_slice.stop - y_slice.start + content_w = x_slice.stop - x_slice.start + + # Convert normalized coordinates relative to pixel space + x1 = int(xtl * content_w) + x_slice.start + y1 = int(ytl * content_h) + y_slice.start + x2 = int(xbr * content_w) + x_slice.start + y2 = int(ybr * content_h) + y_slice.start + + # Create mask for the inner rectangle + mask = np.zeros((H, W), dtype=np.float32) + pts = np.array([[x1, y1], [x2, y1], [x2, y2], [x1, y2]], dtype=np.int32) + cv2.fillPoly(mask, [pts], color=1.0) + + # Create a smooth blend mask around the target rectangle using distance transform + content_mask = np.zeros((H, W), dtype=np.uint8) + content_mask[y_slice, x_slice] = 1 + dist = cv2.distanceTransform(1 - (mask > 0).astype(np.uint8) * content_mask, cv2.DIST_L2, 3) + blend_radius = int(min(content_h, content_w) * self.blend_size) + blend_mask = np.clip(1.0 - dist / blend_radius, 0, 1) + blend_mask *= content_mask + blend_mask = blend_mask[..., np.newaxis] + + # Process all frames with same transformations + result = np.zeros_like(video) + for t in range(T): + frame = video[t] + modified = frame.copy() + content = modified[y_slice, x_slice] + + if apply_blur: + content = cv2.GaussianBlur(content, (self.blur_kernel, self.blur_kernel), 0) + + if apply_noise: + background = np.random.randint(0, 256, content.shape, dtype=np.uint8) / 255.0 + content = alpha * content + (1 - alpha) * background + content = np.clip(content, 0, 1) + + modified[y_slice, x_slice] = content + result[t] = frame * blend_mask + modified * (1 - blend_mask) + + return result[0] if is_single_frame else result + + +class VideoNormalize(VideoTransform): + mean: list[float] = Field(..., description="Mean for normalization") + std: list[float] = Field(..., description="Standard deviation for normalization") + + def get_transform(self, mode: Literal["train", "eval"] = "train") -> Callable: + """Get the normalization transform. Same for train and eval mode. + + Args: + mode (Literal["train", "eval"]): The mode to get the transform for. + + Returns: + Callable: The normalization transform. + """ + print("Using VideoNormalize transform") + if self.backend == "torchvision": + return T.Normalize(mean=self.mean, std=self.std) + elif self.backend == "albumentations": + return A.Normalize(mean=self.mean, std=self.std, max_pixel_value=1.0, p=1.0) + else: + raise ValueError(f"Backend {self.backend} not supported") + + def check_input(self, data: dict): + for key in self.apply_to: + assert key in data, f"Key {key} not found in data" + assert isinstance(data[key], torch.Tensor), f"Video {key} is not a torch tensor" + assert data[key].ndim in [4, 5], f"Video {key} must have 4 or 5 dimensions, got {data[key].ndim}" + assert data[key].dtype == torch.float32, f"Video {key} must be float32, got {data[key].dtype}" + assert data[key].min() >= 0.0 and data[key].max() <= 1.0, ( + f"Video {key} must be in [0,1] range before normalization" + ) + +# class VideoTransformLegacy(ModalityTransform): +# def __init__( +# self, +# modality_keys: list[str], +# backend: str = "torchvision", +# crop_cfg: CropConfig | None = None, +# resize_cfg: ResizeConfig | None = None, +# random_rotation_cfg: RandomRotationConfig | None = None, +# horizontal_flip_cfg: HorizontalFlipConfig | None = None, +# grayscale_cfg: GrayscaleConfig | None = None, +# color_jitter_cfg: ColorJitterConfig | None = None, +# strong_vision_aug: bool = False, +# ): +# """ +# Initialize the video transform. +# With the default settings, the input will be (T, H, W, C) where T is the number of frames. +# The output will be (K, T, C, H, W) where K is the number of video keys. + +# Args: +# modality_keys (list[str]): The keys of the modalities to load and transform. +# backend (str): The backend to use for the transformations. The default is "torchvision". +# crop_cfg (CropConfig | None): Configuration for the crop transformation. See CropConfig for more details. +# resize_cfg (ResizeConfig | None): Configuration for the resize transformation. See ResizeConfig for more details. +# random_rotation_cfg (RandomRotationConfig | None): Configuration for the random rotation transformation. See RandomRotationConfig for more details. +# horizontal_flip_cfg (HorizontalFlipConfig | None): Configuration for the horizontal flip transformation. See HorizontalFlipConfig for more details. +# grayscale_cfg (GrayscaleConfig | None): Configuration for the grayscale transformation. See GrayscaleConfig for more details. +# color_jitter_cfg (ColorJitterConfig | None): Configuration for the color jitter transformation. See ColorJitterConfig for more details. +# strong_vision_aug (bool): Whether to apply strong vision augmentation. The default is False. +# """ +# super().__init__(modality_keys) +# self.backend = backend +# self.crop_cfg = crop_cfg +# self.resize_cfg = resize_cfg +# self.random_rotation_cfg = random_rotation_cfg +# self.horizontal_flip_cfg = horizontal_flip_cfg +# self.grayscale_cfg = grayscale_cfg +# self.color_jitter_cfg = color_jitter_cfg +# self.strong_vision_aug = strong_vision_aug +# self.transforms = None + +# def set_metadata(self, dataset_metadata: TrainableDatasetMetadata_V1_1): +# super().set_metadata(dataset_metadata) +# # Get the original height and width +# video_metadata = dataset_metadata.modalities.video +# resolutions = {} +# for key in self.modality_keys: +# split_keys = key.split(".") +# assert len(split_keys) == 2, f"Invalid key: {key}. Expected format: modality.key" +# sub_key = split_keys[1] +# resolutions[key] = video_metadata[sub_key].resolution +# assert ( +# len(set(resolutions.values())) == 1 +# ), f"All video keys must have the same resolution, got: {resolutions}" +# width, height = resolutions[self.modality_keys[0]] + +# transforms = [] +# if self.crop_cfg is not None: +# self.crop_cfg.set_original_height_width(height, width) +# transforms.append(self.crop_cfg.get_transform(self.backend)) +# if self.resize_cfg is not None: +# transforms.append(self.resize_cfg.get_transform(self.backend)) +# if self.random_rotation_cfg is not None: +# transforms.append(self.random_rotation_cfg.get_transform(self.backend)) +# if self.horizontal_flip_cfg is not None: +# transforms.append(self.horizontal_flip_cfg.get_transform(self.backend)) +# if self.grayscale_cfg is not None: +# transforms.append(self.grayscale_cfg.get_transform(self.backend)) +# if self.color_jitter_cfg is not None: +# transforms.append(self.color_jitter_cfg.get_transform(self.backend)) + +# if self.backend == "torchvision": +# if len(transforms) == 0: +# transforms.append(T.Identity()) +# self.transforms = T.Compose(transforms) +# else: +# raise ValueError(f"Backend {self.backend} not supported") + +# if self.strong_vision_aug: +# import kornia.augmentation as K +# from kornia.augmentation import ImageSequential + +# assert ( +# self.backend == "torchvision" +# ), "Temporarily only support torchvision backend for strong augmentation" +# self.strong_transform = ImageSequential( +# K.RandomErasing(p=0.5, scale=(0.005, 0.01), ratio=(0.3, 1.3)), +# K.RandomSaltAndPepperNoise(p=0.5, amount=0.05, salt_vs_pepper=0.5), +# K.RandomCutMixV2(p=0.5, num_mix=1, cut_size=(0.98, 1.0)), +# random_apply=1, +# keepdim=True, +# same_on_batch=True, +# ) + +# def __call__(self, data: dict) -> dict[str, torch.Tensor | np.ndarray | Image.Image]: +# # Batch frames along the first dimension +# frames = [data[key] for key in self.modality_keys] # view x [T, H, W, C] +# n_view, n_frames = len(frames), len(frames[0]) +# frames = np.concatenate(frames, 0) # [view*T, H, W, C] + +# if self.backend == "torchvision": +# transformed_frames = self.transform_torchvision(frames) +# else: +# raise ValueError(f"Backend {self.backend} not supported") + +# # De-batch the frames +# transformed_frames = np.array(transformed_frames) # [view*T, H, W, C] +# H, W, C = transformed_frames.shape[-3:] +# transformed_frames = { +# key: x +# for key, x in zip( +# self.modality_keys, transformed_frames.reshape(n_view, n_frames, H, W, C) +# ) +# } + +# return transformed_frames + +# def check_input(self, data: dict): +# for key in self.modality_keys: +# assert key in data, f"Key {key} not found in data" +# video = data[key] +# assert isinstance(video, np.ndarray), f"Video {key} is not a numpy array" +# assert video.ndim == 4, f"Video {key} must have 4 dimensions, got {video.ndim}" +# assert video.dtype == np.uint8, f"Video {key} must have dtype uint8, got {video.dtype}" +# shape = video.shape +# expected_resolution = self.dataset_metadata.modalities.video[key].resolution +# assert ( +# shape[1:3] == expected_resolution +# ), f"Video {key} has invalid shape {shape}, expected {expected_resolution}" + +# def transform_torchvision( +# self, frames: np.ndarray +# ) -> list[torch.Tensor | np.ndarray | Image.Image]: +# """ +# frames: [view * T, H, W, C], np.uint8 +# """ +# if self.transforms is None: +# raise ValueError( +# "Transform is not set. Please call set_metadata() before calling __call__()" +# ) +# # Convert to batched tensor, using ToTensor() is too slow +# frames_tensor = torch.from_numpy(frames).to(torch.float32) / 255.0 +# frames_tensor = frames_tensor.permute(0, 3, 1, 2) # [view * T, C, H, W] +# transformed_frames = self.transforms(frames_tensor) +# if self.strong_vision_aug: +# transformed_frames = self.strong_transform(transformed_frames) + +# to_pil = T.ToPILImage() +# transformed_frames = [to_pil(frame) for frame in transformed_frames] +# return transformed_frames # type: ignore + + +# class IdentityTransform(ModalityTransform): +# def __call__(self, data: dict) -> dict: # type: ignore +# warnings.warn("IdentityTransform is used, further transformations is required.") +# output = {} +# for key in self.modality_keys: +# output[key] = data[key] +# return output diff --git a/groot/vla/experiment/__init__.py b/groot/vla/experiment/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2cf08b47119645e476dc0b3cfcb7ed54fa4aaec1 --- /dev/null +++ b/groot/vla/experiment/__init__.py @@ -0,0 +1 @@ +from groot.vla.experiment.experiment import VLATrainer \ No newline at end of file diff --git a/groot/vla/experiment/base.py b/groot/vla/experiment/base.py new file mode 100644 index 0000000000000000000000000000000000000000..33156ff623e5fd7701ef9c3f1db5c8a9c0df99f2 --- /dev/null +++ b/groot/vla/experiment/base.py @@ -0,0 +1,956 @@ +# Copyright 2024 NVIDIA CORPORATION & AFFILIATES +# +# 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. +# +# SPDX-License-Identifier: Apache-2.0 +# This file is modified from https://github.com/haotian-liu/LLaVA/ + +from abc import ABC +import contextlib +import json +import logging +import os +from pathlib import Path +import shutil +import threading +import time +from typing import Optional +import warnings + +from hydra.utils import instantiate +import numpy as np +from omegaconf import DictConfig, OmegaConf, open_dict +import torch +from torch.profiler import ProfilerActivity, profile +from torch.utils.data import DataLoader, Dataset, Sampler +import transformers +from transformers import TrainerCallback, set_seed +from transformers.trainer import ( + # ALL_LAYERNORM_LAYERS, # ShardedDDPOption, # Removed deprecated import + TRAINER_STATE_NAME, + TrainerState, + get_last_checkpoint, + get_parameter_names, + is_sagemaker_mp_enabled, +) + +import groot.vla.common.utils as U +from groot.vla.data.dataset.lerobot_sharded import ShardedLeRobotMixtureDataset +from groot.vla.data.schema import EmbodimentTag +from groot.vla.data.transform import ComposedModalityTransform +from groot.vla.experiment.trex_wandb_video_callback import TrexWandbVideoCallback +from groot.vla.experiment.utils import ( + compute_grad_accum_to_match_global_bs, + dtype_from_string, + estimate_steps_per_data_epoch, + get_checkpoint_path, + mprint, + safe_save_model_for_hf_trainer, +) +from groot.vla.utils.timer import ContextTimer + +# Fix resume: https://github.com/huggingface/transformers/pull/34632/files +np_core = np.core +allowlist = [np_core.multiarray._reconstruct, np.ndarray, np.dtype] +# numpy >1.25 defines numpy.dtypes.UInt32DType, but below works for +# all versions of numpy +allowlist += [type(np.dtype(np.uint32))] +torch.serialization.add_safe_globals(allowlist) + +# Define LayerNorm classes locally to replace deprecated ALL_LAYERNORM_LAYERS +LAYERNORM_LAYERS = [ + torch.nn.LayerNorm, + torch.nn.GroupNorm, + torch.nn.InstanceNorm1d, + torch.nn.InstanceNorm2d, + torch.nn.InstanceNorm3d, + torch.nn.LocalResponseNorm, + torch.nn.BatchNorm1d, + torch.nn.BatchNorm2d, + torch.nn.BatchNorm3d, + torch.nn.SyncBatchNorm, +] + + +class LossLoggerCallback(TrainerCallback): + """Callback that writes per-step loss metrics to a JSONL file for offline analysis.""" + + def __init__(self, output_path: str): + self.output_path = output_path + + def on_log(self, args, state, control, logs=None, **kwargs): + if not state.is_world_process_zero or logs is None: + return + entry = {"step": state.global_step} + for key, value in logs.items(): + if key in ("loss", "learning_rate") or key.endswith("_avg"): + entry[key] = value + if len(entry) > 1: # more than just "step" + with open(self.output_path, "a") as f: + f.write(json.dumps(entry) + "\n") + + +class CheckpointFormatCallback(TrainerCallback): + """This callback format checkpoint to make them standalone. For now, it copies all config + files to /checkpoint-{step}/experiment_cfg/: + - conf.yaml + - initial_actions.npz + - metadata.json + """ + + def __init__( + self, run_name: str, exp_cfg_dir: Path | None = None, processor_dir: Path | None = None + ): + """ + Args: + run_name: Name of the experiment run + exp_cfg_dir: Path to the directory containing all experiment metadata + """ + self.exp_cfg_dir = exp_cfg_dir + self.processor_dir = processor_dir + + def on_save(self, args, state, control, **kwargs): + """Called after the trainer saves a checkpoint.""" + if state.is_world_process_zero: + checkpoint_dir = Path(args.output_dir) / f"checkpoint-{state.global_step}" + + # Copy experiment config directory if provided + if self.exp_cfg_dir is not None: + exp_cfg_dst = checkpoint_dir / self.exp_cfg_dir.name + if self.exp_cfg_dir.exists(): + print( + f"Copying experiment config directory {self.exp_cfg_dir} to {exp_cfg_dst}" + ) + shutil.copytree(self.exp_cfg_dir, exp_cfg_dst, dirs_exist_ok=True) + + # Copy processor directory if provided + if self.processor_dir is not None: + if self.processor_dir.exists(): + print(f"Copying processor directory {self.processor_dir} to {checkpoint_dir}") + shutil.copytree(self.processor_dir, checkpoint_dir, dirs_exist_ok=True) + + # Copy wandb_config.json if provided + wandb_config_src = Path(args.output_dir) / "wandb_config.json" + wandb_config_dst = checkpoint_dir / "wandb_config.json" + if wandb_config_src.exists(): + print(f"Copying wandb_config.json from {wandb_config_src} to {wandb_config_dst}") + shutil.copy2(wandb_config_src, wandb_config_dst) + + +class ProfCallback(transformers.TrainerCallback): + """Callback to manage PyTorch profiler during training. + + Dynamically starts/stops the profiler within a specified session step window. + After profiling completes, triggers optional S3 upload and removes itself. + + Args: + profile_dir: Directory to save profile traces + upload_callback: Optional callback to trigger S3 upload after profiling + profile_start_step: Session step to start profiling (default: 50) + profile_end_step: Session step to stop profiling + warmup_steps: Number of warmup steps for profiler schedule (default: 1) + active_steps: Number of active profiling steps (default: 5) + trainer: Trainer instance (required for self-removal after profiling) + record_shapes: Record tensor shapes in profiler (default: False) + with_stack: Record Python stack traces (default: True) + profile_memory: Record memory allocation/deallocation (default: False) + """ + + def __init__( + self, + profile_dir, + upload_callback=None, + profile_start_step=50, + profile_end_step=55, + warmup_steps=1, + active_steps=5, + trainer=None, + record_shapes=False, + with_stack=True, + profile_memory=False, + ): + self.profile_dir = profile_dir + self.upload_callback = upload_callback + self.profile_start_step = profile_start_step + self.profile_end_step = profile_end_step + self.warmup_steps = warmup_steps + self.active_steps = active_steps + self.trainer = trainer + self.record_shapes = record_shapes + self.with_stack = with_stack + self.profile_memory = profile_memory + self.upload_triggered = False + self.starting_global_step = None + self.session_step = 0 + self.prof = None + self.profiling_active = False + self.profiling_complete = False + self.removed_from_trainer = False + + def on_step_begin(self, args, state, control, **kwargs): + # Remove callback after upload triggered to eliminate all overhead + if self.profiling_complete and self.upload_triggered and not self.removed_from_trainer: + if self.trainer is not None and hasattr(self.trainer, "callback_handler"): + try: + self.trainer.callback_handler.callbacks.remove(self) + self.removed_from_trainer = True + logging.info( + f"Removed ProfCallback from trainer at global step {state.global_step}" + ) + except (ValueError, AttributeError) as e: + logging.warning(f"Failed to remove ProfCallback: {e}") + return + + # Early return if profiling already complete + if self.profiling_complete: + return + + # Record starting global step on first call + if self.starting_global_step is None: + self.starting_global_step = state.global_step + + # Calculate session step + self.session_step = state.global_step - self.starting_global_step + + # Start profiler when we reach the profiling window + if self.session_step == self.profile_start_step and self.prof is None: + logging.info( + f"Starting profiler at global step {state.global_step} (session step {self.session_step})" + ) + self.prof = torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + schedule=torch.profiler.schedule( + skip_first=0, + wait=0, + warmup=self.warmup_steps, + active=self.active_steps, + repeat=1, + ), + profile_memory=self.profile_memory, + with_stack=self.with_stack, + record_shapes=self.record_shapes, + on_trace_ready=torch.profiler.tensorboard_trace_handler(str(self.profile_dir)), + ) + self.prof.__enter__() + self.profiling_active = True + + def on_step_end(self, args, state, control, **kwargs): + # Early return if profiling already complete + if self.profiling_complete: + return + + # Recalculate session_step to ensure accuracy + if self.starting_global_step is not None: + self.session_step = state.global_step - self.starting_global_step + + # Step profiler if active + if self.profiling_active and self.prof is not None: + self.prof.step() + + # Stop profiler when we reach the end of profiling window + if self.session_step == self.profile_end_step and self.prof is not None: + self.prof.__exit__(None, None, None) + self.profiling_active = False + + # Explicitly release profiler resources to minimize CUPTI overhead + # Combined with TEARDOWN_CUPTI=1 env var for full cleanup + del self.prof + self.prof = None + + # Force CUDA synchronization to ensure profiler cleanup completes + if torch.cuda.is_available(): + torch.cuda.synchronize() + + self.profiling_complete = True + logging.info( + f"Profiler stopped and resources released at global step {state.global_step} " + f"(session step {self.session_step})" + ) + + # Trigger upload if callback provided + if self.upload_callback: + logging.info(f"Triggering upload at global step {state.global_step}...") + self.upload_callback() + + # Mark as ready for callback removal + self.upload_triggered = True + + +class BaseSampler(Sampler): + """Sampler for dataset, which enables `set_epoch` for Dataset. + `set_epoch` will be called by huggingface Trainer at the end of each epoch. + `shuffle` is also supported for training set shuffling + """ + + def __init__(self, data_source: Dataset, shuffle: bool = False, seed: int = 0): + self.data_source = data_source + self.shuffle = shuffle + self.seed = seed + self.epoch = 0 + + def __iter__(self): + if self.shuffle: + g = torch.Generator() + g.manual_seed(self.seed + self.epoch) + # must not add rank here, or randomization will be different for each rank + return iter(torch.randperm(len(self.data_source), generator=g).tolist()) + return iter(range(len(self.data_source))) + + def set_epoch(self, epoch): + self.epoch = epoch + if hasattr(self.data_source, "set_epoch"): + # this is important for dataset + self.data_source.set_epoch(epoch) + + def __len__(self): + return len(self.data_source) + + +class BaseTrainer(transformers.Trainer): + + def __init__(self, **kwargs): + # Increase the cache size limit for torch._dynamo to + # accommodate videos with different numbers of frames. + torch._dynamo.config.cache_size_limit = 1000 + + self.compute_dtype = kwargs.pop("compute_dtype") + self.output_dir = kwargs.pop("output_dir") + self.timer = ContextTimer(self) + + self.world_size = int(os.environ.get("WORLD_SIZE", "1")) + self.local_rank = int(os.environ.get("LOCAL_RANK", "0")) + self.global_rank = int(os.environ.get("RANK", "0")) + self.node_rank = int(os.environ.get("NODE_RANK", "0")) + + # Get distributed info + self.current_step = 0 + + # Profiling (legacy per-step profiling) + self.enable_profiling = kwargs.pop("enable_profiling", False) + self.profiling_steps = kwargs.pop("profiling_steps", 5) + # Pop new ProfCallback config options (handled in create_trainer, not here) + kwargs.pop("enable_prof_callback", None) + kwargs.pop("profile_start_step", None) + kwargs.pop("profile_warmup_steps", None) + kwargs.pop("profile_active_steps", None) + kwargs.pop("profile_record_shapes", None) + kwargs.pop("profile_with_stack", None) + kwargs.pop("profile_memory", None) + kwargs.pop("msc_profile_url", None) + kwargs.pop("profile_delete_after_upload", None) + if self.enable_profiling: + # Setup profiling directories + self.profile_dir = Path(self.output_dir) / "profiling" + self.memory_profile_dir = self.profile_dir / "memory" + self.torch_profile_dir = self.profile_dir / "torch" + + self.memory_profile_dir.mkdir(exist_ok=True, parents=True) + self.torch_profile_dir.mkdir(exist_ok=True, parents=True) + + # Start recording the memory history. + torch.cuda.memory._record_memory_history(max_entries=100000) + + super().__init__(**kwargs) + + self.loss_queues = {} + self.loss_queue_size = 10 + + def _get_train_sampler(self): + return BaseSampler(self.train_dataset, shuffle=True, seed=self.args.seed) + + def _get_eval_sampler(self, eval_dataset): + return BaseSampler(eval_dataset, shuffle=False) + + def training_step(self, model, inputs, num_items_in_batch=None): + enable_profile = self.enable_profiling and self.current_step % self.profiling_steps == 0 + if enable_profile: + profile_context = profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + record_shapes=True, + with_stack=True, + ) + else: + profile_context = contextlib.nullcontext() + + start_time = time.time() + + with self.timer.with_label("training_step"), profile_context as prof: + output = super().training_step(model, inputs) + + time_taken = time.time() - start_time + print( + f"Rank {self.global_rank} time taken for training_step {self.current_step}: {time_taken:.2f} seconds" + ) + + if enable_profile: + trace_path = f"{self.torch_profile_dir}/trace_rank_{self.global_rank}_step_{self.current_step}.json.gz" + print(f"Rank {self.global_rank} exporting torch profile to {trace_path}") + prof.export_chrome_trace(trace_path) + + snapshot_path = f"{self.memory_profile_dir}/memory_snapshot_rank_{self.global_rank}_step_{self.current_step}.pickle" + print(f"Rank {self.global_rank} dumping memory snapshot to {snapshot_path}") + torch.cuda.memory._dump_snapshot(snapshot_path) + + self.current_step += 1 + return output + + def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): + with self.timer.with_label("model_forward"): + outputs = model(inputs) + ### For additional losses/diagnostics, track and log their moving averages + for key, value in outputs.items(): + if ( + key.endswith("_loss") + or key in ("tactile_perplexity", "tactile_active_codes") + ) and key != "loss": + # Initialize queue if not exists + if key not in self.loss_queues: + self.loss_queues[key] = [] + + # Add current loss value to queue + current_value = value.item() if torch.is_tensor(value) else value + self.loss_queues[key].append(current_value) + + # Keep only last N values + if len(self.loss_queues[key]) > self.loss_queue_size: + self.loss_queues[key].pop(0) + + # Log average every 10 steps + if self.current_step % self.loss_queue_size == 0: + avg_loss = sum(self.loss_queues[key]) / len(self.loss_queues[key]) + self.log({f"{key}_avg": avg_loss}) + + loss = outputs["loss"] + + return (loss, outputs) if return_outputs else loss + + def create_optimizer(self): + """ + Setup the optimizer. + + We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the + Trainer's init through `optimizers`, or subclass and override this method in a subclass. + """ + if is_sagemaker_mp_enabled(): + return super().create_optimizer() + + opt_model = self.model + + if self.optimizer is None: + decay_parameters = get_parameter_names(opt_model, LAYERNORM_LAYERS) + decay_parameters = [name for name in decay_parameters if "bias" not in name] + optimizer_grouped_parameters = [ + { + "params": [ + p + for n, p in opt_model.named_parameters() + if (n in decay_parameters and p.requires_grad) + ], + "weight_decay": self.args.weight_decay, + }, + { + "params": [ + p + for n, p in opt_model.named_parameters() + if (n not in decay_parameters and p.requires_grad) + ], + "weight_decay": 0.0, + }, + ] + + optimizer_cls, optimizer_kwargs = transformers.Trainer.get_optimizer_cls_and_kwargs( + self.args + ) + self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs) + + # DeepSpeed CPU Adam (ZeRO offload) expects 'bias_correction' in each param group. + # HuggingFace Trainer's AdamW does not set it, causing KeyError in cpu_adam.step(). + if getattr(self.args, "deepspeed", None): + for group in self.optimizer.param_groups: + group.setdefault("bias_correction", True) + + return self.optimizer + + def save_model(self, output_dir: Optional[str], _internal_call: bool): + + ## save tuned model separately + if self.is_deepspeed_enabled: + state_dict = self.accelerator.get_state_dict(self.deepspeed) + else: + state_dict = self.model.state_dict() + + if self.base_cfg.save_lora_only: + if hasattr(self.model, "get_trainable_checkpoint_state_dict"): + state_dict = self.model.get_trainable_checkpoint_state_dict( + state_dict + ) + else: + # Save only the trainable parameters. + train_key = [ + k for k, v in self.model.named_parameters() if v.requires_grad + ] + state_dict = { + k: v + for k, v in self.model.state_dict().items() + if k in train_key + } + + if self.args.should_save: + ret = self.model.save_pretrained(output_dir, state_dict=state_dict) + + # can separately save the VLM model for downstream evalualtion + if self.base_cfg.save_llm: + llm_output_dir = os.path.join(output_dir, "llm") + self.model.backbone.model.save_pretrained(llm_output_dir) + + if self.base_cfg.save_value_model: + assert hasattr( + self.model.action_head, "value_model" + ), f"Value model not found in action head: {type(self.model.action_head)}" + value_model_output_dir = os.path.join(output_dir, "value_model") + self.model.action_head.value_model.save_pretrained(value_model_output_dir) + + return ret + + def train( + self, + resume_from_checkpoint=None, + trial=None, + ignore_keys_for_eval=None, + **kwargs, + ): + """Correctly set self.state from checkpoint so get_train_dataloader can read from it.""" + if resume_from_checkpoint is False: + resume_from_checkpoint = None + + if isinstance(resume_from_checkpoint, bool) and resume_from_checkpoint: + resume_from_checkpoint = get_last_checkpoint(self.args.output_dir) + if resume_from_checkpoint is None: + raise ValueError( + f"No valid checkpoint found in output directory ({self.args.output_dir})" + ) + + if resume_from_checkpoint is not None: + # In case of repeating the find_executable_batch_size, set `self._train_batch_size` properly + self.state = TrainerState.load_from_json( + os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME) + ) + return super().train(resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs) + + def get_train_dataloader(self) -> DataLoader: + """ + Returns the training [`~torch.utils.data.DataLoader`]. + + Will use no sampler if `train_dataset` does not implement `__len__`, a random sampler (adapted to distributed + training if necessary) otherwise. + + Subclass and override this method if you want to inject some custom behavior. + """ + if self.train_dataset is None: + raise ValueError("Trainer: training requires a train_dataset.") + + train_dataset = self.train_dataset + if not isinstance(train_dataset, (ShardedLeRobotMixtureDataset)): + return super().get_train_dataloader() + + # During resume, don't skip the data + self.args.ignore_data_skip = True + curr_global_step = self.state.global_step + print(f"Current global step: {curr_global_step}") + if curr_global_step > 0: + new_seed = train_dataset.seed + curr_global_step + train_dataset.reset_seed(new_seed) + print( + f"Resetting seed to {new_seed}. Please note that this will make the experiment non-reproducible." + ) + + print("Creating custom train dataloader") + # Handle the case where the dataset is an IterableDataset + data_collator = self.data_collator + data_collator = self._get_collator_with_removed_columns( + data_collator, description="training" + ) + + dataloader_params = { + "batch_size": self._train_batch_size, + "collate_fn": data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + } + # persistent_workers is only valid when num_workers > 0 (PyTorch raises otherwise) + if self.args.dataloader_num_workers > 0: + dataloader_params["persistent_workers"] = self.args.dataloader_persistent_workers + + return DataLoader(train_dataset, **dataloader_params) + + +class BaseExperiment(ABC): + def __init__(self, cfg: DictConfig): + # assert cfg.save_steps == 500, "save_steps must be 500 for standarized evaluation" + assert cfg.max_steps > 0, "max_steps must be > 0 for standarized evaluation" + assert cfg.save_total_limit >= 5, "save_total_limit must be >= 5 for standarized evaluation" + + if cfg.load_from_yaml is not None: + # Override the default config with the loaded config. + loaded_cfg = OmegaConf.load(cfg.load_from_yaml) + cfg = loaded_cfg # overwrite + + # Check if evaluation transforms are valid. + assert cfg.transforms is not None, "Evaluation transforms are not provided." + for tag, transform_cfg in cfg.transforms.items(): + try: + # Check if the tag is a valid EmbodimentTag + _ = EmbodimentTag(tag) + # Check if the transform is a valid ComposedModalityTransform + transform = instantiate(transform_cfg) + assert isinstance(transform, ComposedModalityTransform), f"{transform=}" + except Exception as e: + raise ValueError(f"Evaluation transform {tag} is invalid: {e}") + + # Instantiate the training arguments. + cfg.training_args.output_dir = cfg.training_args.output_dir.rstrip("/") + cfg.training_args.run_name = cfg.training_args.output_dir.split("/")[-1] + print(f"Run name: {cfg.training_args.run_name}") + training_args = instantiate(cfg.training_args) + set_seed(training_args.seed) + + # Set the environment variables for wandb. + if "WANDB_PROJECT" not in os.environ: + os.environ["WANDB_PROJECT"] = cfg.wandb_project + if "WANDB_RUN_ID" not in os.environ: + runtime_id = os.environ.get("RUNTIME_ID", None) + """If a RUNTIME_ID is available in the environment, we use it as the wandb id, + which will allow to display the evaluation results and the training results + in the same wandb run. Otherwise, we create a new run.""" + if runtime_id: + os.environ["WANDB_RUN_ID"] = runtime_id + os.environ["WANDB_DIR"] = training_args.output_dir + + # Create the experiment config directory. + output_dir = Path(training_args.output_dir) + exp_cfg_dir = output_dir / "experiment_cfg" + exp_cfg_dir.mkdir(parents=True, exist_ok=True) + OmegaConf.save(cfg, exp_cfg_dir / "conf.yaml", resolve=True) + + wandb_config_file = output_dir / "wandb_config.json" + with open(wandb_config_file, "w") as f: + json.dump( + { + "project": os.environ.get("WANDB_PROJECT", ""), + "run_id": os.environ.get("WANDB_RUN_ID", ""), + }, + f, + ) + + # Check if we are resuming training. + resume_path, continue_training = get_checkpoint_path(training_args.output_dir) + if not continue_training: + print(f"Models is ready under {training_args.output_dir}. Skip training.") + exit(0) + if resume_path: + print(f"Resuming training from {resume_path}") + resume_from_checkpoint = True + else: + # First time training. + resume_from_checkpoint = False + + # Instantiate the model. + model = self.create_model(cfg, training_args) + + if hasattr(model.action_head, "max_steps"): + model.action_head.max_steps = cfg.max_steps + + # Make sure model_dtype and training_args dtype are compatible. + compute_dtype = dtype_from_string(model.config.model_dtype) + + # Create the train dataset. + # Dump the metadata; necessary for policy to normalize the input and unnormalize the output + train_dataset = self.create_train_dataset(cfg, model) + print("Using dataset:") + print(train_dataset) + assert ( + train_dataset.merged_metadata is not None + ), "You must set metadata_config.merge=true in order to save the metadata." + + metadata_save_path = exp_cfg_dir / "metadata.json" + U.json_dump( + {k: v.model_dump(mode="json") for k, v in train_dataset.merged_metadata.items()}, + metadata_save_path, + indent=4, + ) + print("Successfully dumped metadata") + + val_dataset = self.create_val_dataset(cfg, model) + data_collator = self.create_data_collator(cfg, model) + trainer = self.create_trainer( + cfg=cfg, + exp_cfg_dir=exp_cfg_dir, + model=model, + training_args=training_args, + train_dataset=train_dataset, + val_dataset=val_dataset, + data_collator=data_collator, + compute_dtype=compute_dtype, + ) + self.cfg = cfg + self.exp_cfg_dir = exp_cfg_dir + self.training_args = training_args + self.resume_from_checkpoint = resume_from_checkpoint + self.train_dataset = train_dataset + self.trainer = trainer + + def create_model(self, cfg, training_args): + model = instantiate(cfg.model) + + if cfg.pretrained_model_path is not None: + mprint(f"Loading pretrained weights from: {cfg.pretrained_model_path}") + import json, gc + from safetensors.torch import load_file + + ckpt_dir = cfg.pretrained_model_path + safetensors_index_path = os.path.join(ckpt_dir, "model.safetensors.index.json") + safetensors_path = os.path.join(ckpt_dir, "model.safetensors") + + if os.path.exists(safetensors_index_path): + with open(safetensors_index_path, 'r') as f: + index = json.load(f) + for shard_file in sorted(set(index["weight_map"].values())): + shard_path = os.path.join(ckpt_dir, shard_file) + mprint(f"Loading shard: {shard_path}") + shard_state_dict = load_file(shard_path) + model.load_state_dict(shard_state_dict, strict=False) + del shard_state_dict + gc.collect() + elif os.path.exists(safetensors_path): + state_dict = load_file(safetensors_path) + model.load_state_dict(state_dict, strict=False) + else: + raise FileNotFoundError( + f"No weights found at '{ckpt_dir}'. " + "Expected 'model.safetensors' or 'model.safetensors.index.json'." + ) + + if (hasattr(model, 'action_head') + and hasattr(model.action_head, 'inject_lora_after_loading') + and model.action_head.config.defer_lora_injection): + model.action_head.inject_lora_after_loading() + + mprint("Successfully loaded pretrained weights") + + model.config.resume_path = model.config._name_or_path = training_args.output_dir + mprint(f"{model}\n") + return model + + def create_train_dataset(self, cfg, model): + assert torch.distributed.is_initialized() + train_dataset = instantiate(cfg.train_dataset) + return train_dataset + + def create_val_dataset(self, cfg, model): + return None + + def create_data_collator(self, cfg, model): + return instantiate(cfg.data_collator) + + def create_trainer( + self, + cfg, + exp_cfg_dir, + model, + training_args, + train_dataset, + val_dataset, + data_collator, + compute_dtype, + ): + # Set the gradient accumulation steps. + if cfg.global_batch_size is not None: + global_bs = cfg.global_batch_size + bs = training_args.per_device_train_batch_size + grad_acc = compute_grad_accum_to_match_global_bs(global_bs, bs) + training_args.gradient_accumulation_steps = grad_acc + print( + f"Set global batch size to {global_bs}, set gradient accumulation steps to {grad_acc}" + ) + elif cfg.raise_error_if_global_batch_size_not_set: + raise ValueError( + "global_batch_size is not set. To ensure the scripts can be reproduced regardless of the number of nodes used, please set this." + ) + else: + warnings.warn( + "global_batch_size is not set. This is fine for debugging, but please set this for real experiments." + ) + + # Instantiate the partial trainer. + import torch.distributed as dist + + world_size = dist.get_world_size() if dist.is_initialized() else 1 + data_root = cfg.get("trex_data_root") or cfg.get("wandb_video_dataset_root") + steps_per_data_epoch = estimate_steps_per_data_epoch( + data_root, + world_size=world_size, + per_device_batch_size=training_args.per_device_train_batch_size, + ) + if steps_per_data_epoch > 0: + mprint( + "Dataset epoch logging: " + f"steps_per_data_epoch={steps_per_data_epoch:.1f} " + f"(total_frames / global_batch)" + ) + + trainer_partial = instantiate( + cfg.trainer, + model=model, + output_dir=training_args.output_dir, + train_dataset=train_dataset, + eval_dataset=val_dataset, + compute_dtype=compute_dtype, + steps_per_data_epoch=steps_per_data_epoch, + ) + + # Fully instantiate the trainer with dataclasses instances. + trainer = trainer_partial(data_collator=data_collator, args=training_args) + trainer.base_cfg = cfg + train_dl_len = len(trainer.get_train_dataloader()) + eval_dl_len = ( + len(trainer.get_eval_dataloader()) if val_dataset is not None else "no eval dataloader" + ) + + # Save the total training steps in the config. + with open_dict(cfg): + cfg.total_training_steps = train_dl_len * cfg.training_args.num_train_epochs + + # Save config. + OmegaConf.save(cfg, exp_cfg_dir / "conf.yaml", resolve=True) + + run_name = cfg.training_args.get("run_name", None) + ckpt_format_callback = CheckpointFormatCallback(run_name=run_name, exp_cfg_dir=exp_cfg_dir) + trainer.add_callback(ckpt_format_callback) + + loss_log_path = str(Path(training_args.output_dir) / "loss_log.jsonl") + trainer.add_callback(LossLoggerCallback(output_path=loss_log_path)) + + if cfg.get("enable_wandb_video_reconstruction", False): + dataset_root = cfg.get("trex_data_root") or cfg.get("wandb_video_dataset_root") + if not dataset_root: + raise ValueError( + "enable_wandb_video_reconstruction=true requires trex_data_root " + "or wandb_video_dataset_root" + ) + metadata_path = exp_cfg_dir / "metadata.json" + if not metadata_path.exists(): + raise FileNotFoundError( + f"Missing {metadata_path}; cannot build eval transform for wandb video callback" + ) + trainer.add_callback( + TrexWandbVideoCallback( + dataset_root=str(dataset_root), + metadata_path=str(metadata_path), + transform_cfg=cfg.transforms["trex"], + output_dir=( + str(training_args.output_dir) + if bool(cfg.get("wandb_video_save_local", True)) + else None + ), + episode=int(cfg.get("wandb_video_reconstruction_episode", 0)), + num_chunks=int(cfg.get("wandb_video_reconstruction_num_chunks", 4)), + every_n_steps=int(cfg.get("wandb_video_reconstruction_steps", 100)), + fps=int(cfg.get("wandb_video_reconstruction_fps", 5)), + eval_bf16=bool(cfg.get("eval_bf16", False)), + use_dataset_prompt=bool(cfg.get("wandb_video_use_dataset_prompt", True)), + prompt=str(cfg.get("wandb_video_prompt", "perform the task")), + overlay_tracks=bool(cfg.get("wandb_video_overlay_tracks", False)), + save_tracks=bool(cfg.get("wandb_video_save_tracks", False)), + track_trail_steps=int(cfg.get("wandb_video_track_trail_steps", 8)), + reconstruction_inference_steps=int( + cfg.get("wandb_video_reconstruction_inference_steps", 1) + ), + start_chunk_index=int( + cfg.get("wandb_video_start_chunk_index", 0) + ), + compare_tracks_on_gt=bool( + cfg.get("wandb_video_compare_tracks_on_gt", True) + ), + ) + ) + mprint( + "Enabled TrexWandbVideoCallback: " + f"every {cfg.get('wandb_video_reconstruction_steps', 100)} steps, " + f"episode {cfg.get('wandb_video_reconstruction_episode', 0)}, " + f"track overlay={cfg.get('wandb_video_overlay_tracks', False)}, " + f"save local={cfg.get('wandb_video_save_local', True)}" + ) + + # Add profiling callback (local profiling only, no S3 upload) + # Local: {output_dir}/profiling/rank_{id}/*.pt.trace.json + if cfg.trainer.get("enable_prof_callback", False): + output_dir = Path(training_args.output_dir) + global_rank = int(os.environ.get("RANK", "0")) + + # Get profiling configuration from trainer config + profile_start_step = cfg.trainer.get("profile_start_step", 50) + profile_warmup_steps = cfg.trainer.get("profile_warmup_steps", 1) + profile_active_steps = cfg.trainer.get("profile_active_steps", 5) + profile_record_shapes = cfg.trainer.get("profile_record_shapes", False) + profile_with_stack = cfg.trainer.get( + "profile_with_stack", False + ) # Default False to match omni (stack traces add significant file size) + profile_memory = cfg.trainer.get("profile_memory", False) + + # Calculate end step + profile_end_step = profile_start_step + profile_warmup_steps + profile_active_steps - 1 + + # Setup profile directory with rank subdirectory: {output_dir}/profiling/rank_{id}/ + profile_dir = output_dir / "profiling" / f"rank_{global_rank}" + profile_dir.mkdir(parents=True, exist_ok=True) + + mprint( + f"Profiling enabled: steps {profile_start_step}-{profile_end_step}, " + f"saving to {profile_dir}" + ) + + # Add ProfCallback + trainer.add_callback( + ProfCallback( + profile_dir=profile_dir, + upload_callback=None, + profile_start_step=profile_start_step, + profile_end_step=profile_end_step, + warmup_steps=profile_warmup_steps, + active_steps=profile_active_steps, + trainer=trainer, + record_shapes=profile_record_shapes, + with_stack=profile_with_stack, + profile_memory=profile_memory, + ) + ) + + mprint( + f"train dataloader length: {train_dl_len}\n" + f"eval dataloader length: {eval_dl_len}\n" + f"train dataset length: {len(trainer.train_dataset)}\n" + f"GPU memory before training: {torch.cuda.memory_allocated() / 1024 / 1024 / 1024} GB", + flush=True, + ) + return trainer + + def train(self): + # Start training. + self.trainer.train(resume_from_checkpoint=self.resume_from_checkpoint) + self.trainer.save_state() + safe_save_model_for_hf_trainer( + trainer=self.trainer, output_dir=self.training_args.output_dir + ) diff --git a/groot/vla/experiment/experiment.py b/groot/vla/experiment/experiment.py new file mode 100644 index 0000000000000000000000000000000000000000..faa25f2f8e97548140fbea9cd454d94327ed27f9 --- /dev/null +++ b/groot/vla/experiment/experiment.py @@ -0,0 +1,163 @@ +import logging +import os +import time +from typing import Optional + +import hydra +import numpy as np +from omegaconf import DictConfig +import torch + +from groot.vla.experiment.base import BaseExperiment, BaseTrainer +from groot.vla.utils.action_args_override_utils import apply_action_overrides + +logger = logging.getLogger(__name__) + + +INITIAL_ACTIONS_FILENAME = "initial_actions.npz" + + +class ForceRestart(ValueError): + pass + + +class VLATrainer(BaseTrainer): + def __init__(self, **kwargs): + self.benchmark_time = kwargs.pop("benchmark_time", False) + self.step_timer = None + self.num_trials = kwargs.pop("num_trials", 10) + self.curr_trial = 0 + self.all_times = [] + self.start_time = time.time() + self.restart_max_seconds = kwargs.pop("restart_max_seconds", 0) + self.steps_per_data_epoch = kwargs.pop("steps_per_data_epoch", 0) + import torch.distributed as dist + + self.rank = dist.get_rank() + + self.micro_global_step = 0 + + super().__init__(**kwargs) + + def log(self, logs, start_time=None): + """Use dataset-sized epochs for IterableDataset + max_steps training.""" + from transformers.trainer_utils import speed_metrics + + if self.steps_per_data_epoch > 0: + logs["train/data_epoch"] = self.state.global_step / self.steps_per_data_epoch + + if self.state.epoch is not None: + logs["epoch"] = ( + logs["train/data_epoch"] + if self.steps_per_data_epoch > 0 + else self.state.epoch + ) + + if self.args.include_num_input_tokens_seen: + logs["num_input_tokens_seen"] = self.state.num_input_tokens_seen + if start_time is not None: + speed_metrics( + "train", + start_time, + num_tokens=self.state.num_input_tokens_seen, + ) + + output = {**logs, **{"step": self.state.global_step}} + self.state.log_history.append(output) + self.control = self.callback_handler.on_log(self.args, self.state, self.control, logs) + + def training_step(self, model, inputs, *args, **kwargs): + self.micro_global_step += 1 + + if hasattr(self.model.action_head, "global_step"): + self.model.action_head.global_step = self.state.global_step + + if self.benchmark_time: + if self.state.global_step % 100 == 0: + if self.step_timer is not None: + elapsed_time = time.time() - self.step_timer + self.all_times.append(elapsed_time) + self.curr_trial += 1 + self.step_timer = time.time() + if self.curr_trial >= self.num_trials: + exit(0) + if self.state.global_step % self.state.save_steps == 1: + if self.restart_max_seconds > 0: + cur_time = time.time() + if (cur_time - self.start_time) > self.restart_max_seconds: + raise ForceRestart(f"Exceeded time limit {self.restart_max_seconds} seconds") + loss_dict = super().training_step(model, inputs, *args, **kwargs) + return loss_dict + + +class VLATrainerInferenceBenchmark(VLATrainer): + + def compute_loss(self, model, inputs, return_outputs=False): + + warmup_steps = 100 + measure_steps = 100 + + model.eval() + + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + with torch.inference_mode(): + for i in range(warmup_steps): + action = model.module.get_action(inputs) + action.keys() + + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + torch.cuda.synchronize() + start_event.record() + + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + with torch.inference_mode(): + for i in range(measure_steps): + action = model.module.get_action(inputs) + action.keys() + + end_event.record() + torch.cuda.synchronize() + elapsed_time = start_event.elapsed_time(end_event) + + time_per_step = elapsed_time / measure_steps + exit() + + +class VLAExperiment(BaseExperiment): + + def __init__(self, cfg: DictConfig): + super().__init__(cfg) + # Dump the initial actions + if hasattr(self.train_dataset, "get_initial_actions"): + # We only dump the initial actions for the real robot dataset + # Sim dataset doesn't have this function + """ + initial_actions: list[dict[str, dict[str, np.ndarray]]] + 0: (the first dataset) + trajectory_name: + action_key: + action: np.ndarray + 1: (the second dataset) + ... + """ + initial_actions = self.train_dataset.get_initial_actions() + if len(initial_actions) > 0: + initial_actions_path = self.exp_cfg_dir / INITIAL_ACTIONS_FILENAME + np.savez(str(initial_actions_path), initial_actions) + print("Successfully dumped initial actions") + else: + print("No initial actions to dump") + + +@hydra.main(config_path="../configs", config_name="conf", version_base=None) +def main(cfg): + # Automatically update action dim and action horizon keys if specified in the config + cfg = apply_action_overrides(cfg) + + experiment = VLAExperiment(cfg) + experiment.train() + + +if __name__ == "__main__": + main() diff --git a/groot/vla/experiment/trex_eval_utils.py b/groot/vla/experiment/trex_eval_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c02ce08d01c7b5d763004d12ad75e589cb2b02e9 --- /dev/null +++ b/groot/vla/experiment/trex_eval_utils.py @@ -0,0 +1,251 @@ +"""Shared T-Rex causal video inference helpers (training callback + offline scripts).""" + +from __future__ import annotations + +import glob +import os +from typing import TYPE_CHECKING + +import cv2 +import numpy as np +import pyarrow.parquet as pq +import torch +from einops import rearrange +from groot.vla.model.n1_5.sim_policy import unsqueeze_dict_values + +if TYPE_CHECKING: + from groot.vla.data.transform import ComposedModalityTransform + +VIDEO_KEYS = [ + "video.head_left", + "video.left_wrist", + "video.right_wrist", +] +STATE_KEYS = { + "state.left_arm": (0, 7), + "state.left_hand": (7, 29), + "state.right_arm": (29, 36), + "state.right_hand": (36, 58), +} +VIDEO_FOLDERS = { + "video.head_left": "observation.images.head_left", + "video.left_wrist": "observation.images.left_wrist", + "video.right_wrist": "observation.images.right_wrist", +} + +RELATIVE_OFFSETS = [-23, -16, -8, 0] +ACTION_HORIZON = 24 + + +def unwrap_model(model: torch.nn.Module) -> torch.nn.Module: + return model.module if hasattr(model, "module") else model + + +def get_expected_video_resolution(model: torch.nn.Module) -> tuple[int, int]: + cfg = model.action_head.config + target_h = getattr(cfg, "target_video_height", None) + target_w = getattr(cfg, "target_video_width", None) + if target_h is not None and target_w is not None: + return int(target_h), int(target_w) + return 160, 320 + + +def resize_frames(frames: np.ndarray, target_h: int, target_w: int) -> np.ndarray: + if frames.ndim == 3: + if (frames.shape[0], frames.shape[1]) == (target_h, target_w): + return frames + return cv2.resize(frames, (target_w, target_h), interpolation=cv2.INTER_LINEAR) + return np.stack( + [cv2.resize(f, (target_w, target_h), interpolation=cv2.INTER_LINEAR) for f in frames], + axis=0, + ) + + +class TrexEpisode: + """One T-Rex episode from LeRobot v2 layout.""" + + def __init__(self, dataset_root: str, episode_index: int): + pq_path = os.path.join( + dataset_root, + "data", + f"chunk-{episode_index // 1000:03d}", + f"episode_{episode_index:06d}.parquet", + ) + if not os.path.isfile(pq_path): + raise FileNotFoundError(pq_path) + self.table = pq.read_table(pq_path) + self.length = self.table.num_rows + self.episode_index = episode_index + self.root = dataset_root + self._frame_cache: dict[str, dict[int, np.ndarray]] = {} + + self.video_dirs: dict[str, str] = {} + for key, folder in VIDEO_FOLDERS.items(): + pattern = os.path.join( + dataset_root, + "videos", + "**", + folder, + f"episode_{episode_index:06d}.mp4", + ) + hits = sorted(glob.glob(pattern, recursive=True)) + if not hits: + raise FileNotFoundError(f"No video for {key}: {pattern}") + self.video_dirs[key] = hits[0] + self._frame_cache[key] = {} + + def get_task(self, row: int) -> str: + try: + return str(self.table.column("annotation.task")[row].as_py()) + except Exception: + return "" + + def get_state(self, row: int) -> np.ndarray: + return np.array(self.table.column("observation.state")[row].as_py(), dtype=np.float64) + + def get_frames( + self, rows: list[int] | np.ndarray, video_key: str + ) -> np.ndarray: + indices = np.asarray(rows, dtype=np.int64).reshape(-1) + if video_key not in self.video_dirs: + raise KeyError(video_key) + if indices.size == 0: + raise ValueError("rows must not be empty") + if indices.min() < 0 or indices.max() >= self.length: + raise IndexError("video frame index outside episode") + cache = self._frame_cache[video_key] + missing = sorted(set(indices.tolist()) - set(cache)) + if missing: + cap = cv2.VideoCapture(self.video_dirs[video_key]) + if not cap.isOpened(): + raise RuntimeError(f"Failed to open {self.video_dirs[video_key]}") + previous = -2 + try: + for row in missing: + if row != previous + 1: + cap.set(cv2.CAP_PROP_POS_FRAMES, row) + ok, frame = cap.read() + if not ok: + raise RuntimeError( + f"Failed frame {row} from {self.video_dirs[video_key]}" + ) + cache[row] = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + previous = row + finally: + cap.release() + return np.stack([cache[int(row)] for row in indices]) + + def get_frame(self, row: int, video_key: str) -> np.ndarray: + return self.get_frames([row], video_key)[0] + + +def build_chunk_schedule(total_frames: int, num_chunks: int) -> list[list[int]]: + chunks: list[list[int]] = [] + anchor = 0 + for i in range(num_chunks): + if i == 0: + indices = [0] + else: + indices = [max(anchor + off, 0) for off in RELATIVE_OFFSETS] + if indices[-1] >= total_frames: + break + chunks.append(indices) + anchor += ACTION_HORIZON + return chunks + + +def build_obs( + episode: TrexEpisode, + frame_indices: list[int], + prompt: str, + video_height: int, + video_width: int, + *, + pre_resize: bool = True, +) -> dict: + obs: dict = {} + anchor = frame_indices[-1] + state = episode.get_state(anchor) + + for key in VIDEO_KEYS: + frames = episode.get_frames(frame_indices, key) + frames = frames.astype(np.uint8) + if pre_resize: + frames = resize_frames(frames, video_height, video_width) + if len(frame_indices) == 1: + obs[key] = frames[0:1] # (1, H, W, 3); transforms require a time dimension + else: + obs[key] = frames # (T, H, W, 3) + + for key, (start, end) in STATE_KEYS.items(): + obs[key] = state[start:end].reshape(1, -1).astype(np.float64) + + obs["annotation.task"] = prompt + return obs + + +def reset_action_head_causal_state(action_head) -> None: + action_head.current_start_frame = 0 + action_head.kv_cache1 = None + action_head.kv_cache_neg = None + if hasattr(action_head, "last_language"): + action_head.last_language = None + + +def decode_video_latents(model: torch.nn.Module, video_chunks: list[torch.Tensor]) -> np.ndarray: + if not video_chunks: + raise ValueError("No video chunks to decode") + cat = torch.cat(video_chunks, dim=2) + ah = unwrap_model(model).action_head + frames = ah.vae.decode( + cat, + tiled=ah.tiled, + tile_size=(ah.tile_size_height, ah.tile_size_width), + tile_stride=(ah.tile_stride_height, ah.tile_stride_width), + ) + frames = rearrange(frames, "B C T H W -> B T H W C")[0] + return ((frames.float() + 1) * 127.5).clip(0, 255).cpu().numpy().astype(np.uint8) + + +def run_trex_causal_video_prediction( + model: torch.nn.Module, + eval_transform: ComposedModalityTransform, + episode: TrexEpisode, + *, + prompt: str, + num_chunks: int, + eval_bf16: bool = True, + pre_resize_video: bool = False, +) -> np.ndarray: + """Run causal video reconstruction and return decoded uint8 frames (T, H, W, 3).""" + unwrapped = unwrap_model(model) + video_height, video_width = get_expected_video_resolution(unwrapped) + schedule = build_chunk_schedule(episode.length, num_chunks) + if not schedule: + raise RuntimeError(f"Episode {episode.episode_index} is too short for causal inference") + + reset_action_head_causal_state(unwrapped.action_head) + video_chunks: list[torch.Tensor] = [] + + for frame_indices in schedule: + obs = build_obs( + episode, + frame_indices, + prompt, + video_height, + video_width, + pre_resize=pre_resize_video, + ) + # Match GrootSimPolicy inference: add batch dim before eval_transform so + # DreamTransform.apply_batch splits state/video consistently. + obs = unsqueeze_dict_values(obs) + normalized_input = eval_transform(obs) + for key, value in list(normalized_input.items()): + if torch.is_tensor(value) and value.dtype == torch.float32 and eval_bf16: + normalized_input[key] = value.to(dtype=torch.bfloat16) + + with torch.inference_mode(): + model_pred = unwrapped.lazy_joint_video_action_causal(normalized_input) + video_chunks.append(model_pred["video_pred"]) + + return decode_video_latents(model, video_chunks) diff --git a/groot/vla/experiment/trex_track_force_eval_utils.py b/groot/vla/experiment/trex_track_force_eval_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..6d2042b4684305b2d5e780fa6908e9136a6ac094 --- /dev/null +++ b/groot/vla/experiment/trex_track_force_eval_utils.py @@ -0,0 +1,726 @@ +"""Periodic T-Rex Track-Force video and track visualization helpers.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import cv2 +from einops import rearrange +import numpy as np +import torch + +from groot.vla.data.transform import ComposedModalityTransform +from groot.vla.experiment.trex_eval_utils import TrexEpisode, unwrap_model +from groot.vla.model.n1_5.sim_policy import unsqueeze_dict_values +from groot.vla.model.trex_track_force.dataset import ( + nearest_timestamp_indices, + uniform_target_times, +) +from groot.vla.model.trex_track_force.runtime import TrexRuntimeStatistics +from groot.vla.model.trex_track_force.track import ( + NUM_TRACK_POINTS, + TRACK_HORIZON, + TRACK_SEGMENTS, + TrackView, +) + +ACTION_RATE_HZ = 20.0 +VIDEO_RATE_HZ = 10.0 +TACTILE_RATE_HZ = 5.0 +ACTION_HORIZON = 16 +VIDEO_HISTORY_FRAMES = 9 +VIDEO_FUTURE_FRAMES = 8 +FORCE_HISTORY_FRAMES = 16 +VIDEO_KEYS = ( + "video.head_left", + "video.left_wrist", + "video.right_wrist", +) + +_SEGMENT_RGB = ( + (0, 255, 0), + (255, 128, 0), + (255, 255, 0), + (255, 0, 255), + (255, 255, 0), + (0, 170, 255), +) + + +@dataclass(frozen=True) +class TrexTrackForceVisualization: + video_frames: np.ndarray + predicted_track_xy: np.ndarray + target_track_xy: np.ndarray + target_track_visibility: np.ndarray + anchor_timestamps: np.ndarray + frame_track_indices: np.ndarray + conditioning_frame: np.ndarray | None = None + conditioning_track_xy: np.ndarray | None = None + conditioning_track_visibility: np.ndarray | None = None + gt_video_frames: np.ndarray | None = None + + +def _column_array( + episode: TrexEpisode, column: str, *, dtype: np.dtype +) -> np.ndarray: + if column not in episode.table.column_names: + raise KeyError(f"episode is missing required column {column!r}") + values = np.asarray(episode.table.column(column).to_pylist(), dtype=dtype) + if not np.isfinite(values).all(): + raise ValueError(f"episode column {column!r} contains NaN/Inf") + return values + + +def _sample_indices( + timestamps: np.ndarray, + anchor_timestamp: float, + offsets: range, + rate_hz: float, +): + targets = uniform_target_times(anchor_timestamp, offsets, rate_hz) + return nearest_timestamp_indices(timestamps, targets) + + +def _read_video_frames( + episode: TrexEpisode, video_key: str, indices: np.ndarray +) -> np.ndarray: + return episode.get_frames(indices, video_key).astype(np.uint8) + + +def _decode_future_video( + model: torch.nn.Module, + conditioning_latents: torch.Tensor, + future_latents: torch.Tensor, +) -> np.ndarray: + action_head = unwrap_model(model).action_head + config = action_head.config + tiled = bool(getattr(config, "tiled", False)) + tile_size = ( + int(getattr(config, "tile_size_height", 34)), + int(getattr(config, "tile_size_width", 34)), + ) + tile_stride = ( + int(getattr(config, "tile_stride_height", 18)), + int(getattr(config, "tile_stride_width", 16)), + ) + latent_sequence = torch.cat( + (conditioning_latents[:, :, -1:], future_latents), + dim=2, + ) + decoded = action_head.vae.decode( + latent_sequence, + tiled=tiled, + tile_size=tile_size, + tile_stride=tile_stride, + ) + frames = rearrange(decoded, "B C T H W -> B T H W C")[0] + frames = ( + (frames.float() + 1.0) * 127.5 + ).clamp(0, 255).cpu().numpy().astype(np.uint8) + if frames.shape[0] < VIDEO_FUTURE_FRAMES: + raise ValueError( + f"VAE decoded {frames.shape[0]} frames; need {VIDEO_FUTURE_FRAMES}" + ) + return frames[-VIDEO_FUTURE_FRAMES:] + + +def _build_chunk_input( + *, + episode: TrexEpisode, + timestamps: np.ndarray, + state_eef62: np.ndarray, + track_xy: np.ndarray, + track_visibility: np.ndarray, + force: np.ndarray, + statistics: TrexRuntimeStatistics, + anchor_timestamp: float, + prompt: str, +) -> tuple[dict, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + state_selection = _sample_indices( + timestamps, anchor_timestamp, range(1), ACTION_RATE_HZ + ) + video_history = _sample_indices( + timestamps, + anchor_timestamp, + range(-(VIDEO_HISTORY_FRAMES - 1), 1), + VIDEO_RATE_HZ, + ) + track_history = _sample_indices( + timestamps, + anchor_timestamp, + range(-(FORCE_HISTORY_FRAMES - 1), 1), + ACTION_RATE_HZ, + ) + track_future = _sample_indices( + timestamps, + anchor_timestamp, + range(TRACK_HORIZON), + ACTION_RATE_HZ, + ) + force_history = _sample_indices( + timestamps, + anchor_timestamp, + range(-(FORCE_HISTORY_FRAMES - 1), 1), + TACTILE_RATE_HZ, + ) + if track_future.padding_mask.any(): + raise ValueError("evaluation chunk does not have a complete future track") + + past_visibility = track_visibility[track_history.indices].copy() + past_visibility[track_history.padding_mask] = 0.0 + normalized_force = statistics.normalize_force( + force[force_history.indices].reshape(-1, 10, 6) + ) + normalized_force[force_history.padding_mask] = 0.0 + + raw: dict[str, object] = { + key: _read_video_frames(episode, key, video_history.indices) + for key in VIDEO_KEYS + } + raw.update( + { + "state.eef62": state_eef62[state_selection.indices], + "track_past_xy": track_xy[track_history.indices][None], + "track_past_visibility": past_visibility[None], + "track_future_xy": track_xy[track_future.indices][None], + "track_future_visibility": track_visibility[ + track_future.indices + ][None], + "current_force": normalized_force[-1:], + "force_history": normalized_force[None], + "force_history_padding_mask": force_history.padding_mask[None], + "annotation.task": prompt, + } + ) + return ( + unsqueeze_dict_values(raw), + track_xy[track_future.indices], + track_visibility[track_future.indices], + track_xy[track_history.indices[-1]], + track_visibility[track_history.indices[-1]], + ) + + +def _read_gt_future_video_strip( + episode: TrexEpisode, + timestamps: np.ndarray, + anchor_timestamp: float, + *, + height: int, + width: int, +) -> np.ndarray: + """Load GT future RGB frames and pack them into the model 2x2 canvas.""" + + video_future = _sample_indices( + timestamps, + anchor_timestamp, + range(1, VIDEO_FUTURE_FRAMES + 1), + VIDEO_RATE_HZ, + ) + if video_future.padding_mask.any(): + raise ValueError("evaluation chunk does not have complete GT future video") + panels = [ + _read_video_frames(episode, key, video_future.indices) for key in VIDEO_KEYS + ] + resized = [] + for panel in panels: + frames = np.empty((panel.shape[0], height, width, 3), dtype=np.uint8) + for index, frame in enumerate(panel): + frames[index] = cv2.resize( + frame, (width, height), interpolation=cv2.INTER_LINEAR + ) + resized.append(frames) + # Match the model canvas: [head, left_wrist; right_wrist, black]. + black = np.zeros_like(resized[0]) + top = np.concatenate((resized[0], resized[2]), axis=2) + bottom = np.concatenate((resized[1], black), axis=2) + return np.concatenate((top, bottom), axis=1) + + +def run_trex_track_force_prediction( + model: torch.nn.Module, + eval_transform: ComposedModalityTransform, + episode: TrexEpisode, + *, + prompt: str, + num_chunks: int, + eval_bf16: bool = True, + coarse_steps: int = 1, + seed: int = 42, + start_chunk_index: int = 0, +) -> TrexTrackForceVisualization: + """Predict fixed GT-conditioned chunks and return video plus raw tracks.""" + + if num_chunks <= 0: + raise ValueError("num_chunks must be positive") + if coarse_steps <= 0: + raise ValueError("coarse_steps must be positive") + if start_chunk_index < 0: + raise ValueError("start_chunk_index must be non-negative") + unwrapped = unwrap_model(model) + statistics = TrexRuntimeStatistics.from_dataset(episode.root) + timestamps = _column_array(episode, "timestamp", dtype=np.float64) + state_eef62 = _column_array( + episode, "observation.state_eef62", dtype=np.float32 + ) + track_xy = _column_array( + episode, "observation.track_xy", dtype=np.float32 + ) + track_visibility = _column_array( + episode, "observation.track_visibility", dtype=np.float32 + ) + force = _column_array( + episode, "observation.tactile_force", dtype=np.float32 + ) + + generated_chunks: list[np.ndarray] = [] + gt_chunks: list[np.ndarray] = [] + predicted_tracks: list[np.ndarray] = [] + target_tracks: list[np.ndarray] = [] + target_visibility: list[np.ndarray] = [] + anchor_timestamps: list[float] = [] + frame_track_indices: list[int] = [] + conditioning_tracks: list[np.ndarray] = [] + conditioning_visibility: list[np.ndarray] = [] + conditioning_frame: np.ndarray | None = None + canvas_height = canvas_width = None + + for local_index in range(num_chunks): + chunk_index = start_chunk_index + local_index + anchor_timestamp = float( + timestamps[0] + chunk_index * ACTION_HORIZON / ACTION_RATE_HZ + ) + if anchor_timestamp + TRACK_HORIZON / ACTION_RATE_HZ > timestamps[-1] + 1e-9: + break + ( + raw_input, + target_xy, + target_vis, + conditioning_xy, + conditioning_vis, + ) = _build_chunk_input( + episode=episode, + timestamps=timestamps, + state_eef62=state_eef62, + track_xy=track_xy, + track_visibility=track_visibility, + force=force, + statistics=statistics, + anchor_timestamp=anchor_timestamp, + prompt=prompt, + ) + normalized_input = eval_transform(raw_input) + if conditioning_frame is None: + history_images = normalized_input["history_images"] + if torch.is_tensor(history_images): + history_images = history_images.detach().cpu().numpy() + conditioning_frame = np.asarray(history_images[0, -1], dtype=np.uint8) + for key, value in list(normalized_input.items()): + if ( + torch.is_tensor(value) + and value.dtype == torch.float32 + and eval_bf16 + ): + normalized_input[key] = value.to(dtype=torch.bfloat16) + + with torch.inference_mode(): + prediction = unwrapped.get_visualization( + normalized_input, + coarse_steps=coarse_steps, + seed=seed + chunk_index, + ) + if "conditioning_video_latents" not in prediction: + raise KeyError("sampler did not return conditioning video latents") + generated = _decode_future_video( + unwrapped, + prediction["conditioning_video_latents"], + prediction["video_latents_pred"], + ) + generated_chunks.append(generated) + if canvas_height is None: + canvas_height, canvas_width = generated.shape[1:3] + panel_height = canvas_height // 2 + panel_width = canvas_width // 2 + gt_chunks.append( + _read_gt_future_video_strip( + episode, + timestamps, + anchor_timestamp, + height=panel_height, + width=panel_width, + ) + ) + predicted = ( + prediction["track_pred"][0] + .transpose(0, 1) + .float() + .cpu() + .numpy() + ) + # The model predicts anchor-relative motion and restores absolute XY. + # Only overwrite t=0 to remove BF16 round-off; never translate the + # entire predicted trajectory after inference. + predicted[0] = conditioning_xy + predicted_tracks.append(predicted) + target_tracks.append(target_xy) + target_visibility.append(target_vis) + conditioning_tracks.append(conditioning_xy) + conditioning_visibility.append(conditioning_vis) + anchor_timestamps.append(anchor_timestamp) + track_base = local_index * TRACK_HORIZON + frame_track_indices.extend( + track_base + min(2 * (index + 1), TRACK_HORIZON - 1) + for index in range(VIDEO_FUTURE_FRAMES) + ) + + if not generated_chunks: + raise RuntimeError( + f"episode {episode.episode_index} has no complete visualization chunk" + ) + generated_frames = np.concatenate(generated_chunks, axis=0) + gt_frames = np.concatenate(gt_chunks, axis=0) + assert conditioning_frame is not None + conditioning_frame = cv2.resize( + conditioning_frame, + (generated_frames.shape[2], generated_frames.shape[1]), + interpolation=cv2.INTER_LINEAR, + ) + return TrexTrackForceVisualization( + video_frames=np.concatenate( + (conditioning_frame[None], generated_frames), axis=0 + ), + predicted_track_xy=np.stack(predicted_tracks), + target_track_xy=np.stack(target_tracks), + target_track_visibility=np.stack(target_visibility), + anchor_timestamps=np.asarray(anchor_timestamps, dtype=np.float64), + frame_track_indices=np.asarray( + [-1, *frame_track_indices], dtype=np.int64 + ), + conditioning_frame=conditioning_frame, + conditioning_track_xy=np.stack(conditioning_tracks), + conditioning_track_visibility=np.stack(conditioning_visibility), + gt_video_frames=np.concatenate( + (conditioning_frame[None], gt_frames), axis=0 + ), + ) + + +def _three_view_strip(frames: np.ndarray) -> np.ndarray: + """Extract the three valid cells from the model's 2x2 canvas.""" + + height, width = frames.shape[1:3] + if height % 2 or width % 2: + raise ValueError("2x2 canvas dimensions must be even") + panel_height, panel_width = height // 2, width // 2 + return np.concatenate( + ( + frames[:, :panel_height, :panel_width], + frames[:, panel_height:, :panel_width], + frames[:, :panel_height, panel_width:], + ), + axis=2, + ) + + +def _project_tracks_to_three_view_strip( + tracks: np.ndarray, frame_height: int, frame_width: int +) -> np.ndarray: + projected = np.asarray(tracks, dtype=np.float32).copy() + panel_height = float(frame_height) + panel_width = frame_width / 3.0 + for segment in TRACK_SEGMENTS: + values = projected[:, segment.start : segment.stop] + values[..., 0] = np.clip(values[..., 0], 0.0, 1.0) * panel_width + values[..., 1] = np.clip(values[..., 1], 0.0, 1.0) * panel_height + if segment.view == TrackView.LEFT_WRIST: + values[..., 0] += panel_width + elif segment.view == TrackView.RIGHT_WRIST: + values[..., 0] += 2.0 * panel_width + return projected + + +def overlay_track_motion( + visualization: TrexTrackForceVisualization, + *, + trail_steps: int = 8, +) -> np.ndarray: + """Overlay predicted 20 Hz track motion on generated 10 Hz grid video.""" + + frames = np.asarray(visualization.video_frames, dtype=np.uint8).copy() + if frames.ndim != 4 or frames.shape[-1] != 3: + raise ValueError("video frames must be [T,H,W,3]") + if frames.shape[0] != visualization.frame_track_indices.shape[0]: + raise ValueError("video and frame-track index counts do not match") + frames = _three_view_strip(frames) + tracks = visualization.predicted_track_xy.reshape( + -1, NUM_TRACK_POINTS, 2 + ) + visibility = np.asarray( + visualization.target_track_visibility + ).reshape(-1, NUM_TRACK_POINTS) + if visibility.shape[:2] != tracks.shape[:2]: + raise ValueError("target track visibility does not match predicted tracks") + projected = _project_tracks_to_three_view_strip( + tracks, frames.shape[1], frames.shape[2] + ) + conditioning_projected = None + if visualization.conditioning_track_xy is not None: + conditioning_projected = _project_tracks_to_three_view_strip( + visualization.conditioning_track_xy[:1], + frames.shape[1], + frames.shape[2], + )[0] + trail_steps = max(int(trail_steps), 0) + + for frame_index, track_index_value in enumerate( + visualization.frame_track_indices + ): + track_index = int(track_index_value) + if track_index < 0: + if conditioning_projected is None: + raise ValueError( + "GT conditioning frame requires conditioning track xy" + ) + conditioning_visibility = ( + visualization.conditioning_track_visibility + ) + visible = ( + np.ones(NUM_TRACK_POINTS, dtype=bool) + if conditioning_visibility is None + else np.asarray(conditioning_visibility[0]) > 0 + ) + frame = np.ascontiguousarray(frames[frame_index]) + for segment, color in zip(TRACK_SEGMENTS, _SEGMENT_RGB): + for point_index in range(segment.start, segment.stop): + if not visible[point_index]: + continue + point = conditioning_projected[point_index] + cv2.circle( + frame, + tuple(np.rint(point).astype(int)), + 2, + color, + -1, + cv2.LINE_AA, + ) + cv2.putText( + frame, + "GT conditioning tracks", + (8, 16), + cv2.FONT_HERSHEY_SIMPLEX, + 0.42, + (255, 255, 255), + 1, + cv2.LINE_AA, + ) + frames[frame_index] = frame + continue + chunk_start = (track_index // TRACK_HORIZON) * TRACK_HORIZON + trail_start = max(chunk_start, track_index - trail_steps) + frame = np.ascontiguousarray(frames[frame_index]) + for segment, color in zip(TRACK_SEGMENTS, _SEGMENT_RGB): + for point_index in range(segment.start, segment.stop): + path = projected[ + trail_start : track_index + 1, point_index + ] + path_visible = ( + visibility[trail_start : track_index + 1, point_index] > 0 + ) + valid = np.isfinite(path).all(axis=-1) & path_visible + valid_indices = np.flatnonzero(valid) + if valid_indices.size >= 2: + split_at = np.flatnonzero(np.diff(valid_indices) > 1) + 1 + for segment_indices in np.split(valid_indices, split_at): + if segment_indices.size < 2: + continue + cv2.polylines( + frame, + [ + np.rint(path[segment_indices]).astype( + np.int32 + ) + ], + False, + color, + 1, + cv2.LINE_AA, + ) + point = projected[track_index, point_index] + if ( + visibility[track_index, point_index] > 0 + and np.isfinite(point).all() + ): + cv2.circle( + frame, + tuple(np.rint(point).astype(int)), + 2, + color, + -1, + cv2.LINE_AA, + ) + cv2.putText( + frame, + "predicted track motion", + (8, 16), + cv2.FONT_HERSHEY_SIMPLEX, + 0.42, + (255, 255, 255), + 1, + cv2.LINE_AA, + ) + frames[frame_index] = frame + return frames + + +def _draw_track_paths( + frame: np.ndarray, + projected: np.ndarray, + visibility: np.ndarray, + *, + track_index: int, + trail_steps: int, + colors: tuple[tuple[int, int, int], ...], + radius: int = 2, +) -> None: + chunk_start = (track_index // TRACK_HORIZON) * TRACK_HORIZON + trail_start = max(chunk_start, track_index - trail_steps) + for segment_index, segment in enumerate(TRACK_SEGMENTS): + color = colors[segment_index % len(colors)] + for point_index in range(segment.start, segment.stop): + path = projected[trail_start : track_index + 1, point_index] + path_visible = visibility[trail_start : track_index + 1, point_index] > 0 + valid = np.isfinite(path).all(axis=-1) & path_visible + valid_indices = np.flatnonzero(valid) + if valid_indices.size >= 2: + split_at = np.flatnonzero(np.diff(valid_indices) > 1) + 1 + for segment_indices in np.split(valid_indices, split_at): + if segment_indices.size < 2: + continue + cv2.polylines( + frame, + [np.rint(path[segment_indices]).astype(np.int32)], + False, + color, + 1, + cv2.LINE_AA, + ) + point = projected[track_index, point_index] + if visibility[track_index, point_index] > 0 and np.isfinite(point).all(): + cv2.circle( + frame, + tuple(np.rint(point).astype(int)), + radius, + color, + -1, + cv2.LINE_AA, + ) + + +def overlay_track_gt_vs_pred( + visualization: TrexTrackForceVisualization, + *, + trail_steps: int = 2, +) -> np.ndarray: + """Side-by-side GT vs predicted tracks drawn on GT video (not generated video).""" + + if visualization.gt_video_frames is None: + raise ValueError("visualization is missing gt_video_frames") + frames = np.asarray(visualization.gt_video_frames, dtype=np.uint8).copy() + if frames.ndim != 4 or frames.shape[-1] != 3: + raise ValueError("gt video frames must be [T,H,W,3]") + if frames.shape[0] != visualization.frame_track_indices.shape[0]: + raise ValueError("gt video and frame-track index counts do not match") + frames = _three_view_strip(frames) + pred_tracks = visualization.predicted_track_xy.reshape(-1, NUM_TRACK_POINTS, 2) + gt_tracks = visualization.target_track_xy.reshape(-1, NUM_TRACK_POINTS, 2) + visibility = np.asarray(visualization.target_track_visibility).reshape( + -1, NUM_TRACK_POINTS + ) + pred_projected = _project_tracks_to_three_view_strip( + pred_tracks, frames.shape[1], frames.shape[2] + ) + gt_projected = _project_tracks_to_three_view_strip( + gt_tracks, frames.shape[1], frames.shape[2] + ) + trail_steps = max(int(trail_steps), 0) + gt_colors = tuple((220, 220, 220) for _ in TRACK_SEGMENTS) + output_frames = [] + for frame_index, track_index_value in enumerate( + visualization.frame_track_indices + ): + track_index = int(track_index_value) + base = np.ascontiguousarray(frames[frame_index]) + if track_index < 0: + left = base.copy() + right = base.copy() + cv2.putText( + left, + "GT tracks", + (8, 16), + cv2.FONT_HERSHEY_SIMPLEX, + 0.42, + (255, 255, 255), + 1, + cv2.LINE_AA, + ) + cv2.putText( + right, + "predicted tracks", + (8, 16), + cv2.FONT_HERSHEY_SIMPLEX, + 0.42, + (255, 255, 255), + 1, + cv2.LINE_AA, + ) + output_frames.append(np.concatenate((left, right), axis=1)) + continue + left = base.copy() + right = base.copy() + _draw_track_paths( + left, + gt_projected, + visibility, + track_index=track_index, + trail_steps=trail_steps, + colors=gt_colors, + ) + _draw_track_paths( + right, + pred_projected, + visibility, + track_index=track_index, + trail_steps=trail_steps, + colors=_SEGMENT_RGB, + ) + cv2.putText( + left, + "GT on GT video", + (8, 16), + cv2.FONT_HERSHEY_SIMPLEX, + 0.42, + (255, 255, 255), + 1, + cv2.LINE_AA, + ) + cv2.putText( + right, + "pred on GT video", + (8, 16), + cv2.FONT_HERSHEY_SIMPLEX, + 0.42, + (255, 255, 255), + 1, + cv2.LINE_AA, + ) + output_frames.append(np.concatenate((left, right), axis=1)) + return np.stack(output_frames, axis=0) + + +__all__ = [ + "TrexTrackForceVisualization", + "overlay_track_gt_vs_pred", + "overlay_track_motion", + "run_trex_track_force_prediction", +] diff --git a/groot/vla/experiment/trex_wandb_video_callback.py b/groot/vla/experiment/trex_wandb_video_callback.py new file mode 100644 index 0000000000000000000000000000000000000000..ff81f21525d7d8b8cfdaa90272c5234ac3258004 --- /dev/null +++ b/groot/vla/experiment/trex_wandb_video_callback.py @@ -0,0 +1,363 @@ +"""Log T-Rex causal video reconstructions to Weights & Biases during training.""" + +from __future__ import annotations + +import json +import logging +import tempfile +import time +from pathlib import Path + +import imageio +import numpy as np +import torch +import torch.distributed as dist +from hydra.utils import instantiate +from transformers import TrainerCallback + +from groot.vla.data.schema import DatasetMetadata, EmbodimentTag +from groot.vla.data.transform import ComposedModalityTransform +from groot.vla.experiment.trex_eval_utils import TrexEpisode, run_trex_causal_video_prediction +from groot.vla.experiment.trex_track_force_eval_utils import ( + overlay_track_gt_vs_pred, + overlay_track_motion, + run_trex_track_force_prediction, +) +from groot.vla.experiment.utils import mprint + +logger = logging.getLogger(__name__) + + +class TrexWandbVideoCallback(TrainerCallback): + """Periodically run causal video inference on a fixed episode and log to wandb.""" + + TRAIN_STEP_METRIC = "train/global_step" + VIDEO_METRIC = "eval/predicted_video" + TRACK_VIDEO_METRIC = "eval/predicted_track_video" + TRACK_COMPARE_METRIC = "eval/track_gt_vs_pred" + + def __init__( + self, + *, + dataset_root: str, + metadata_path: str, + transform_cfg, + output_dir: str | None = None, + episode: int = 0, + num_chunks: int = 4, + every_n_steps: int = 100, + fps: int = 5, + eval_bf16: bool = True, + use_dataset_prompt: bool = True, + prompt: str = "perform the task", + overlay_tracks: bool = False, + save_tracks: bool = False, + track_trail_steps: int = 8, + reconstruction_inference_steps: int = 1, + start_chunk_index: int = 0, + compare_tracks_on_gt: bool = True, + ) -> None: + if reconstruction_inference_steps <= 0: + raise ValueError("reconstruction_inference_steps must be positive") + if start_chunk_index < 0: + raise ValueError("start_chunk_index must be non-negative") + self.dataset_root = dataset_root + self.episode_index = episode + self.num_chunks = num_chunks + self.every_n_steps = every_n_steps + self.fps = fps + self.eval_bf16 = eval_bf16 + self.use_dataset_prompt = use_dataset_prompt + self.prompt = prompt + self.overlay_tracks = overlay_tracks + self.save_tracks = save_tracks + self.track_trail_steps = track_trail_steps + self.reconstruction_inference_steps = reconstruction_inference_steps + self.start_chunk_index = start_chunk_index + self.compare_tracks_on_gt = compare_tracks_on_gt + self.eval_video_dir = ( + Path(output_dir) / "eval_videos" if output_dir is not None else None + ) + self.eval_track_video_dir = ( + Path(output_dir) / "eval_track_videos" + if output_dir is not None and overlay_tracks + else None + ) + self.eval_track_dir = ( + Path(output_dir) / "eval_tracks" + if output_dir is not None and save_tracks + else None + ) + self._episode: TrexEpisode | None = None + self._eval_transform = self._build_eval_transform(metadata_path, transform_cfg) + + @staticmethod + def _build_eval_transform(metadata_path: str, transform_cfg) -> ComposedModalityTransform: + with open(metadata_path, "r", encoding="utf-8") as f: + metadatas = json.load(f) + metadata = DatasetMetadata.model_validate(metadatas[EmbodimentTag.TREX.value]) + eval_transform = instantiate(transform_cfg) + if not isinstance(eval_transform, ComposedModalityTransform): + raise TypeError(f"Expected ComposedModalityTransform, got {type(eval_transform)}") + eval_transform.set_metadata(metadata) + eval_transform.eval() + return eval_transform + + def _get_episode(self) -> TrexEpisode: + if self._episode is None: + self._episode = TrexEpisode(self.dataset_root, self.episode_index) + return self._episode + + @staticmethod + def _distributed_barrier() -> None: + if dist.is_available() and dist.is_initialized(): + dist.barrier() + + @staticmethod + def _setup_wandb_train_step_axis() -> None: + import wandb + + if wandb.run is None: + return + wandb.define_metric(TrexWandbVideoCallback.TRAIN_STEP_METRIC) + wandb.define_metric( + TrexWandbVideoCallback.VIDEO_METRIC, + step_metric=TrexWandbVideoCallback.TRAIN_STEP_METRIC, + ) + wandb.define_metric( + TrexWandbVideoCallback.TRACK_VIDEO_METRIC, + step_metric=TrexWandbVideoCallback.TRAIN_STEP_METRIC, + ) + wandb.define_metric( + TrexWandbVideoCallback.TRACK_COMPARE_METRIC, + step_metric=TrexWandbVideoCallback.TRAIN_STEP_METRIC, + ) + + def on_train_begin(self, args, state, control, **kwargs): + if state.is_world_process_zero: + self._setup_wandb_train_step_axis() + if self.eval_video_dir is not None: + self.eval_video_dir.mkdir(parents=True, exist_ok=True) + if self.eval_track_video_dir is not None: + self.eval_track_video_dir.mkdir(parents=True, exist_ok=True) + if self.eval_track_dir is not None: + self.eval_track_dir.mkdir(parents=True, exist_ok=True) + + def on_step_end(self, args, state, control, model=None, **kwargs): + if state.global_step <= 0 or state.global_step % self.every_n_steps != 0: + return + + # Keep all ranks in sync while rank 0 runs eval inference. + self._distributed_barrier() + + if not state.is_world_process_zero: + self._distributed_barrier() + return + + if model is None: + mprint("TrexWandbVideoCallback skipped: model is None") + self._distributed_barrier() + return + + try: + import wandb + except ImportError: + mprint("wandb is not installed; skipping video reconstruction callback") + self._distributed_barrier() + return + + if wandb.run is None: + mprint("wandb.run is None; skipping video reconstruction callback") + self._distributed_barrier() + return + + self._setup_wandb_train_step_axis() + + was_training = model.training + model.eval() + try: + callback_start = time.perf_counter() + # Avoid torch.compile recompiles during mid-training eval. + torch._dynamo.config.disable = True + + episode = self._get_episode() + prompt = episode.get_task(0) if self.use_dataset_prompt else self.prompt + mprint( + "Starting T-Rex training visualization: " + f"chunks={self.num_chunks}, " + f"start_chunk={self.start_chunk_index}, " + f"coarse_steps={self.reconstruction_inference_steps}" + ) + with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=self.eval_bf16): + track_visualization = None + compare_frames = None + if self.overlay_tracks: + track_visualization = run_trex_track_force_prediction( + model, + self._eval_transform, + episode, + prompt=prompt, + num_chunks=self.num_chunks, + eval_bf16=self.eval_bf16, + coarse_steps=self.reconstruction_inference_steps, + start_chunk_index=self.start_chunk_index, + ) + pred_frames = track_visualization.video_frames + track_frames = overlay_track_motion( + track_visualization, + trail_steps=self.track_trail_steps, + ) + if self.compare_tracks_on_gt: + compare_frames = overlay_track_gt_vs_pred( + track_visualization, + trail_steps=self.track_trail_steps, + ) + else: + track_frames = None + pred_frames = run_trex_causal_video_prediction( + model, + self._eval_transform, + episode, + prompt=prompt, + num_chunks=self.num_chunks, + eval_bf16=self.eval_bf16, + # Feed native dataset resolution; eval_transform resizes to 160x320. + pre_resize_video=False, + ) + train_step = int(state.global_step) + track_path = None + if track_visualization is not None and self.eval_track_dir is not None: + track_path = self.eval_track_dir / f"train_step_{train_step:06d}.npz" + np.savez_compressed( + track_path, + predicted_track_xy=track_visualization.predicted_track_xy, + target_track_xy=track_visualization.target_track_xy, + target_track_visibility=track_visualization.target_track_visibility, + anchor_timestamps=track_visualization.anchor_timestamps, + frame_track_indices=track_visualization.frame_track_indices, + conditioning_track_xy=track_visualization.conditioning_track_xy, + conditioning_track_visibility=( + track_visualization.conditioning_track_visibility + ), + ) + if self.eval_video_dir is not None: + video_path = self.eval_video_dir / f"train_step_{train_step:06d}.mp4" + else: + with tempfile.NamedTemporaryFile( + suffix=f"_train_step_{train_step:06d}.mp4", + delete=False, + ) as tmp: + video_path = Path(tmp.name) + imageio.mimsave( + str(video_path), + list(pred_frames), + fps=self.fps, + format="FFMPEG", + codec="libx264", + pixelformat="yuv420p", + ) + track_video_path = None + if track_frames is not None: + if self.eval_track_video_dir is not None: + track_video_path = ( + self.eval_track_video_dir + / f"train_step_{train_step:06d}.mp4" + ) + else: + with tempfile.NamedTemporaryFile( + suffix=f"_track_train_step_{train_step:06d}.mp4", + delete=False, + ) as tmp: + track_video_path = Path(tmp.name) + imageio.mimsave( + str(track_video_path), + list(track_frames), + fps=self.fps, + format="FFMPEG", + codec="libx264", + pixelformat="yuv420p", + ) + compare_video_path = None + if compare_frames is not None: + if self.eval_track_video_dir is not None: + compare_video_path = ( + self.eval_track_video_dir + / f"train_step_{train_step:06d}_gt_vs_pred.mp4" + ) + else: + with tempfile.NamedTemporaryFile( + suffix=f"_track_compare_train_step_{train_step:06d}.mp4", + delete=False, + ) as tmp: + compare_video_path = Path(tmp.name) + imageio.mimsave( + str(compare_video_path), + list(compare_frames), + fps=self.fps, + format="FFMPEG", + codec="libx264", + pixelformat="yuv420p", + ) + + mprint( + "Finished visualization inference/local encoding in " + f"{time.perf_counter() - callback_start:.1f}s; uploading to wandb" + ) + wandb_payload = { + self.TRAIN_STEP_METRIC: train_step, + self.VIDEO_METRIC: wandb.Video( + str(video_path), + format="mp4", + caption=( + f"train_step={train_step}, episode={self.episode_index}, " + f"prompt={prompt!r}" + ), + ), + } + if track_video_path is not None: + wandb_payload[self.TRACK_VIDEO_METRIC] = wandb.Video( + str(track_video_path), + format="mp4", + caption=( + f"train_step={train_step}, episode={self.episode_index}, " + f"prompt={prompt!r}, predicted track motion" + ), + ) + if compare_video_path is not None: + wandb_payload[self.TRACK_COMPARE_METRIC] = wandb.Video( + str(compare_video_path), + format="mp4", + caption=( + f"train_step={train_step}, episode={self.episode_index}, " + f"start_chunk={self.start_chunk_index}, " + f"GT vs pred tracks on GT video" + ), + ) + wandb.log(wandb_payload) + if self.eval_video_dir is None: + video_path.unlink(missing_ok=True) + if ( + track_video_path is not None + and self.eval_track_video_dir is None + ): + track_video_path.unlink(missing_ok=True) + if ( + compare_video_path is not None + and self.eval_track_video_dir is None + ): + compare_video_path.unlink(missing_ok=True) + mprint( + f"Logged T-Rex predicted video at train_step={train_step} " + f"({len(pred_frames)} frames, wandb key={self.VIDEO_METRIC}, " + f"x-axis={self.TRAIN_STEP_METRIC}, local={video_path}, " + f"track_video={track_video_path}, compare={compare_video_path}, " + f"tracks={track_path}, " + f"elapsed={time.perf_counter() - callback_start:.1f}s)" + ) + except Exception as exc: + mprint(f"TrexWandbVideoCallback failed at step {state.global_step}: {exc}") + logger.exception("TrexWandbVideoCallback failed at step %s", state.global_step) + finally: + if was_training: + model.train() + self._distributed_barrier() diff --git a/groot/vla/experiment/utils.py b/groot/vla/experiment/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f49b6084a439ad5c59be6a5ae9fbde01a5d6fb19 --- /dev/null +++ b/groot/vla/experiment/utils.py @@ -0,0 +1,187 @@ +""" +Originally trinity.train.utils +""" + +from dataclasses import dataclass +import json +import os +import pathlib +from pathlib import Path +import re +import shutil + +import torch +import torch.nn as nn +from transformers import PretrainedConfig, Trainer + + +def dtype_from_string(dtype_str): + if dtype_str == "bfloat16": + return torch.bfloat16 + elif dtype_str == "float16": + return torch.float16 + elif dtype_str == "float32": + return torch.float32 + else: + raise ValueError(f"Unsupported dtype_str {dtype_str}") + + +def rprint(*args, **kwargs): + rank = int(os.environ.get("RANK", 0)) + world_size = int(os.environ.get("WORLD_SIZE", 1)) + if world_size > 1: + return print(f"[dist-{rank}-of-{world_size}]", *args, **kwargs) + else: + return print(*args, **kwargs) + + +def mprint(*args, **kwargs): + rank = int(os.environ.get("RANK", 0)) + world_size = int(os.environ.get("WORLD_SIZE", 1)) + if world_size > 1: + if rank == 0: + return print(f"[dist-{rank}-of-{world_size}]", *args, **kwargs) + else: + return + else: + return print(*args, **kwargs) + + +def is_local(model_name_or_path: str) -> bool: + return os.path.isdir(model_name_or_path) + + +def get_checkpoint_path(output_dir: str, checkpoint_prefix: str = "checkpoint") -> str | None: + output_dir = os.path.abspath(output_dir) + pathlib_dir = pathlib.Path(output_dir) + + if list(pathlib_dir.glob("config.json")): + # training has been finished + return output_dir, False + else: + try: + ordering_and_checkpoint_path = [] + glob_checkpoints = [ + str(x) + for x in pathlib.Path(output_dir).glob(f"{checkpoint_prefix}-*") + if os.path.isdir(x) + ] + for path in glob_checkpoints: + regex_match = re.match(f".*{checkpoint_prefix}-([0-9]+)", path) + if regex_match is not None and regex_match.groups() is not None: + ordering_and_checkpoint_path.append((int(regex_match.groups()[0]), path)) + checkpoints_sorted = sorted(ordering_and_checkpoint_path) + return checkpoints_sorted[-1][1], True + except IndexError: + return None, True + + +def prepare_config_for_training( + config: PretrainedConfig, model_args: dataclass, training_args: dataclass, data_args: dataclass +) -> None: + ## set default dtype + # config.model_dtype = "bfloat16" if training_args.bf16 else "float16" + + ## set tuning modules + config.tune_language_model = training_args.tune_language_model + config.tune_vision_tower = training_args.tune_vision_tower + config.tune_mm_projector = training_args.tune_mm_projector + + +def safe_save_model_for_hf_trainer(trainer: Trainer, output_dir: str): + """Collects the state dict and dump to disk.""" + if trainer.deepspeed: + torch.cuda.synchronize() + trainer.save_model(output_dir, _internal_call=True) + return + + state_dict = trainer.model.state_dict() + if trainer.args.should_save: + cpu_state_dict = {key: value.cpu() for key, value in state_dict.items()} + del state_dict + trainer._save(output_dir, state_dict=cpu_state_dict) # noqa + + +def compute_grad_accum_to_match_global_bs(global_bs: int, bs: int): + num_devices = torch.distributed.get_world_size() + per_step_bs = bs * num_devices + assert global_bs % per_step_bs == 0, f"{global_bs=}, {per_step_bs=}" + num_grad_accum = global_bs // per_step_bs + return num_grad_accum + + +def estimate_steps_per_data_epoch( + data_root: str | os.PathLike[str] | None, + *, + world_size: int, + per_device_batch_size: int, +) -> float: + """Approximate optimizer steps needed for one pass over all dataset frames.""" + if not data_root: + return 0.0 + info_path = Path(data_root) / "meta" / "info.json" + if not info_path.exists(): + return 0.0 + with open(info_path, encoding="utf-8") as f: + info = json.load(f) + total_frames = int(info.get("total_frames", 0)) + if total_frames <= 0: + return 0.0 + samples_per_step = max(per_device_batch_size * world_size, 1) + return total_frames / samples_per_step + + +def get_training_param_info(model): + module_states = dict() + for module_name, module in model.named_children(): + key = f"{module_name}({module.__class__.__name__})" + if all([p.requires_grad for p in module.parameters()]): + module_states[key] = "true" + elif all([not p.requires_grad for p in module.parameters()]): + module_states[key] = "false" + else: + module_states[key] = get_training_param_info(module) + + return module_states + + +def get_param_count_tree(model: nn.Module): + """ + Calculate parameters for the model, structure them as a nested dictionary, + and save the result as a formatted JSON file. + """ + + def format_param_count(count: int) -> str: + """Format the count as a string in millions, e.g. 11M or 5.5M.""" + count_in_millions = count / 1e6 + # If the value is an integer, display without decimal places + if count_in_millions.is_integer(): + return f"{int(count_in_millions)}M" + else: + return f"{count_in_millions:.2f}M" + + def module_to_dict(module: nn.Module, module_name: str) -> dict: + """ + Recursively convert a module and its children into a nested dictionary. + The key is formatted as "module_name (ClassName, param_count)". + """ + + total_count = sum(p.numel() for p in module.parameters()) + formatted_total = format_param_count(total_count) + key = f"{module_name} ({module.__class__.__name__}, {formatted_total})" + + # Get immediate children modules + children = list(module.named_children()) + if children: + nested = {} + for child_name, child_module in children: + # Recursively convert child modules + nested.update(module_to_dict(child_module, child_name)) + return {key: nested} + else: + # Leaf module: return empty dict as value + return {key: {}} + + # Start from the top-level module (you can change the name "model" as needed) + nested_dict = module_to_dict(model, "model") + return nested_dict diff --git a/groot/vla/model/__init__.py b/groot/vla/model/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/model/dreamzero/__init__.py b/groot/vla/model/dreamzero/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/model/dreamzero/action_head/__init__.py b/groot/vla/model/dreamzero/action_head/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/model/dreamzero/action_head/wan_flow_matching_action_tf.py b/groot/vla/model/dreamzero/action_head/wan_flow_matching_action_tf.py new file mode 100644 index 0000000000000000000000000000000000000000..38a4099989fb601378011fd7589855119bfb6a4e --- /dev/null +++ b/groot/vla/model/dreamzero/action_head/wan_flow_matching_action_tf.py @@ -0,0 +1,1402 @@ +from dataclasses import dataclass, field +import logging +import time +from typing import TypeAlias, cast +import os + +from accelerate import load_checkpoint_and_dispatch + +from einops import rearrange +from hydra.utils import instantiate +from peft import LoraConfig, get_peft_model +import torch +from torch import nn +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh +from safetensors.torch import load_file +import json +from huggingface_hub import hf_hub_download + + +logger = logging.getLogger(__name__) + +WAN_HF_REPO_ID = "Wan-AI/Wan2.1-I2V-14B-480P" +WAN22_HF_REPO_ID = "Wan-AI/Wan2.2-TI2V-5B" + + +def hf_download(filename: str, repo_id: str = WAN_HF_REPO_ID) -> str: + """Download a file from the specified HuggingFace repo to HF cache.""" + path = hf_hub_download(repo_id=repo_id, filename=filename) + return path + + +def ensure_file(path: str | None, hf_filename: str, repo_id: str = WAN_HF_REPO_ID) -> str: + """Return a valid local path: use `path` if it exists, otherwise download from HuggingFace.""" + if path is not None and os.path.exists(path): + return path + return hf_download(hf_filename, repo_id) + +from torch.distributions import Beta +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh +from torchvision.transforms import v2 +from transformers import PretrainedConfig +from transformers.feature_extraction_utils import BatchFeature + +from groot.vla.model.n1_5.action_head.base_action_head import ActionHead +from groot.vla.model.dreamzero.modules.flow_match_scheduler import FlowMatchScheduler +from groot.vla.model.dreamzero.modules.vram_management import enable_vram_management, AutoWrappedModule, AutoWrappedLinear +from groot.vla.model.dreamzero.modules.wan_video_text_encoder import T5RelativeEmbedding, T5LayerNorm +from groot.vla.model.dreamzero.modules.flow_unipc_multistep_scheduler import FlowUniPCMultistepScheduler + + +KVCacheType: TypeAlias = torch.Tensor + +@dataclass +class WANPolicyHeadConfig(PretrainedConfig): + add_pos_embed: bool = field( + default=True, metadata={"help": "Whether to add positional embedding"} + ) + model_dtype: str = field(default="float32", metadata={"help": "Model data type."}) + 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."} + ) + backbone_embedding_dim: int = field( + default=1536, metadata={"help": "Backbone embedding channel dimension."} + ) + tiled: bool = field(default=True, metadata={"help": "Whether to use tiled input."}) + tile_size_height: int = field(default=34, metadata={"help": "Tile size height."}) + tile_size_width: int = field(default=34, metadata={"help": "Tile size width."}) + tile_stride_height: int = field(default=18, metadata={"help": "Tile stride height."}) + tile_stride_width: int = field(default=16, metadata={"help": "Tile stride width."}) + num_frame_per_block: int = field(default=1, metadata={"help": "Number of frames per block."}) + # Target video (H, W) for Wan22 resize. When set, videos are resized to this before VAE so latent + # spatial size matches. Use height/width divisible by 32 for WanVideoVAE38 (16x) so latent H,W are even. + target_video_height: int | None = field(default=None, metadata={"help": "Target video height for resize (e.g. 160 for even latent with VAE38)."}) + target_video_width: int | None = field(default=None, metadata={"help": "Target video width for resize (e.g. 320)."}) + + lora_rank: int = field(default=4, metadata={"help": "LoRA rank."}) + lora_alpha: int = field(default=4, metadata={"help": "LoRA alpha."}) + lora_target_modules: str = field(default="q,k,v,o,ffn.0,ffn.2") + init_lora_weights: str = field(default="kaiming", metadata={"help": "LoRA initialization method."}) + train_architecture: str= field(default="lora", metadata={"help": "Train architecture."}) + skip_component_loading: bool = field(default=False, metadata={"help": "Skip loading individual component weights (used when loading from full pretrained model)."}) + + use_gradient_checkpointing: bool = field(default=True, metadata={"help": "Whether to use gradient checkpointing."}) + qformer_cfg: dict = field(default=None, metadata={"help": "Qformer configuration."}) + 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."} + ) + # High noise emphasis for BASE (coupled) training - applies Beta distribution to BOTH video and action together + use_high_noise_emphasis: bool = field( + default=False, metadata={"help": "Use Beta distribution for noise sampling (biases BOTH video and action towards high noise levels together)."} + ) + high_noise_beta_alpha: float = field( + default=3.0, metadata={"help": "Beta alpha for high noise emphasis. Beta(3,1): mean=0.75, Beta(5,1): mean=0.83. Higher = more high noise bias."} + ) + # Decoupled noise sampling config for training-inference alignment + # When enabled: video uses Beta(alpha,beta) biased towards high noise, action uses independent uniform + decouple_video_action_noise: bool = field( + default=False, metadata={"help": "Decouple video/action noise: video uses Beta distribution (high noise bias), action uses independent uniform."} + ) + video_noise_beta_alpha: float = field( + default=3.0, metadata={"help": "Beta alpha for video noise. Beta(3,1): mean=0.75, Beta(5,1): mean=0.83. Higher alpha = more bias to high noise."} + ) + video_noise_beta_beta: float = field( + default=1.0, metadata={"help": "Beta beta for video noise. Keep at 1.0."} + ) + # Decoupled inference config - allows video to stay noisy while action fully denoises + decouple_inference_noise: bool = field( + default=False, metadata={"help": "Use decoupled noise schedules during inference (video stays noisy, action fully denoises)."} + ) + video_inference_final_noise: float = field( + default=0.8, metadata={"help": "Final noise level for video during decoupled inference (0.0-1.0). E.g., 0.8 means video ends at 80% noise."} + ) + 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) + defer_lora_injection: bool = field(default=False, metadata={"help": "Defer LoRA injection until after loading pretrained weights."}) + + vl_self_attention_cfg: dict = field(default=None) + text_encoder_cfg: dict = field(default=None) + image_encoder_cfg: dict = field(default=None) + vae_cfg: dict = field(default=None) + + def __init__(self, **kwargs): + super().__init__(**kwargs) + for key, value in kwargs.items(): + setattr(self, key, value) + + +class WANPolicyHead(ActionHead): + config_class = WANPolicyHeadConfig + supports_gradient_checkpointing = True + + def __init__( + self, + config: WANPolicyHeadConfig, + ): + super().__init__() + self.tiled = config.tiled + self.tile_size_height = config.tile_size_height + self.tile_size_width = config.tile_size_width + self.tile_stride_height = config.tile_stride_height + self.tile_stride_width = config.tile_stride_width + self.num_frame_per_block = config.num_frame_per_block + self.hidden_size = config.hidden_size + self.num_frames = config.num_frames + self.text_encoder = instantiate(config.text_encoder_cfg) + self.image_encoder = instantiate(config.image_encoder_cfg) + self.vae = instantiate(config.vae_cfg) + self.scheduler = FlowMatchScheduler(shift=5, sigma_min=0.0, extra_one_step=True) + self.model_names = ['text_encoder'] + + self.num_inference_steps = 16 + self.seed = 1140 + self.cfg_scale = 5.0 + self.denoising_strength = 1.0 + self.sigma_shift = 5.0 + self.kv_cache1: KVCacheType | None = None + self.kv_cache_neg: KVCacheType | None = None + self.crossattn_cache: KVCacheType | None = None + self.crossattn_cache_neg: KVCacheType | None = None + + self.global_step = 0 + self.max_steps = 0 + self.lora_rank = config.lora_rank + self.lora_alpha = config.lora_alpha + self.lora_target_modules = config.lora_target_modules + self.init_lora_weights = config.init_lora_weights + self.train_architecture = config.train_architecture + self.clip_feas = None + self.ys = None + self.current_start_frame = 0 + self.language = None + + self.ip_rank = 0 + self.ip_size = 1 + self.ip_group = None + self.trt_engine = None + + self._device = "cuda" + self.dynamic_cache_schedule = os.getenv("DYNAMIC_CACHE_SCHEDULE", "False").lower() == "true" + + + num_dit_steps = 8 + if os.getenv("NUM_DIT_STEPS") is not None: + num_dit_steps = int(os.getenv("NUM_DIT_STEPS")) + if num_dit_steps == 5: + self.dit_step_mask = [True, True, True, False, False, False, False, True, False, False, False, False, True, False, False, False] + elif num_dit_steps == 6: + self.dit_step_mask = [True, True, False, False, False, True, False, False, False, False, True, False, False, False, True, True] + elif num_dit_steps == 7: + self.dit_step_mask = [True, True, True, False, False, False, True, False, False, False, True, False, False, False, True, True] + elif num_dit_steps == 8: + self.dit_step_mask = [True, True, True, False, False, False, True, False, False, False, True, False, False, True, True, True] + else: + self.dit_step_mask = [True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True] + assert self.dit_step_mask[0] == True, "first step must be True" + + self.normalize_video = v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) + + + self.use_gradient_checkpointing = config.use_gradient_checkpointing + if self.training: + self.scheduler.set_timesteps(1000, training=True) + + + self.input_embedding_dim = config.input_embedding_dim + + self.cpu_offload = False + + self.model = instantiate(config.diffusion_model_cfg) + self.action_dim = config.action_dim + self.action_horizon = config.action_horizon + self.num_inference_timesteps = config.num_inference_timesteps + + text_enc_path = ensure_file( + self.text_encoder.text_encoder_pretrained_path, + "models_t5_umt5-xxl-enc-bf16.pth", + ) + self.text_encoder.load_state_dict(torch.load(text_enc_path, map_location='cpu')) + + img_enc_path = ensure_file( + self.image_encoder.image_encoder_pretrained_path, + "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth", + ) + self.image_encoder.model.load_state_dict(torch.load(img_enc_path, map_location='cpu'), strict=False) + + # Wan2.2 (WanVideoVAE38, z_dim=48) uses Wan2.2_VAE.pth; Wan2.1 uses Wan2.1_VAE.pth + vae_hf_filename = "Wan2.2_VAE.pth" if getattr(self.vae, "z_dim", 16) == 48 else "Wan2.1_VAE.pth" + vae_repo_id = WAN22_HF_REPO_ID if getattr(self.vae, "z_dim", 16) == 48 else WAN_HF_REPO_ID + vae_path = ensure_file( + self.vae.vae_pretrained_path, + vae_hf_filename, + repo_id=vae_repo_id, + ) + self.vae.model.load_state_dict(torch.load(vae_path, map_location='cpu')) + + if not config.skip_component_loading: + dit_dir = self.model.diffusion_model_pretrained_path + # Wan2.2 (in_dim=48) uses Wan2.2-TI2V-5B repo; Wan2.1 uses Wan2.1-I2V-14B-480P + dit_repo_id = WAN22_HF_REPO_ID if getattr(self.model, "in_dim", 16) == 48 else WAN_HF_REPO_ID + if dit_dir is None or not os.path.isdir(dit_dir): + index_path = hf_hub_download(repo_id=dit_repo_id, filename="diffusion_pytorch_model.safetensors.index.json") + dit_dir = os.path.dirname(index_path) + with open(index_path, 'r') as f: + index = json.load(f) + for shard_file in set(index["weight_map"].values()): + hf_hub_download(repo_id=dit_repo_id, filename=shard_file) + + if dit_dir is not None: + safetensors_path = os.path.join(dit_dir, "diffusion_pytorch_model.safetensors") + safetensors_index_path = os.path.join(dit_dir, "diffusion_pytorch_model.safetensors.index.json") + state_dict = {} + + if os.path.exists(safetensors_index_path): + # Handle sharded safetensors + print(f"Loading sharded safetensors using index: {safetensors_index_path}") + + with open(safetensors_index_path, 'r') as f: + index = json.load(f) + + # Load each shard + for shard_file in set(index["weight_map"].values()): + shard_path = os.path.join(dit_dir, shard_file) + print(f"Loading shard: {shard_path}") + shard_state_dict = load_file(shard_path) + state_dict.update(shard_state_dict) + + elif os.path.exists(safetensors_path): + # Handle single safetensors file + print(f"Loading weights from safetensors: {safetensors_path}") + state_dict = load_file(safetensors_path) + + else: + raise ValueError(f"No safetensors file found at {safetensors_path} or {safetensors_index_path}") + + missing_keys, unexpected_keys = self.model.load_state_dict(state_dict, strict=False) + + if missing_keys: + print(f"Missing keys when loading pretrained weights: {missing_keys}") + if unexpected_keys: + print(f"Unexpected keys when loading pretrained weights: {unexpected_keys}") + + print("Successfully loaded pretrained weights") + else: + print("Skipping individual component loading (loading from full pretrained model)") + self.beta_dist = Beta(config.noise_beta_alpha, config.noise_beta_beta) + # Video noise Beta distribution (biased towards high noise levels when enabled) + self.video_beta_dist = Beta(config.video_noise_beta_alpha, config.video_noise_beta_beta) + # High noise emphasis Beta distribution for coupled training (applies to both video and action) + self.high_noise_beta_dist = Beta(config.high_noise_beta_alpha, 1.0) + # self.num_timestep_buckets = config.num_timestep_buckets + self.config = config + self._noise_logged = False + self.defer_lora_injection = config.defer_lora_injection + print("defer_lora_injection@@", self.defer_lora_injection) + self.set_trainable_parameters(config.tune_projector, config.tune_diffusion_model) + + def set_trainable_parameters(self, tune_projector: bool, tune_diffusion_model: bool): + self.tune_projector = tune_projector + self.tune_diffusion_model = tune_diffusion_model + for p in self.parameters(): + p.requires_grad = True + if not tune_diffusion_model: + self.model.requires_grad_(False) + print(f"Tune action head projector: {self.tune_projector}") + print(f"Tune action head diffusion model: {self.tune_diffusion_model}") + # Check if any parameters are still trainable. If not, print a warning. + if not tune_projector and not tune_diffusion_model: + for name, p in self.named_parameters(): + if p.requires_grad: + print(f"Action head trainable parameter: {name}") + if not any(p.requires_grad for p in self.parameters()): + print("Warning: No action head trainable parameters found.") + + if self.train_architecture == "lora" and not self.defer_lora_injection: + print("Adding LoRA to model") + for p in self.parameters(): + p.requires_grad = False + self.model = self.add_lora_to_model( + self.model, + lora_rank=self.lora_rank, + lora_alpha=self.lora_alpha, + lora_target_modules=self.lora_target_modules, + init_lora_weights=self.init_lora_weights, + ) + self.model.state_encoder.requires_grad_(True) + self.model.action_encoder.requires_grad_(True) + self.model.action_decoder.requires_grad_(True) + elif self.train_architecture == "lora" and self.defer_lora_injection: + print("Deferring LoRA injection until after pretrained weights are loaded") + else: + self.print_trainable_params() + + self.text_encoder.requires_grad_(False) + self.image_encoder.requires_grad_(False) + self.vae.requires_grad_(False) + if not self.defer_lora_injection: + self.print_trainable_params() + + + def print_trainable_params(self): + """Print trainable parameters of the diffusion model.""" + trainable_params = [] + total_params = 0 + trainable_total = 0 + + for name, param in self.model.named_parameters(): + total_params += param.numel() + if param.requires_grad: + trainable_params.append(name) + trainable_total += param.numel() + + print(f"Total parameters in diffusion model: {total_params:,}") + print(f"Trainable parameters in diffusion model: {trainable_total:,}") + # print(trainable_params) + + + def inject_lora_after_loading(self): + """ + Inject LoRA adapters after pretrained weights have been loaded. + This should be called when defer_lora_injection=True. + """ + if self.train_architecture == "lora": + print("Injecting LoRA after loading pretrained weights") + for p in self.parameters(): + p.requires_grad = False + self.model = self.add_lora_to_model( + self.model, + lora_rank=self.lora_rank, + lora_alpha=self.lora_alpha, + lora_target_modules=self.lora_target_modules, + init_lora_weights=self.init_lora_weights, + ) + self.model.state_encoder.requires_grad_(True) + self.model.action_encoder.requires_grad_(True) + self.model.action_decoder.requires_grad_(True) + # self.model.registers.requires_grad_(True) + # self.model.time_modality_projection.requires_grad_(True) + + self.text_encoder.requires_grad_(False) + self.image_encoder.requires_grad_(False) + self.vae.requires_grad_(False) + self.print_trainable_params() + else: + print("LoRA injection not needed (train_architecture != 'lora')") + + def set_frozen_modules_to_eval_mode(self): + """ + Huggingface will call model.train() at each training_step. To ensure + the expected behaviors for modules like dropout, batchnorm, etc., we + need to call model.eval() for the frozen modules. + """ + if self.training: + if not self.tune_diffusion_model: + self.model.eval() + self.text_encoder.eval() + self.image_encoder.eval() + self.vae.eval() + + + def enable_vram_management(self, num_persistent_param_in_dit=None): + dtype = next(iter(self.text_encoder.parameters())).dtype + enable_vram_management( + self.text_encoder, + module_map = { + torch.nn.Linear: AutoWrappedLinear, + torch.nn.Embedding: AutoWrappedModule, + T5RelativeEmbedding: AutoWrappedModule, + T5LayerNorm: AutoWrappedModule, + }, + module_config = dict( + offload_dtype=dtype, + offload_device="cpu", + onload_dtype=dtype, + onload_device="cpu", + computation_dtype=self.dtype, + computation_device='cuda', + ), + ) + + self.cpu_offload = True + + def load_models_to_device(self, loadmodel_names=[]): + # only load models to device if cpu_offload is enabled + if not self.cpu_offload: + return + # offload the unneeded models to cpu + for model_name in self.model_names: + if model_name not in loadmodel_names: + model = getattr(self, model_name) + if model is not None: + if hasattr(model, "vram_management_enabled") and model.vram_management_enabled: + print("offloadd") + for module in model.modules(): + if hasattr(module, "offload"): + # print("offload", module) + module.offload() + else: + print("tocpu") + model.cpu() + # load the needed models to device + for model_name in loadmodel_names: + model = getattr(self, model_name) + if model is not None: + if hasattr(model, "vram_management_enabled") and model.vram_management_enabled: + print("onload") + for module in model.modules(): + if hasattr(module, "onload"): + # print("onload", module) + module.onload() + else: + print("togpu") + model.to(self._device) + # fresh the cuda cache + torch.cuda.empty_cache() + + def _create_kv_caches( + self, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + frame_seqlen: int, + ) -> tuple[KVCacheType, KVCacheType]: + """ + Initialize a Per-GPU KV cache for the Wan model. + Use the model's num_heads and head_dim (5B has 24 heads, 14B has 40). + """ + num_heads = self.model.num_heads + head_dim = self.model.dim // num_heads + kv_cache1: KVCacheType = [] + kv_cache_neg: KVCacheType = [] + for _ in range(self.model.num_layers): + kv_cache1.append( + torch.zeros([2, batch_size, 0, num_heads, head_dim], dtype=dtype, device=device), + ) + kv_cache_neg.append( + torch.zeros([2, batch_size, 0, num_heads, head_dim], dtype=dtype, device=device), + ) + + return kv_cache1, kv_cache_neg + + def _create_crossattn_caches( + self, batch_size: int, dtype: torch.dtype, device: torch.device, + ) -> tuple[KVCacheType, KVCacheType]: + """ + Initialize a Per-GPU cross-attention cache for the Wan model. + Use the model's num_heads and head_dim (5B has 24 heads, 14B has 40). + """ + num_heads = self.model.num_heads + head_dim = self.model.dim // num_heads + crossattn_cache: KVCacheType = [] + crossattn_cache_neg: KVCacheType = [] + + for _ in range(self.model.num_layers): + crossattn_cache.append( + torch.zeros([2, batch_size, 512, num_heads, head_dim], dtype=dtype, device=device), + ) + crossattn_cache_neg.append( + torch.zeros([2, batch_size, 512, num_heads, head_dim], dtype=dtype, device=device), + ) + + return crossattn_cache, crossattn_cache_neg + + def sample_time(self, batch_size, device, dtype): + sample = self.beta_dist.sample([batch_size]).to(device, dtype=dtype) + return (self.config.noise_s - sample) / self.config.noise_s + + def prepare_input(self, batch: dict) -> BatchFeature: + return BatchFeature(data=batch) + + def preprocess_image(self, image): + image = (image * (2 / 255) - 1).permute(0, 1, 4, 2, 3) + return image + + def encode_prompt(self, input_ids, attention_mask): + seq_lens = attention_mask.gt(0).sum(dim=1).long() + prompt_emb = self.text_encoder(input_ids, attention_mask) + prompt_emb = prompt_emb.clone().to(dtype=torch.bfloat16) + for i, v in enumerate(seq_lens): + prompt_emb[:, v:] = 0 + return prompt_emb + + def _ensure_vae_on_device(self, ref_tensor): + """Lazily move the VAE to the correct device/dtype on first use.""" + if not getattr(self, '_vae_device_ready', False): + self.vae.to(device=ref_tensor.device, dtype=torch.bfloat16) + self.vae.eval() + self._vae_device_ready = True + + def encode_video(self, input_video, tiled=True, tile_size=(34, 34), tile_stride=(18, 16)): + self._ensure_vae_on_device(input_video) + with torch.no_grad(): + latents = self.vae.encode(input_video, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) + return latents + + def encode_image(self, image, num_frames, height, width): + with torch.amp.autocast(dtype=torch.bfloat16, device_type=torch.device(self._device).type): + batch_size = image.shape[0] + clip_context = self.image_encoder.encode_image(image) + image_input = image.transpose(1, 2) + image_zeros = torch.zeros(batch_size, 3, num_frames-1, height, width, dtype=torch.bfloat16, device=self._device) + self._ensure_vae_on_device(image_input) + with torch.no_grad(): + y = self.vae.encode(torch.concat([image_input, image_zeros], dim=2)) + # Build mask to match VAE output shape (VAE may use different spatial downsampling, e.g. WanVideoVAE38 uses patch_size=2 -> height/16) + # y shape is B * 16 * (1+(T-1)/4) * H_latent * W_latent + num_t = y.shape[2] + h_latent, w_latent = y.shape[3], y.shape[4] + msk = torch.zeros(batch_size, 4, num_t, h_latent, w_latent, dtype=y.dtype, device=self._device) + msk[:, :, 0:1, :, :] = 1 + new_image = y[:, :, 0:1] + # concat: B * (4+16) * (1+(T-1)/4) * H_latent * W_latent + y = torch.concat([msk, y], dim=1) + return clip_context, y, new_image + + def prepare_extra_input(self, latents=None): + return {} + + def add_lora_to_model(self, model, lora_rank=4, lora_alpha=4, lora_target_modules="q,k,v,o,ffn.0,ffn.2", init_lora_weights="kaiming") -> nn.Module: + # Add LoRA to UNet + self.lora_alpha = lora_alpha + if init_lora_weights == "kaiming": + init_lora_weights = True + + lora_config = LoraConfig( + r=lora_rank, + lora_alpha=lora_alpha, + init_lora_weights=init_lora_weights, + target_modules=lora_target_modules.split(","), + ) + model = get_peft_model(model, lora_config) + for param in model.parameters(): + param.data = param.to(torch.float32) + return model + + def forward(self, backbone_output: BatchFeature, action_input: BatchFeature) -> BatchFeature: + # Set frozen modules to eval + self.set_frozen_modules_to_eval_mode() + + data = action_input + # Get embodiment ID. + embodiment_id = action_input.embodiment_id + # print("embodiment_id", embodiment_id) + has_real_action = action_input.has_real_action + action_mask = action_input.action_mask + + state_features = action_input.state + + actions = action_input.action + # assert the values of action is in between -1 and 1 + if actions.numel() > 0: + assert actions.min() >= -1.0 and actions.max() <= 1.0, "actions must be in [-1,1] range" + videos = data["images"] + + videos = rearrange(videos, "b t h w c -> b c t h w") + print("videos", videos.shape) + + + if videos.dtype == torch.uint8: + videos = videos.float() / 255.0 + b, c, t, h, w = videos.shape + videos = videos.permute(0, 2, 1, 3, 4) # [b, t, c, h, w] + videos = videos.reshape(b * t, c, h, w) + videos = self.normalize_video(videos) + videos = videos.reshape(b, t, c, h, w).permute(0, 2, 1, 3, 4) # back to [b, c, t, h, w] + assert videos.min() >= -1.0 and videos.max() <= 1.0, "videos must be in [-1,1] range" + videos = videos.to(dtype=self.dtype) + + # shape of B * max_length * dim + prompt_embs = self.encode_prompt(data["text"], data["text_attention_mask"]) + + # Wan 5B: resize to target resolution so latent tokens/frame matches DiT. Use config target when set + # (e.g. 160x320 so latent is 10x20 with VAE38 16x → even H,W, no crop in dynamics loss); else 176x320. + target_h = getattr(self.config, "target_video_height", None) + target_w = getattr(self.config, "target_video_width", None) + if target_h is None or target_w is None: + if getattr(self.model, "frame_seqlen", None) in (50, 55): + target_h, target_w = 176, 320 + else: + target_h, target_w = None, None + if target_h is not None and target_w is not None: + _, _, _, h, w = videos.shape + if (h, w) != (target_h, target_w): + b, c, t, _, _ = videos.shape + videos = torch.nn.functional.interpolate( + videos.reshape(b * t, c, h, w), + size=(target_h, target_w), + mode="bilinear", + align_corners=False, + ).reshape(b, c, t, target_h, target_w) + + latents = self.encode_video(videos, self.tiled, (self.tile_size_height, self.tile_size_width), (self.tile_stride_height, self.tile_stride_width)) + + # print("latents shape", latents.shape, self.dtype) + _, _, num_frames, height, width = videos.shape + image = videos[:, :, :1].transpose(1, 2) + + clip_feas, ys, _ = self.encode_image(image, num_frames, height, width) + + latents = latents.to(self._device) + clip_feas = clip_feas.to(self._device) + ys = ys.to(self._device) + prompt_embs = prompt_embs.to(self._device) + + # Loss + noise = torch.randn_like(latents) + + # specific to autoregressive + noise = noise.transpose(1, 2) + latents = latents.transpose(1, 2) + + # ============ VIDEO TIMESTEP SAMPLING ============ + if self.config.decouple_video_action_noise: + # Decoupled mode: sample video from Beta distribution biased towards HIGH noise + video_noise_ratio = self.video_beta_dist.sample([noise.shape[0], noise.shape[1]]) + timestep_id = ((1.0 - video_noise_ratio) * self.scheduler.num_train_timesteps).long() + timestep_id = torch.clamp(timestep_id, 0, self.scheduler.num_train_timesteps - 1) + noise_mode = "DECOUPLED" + elif self.config.use_high_noise_emphasis: + # High noise emphasis mode (coupled): BOTH video and action use Beta distribution + noise_ratio = self.high_noise_beta_dist.sample([noise.shape[0], noise.shape[1]]) + timestep_id = ((1.0 - noise_ratio) * self.scheduler.num_train_timesteps).long() + timestep_id = torch.clamp(timestep_id, 0, self.scheduler.num_train_timesteps - 1) + noise_mode = "HIGH_NOISE_EMPHASIS" + else: + # Original: uniform sampling over full range + timestep_id = torch.randint(0, self.scheduler.num_train_timesteps, (noise.shape[0], noise.shape[1])) + noise_mode = "STANDARD" + + timestep_id_block = timestep_id[:, 1:].reshape( + timestep_id.shape[0], -1, self.num_frame_per_block) + timestep_id_block[:, :, 1:] = timestep_id_block[:, :, 0:1] + + if actions.numel() > 0: + noise_action = torch.randn_like(actions) + assert actions.shape[1] / (noise.shape[1]-1) == (self.model.num_action_per_block // self.num_frame_per_block), f"actions.shape, {actions.shape}, noise.shape, {noise.shape}, video.shape, {videos.shape}, latents.shape, {latents.shape}" + assert (noise.shape[1]-1) / state_features.shape[1] == (self.num_frame_per_block // self.model.num_state_per_block), f"state_features.shape, {state_features.shape}, noise.shape, {noise.shape}, video.shape, {videos.shape}, latents.shape, {latents.shape}" + + # ============ ACTION TIMESTEP SAMPLING ============ + if self.config.decouple_video_action_noise: + # Decoupled: sample action timestep independently with full range + timestep_action_id = torch.randint( + 0, + self.scheduler.num_train_timesteps, + (actions.shape[0], actions.shape[1]) + ) + action_mode = "INDEPENDENT" + else: + # Original coupled: action timestep derived from video timestep + timestep_action_id = timestep_id_block.repeat(1, 1, actions.shape[1]//(noise.shape[1]-1)) + timestep_action_id = timestep_action_id.reshape(timestep_action_id.shape[0], -1) + action_mode = "COUPLED" + + # Log noise mode once + if not self._noise_logged: + video_mean = timestep_id.float().mean().item() + action_mean = timestep_action_id.float().mean().item() + if noise_mode == "DECOUPLED": + print(f"[NOISE] Mode={noise_mode} | Video: Beta({self.config.video_noise_beta_alpha},1) mean_t={video_mean:.0f} | Action: {action_mode} Uniform mean_t={action_mean:.0f}") + elif noise_mode == "HIGH_NOISE_EMPHASIS": + print(f"[NOISE] Mode={noise_mode} | Video+Action: Beta({self.config.high_noise_beta_alpha},1) mean_t={video_mean:.0f} | Action: {action_mode}") + else: + print(f"[NOISE] Mode={noise_mode} | Video+Action: Uniform mean_t={video_mean:.0f} | Action: {action_mode}") + self._noise_logged = True + else: + noise_action = None + timestep_action_id = None + + timestep_id_block = timestep_id_block.reshape(timestep_id_block.shape[0], -1) + timestep_id = torch.concat([timestep_id[:, :1], timestep_id_block], dim=1) + _, num_frames, num_channels, height, width = noise.shape + # DiT patch_embedding uses stride (1,2,2), so sequence length is num_frames * (H//2) * (W//2) + tokens_per_frame = (height // 2) * (width // 2) + seq_len = num_frames * tokens_per_frame + + timestep = self.scheduler.timesteps[timestep_id].to(self._device) + noisy_latents = self.scheduler.add_noise(latents.flatten(0, 1), noise.flatten(0, 1), timestep.flatten(0, 1)).unflatten(0, (noise.shape[0], noise.shape[1])) + training_target = self.scheduler.training_target(latents, noise, timestep).transpose(1, 2) + + if actions.numel() > 0: + timestep_action = self.scheduler.timesteps[timestep_action_id].to(self._device) + noisy_actions = self.scheduler.add_noise( + actions.flatten(0, 1), + noise_action.flatten(0, 1), + timestep_action.flatten(0, 1), + ).unflatten(0, (noise_action.shape[0], noise_action.shape[1])) + training_target_action = self.scheduler.training_target(actions, noise_action, timestep_action) + else: + timestep_action = None + noisy_actions = None + training_target_action = None + + # Compute loss + with torch.amp.autocast(dtype=torch.bfloat16, device_type=torch.device(self._device).type): + if actions.numel() > 0: + video_noise_pred, action_noise_pred = self.model( + noisy_latents.transpose(1, 2), timestep=timestep, clip_feature=clip_feas, y=ys, context=prompt_embs, seq_len=seq_len, + state=state_features, embodiment_id=embodiment_id, + action=noisy_actions, timestep_action=timestep_action, + clean_x=latents.transpose(1, 2), + ) + else: + video_noise_pred, action_noise_pred = self.model( + noisy_latents.transpose(1, 2), timestep=timestep, timestep_action=timestep_action, + clip_feature=clip_feas, y=ys, context=prompt_embs, seq_len=seq_len, + state=state_features, embodiment_id=embodiment_id, + clean_x=latents.transpose(1, 2), + ) + + # Per-sample dynamics loss + # DiT patch_embedding uses stride (1,2,2), so output spatial size can be smaller than + # latent when H or W is odd (e.g. latent 11x20 -> model output 10x20). Crop target to match. + if training_target.shape != video_noise_pred.shape: + training_target = training_target[ + ..., : video_noise_pred.shape[3], : video_noise_pred.shape[4] + ] + dynamics_loss_per_sample = torch.nn.functional.mse_loss( + video_noise_pred.float(), training_target.float(), reduction='none' + ).mean(dim=(1,3,4)) # shape: [B, ...] + + weight_dynamics = dynamics_loss_per_sample * self.scheduler.training_weight(timestep.flatten(0, 1)).unflatten(0, (noise.shape[0], noise.shape[1])).to(self._device) + weighted_dynamics_loss = weight_dynamics.mean() + + if actions.numel() > 0: + action_loss_per_sample = torch.nn.functional.mse_loss( + action_noise_pred.float(), training_target_action.float(), reduction='none' + ) * action_mask # shape: [B, ...] + action_loss_per_sample = has_real_action[:, None].float() * action_loss_per_sample # apply has_real_action + weight_action = action_loss_per_sample.mean(dim=2) * self.scheduler.training_weight( + timestep_action.flatten(0, 1), + ).unflatten(0, (noise_action.shape[0], noise_action.shape[1])).to(self._device) + weighted_action_loss = weight_action.mean() + loss = weighted_dynamics_loss + weighted_action_loss + else: + weighted_action_loss = torch.tensor(0.0, device=self._device) + loss = weighted_dynamics_loss + # loss = dynamics_loss_per_sample.mean() + + # Record log + output_dict = { + "loss": loss, + "dynamics_loss": weighted_dynamics_loss, + "action_loss": weighted_action_loss, + } + + return BatchFeature(data=output_dict) + + def generate_noise(self, shape, seed=None, device="cpu", dtype=torch.float16): + generator = None if seed is None else torch.Generator(device).manual_seed(seed) + noise = torch.randn(shape, generator=generator, device=device, dtype=dtype) + return noise + + def _get_caches( + self, kv_caches_input: list[KVCacheType], + ) -> list[KVCacheType]: + if self.ip_size > 1: + assert self.cfg_scale != 1.0, "cfg_scale must be != 1.0 when ip_size > 1" + assert len(kv_caches_input) == 2 + if self.ip_rank == 0: + kv_caches = [kv_caches_input[0]] + else: + kv_caches = [kv_caches_input[1]] + else: + assert len(kv_caches_input) <= 2 + kv_caches = [kv_caches_input[0]] + if self.cfg_scale != 1.0: + kv_caches.append(kv_caches_input[1]) + return kv_caches + + def _prepare_text_inputs(self, data: BatchFeature) -> list[tuple[torch.Tensor, torch.Tensor]]: + + if self.ip_size > 1: + assert self.cfg_scale != 1.0, "cfg_scale must be != 1.0 when ip_size > 1" + if self.ip_rank == 0: + text_inputs = [(data["text"], data["text_attention_mask"])] + else: + text_inputs = [(data["text_negative"], data["text_attention_mask_negative"])] + else: + text_inputs = [(data["text"], data["text_attention_mask"])] + if self.cfg_scale != 1.0: + text_inputs.append((data["text_negative"], data["text_attention_mask_negative"])) + return text_inputs + + + def _run_diffusion_steps( + self, + noisy_input: torch.Tensor, + timestep: torch.Tensor, + action: torch.Tensor, + timestep_action: torch.Tensor, + state: torch.Tensor, + embodiment_id: torch.Tensor, + context: torch.Tensor, + seq_len: int, + y: torch.Tensor, + clip_feature: torch.Tensor, + kv_caches: list[KVCacheType], + crossattn_caches: list[KVCacheType], + kv_cache_metadata: dict[str, bool | int], + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + predictions = [] + for index, prompt_emb in enumerate(context): + kv_cache = kv_caches[index] + crossattn_cache = crossattn_caches[index] + if not kv_cache_metadata["update_kv_cache"] and getattr(self, "trt_engine", None) is not None: + obs_noise_pred, action_noise_pred = self.trt_engine( + noisy_input, + timestep, + action=action, + timestep_action=timestep_action, + state=state, + context=prompt_emb, + y=y, + clip_feature=clip_feature, + kv_cache=kv_cache, + ) + else: + obs_noise_pred, action_noise_pred, updated_kv_caches = self.model( + noisy_input, + timestep, + action=action, + timestep_action=timestep_action, + state=state, + embodiment_id=embodiment_id, + context=prompt_emb, + seq_len=seq_len, + y=y, + clip_feature=clip_feature, + kv_cache=kv_cache, + crossattn_cache=crossattn_cache, + current_start_frame=kv_cache_metadata["start_frame"], + ) + if kv_cache_metadata["update_kv_cache"]: + for block_index, updated_kv_cache in enumerate(updated_kv_caches): + kv_cache[block_index] = updated_kv_cache.clone() + obs_noise_pred = obs_noise_pred.clone() + if action_noise_pred is not None: + action_noise_pred = action_noise_pred.clone() + else: + action_noise_pred = torch.tensor(0.0, device=obs_noise_pred.device) # dummy action noise prediction + predictions.append((obs_noise_pred, action_noise_pred)) + return self._exchange_predictions(predictions) + + def _exchange_predictions( + self, + predictions: list[tuple[torch.Tensor, torch.Tensor]], + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + if self.ip_size == 1: + return predictions + + assert len(predictions) == 1 + my_predictions = list(predictions[0]) + + other_predictions = [torch.empty_like(pred) for pred in my_predictions] + + send_ops = [ + dist.P2POp(op=dist.isend, tensor=pred, group_peer=(self.ip_rank + 1) % self.ip_size, group=self.ip_group) + for pred in my_predictions + ] + recv_ops = [ + dist.P2POp(op=dist.irecv, tensor=other_pred, group_peer=(self.ip_rank + 1) % self.ip_size, group=self.ip_group) + for other_pred in other_predictions + ] + ops = send_ops + recv_ops + + reqs = dist.batch_isend_irecv(ops) + for req in reqs: + req.wait() + + output_predictions: list[tuple[torch.Tensor, torch.Tensor] | None] = [None for _ in range(self.ip_size)] + output_predictions[self.ip_rank] = tuple(my_predictions) + output_predictions[(self.ip_rank + 1) % self.ip_size] = tuple(other_predictions) + assert all(isinstance(pred, tuple) for pred in output_predictions) + return cast(list[tuple[torch.Tensor, torch.Tensor]], output_predictions) + + def should_run_model(self, index, current_timestep, prev_predictions): + + if not self.dynamic_cache_schedule: + return self.dit_step_mask[index] + + # Always run first 2 steps to establish history + if len(prev_predictions) < 2: + return True + + if self.skip_countdown > 1: + self.skip_countdown -= 1 + return False + elif self.skip_countdown == 1: + self.skip_countdown = 0 + return True + + v_last = prev_predictions[-1][1].flatten(1).float() + v_prev = prev_predictions[-2][1].flatten(1).float() + sim = torch.nn.functional.cosine_similarity(v_last, v_prev, dim=1).mean() + + thresholds = [0.95, 0.93] + countdowns = [4, 2] + + for threshold, countdown in zip(thresholds, countdowns): + if sim > threshold: + self.skip_countdown = countdown + return False + + return True + + def lazy_joint_video_action(self, backbone_output: BatchFeature, action_input: BatchFeature, latent_video: torch.Tensor | None = None) -> BatchFeature: + start_time = time.perf_counter() + + # Tracking time taken on GPU for various operations. + start_text_encoder_event = torch.cuda.Event(enable_timing=True) + end_text_encoder_event = torch.cuda.Event(enable_timing=True) + start_image_encoder_event = torch.cuda.Event(enable_timing=True) + end_image_encoder_event = torch.cuda.Event(enable_timing=True) + start_vae_event = torch.cuda.Event(enable_timing=True) + end_vae_event = torch.cuda.Event(enable_timing=True) + start_kv_event = torch.cuda.Event(enable_timing=True) + end_kv_event = torch.cuda.Event(enable_timing=True) + start_diffusion_events = [torch.cuda.Event(enable_timing=True) for _ in range(self.num_inference_steps)] + end_diffusion_events = [torch.cuda.Event(enable_timing=True) for _ in range(self.num_inference_steps)] + + self.set_frozen_modules_to_eval_mode() + data = action_input + + videos = data["images"] + + embodiment_id = action_input.embodiment_id + state_features = action_input.state + + videos = rearrange(videos, "b t h w c -> b c t h w") + + if videos.dtype == torch.uint8: + videos = videos.float() / 255.0 + videos = videos.to(dtype=self.dtype) + b, c, t, h, w = videos.shape + videos = videos.permute(0, 2, 1, 3, 4) # [b, t, c, h, w] + videos = videos.reshape(b * t, c, h, w) + videos = self.normalize_video(videos) + videos = videos.reshape(b, t, c, h, w).permute(0, 2, 1, 3, 4) # back to [b, c, t, h, w] + assert videos.min() >= -1.0 and videos.max() <= 1.0, "videos must be in [-1,1] range" + videos = videos.to(dtype=self.dtype) + + state_features = state_features.to(dtype=torch.bfloat16) + videos = videos.to(dtype=torch.bfloat16) + + # Wan 5B: same as training — resize to target resolution so latent matches DiT + target_h = getattr(self.config, "target_video_height", None) + target_w = getattr(self.config, "target_video_width", None) + if target_h is None or target_w is None: + if getattr(self.model, "frame_seqlen", None) in (50, 55): + target_h, target_w = 176, 320 + else: + target_h, target_w = None, None + if target_h is not None and target_w is not None: + _, _, _, h, w = videos.shape + if (h, w) != (target_h, target_w): + b, c, t, _, _ = videos.shape + videos = torch.nn.functional.interpolate( + videos.reshape(b * t, c, h, w), + size=(target_h, target_w), + mode="bilinear", + align_corners=False, + ).reshape(b, c, t, target_h, target_w) + + if self.language is None: + print("language is None, reset current_start_frame to 0") + self.language = data["text"] + self.current_start_frame = 0 + elif not torch.equal(self.language, data["text"]): + print("language changed, reset current_start_frame to 0") + self.current_start_frame = 0 + self.language = data["text"] + elif videos.shape[2] == 1: + print("videos.shape[2] == 1, reset current_start_frame to 0") + self.current_start_frame = 0 + elif self.current_start_frame >= self.model.local_attn_size: + print("current_start_frame >= local_attn_size, reset current_start_frame to 0") + self.current_start_frame = 0 + + if self.ip_rank == 0: + print("videos shape", videos.shape, self.num_frames) + + start_text_encoder_event.record() + + text_inputs = self._prepare_text_inputs(data) + prompt_embs = [self.encode_prompt(text, attention_mask) for text, attention_mask in text_inputs] + + end_text_encoder_event.record() + + start_image_encoder_event.record() + + _, _, num_frames, height, width = videos.shape + if videos.shape[2] == 4 or videos.shape[2] == 9: + # special case for real-world eval where language is updated + image = videos[:, :, -1:].transpose(1, 2) + else: + image = videos[:, :, :1].transpose(1, 2) + + if self.current_start_frame == 0: + clip_feas, ys, image = self.encode_image(image, self.num_frames, height, width) + self.clip_feas = clip_feas.to(dtype=image.dtype) + self.ys = ys.to(dtype=image.dtype) + + assert self.clip_feas is not None and self.ys is not None, "clip_feas and ys must be set" + + end_image_encoder_event.record() + + start_vae_event.record() + + if latent_video is not None and self.current_start_frame != 0: + image = latent_video + if self.ip_rank == 0: + print("image shape@@", image.shape) + elif self.current_start_frame != 0: + # this is for real world execution + if (videos.shape[2] - 1) // 4 == self.num_frame_per_block: + print("no further action") + elif videos.shape[2] // 4 != self.num_frame_per_block: + # Repeating videos along dim 2. + repeat_factor = self.num_frame_per_block // (videos.shape[2] // 4) + videos = torch.repeat_interleave(videos, repeat_factor, dim=2) + + first_frame = videos[:, :, 0:1] # Extract first frame + videos = torch.cat([first_frame, videos], dim=2) + else: + first_frame = videos[:, :, 0:1] # Extract first frame + videos = torch.cat([first_frame, videos], dim=2) + + image = self.vae.encode( + videos, + tiled=self.tiled, + tile_size=(self.tile_size_height, self.tile_size_width), + tile_stride=(self.tile_stride_height, self.tile_stride_width), + ) + + end_vae_event.record() + + noise_obs = self.generate_noise((image.shape[0], image.shape[1], self.num_frame_per_block, image.shape[3], image.shape[4]), seed=self.seed, device='cuda', dtype=torch.bfloat16) + noise_action = self.generate_noise((image.shape[0], self.action_horizon, self.model.action_dim), seed=self.seed, device='cuda', dtype=torch.bfloat16) + batch_size, num_channels, num_frames, height, width = noise_obs.shape + ######### Generate video ######### + # DiT patch_embedding uses stride (1,2,2), so tokens per frame = (H//2)*(W//2) + tokens_per_frame = (height // 2) * (width // 2) + frame_seqlen = tokens_per_frame + seq_len = num_frames * frame_seqlen + + image = image.transpose(1, 2) + noise_obs = noise_obs.transpose(1, 2) + + if self.current_start_frame == 0: + # Reinitialize KV cache and crossattn cache for each new sequence. + self.kv_cache1, self.kv_cache_neg = self._create_kv_caches( + batch_size=batch_size, + dtype=noise_obs.dtype, + device=noise_obs.device, + frame_seqlen=frame_seqlen, + ) + self.crossattn_cache, self.crossattn_cache_neg = self._create_crossattn_caches( + batch_size=batch_size, + dtype=noise_obs.dtype, + device=noise_obs.device, + ) + + assert self.kv_cache1 is not None + assert self.kv_cache_neg is not None + assert self.crossattn_cache is not None + assert self.crossattn_cache_neg is not None + kv_caches = self._get_caches( + [self.kv_cache1, self.kv_cache_neg], + ) + crossattn_caches = self._get_caches( + [self.crossattn_cache, self.crossattn_cache_neg], + ) + + start_kv_event.record() + + if self.current_start_frame == 0: + timestep = torch.ones([batch_size, 1], device=noise_obs.device, dtype=torch.int64) * 0 + self._run_diffusion_steps( + noisy_input=image.transpose(1, 2), + timestep=timestep * 0, + action=None, + timestep_action=None, + state=None, + embodiment_id=None, + context=prompt_embs, + seq_len=frame_seqlen, + y=self.ys[:, :, 0:1], + clip_feature=self.clip_feas, + kv_caches=kv_caches, + crossattn_caches=crossattn_caches, + kv_cache_metadata=dict( + start_frame=0, + update_kv_cache=True, + ), + ) + self.current_start_frame += 1 + + timestep = torch.ones([batch_size, self.num_frame_per_block], device=noise_obs.device, dtype=torch.int64) * 0 + + if self.current_start_frame != 1: + current_ref_latents = image[:, -self.num_frame_per_block:] + if self.current_start_frame <= self.ys.shape[2]: + y = self.ys[:, :, self.current_start_frame - self.num_frame_per_block : self.current_start_frame] + else: + y = self.ys[:, :, -self.num_frame_per_block:] + self._run_diffusion_steps( + noisy_input=current_ref_latents.transpose(1, 2), + timestep=timestep * 0, + action=None, + timestep_action=None, + state=None, + embodiment_id=None, + context=prompt_embs, + seq_len=seq_len, + y=y, + clip_feature=self.clip_feas, + kv_caches=kv_caches, + crossattn_caches=crossattn_caches, + kv_cache_metadata=dict( + start_frame=self.current_start_frame - self.num_frame_per_block, + update_kv_cache=True, + ), + ) + + end_kv_event.record() + + noisy_input = noise_obs + noisy_input_action = noise_action + + # Step 3.1: Spatial denoising loop + + sample_scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=self.scheduler.num_train_timesteps, + shift=1, + use_dynamic_shifting=False) + sample_scheduler_action = FlowUniPCMultistepScheduler( + num_train_timesteps=self.scheduler.num_train_timesteps, + shift=1, + use_dynamic_shifting=False) + sample_scheduler.set_timesteps( + self.num_inference_steps, device=noise_obs.device, shift=self.sigma_shift) + sample_scheduler_action.set_timesteps( + self.num_inference_steps, device=noise_obs.device, shift=self.sigma_shift) + + # Decoupled inference: video sigmas end at video_final_noise instead of 0 + # This rescales the schedule so video still takes all denoising steps, + # but ends at a higher noise level (e.g., 1.0 → 0.9 → 0.8 instead of 1.0 → 0.5 → 0.0) + if self.config.decouple_inference_noise: + video_final_noise = self.config.video_inference_final_noise + # Rescale video sigmas: map [sigma_max, 0] -> [sigma_max, video_final_noise] + sigma_max = sample_scheduler.sigmas[0].item() + sample_scheduler.sigmas = sample_scheduler.sigmas * (sigma_max - video_final_noise) / sigma_max + video_final_noise + sample_scheduler.timesteps = (sample_scheduler.sigmas[:-1] * 1000).to(torch.int64) + if self.ip_rank == 0: + print(f"Decoupled inference: video sigmas {sigma_max:.3f} -> {sample_scheduler.sigmas[-1].item():.3f}") + + start_diffusion_events = [torch.cuda.Event(enable_timing=True) for _ in sample_scheduler.timesteps] + end_diffusion_events = [torch.cuda.Event(enable_timing=True) for _ in sample_scheduler.timesteps] + prev_predictions = [] + self.skip_countdown = 0 + dit_compute_steps = 0 + for index, current_timestep in enumerate(sample_scheduler.timesteps): + start_diffusion_events[index].record() + + # Get timesteps from respective schedulers + action_timestep = sample_scheduler_action.timesteps[index] + video_timestep = sample_scheduler.timesteps[index] # Already rescaled if decoupled + + # set current timestep + timestep = torch.ones( + [batch_size, self.num_frame_per_block], + device=noise_obs.device, + dtype=torch.int64, + ) * video_timestep + timestep_action = torch.ones( + [batch_size, self.action_horizon], + device=noise_obs.device, + dtype=torch.int64, + ) * action_timestep + + # check if we need to run the DIT step + should_run_model = self.should_run_model(index, current_timestep, prev_predictions) + if should_run_model: + dit_compute_steps += 1 + if self.current_start_frame + self.num_frame_per_block <= self.ys.shape[2]: + y = self.ys[:, :, self.current_start_frame : self.current_start_frame + self.num_frame_per_block] + else: + y = self.ys[:, :, -self.num_frame_per_block:] + predictions = self._run_diffusion_steps( + noisy_input=noisy_input.transpose(1, 2), + timestep=timestep, + action=noisy_input_action, + timestep_action=timestep_action, + state=state_features, + embodiment_id=embodiment_id, + context=prompt_embs, + seq_len=seq_len, + y=y, + clip_feature=self.clip_feas, + kv_caches=kv_caches, + crossattn_caches=crossattn_caches, + kv_cache_metadata=dict( + start_frame=self.current_start_frame, + update_kv_cache=False, + ), + ) + flow_pred_cond, flow_pred_cond_action = predictions[0] + flow_pred_uncond, flow_pred_uncond_action = predictions[1] + + flow_pred = flow_pred_uncond + self.cfg_scale * (flow_pred_cond - flow_pred_uncond) + prev_predictions.append((current_timestep, flow_pred, flow_pred_cond_action)) + max_cache_size = 2 + if len(prev_predictions) > max_cache_size: + prev_predictions.pop(0) + + else: + assert len(prev_predictions) > 0, "prev_predictions must be set when skipping" + _, flow_pred, flow_pred_cond_action = prev_predictions[-1] + + end_diffusion_events[index].record() + + # Video: denoising step (uses rescaled schedule if decoupled) + noisy_input = sample_scheduler.step( + model_output=flow_pred.transpose(1, 2), + timestep=video_timestep, + sample=noisy_input, + step_index=index, + return_dict=False, + )[0] + + # Action: always fully denoises with standard schedule (1000->0) + noisy_input_action = sample_scheduler_action.step( + model_output=flow_pred_cond_action, + timestep=action_timestep, + sample=noisy_input_action, + step_index=index, + return_dict=False, + )[0] + + latents = noisy_input + latents_action = noisy_input_action + output = latents + + if self.current_start_frame == 1: + output = torch.cat([image, output], dim=1) + self.current_start_frame += self.num_frame_per_block + + # Do torch.cuda.synchronize() to ensure all operations are completed before timing. + # This isn't expected to affect inference performance since it's at the end of an inference step. + torch.cuda.synchronize() + + total_time = time.perf_counter() - start_time + text_encoder_time = start_text_encoder_event.elapsed_time(end_text_encoder_event) / 1000 + image_encoder_time = start_image_encoder_event.elapsed_time(end_image_encoder_event) / 1000 + vae_time = start_vae_event.elapsed_time(end_vae_event) / 1000 + kv_creation_time = start_kv_event.elapsed_time(end_kv_event) / 1000 + diffusion_times = [s.elapsed_time(e) for s, e in zip(start_diffusion_events, end_diffusion_events)] + diffusion_time = sum(diffusion_times) / 1000 + scheduler_time = total_time - kv_creation_time - diffusion_time - text_encoder_time - image_encoder_time - vae_time + + if self.ip_rank == 0: + print(f"Time taken: Total {total_time:.2f} seconds, " + f"Text Encoder {text_encoder_time:.2f} seconds, " + f"Image Encoder {image_encoder_time:.2f} seconds, " + f"VAE {vae_time:.2f} seconds, " + f"KV Cache Creation {kv_creation_time:.2f} seconds, " + f"Diffusion {diffusion_time:.2f} seconds, " + f"DIT Compute Steps {dit_compute_steps} steps, " + f"Scheduler {scheduler_time:.2f} seconds") + + return BatchFeature(data={"action_pred": latents_action, "video_pred": output.transpose(1, 2)}) + + def cache_predict_order1(self, current_timestep, timestep_1, f1, timestep_2, f2): + h_curr = current_timestep - timestep_1 + h_past = timestep_1 - timestep_2 + + v_prime = (f1 - f2) / h_past + + # Prediction + damping_factor = 0.25 + flow_pred = f1 + (v_prime * h_curr) * damping_factor + return flow_pred + + def post_initialize(self): + # Move models to the cuda device and set the dtype to bfloat16. + print("Moving models to the cuda device and setting the dtype to bfloat16.") + self.model.to(device=self._device, dtype=torch.bfloat16) + self.text_encoder.to(device=self._device, dtype=torch.bfloat16) + self.image_encoder.to(device=self._device, dtype=torch.bfloat16) + self.vae.to(device=self._device, dtype=torch.bfloat16) + import os + ENABLE_TENSORRT = os.getenv("ENABLE_TENSORRT", "False").lower() == "true" + LOAD_TRT_ENGINE = os.getenv("LOAD_TRT_ENGINE", None) + + # Torch compile the modules. Skip _forward_blocks: Dynamo with fullgraph can fail on + # shape variation (e.g. x [1,50,C] vs e [1,200,C]); the block aligns e to x at runtime. + if not ENABLE_TENSORRT: + print("Torch compiling the TextEncoder, ImageEncoder, and VAE modules (Wan _forward_blocks not compiled).") + + self.text_encoder.forward = torch.compile( + mode="reduce-overhead", fullgraph=True, dynamic=False, + )(self.text_encoder.forward) + + self.image_encoder.model.visual.forward = torch.compile( + mode="reduce-overhead", fullgraph=True, dynamic=False, + )(self.image_encoder.model.visual.forward) + + self.vae.model.encode = torch.compile( + mode="reduce-overhead", fullgraph=True, dynamic=False, + )(self.vae.model.encode) + + self.trt_engine = None + if LOAD_TRT_ENGINE is not None: + print(f"Loading TRT engine from {LOAD_TRT_ENGINE}") + import groot.control.tensorrt_utils as trt_utils + model_path = LOAD_TRT_ENGINE + self.trt_engine = trt_utils.load_tensorrt_engine(model_path, model_type="ar_14B") + + def parallelize(self, device_mesh: DeviceMesh) -> None: + ip_mesh = device_mesh["ip"] + self.ip_rank = ip_mesh.get_local_rank() + self.ip_size = ip_mesh.size() + self.ip_group = ip_mesh.get_group() + + assert self.ip_size == 1 or self.ip_size == 2, "ip_size must be 1 or 2" + assert self.ip_rank >= 0 and self.ip_rank < self.ip_size, "ip_rank must be in [0, ip_size)" + + @property + def device(self): + return next(iter(self.parameters())).device + + @property + def dtype(self): + return next(iter(self.parameters())).dtype diff --git a/groot/vla/model/dreamzero/backbone/__init__.py b/groot/vla/model/dreamzero/backbone/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/model/dreamzero/backbone/base_backbone.py b/groot/vla/model/dreamzero/backbone/base_backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..b768bcdef5fa0e6b08ee4634ec2e144e549edbb7 --- /dev/null +++ b/groot/vla/model/dreamzero/backbone/base_backbone.py @@ -0,0 +1,16 @@ +from abc import ABC, abstractmethod + +from torch import nn +from transformers.feature_extraction_utils import BatchFeature + + +class Backbone(ABC, nn.Module): + def __init__(self): + super(Backbone, self).__init__() + + @abstractmethod + def forward(self, backbone_input: BatchFeature) -> BatchFeature: + pass + + def prepare_input(self, batch: dict) -> BatchFeature: + pass diff --git a/groot/vla/model/dreamzero/backbone/identity.py b/groot/vla/model/dreamzero/backbone/identity.py new file mode 100644 index 0000000000000000000000000000000000000000..d63ad24139b9b796f0244025618f472e2fe0235a --- /dev/null +++ b/groot/vla/model/dreamzero/backbone/identity.py @@ -0,0 +1,49 @@ +import torch +from transformers.feature_extraction_utils import BatchFeature + +from groot.vla.model.dreamzero.backbone.base_backbone import Backbone + + +class IdentityBackbone(Backbone): + """ + This class allows pretraining the action head without depending on any backbone. + That's why it's called "identity" — it preserves the action head to be a standalone trainable model. + """ + + def set_trainable_parameters(self, **kwargs): + return + + def forward(self, backbone_input: BatchFeature) -> BatchFeature: + backbone_input_first_value = next(iter(backbone_input.values())) + B = backbone_input_first_value.shape[0] + + backbone_features = torch.empty( + B, 1, 0, dtype=torch.float32, device=backbone_input_first_value.device + ) + output_dict = { + "backbone_features": backbone_features, + } + + return BatchFeature(data=output_dict) + + def prepare_input(self, batch: dict) -> BatchFeature: + """ + Args: + batch: dict + Must contain at least one key-value pair to inform the batch size. + Expects the first dimension to be the batch size. See `forward`. + """ + if "action" in batch: + return BatchFeature(data={"action": batch["action"]}) + else: + # at inference time, we have to use either state or video + if "state" in batch: + return BatchFeature(data={"state": batch["state"]}) + elif "video" in batch: + # For video, it's tricky because it's a numpy array, which isn't compatible with BatchFeature's `to` method + # So instead, we make it a tensor and return it + video = batch["video"] + video_tensor = torch.from_numpy(video) + return BatchFeature(data={"video": video_tensor}) + else: + return BatchFeature(data=batch) diff --git a/groot/vla/model/dreamzero/base_vla.py b/groot/vla/model/dreamzero/base_vla.py new file mode 100644 index 0000000000000000000000000000000000000000..274b8640f44b2ac06b7fe8f60f43b4ce8dd51a28 --- /dev/null +++ b/groot/vla/model/dreamzero/base_vla.py @@ -0,0 +1,614 @@ +from dataclasses import dataclass, field +from typing import Tuple + +from hydra.utils import instantiate +import numpy as np +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh + +from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel +from transformers.feature_extraction_utils import BatchFeature +import tree + +BACKBONE_FEATURE_KEY = "backbone_features" +ACTION_KEY = "action_pred" +LOSS_KEY = "loss" +ERROR_MSG = "Error: unexpected input/output" +N_COLOR_CHANNELS = 3 + + +@dataclass +class VLAConfig(PretrainedConfig): + model_type = "vla" + backbone_cfg: PretrainedConfig = field( + default=None, metadata={"help": "Backbone configuration."} + ) + + action_head_cfg: PretrainedConfig = field( + default=None, metadata={"help": "Action head configuration."} + ) + + action_horizon: int = field(default=None, metadata={"help": "Action horizon."}) + + action_dim: int = field(default=None, metadata={"help": "Action dimension."}) + compute_dtype: str = field(default="float32", metadata={"help": "Compute dtype."}) + + def __init__(self, **kwargs): + super().__init__(**kwargs) + for key, value in kwargs.items(): + setattr(self, key, value) + + +class VLA(PreTrainedModel): + supports_gradient_checkpointing = True + config_class = VLAConfig + """ + we expect the backbone output to have a key 'backbone_features' with shape (batch_size, n, hidden_size) + here n is variable and can be e.g. time, 1 or user specified + we expect the action head output to have a key 'action_pred' with shape (batch_size, time, action_dim) during inference time + we expect these to have type BatchFeature, and they can of course have many other user specified keys too + see discussion at https://nvidia.slack.com/archives/C07T1V7L886/p1732550624654139 + """ + + def __init__( + self, + config: VLAConfig, + ): + assert isinstance(config.backbone_cfg, dict) + assert isinstance(config.action_head_cfg, dict) + super().__init__(config) + self.backbone = instantiate(config.backbone_cfg) + self.action_head = instantiate(config.action_head_cfg) + self.action_horizon = config.action_horizon + self.action_dim = config.action_dim + self.compute_dtype = config.compute_dtype + + self.rank = dist.get_rank() if dist.is_initialized() else 0 + + def validate_inputs(self, inputs): + detected_error = False + error_msg = ERROR_MSG + if "action" in inputs: + action = inputs["action"] + type_ok = isinstance(action, torch.Tensor) + shape_ok = ( + len(action.shape) == 3 + and action.shape[1] % self.action_horizon == 0 + and action.shape[2] == self.action_dim + ) + if not type_ok: + error_msg += f"\n{action.dtype=}" + detected_error = True + if not shape_ok: + error_msg += f"\n{action.shape=}" + detected_error = True + + if "video" in inputs: + video = inputs["video"] + type_ok = isinstance(video, np.ndarray) + dtype_ok = video.dtype == np.uint8 + shape_ok = len(video.shape) == 6 and video.shape[3] == N_COLOR_CHANNELS + if not type_ok: + error_msg += f"\n{type(video)=}" + detected_error = True + if not dtype_ok: + error_msg += f"\n{video.dtype=}" + detected_error = True + if not shape_ok: + error_msg += f"\n{video.shape=}" + detected_error = True + + if detected_error: + raise ValueError(error_msg) + + def validate_data(self, action_head_outputs, backbone_outputs, is_training): + + fail_backbone = ( + not isinstance(backbone_outputs, BatchFeature) + or BACKBONE_FEATURE_KEY not in backbone_outputs + ) + + if fail_backbone: + error_msg = ERROR_MSG + error_msg += f"\n{isinstance(backbone_outputs, BatchFeature)=}" + error_msg += f"\n{BACKBONE_FEATURE_KEY in backbone_outputs=}" + error_msg += f"\n{backbone_outputs[BACKBONE_FEATURE_KEY].shape=}" + raise ValueError(error_msg) + + fail_action_head = (not isinstance(action_head_outputs, BatchFeature)) or not ( + ( + LOSS_KEY in action_head_outputs and is_training + ) # there might not be an action prediction during training + or ( + ACTION_KEY in action_head_outputs + and action_head_outputs[ACTION_KEY].shape[1] == self.action_horizon + and action_head_outputs[ACTION_KEY].shape[2] == self.action_dim + ) + ) + + if fail_action_head: + error_msg = ERROR_MSG + error_msg += f"\n{isinstance(action_head_outputs, BatchFeature)=}" + error_msg += f"\n{LOSS_KEY in action_head_outputs=}" + error_msg += f"\n{action_head_outputs[ACTION_KEY].shape=}" + error_msg += f"\n{self.action_horizon=}" + error_msg += f"\n{self.action_dim=}" + raise ValueError(error_msg) + + def forward( + self, + inputs: dict, + ) -> BatchFeature: + + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + action_head_outputs = self.action_head(backbone_outputs, action_inputs) + + return action_head_outputs + + def get_action( + self, + inputs: dict, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + action_head_outputs = self.action_head.get_action(backbone_outputs, action_inputs) + self.validate_data(action_head_outputs, backbone_outputs, is_training=False) + return action_head_outputs + + def joint_video_action( + self, + inputs: dict, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + action_head_outputs = self.action_head.joint_video_action(backbone_outputs, action_inputs) + self.validate_data(action_head_outputs, backbone_outputs, is_training=False) + return action_head_outputs + + def lazy_joint_video_action( + self, + inputs: dict, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + action_head_outputs = self.action_head.lazy_joint_video_action(backbone_outputs, action_inputs) + self.validate_data(action_head_outputs, backbone_outputs, is_training=False) + return action_head_outputs + + def lazy_joint_video_action_causal( + self, + inputs: dict, + latent_video: torch.Tensor | None = None, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + action_head_outputs = self.action_head.lazy_joint_video_action(backbone_outputs, action_inputs, latent_video=latent_video) + self.validate_data(action_head_outputs, backbone_outputs, is_training=False) + return action_head_outputs + + def lazy_joint_video_action_causal_gt_cond( + self, + inputs: dict, + latent_video: torch.Tensor | None = None, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + + action_head_outputs = self.action_head.lazy_joint_video_action_causal_gt_cond(backbone_outputs, action_inputs, latent_video=latent_video) + self.validate_data(action_head_outputs, backbone_outputs, is_training=False) + return action_head_outputs + + def lazy_joint_video_action_efficient( + self, + inputs: dict, + prompt_embs: torch.Tensor | None = None, + prompt_emb_nega: torch.Tensor | None = None, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + action_head_outputs = self.action_head.lazy_joint_video_action_efficient(backbone_outputs, action_inputs, prompt_embs=prompt_embs, prompt_emb_nega=prompt_emb_nega) + self.validate_data(action_head_outputs, backbone_outputs, is_training=False) + return action_head_outputs + + def gt_video_action_pred( + self, + inputs: dict, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + action_head_outputs = self.action_head.gt_video_action_pred(backbone_outputs, action_inputs) + self.validate_data(action_head_outputs, backbone_outputs, is_training=False) + return action_head_outputs + + def get_language( + self, + inputs: dict, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + # Because the behavior of backbones remains the same for training and inference, we can use `forward` for backbones. + backbone_outputs = self.backbone.generate(backbone_inputs) + return backbone_outputs + + def get_video( + self, + inputs: dict, + ) -> BatchFeature: + _, video_inputs = self.prepare_input(inputs) + video_outputs = self.action_head.get_video(video_inputs) + return video_outputs + + def prepare_input(self, inputs) -> Tuple[BatchFeature, BatchFeature]: + self.validate_inputs(inputs) + backbone_inputs = self.backbone.prepare_input(inputs) + action_inputs = self.action_head.prepare_input(inputs) + + def to_device_with_maybe_dtype(x): + # Only cast to self.compute_dtype if the tensor is floating + if torch.is_floating_point(x): + return x.to(self.device, dtype=self.action_head.dtype) + else: + # Keep original dtype + return x.to(self.device) + + backbone_inputs = tree.map_structure(to_device_with_maybe_dtype, backbone_inputs) + action_inputs = tree.map_structure(to_device_with_maybe_dtype, action_inputs) + return backbone_inputs, action_inputs + + + @classmethod + def from_pretrained_for_tuning( + cls, + pretrained_model_name_or_path: str, + config: VLAConfig = None, # This config will now be USED + device_map: str = "auto", + dtype: torch.dtype = torch.bfloat16, + offload_state_dict: bool = True, + lora_weights_path: str | None = None, + ): + if config is None: + raise ValueError( + "A `config` object must be provided to build the model structure." + ) + + import os + import json + import gc + from safetensors.torch import load_file + + model = cls(config) + + safetensors_path = os.path.join(pretrained_model_name_or_path, "model.safetensors") + safetensors_index_path = os.path.join(pretrained_model_name_or_path, "model.safetensors.index.json") + + if os.path.exists(safetensors_index_path): + with open(safetensors_index_path, 'r') as f: + index = json.load(f) + missing_keys_accum = set() + unexpected_keys_accum = set() + shard_files = sorted(set(index["weight_map"].values())) + for shard_file in shard_files: + shard_path = os.path.join(pretrained_model_name_or_path, shard_file) + print(f"Loading shard: {shard_path}") + shard_state_dict = load_file(shard_path) + missing_keys, unexpected_keys = model.load_state_dict(shard_state_dict, strict=False) + if missing_keys: + missing_keys_accum.update(missing_keys) + if unexpected_keys: + unexpected_keys_accum.update(unexpected_keys) + # Free shard immediately + del shard_state_dict + gc.collect() + if missing_keys_accum: + print(f"Missing keys when loading sharded pretrained weights: {sorted(missing_keys_accum)} ... total={len(missing_keys_accum)}") + if unexpected_keys_accum: + print(f"Unexpected keys when loading sharded pretrained weights: {sorted(unexpected_keys_accum)} ... total={len(unexpected_keys_accum)}") + if not missing_keys_accum and not unexpected_keys_accum: + print("Successfully loaded pretrained base weights (sharded)") + elif os.path.exists(safetensors_path): + # Handle single safetensors file + print(f"Loading weights from safetensors: {safetensors_path}") + state_dict = load_file(safetensors_path) + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) + if missing_keys: + print(f"Missing keys when loading pretrained weights: {missing_keys}") + if unexpected_keys: + print(f"Unexpected keys when loading pretrained weights: {unexpected_keys}") + if not missing_keys and not unexpected_keys: + print("Successfully loaded pretrained base weights") + else: + raise FileNotFoundError( + f"No weights found at '{pretrained_model_name_or_path}'. " + "Expected 'model.safetensors' or 'model.safetensors.index.json'." + ) + + if lora_weights_path is not None: + print(f"Loading LoRA weights from: {lora_weights_path}") + model.load_lora_weight(lora_weights_path) + else: + if hasattr(model, 'action_head') and hasattr(model.action_head, 'inject_lora_after_loading') and model.action_head.config.defer_lora_injection: + print("Injecting LoRA adapters into action_head after loading pretrained weights") + model.action_head.inject_lora_after_loading() + + print(f"{cls}\n") + return model + + @classmethod + def load_lora( + cls, + pretrained_model_name_or_path: str + ): + from safetensors.torch import load_file + import os + import json + print("loading lora@@@@@") + + # Check for different checkpoint formats + safetensors_path = os.path.join(pretrained_model_name_or_path, "model.safetensors") + safetensors_index_path = os.path.join(pretrained_model_name_or_path, "model.safetensors.index.json") + + state_dict = {} + if os.path.exists(safetensors_index_path): + # Handle sharded safetensors + print(f"Loading sharded safetensors using index: {safetensors_index_path}") + + with open(safetensors_index_path, 'r') as f: + index = json.load(f) + + # Load each shard + for shard_file in set(index["weight_map"].values()): + shard_path = os.path.join(pretrained_model_name_or_path, shard_file) + print(f"Loading shard: {shard_path}") + shard_state_dict = load_file(shard_path) + state_dict.update(shard_state_dict) + + elif os.path.exists(safetensors_path): + # Handle single safetensors file + print(f"Loading weights from safetensors: {safetensors_path}") + state_dict.update(load_file(safetensors_path)) + + # Load config + print("loading config@@") + config_path = os.path.join(pretrained_model_name_or_path, "config.json") + with open(config_path, "r") as f: + config_dict = json.load(f) + config = VLAConfig(**config_dict) + print("loading model") + + # Disable defer_lora_injection so LoRA layers are created during init, + # matching the PEFT key hierarchy (base_model.model.*) in the checkpoint. + ah_cfg = config.action_head_cfg + inner = ah_cfg.get('config', ah_cfg) if isinstance(ah_cfg.get('config'), dict) else ah_cfg + if 'defer_lora_injection' in inner: + inner['defer_lora_injection'] = False + print("defer_lora_injection disabled for load_lora") + # Enable component loading so DiT base weights are loaded from pretrained + if 'skip_component_loading' in inner: + inner['skip_component_loading'] = False + print("skip_component_loading disabled for load_lora") + + # Instantiate model (LoRA layers now exist from init) + model = cls(config) + + # Remove .base_layer from keys if present + has_base_layer = any(".base_layer." in key for key in state_dict.keys()) + if has_base_layer: + print("Removing '.base_layer' from state dict keys") + state_dict = {k.replace(".base_layer.", "."): v for k, v in state_dict.items()} + + # Load weights + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) + + if missing_keys: + print(f"Missing keys when loading pretrained weights: {missing_keys}") + if unexpected_keys: + print(f"Unexpected keys when loading pretrained weights: {unexpected_keys}") + + print("Successfully loaded pretrained weights") + + print(f"{cls}\n") + return model + + def load_lora_weight(self, pretrained_model_name_or_path: str): + """Load only LoRA weights from a pretrained model without loading config.""" + from safetensors.torch import load_file + import os + import json + + print(f"Loading LoRA weights from {pretrained_model_name_or_path}") + + # Check for different checkpoint formats + safetensors_path = os.path.join(pretrained_model_name_or_path, "model.safetensors") + safetensors_index_path = os.path.join(pretrained_model_name_or_path, "model.safetensors.index.json") + + state_dict = {} + if os.path.exists(safetensors_index_path): + # Handle sharded safetensors + print(f"Loading sharded safetensors using index: {safetensors_index_path}") + + with open(safetensors_index_path, 'r') as f: + index = json.load(f) + + # Load each shard + for shard_file in set(index["weight_map"].values()): + shard_path = os.path.join(pretrained_model_name_or_path, shard_file) + print(f"Loading shard: {shard_path}") + shard_state_dict = load_file(shard_path) + state_dict.update(shard_state_dict) + + elif os.path.exists(safetensors_path): + # Handle single safetensors file + print(f"Loading weights from safetensors: {safetensors_path}") + state_dict.update(load_file(safetensors_path)) + else: + raise FileNotFoundError(f"No valid checkpoint found at {pretrained_model_name_or_path}") + + print("Loading LoRA weights into existing model") + + def rewrite_lora_state_dict_keys(state_dict, pattern, repl): + new_state_dict = {} + for k, v in state_dict.items(): + new_k = k.replace(pattern, repl) + new_state_dict[new_k] = v + return new_state_dict + + has_target_pattern = any("action_head.model.base_model.model" in key for key in state_dict.keys()) + + if not has_target_pattern: + print("Rewriting LoRA state dict keys from 'action_head.model' to 'action_head.model.base_model.model'") + state_dict = rewrite_lora_state_dict_keys( + state_dict, + pattern="action_head.model", + repl="action_head.model.base_model.model", + ) + else: + print("State dict already has 'action_head.model.base_model.model' pattern, skipping key rewrite") + + # Load only the weights into the existing model + missing_keys, unexpected_keys = self.load_state_dict(state_dict, strict=False) + + print("Successfully loaded LoRA state dict") + + if missing_keys: + print(f"Missing keys when loading LoRA weights: {missing_keys}") + if unexpected_keys: + print(f"Unexpected keys when loading LoRA weights: {unexpected_keys}") + + print("Successfully loaded LoRA weights") + + @classmethod + def from_config_with_lora_weights( + cls, + config: VLAConfig, + pretrained_model_path: str, + ): + """Create VLA model from config and then load LoRA weights from pretrained model.""" + print(f"Creating VLA model from config and loading LoRA weights from {pretrained_model_path}") + + # 1. Create model from config (similar to vla.yaml) + model = cls(config) + print("Model created from config") + + # 2. Load LoRA weights into the created model + model.load_lora_weight(pretrained_model_path) + + return model + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: str, + config: VLAConfig = None + ): + del config + + from safetensors.torch import load_file + import os + import json + print("loading pretrained@@@@@") + # Check for different checkpoint formats + safetensors_path = os.path.join(pretrained_model_name_or_path, "model.safetensors") + safetensors_index_path = os.path.join(pretrained_model_name_or_path, "model.safetensors.index.json") + + state_dict = {} + if os.path.exists(safetensors_index_path): + # Handle sharded safetensors + print(f"Loading sharded safetensors using index: {safetensors_index_path}") + + with open(safetensors_index_path, 'r') as f: + index = json.load(f) + + # Load each shard + for shard_file in set(index["weight_map"].values()): + shard_path = os.path.join(pretrained_model_name_or_path, shard_file) + print(f"Loading shard: {shard_path}") + shard_state_dict = load_file(shard_path) + state_dict.update(shard_state_dict) + + elif os.path.exists(safetensors_path): + # Handle single safetensors file + print(f"Loading weights from safetensors: {safetensors_path}") + state_dict.update(load_file(safetensors_path)) + + # Load config + print("loading config@@") + config_path = os.path.join(pretrained_model_name_or_path, "config.json") + with open(config_path, "r") as f: + config_dict = json.load(f) + config = VLAConfig(**config_dict) + print("loading model") + print("config.action_head_cfg", config.action_head_cfg) + # Always disable defer_lora_injection + # config.action_head_cfg is a dict, and defer_lora_injection is nested in config.action_head_cfg['config'] + if 'config' in config.action_head_cfg and isinstance(config.action_head_cfg['config'], dict): + if 'defer_lora_injection' in config.action_head_cfg['config']: + config.action_head_cfg['config']['defer_lora_injection'] = False + print("config.action_head_cfg['config']['defer_lora_injection'] disabled (set to False)") + elif 'defer_lora_injection' in config.action_head_cfg: + config.action_head_cfg['defer_lora_injection'] = False + print("config.action_head_cfg['defer_lora_injection'] disabled (set to False)") + + # Instantiate model + model = cls(config) + print("model", model) + # Remove .base_layer from keys (e.g., 'action_head.model.base_model.model.blocks.19.self_attn.v.base_layer.bias' -> 'action_head.model.base_model.model.blocks.19.self_attn.v.bias') + has_base_layer = any(".base_layer." in key for key in state_dict.keys()) + if has_base_layer: + print("Removing '.base_layer' from state dict keys") + new_state_dict = {} + for k, v in state_dict.items(): + new_k = k.replace(".base_layer.", ".") + new_state_dict[new_k] = v + state_dict = new_state_dict + + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) + + if missing_keys: + print(f"Missing keys when loading pretrained weights: {missing_keys}") + if unexpected_keys: + print(f"Unexpected keys when loading pretrained weights: {unexpected_keys}") + + print("Successfully loaded pretrained weights") + + print(f"{cls}\n") + return model + + def post_initialize(self): + self.action_head.post_initialize() + + def parallelize(self, device_mesh: DeviceMesh): + self.action_head.parallelize(device_mesh=device_mesh) + + +class CotrainVLA(VLA): + + def forward( + self, + inputs: dict, + ) -> BatchFeature: + if "cotrain" in inputs and inputs["cotrain"]: + return self.backbone.cotrain(inputs) + return super().forward(inputs) + + +def create_vla_with_pretrained_action_head(pretrained_vla_path: str, config: VLAConfig): + # 1. Instantiate a new VLAModel + vla = VLA(config) + + # 2. Load the pretrained VLAModel + pretrained_vla = VLA.from_pretrained(pretrained_vla_path) + + # 3. Replace the action head in the new VLAModel with the pretrained action head + vla.action_head = pretrained_vla.action_head + + # 4. Replace the action head config in the new VLAModel with the pretrained action head config + vla.config.action_head_cfg = pretrained_vla.config.action_head_cfg + + # 5. Return the new VLAModel + return vla + + +# register +AutoConfig.register("vla", VLAConfig) +AutoModel.register(VLAConfig, VLA) diff --git a/groot/vla/model/dreamzero/modules/__init__.py b/groot/vla/model/dreamzero/modules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/model/dreamzero/modules/attention.py b/groot/vla/model/dreamzero/modules/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..0ebc62a64a02b84cbecbc682e5612b0e290a87d4 --- /dev/null +++ b/groot/vla/model/dreamzero/modules/attention.py @@ -0,0 +1,258 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import torch +import os + +try: + import flash_attn_interface + FLASH_ATTN_3_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_3_AVAILABLE = False + +try: + import flash_attn + FLASH_ATTN_2_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_2_AVAILABLE = False + +import warnings + + +__all__ = [ + 'flash_attention', + 'attention', +] + + +def _gpu_supports_flash_attention(): + """FlashAttention requires Ampere (compute capability 8.0) or newer.""" + if not (FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE): + return False + try: + if not torch.cuda.is_available(): + return False + cap = torch.cuda.get_device_capability() + return cap[0] >= 8 + except Exception: + return False + + +def _sdpa_attention_fallback( + q, k, v, + q_lens=None, + k_lens=None, + dropout_p=0., + softmax_scale=None, + q_scale=None, + causal=False, + dtype=torch.bfloat16, +): + """PyTorch SDPA fallback for GPUs that don't support FlashAttention (e.g. pre-Ampere).""" + if q_lens is not None or k_lens is not None: + warnings.warn( + 'Padding mask is disabled when using scaled_dot_product_attention on this GPU. ' + 'It can have a slight impact on quality.' + ) + q = q.transpose(1, 2).to(dtype) + k = k.transpose(1, 2).to(dtype) + v = v.transpose(1, 2).to(dtype) + if q_scale is not None: + q = q * q_scale + if softmax_scale is not None: + q = q * softmax_scale + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=None, is_causal=causal, dropout_p=dropout_p + ) + return out.transpose(1, 2).contiguous() + + +def flash_attention( + q, + k, + v, + q_lens=None, + k_lens=None, + dropout_p=0., + softmax_scale=None, + q_scale=None, + causal=False, + window_size=(-1, -1), + deterministic=False, + dtype=torch.bfloat16, + version=None, +): + """ + q: [B, Lq, Nq, C1]. + k: [B, Lk, Nk, C1]. + v: [B, Lk, Nk, C2]. Nq must be divisible by Nk. + q_lens: [B]. + k_lens: [B]. + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + causal: bool. Whether to apply causal attention mask. + window_size: (left right). If not (-1, -1), apply sliding window local attention. + deterministic: bool. If True, slightly slower and uses more memory. + dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16. + """ + half_dtypes = (torch.float16, torch.bfloat16) + assert dtype in half_dtypes + assert q.device.type == 'cuda' and q.size(-1) <= 256 + + # Use PyTorch SDPA on pre-Ampere GPUs (FlashAttention requires Ampere or newer) + if not _gpu_supports_flash_attention(): + return _sdpa_attention_fallback( + q, k, v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + dtype=dtype, + ) + + # params + b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype + + def half(x): + return x if x.dtype in half_dtypes else x.to(dtype) + + # preprocess query + if q_lens is None: + q = half(q.flatten(0, 1)) + q_lens = torch.tensor( + [lq] * b, dtype=torch.int32).to( + device=q.device, non_blocking=True) + else: + q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)])) + + # preprocess key, value + if k_lens is None: + k = half(k.flatten(0, 1)) + v = half(v.flatten(0, 1)) + k_lens = torch.tensor( + [lk] * b, dtype=torch.int32).to( + device=k.device, non_blocking=True) + else: + k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)])) + v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)])) + + q = q.to(v.dtype) + k = k.to(v.dtype) + + if q_scale is not None: + q = q * q_scale + + if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE: + warnings.warn( + 'Flash attention 3 is not available, use flash attention 2 instead.' + ) + + # Check for TensorRT at runtime, not import time + if os.getenv("ENABLE_TENSORRT", "False").lower() == "true": + # use torch.nn.functional.scaled_dot_product_attention for tensorrt export + + # The input is (s, n, d), but sdpa needs (b, n, s, d). + # We add a batch dimension and transpose. + q = q.unsqueeze(0).transpose(1, 2) + k = k.unsqueeze(0).transpose(1, 2) + v = v.unsqueeze(0).transpose(1, 2) + + # Fix for ONNX export: repeat k and v to match q's batch size in cross-attention + if q.shape[0] != k.shape[0] and k.shape[0] == 1: + k = k.repeat(q.shape[0], 1, 1, 1) + v = v.repeat(q.shape[0], 1, 1, 1) + + attn_mask = None + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p) + + # Transpose back to (b, s, n, d) format. + out = out.transpose(1, 2).contiguous() + return out + + elif (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE: + # Note: dropout_p, window_size are not supported in FA3 now. + x = flash_attn_interface.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum( + 0, dtype=torch.int32).to(q.device, non_blocking=True), + cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum( + 0, dtype=torch.int32).to(q.device, non_blocking=True), + seqused_q=None, + seqused_k=None, + max_seqlen_q=lq, + max_seqlen_k=lk, + softmax_scale=softmax_scale, + causal=causal, + deterministic=deterministic)[0].unflatten(0, (b, lq)) + else: + assert FLASH_ATTN_2_AVAILABLE + x = flash_attn.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum( + 0, dtype=torch.int32).to(q.device, non_blocking=True), + cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum( + 0, dtype=torch.int32).to(q.device, non_blocking=True), + max_seqlen_q=lq, + max_seqlen_k=lk, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic).unflatten(0, (b, lq)) + + # output + return x.type(out_dtype) + + +def attention( + q, + k, + v, + q_lens=None, + k_lens=None, + dropout_p=0., + softmax_scale=None, + q_scale=None, + causal=False, + window_size=(-1, -1), + deterministic=False, + dtype=torch.bfloat16, + fa_version=None, +): + if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE: + return flash_attention( + q=q, + k=k, + v=v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + dtype=dtype, + version=fa_version, + ) + else: + if q_lens is not None or k_lens is not None: + warnings.warn( + 'Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance.' + ) + attn_mask = None + + q = q.transpose(1, 2).to(dtype) + k = k.transpose(1, 2).to(dtype) + v = v.transpose(1, 2).to(dtype) + + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p) + + out = out.transpose(1, 2).contiguous() + return out diff --git a/groot/vla/model/dreamzero/modules/cudnn_attention.py b/groot/vla/model/dreamzero/modules/cudnn_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..c996baa33c4f8212652234a8f6d9ba7645d3c234 --- /dev/null +++ b/groot/vla/model/dreamzero/modules/cudnn_attention.py @@ -0,0 +1,610 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +""" +Simplified wrapper around TransformerEngine's C++ pytorch backend. +This supports torch.compile(fullgraph=True). +Lowers to cudnn ultimately. +Only bf16 / fp16 is supported. +Only BSHD layout is supported. +Currently, tensors are made contiguous -- packed th2d, th3d not supported yet. +""" + +import math +from typing import Any, List, Optional, Tuple, Union + +import torch +import transformer_engine + +_TE_VER = tuple(int(x) for x in transformer_engine.__version__.split(".")[:2]) + + +try: + # transformer_engine >= 2.8.0 + import transformer_engine.pytorch.attention.dot_product_attention.utils as dpa_utils +except ImportError: + # transformer_engine < 2.8.0 + import transformer_engine.pytorch.dot_product_attention.utils as dpa_utils + +import transformer_engine_torch as tex +from transformer_engine.pytorch.constants import ( + TE_DType, +) +from transformer_engine.pytorch.cpp_extensions.fused_attn import ( + AttnBiasType, + AttnMaskType, + FusedAttnBackend, + QKVLayout, +) + +if _TE_VER >= (2, 8): + from transformer_engine.pytorch.cpp_extensions.fused_attn import SoftmaxType + +from transformer_engine.pytorch.utils import get_cudnn_version + +__all__ = ["DotProductAttention"] + + +class DotProductAttention(torch.nn.Module): + def __init__( + self, + num_attention_heads: int, + kv_channels: Union[int, Tuple[int, int]], + num_gqa_groups: Optional[int] = None, + attention_dropout: float = 0.0, + qkv_format: str = "bshd", + attn_mask_type: str = "no_mask", + window_size: Optional[Tuple[int, int]] = None, + ) -> None: + super().__init__() + assert qkv_format == "bshd", "Only bshd layout is supported." + + self.qkv_format = qkv_format + self.attn_mask_type = attn_mask_type + + self.softmax_scale = 1.0 / math.sqrt(kv_channels if isinstance(kv_channels, int) else kv_channels[0]) + self.attention_dropout = attention_dropout + self.softmax_type = "vanilla" + self.window_size = dpa_utils.check_set_window_size(attn_mask_type) + + self.fused_attention = FusedAttention( + self.softmax_scale, + deterministic=False, + attention_dropout=self.attention_dropout, + softmax_type=self.softmax_type, + ) + + def forward( + self, + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: torch.Tensor, + ) -> torch.Tensor: + """ + Dot Product Attention Layer. + """ + + qkv_layout = "bshd_bshd_bshd" + batch_size = query_layer.shape[0] + device = query_layer.device + + def _get_cu_seqlens(max_seqlen: int) -> torch.Tensor: + return torch.arange( + 0, + (batch_size + 1) * max_seqlen, + step=max_seqlen, + dtype=torch.int32, + device=device, + ) + + max_seqlen_q = query_layer.shape[1] + max_seqlen_kv = key_layer.shape[1] + cu_seqlens_q = _get_cu_seqlens(max_seqlen_q) + cu_seqlens_kv = _get_cu_seqlens(max_seqlen_kv) + + return self.fused_attention( + query_layer, + key_layer, + value_layer, + qkv_layout=qkv_layout, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + cu_seqlens_q_padded=cu_seqlens_q, + cu_seqlens_kv_padded=cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + attn_mask_type=self.attn_mask_type, + window_size=self.window_size, + fused_attention_backend=FusedAttnBackend["F16_arbitrary_seqlen"], + ) + + +def prepare_for_saving(*tensors) -> Tuple[list[torch.Tensor | None], list[torch.Tensor | None]]: + """Prepare tensors for saving. Needed because save_for_backward accepts only + torch.Tensor/torch.nn.Parameter types, while we want to be able to save + the internal TensorBase types too.""" + + tensor_list, tensor_objects_list = [], [] + for tensor in tensors: + if tensor is None or isinstance(tensor, torch.Tensor): + tensor_list.append(tensor) + tensor_objects_list.append(None) + else: + t, t_obj = tensor.prepare_for_saving() + tensor_list.extend(t) + tensor_objects_list.append(t_obj) + return tensor_list, tensor_objects_list + + +def restore_from_saved(tensors, saved_tensors) -> Tuple[Any, ...]: + """Recombine the tensor data and metadata during backward pass.""" + tensor_objects = [] + for tensor in tensors: + if tensor is None or isinstance(tensor, torch.Tensor): + tensor_objects.append(saved_tensors[0]) + saved_tensors = saved_tensors[1:] + else: + saved_tensors = tensor.restore_from_saved(saved_tensors) + tensor_objects.append(tensor) + + return tuple(tensor_objects) + + +class FusedAttention(torch.nn.Module): + def __init__( + self, + softmax_scale: float, + attention_dropout: float = 0.0, + deterministic: bool = False, + softmax_type: str = "vanilla", + ) -> None: + super().__init__() + self.softmax_scale = softmax_scale + self.attention_dropout = attention_dropout + self.deterministic = deterministic + self.softmax_type = softmax_type + + def forward( + self, + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: torch.Tensor, + qkv_layout: str, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + cu_seqlens_q_padded: torch.Tensor, + cu_seqlens_kv_padded: torch.Tensor, + max_seqlen_q: int, + max_seqlen_kv: int, + attn_mask_type: str = "causal", + attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None, + window_size: Optional[Tuple[int, int]] = None, + fused_attention_backend: tex.NVTE_Fused_Attn_Backend = tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend, + core_attention_bias_type: str = "no_bias", + core_attention_bias: Optional[torch.Tensor] = None, + fast_zero_fill: bool = True, + quantizers=None, + pad_between_seqs: bool = False, + softmax_offset: torch.Tensor | None = None, + softmax_scale: float = None, + attention_dropout: float = 0.0, + deterministic: bool = False, + softmax_type: str = "vanilla", + ) -> torch.Tensor: + """fused attention fprop""" + + cu_seqlens_q_padded = cu_seqlens_q + cu_seqlens_kv_padded = cu_seqlens_kv + + out_nominal_dtype = query_layer.dtype + output_tensors = fused_attn( + self.training, + max_seqlen_q, + max_seqlen_kv, + cu_seqlens_q, + cu_seqlens_kv, + query_layer, + key_layer, + value_layer, + out_nominal_dtype, + window_size, + core_attention_bias, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + None, # page_table_k, + None, # page_table_v, + self.softmax_scale, + self.attention_dropout if self.training else 0.0, + fast_zero_fill, + qkv_layout, + core_attention_bias_type, + attn_mask_type, + self.softmax_type, + self.deterministic, + softmax_offset, + ) + return output_tensors[0] + + +BACKEND_F16arb_ELTS_PER_THREADS = 16 + + +@torch.library.custom_op("groot::fused_attn", mutates_args=()) +def fused_attn( + is_training: bool, + max_seqlen_q: int, + max_seqlen_kv: int, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + fake_dtype: torch.dtype, + window_size: List[int], + attn_bias: torch.Tensor = None, + cu_seqlens_q_padded: torch.Tensor = None, + cu_seqlens_kv_padded: torch.Tensor = None, + page_table_k: torch.Tensor = None, + page_table_v: torch.Tensor = None, + attn_scale: Optional[float] = None, + dropout: float = 0.0, + fast_zero_fill: bool = True, + qkv_layout: str = "sbh3d", + attn_bias_type: str = "no_bias", + attn_mask_type: str = "padding", + softmax_type: str = "vanilla", + deterministic: bool = False, + softmax_offset: torch.Tensor = None, +) -> List[torch.Tensor]: + assert deterministic is not None + + rng_elts_per_thread = BACKEND_F16arb_ELTS_PER_THREADS + s_quantizer = None + o_quantizer = None + rng_gen = None + + q = q.contiguous() + k = k.contiguous() + v = v.contiguous() + + args = ( + max_seqlen_q, + max_seqlen_kv, + is_training, + attn_scale, + dropout, + fast_zero_fill, + QKVLayout[qkv_layout], + AttnBiasType[attn_bias_type], + AttnMaskType[attn_mask_type], + ) + + if _TE_VER >= (2, 8): + args += (SoftmaxType[softmax_type],) + + args += ( + tuple(window_size), + cu_seqlens_q, + cu_seqlens_kv, + q, + k, + v, + fake_dtype, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + page_table_k, + page_table_v, + s_quantizer, + o_quantizer, + attn_bias, + ) + + if _TE_VER >= (2, 8): + args += (softmax_offset,) + + args += ( + rng_gen, + rng_elts_per_thread, + ) + + if _TE_VER >= (2, 9): + # return_max_logit + args += (False,) + + if _TE_VER >= (2, 10): + # is_cuda_graph + args += (False,) + + output_tensors = tex.fused_attn_fwd(*args) + return output_tensors + + +@fused_attn.register_fake +def _( + is_training: bool, + max_seqlen_q: int, + max_seqlen_kv: int, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + fake_dtype: torch.dtype, + window_size: List[int], + attn_bias: torch.Tensor = None, + cu_seqlens_q_padded: torch.Tensor = None, + cu_seqlens_kv_padded: torch.Tensor = None, + page_table_k: torch.Tensor = None, + page_table_v: torch.Tensor = None, + attn_scale: Optional[float] = None, + dropout: float = 0.0, + fast_zero_fill: bool = True, + qkv_layout: str = "sbh3d", + attn_bias_type: str = "no_bias", + attn_mask_type: str = "padding", + softmax_type: str = "vanilla", + deterministic: bool = False, + softmax_offset: torch.Tensor = None, +) -> List[torch.Tensor]: + return [ + q.new_empty(tuple(q.shape[:-1]) + (v.shape[-1],)), + q.new_empty( + q.shape[0], q.shape[2], q.shape[1], 1, dtype=torch.float32 + ), # these are the softmax outputs from cudnn; will always be float32 + q.new_empty((2,), dtype=torch.int64), + ] + + +def fused_attn_bwd_setup_context(ctx, inputs, output) -> None: + ( + _, + max_seqlen_q, + max_seqlen_kv, + cu_seqlens_q, + cu_seqlens_kv, + q, + k, + v, + _, + window_size, + _, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + _, + _, + attn_scale, + dropout, + fast_zero_fill, + qkv_layout, + attn_bias_type, + attn_mask_type, + softmax_type, + deterministic, + _, + ) = inputs + + out = output[0] + aux_ctx_tensors = output[1:] + qkvo_tensors = (q, k, v, out) + + # assume fwd and bwd always use the same high precision, i.e. torch.float16 or torch.bfloat16 + # used when some tensors are base tensors and loose the "dtype" attribute + ctx.nominal_dtype = q.dtype + + tensors_to_save, tensor_objects = prepare_for_saving( + *qkvo_tensors, + cu_seqlens_q, + cu_seqlens_kv, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + *aux_ctx_tensors, + ) + ctx.save_for_backward(*tensors_to_save) + ctx.tensor_objects = tensor_objects + + ctx.QKV_quantizer = None + ctx.O_quantizer = None + ctx.dQKV_quantizer = None + ctx.dO_quantizer = None + ctx.dP_quantizer = None + ctx.S_quantizer = None + + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_kv = max_seqlen_kv + ctx.attn_scale = attn_scale + ctx.dropout_p = dropout + ctx.fast_zero_fill = fast_zero_fill + ctx.qkv_layout = qkv_layout + ctx.attn_bias_type = attn_bias_type + ctx.attn_mask_type = attn_mask_type + ctx.softmax_type = softmax_type + ctx.window_size = window_size + ctx.fused_attention_backend = FusedAttnBackend["F16_arbitrary_seqlen"] + ctx.deterministic = deterministic + + +@torch.library.custom_op("groot::fused_attn_bwd_op", mutates_args=()) +def fused_attn_bwd_op( + max_seqlen_q: int, + max_seqlen_kv: int, + attn_scale: float, + dropout: float, + fast_zero_fill: bool, + qkv_layout: str, + attn_bias_type: str, + attn_mask_type: str, + softmax_type: str, + window_size: List[int], + deterministic: bool, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + d_out: torch.Tensor, + dqkv_nominal_dtype: torch.dtype, + dqkv_te_dtype: torch.dtype, + aux_ctx_tensors: List[torch.Tensor], + cu_seqlens_q_padded: torch.Tensor, + cu_seqlens_kv_padded: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + args = ( + max_seqlen_q, + max_seqlen_kv, + attn_scale, + dropout, + fast_zero_fill, + QKVLayout[qkv_layout], + AttnBiasType[attn_bias_type], + AttnMaskType[attn_mask_type], + ) + + if _TE_VER >= (2, 8): + args += (SoftmaxType[softmax_type],) + + args += ( + window_size, + deterministic, + cu_seqlens_q, + cu_seqlens_kv, + q, + k, + v, + out, + d_out, + dqkv_nominal_dtype, + TE_DType[dqkv_te_dtype], + aux_ctx_tensors, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + None, # s_quantizer, + None, # dp_quantizer, + None, # dqkv_quantizer, + ) + + if _TE_VER >= (2, 10): + # is_cuda_graph + args += (False,) + + dq, dk, dv, *rest = tex.fused_attn_bwd(*args) + return dq, dk, dv + + +@fused_attn_bwd_op.register_fake +def _( + max_seqlen_q: int, + max_seqlen_kv: int, + attn_scale: float, + dropout: float, + fast_zero_fill: bool, + qkv_layout: str, + attn_bias_type: str, + attn_mask_type: str, + softmax_type: str, + window_size: List[int], + deterministic: bool, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + d_out: torch.Tensor, + dqkv_nominal_dtype: torch.dtype, + dqkv_te_dtype: torch.dtype, + aux_ctx_tensors: List[torch.Tensor], + cu_seqlens_q_padded: torch.Tensor, + cu_seqlens_kv_padded: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + + +def fused_attn_bwd_impl(ctx, grad): + d_out, _, _ = grad + d_out = d_out.contiguous() + + ( + q, + k, + v, + out, + cu_seqlens_q, + cu_seqlens_kv, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + *other_tensors, + ) = restore_from_saved(ctx.tensor_objects, ctx.saved_tensors) + + aux_ctx_tensors = other_tensors + + if not aux_ctx_tensors[0].is_contiguous(): + aux_ctx_tensors[0] = aux_ctx_tensors[0].contiguous() + + with torch.cuda.nvtx.range("FusedAttnFunc.backward"): + assert ctx.fused_attention_backend != FusedAttnBackend["No_Backend"], ( + "Fused attention does not support this input combination." + ) + + # get nominal data type of dq, dk, dv + # FP16/BF16 attention: torch.float16 or torch.bfloat16 + dqkv_nominal_dtype = ctx.nominal_dtype + + # q, k, v, out, d_out, dq, dk, dv: torch.Tensor; torch.float16 or torch.bfloat16 + dq, dk, dv = fused_attn_bwd_op( + ctx.max_seqlen_q, + ctx.max_seqlen_kv, + ctx.attn_scale, + ctx.dropout_p, + ctx.fast_zero_fill, + ctx.qkv_layout, + ctx.attn_bias_type, + ctx.attn_mask_type, + ctx.softmax_type, + ctx.window_size, + ctx.deterministic, + cu_seqlens_q, + cu_seqlens_kv, + q, + k, + v, + out, + d_out, + dqkv_nominal_dtype, + d_out.dtype, + aux_ctx_tensors, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + ) + + output = ( + None, # is_training + None, # max_seqlen_q + None, # max_seqlen_kv + None, # cu_seqlens_q + None, # cu_seqlens_kv + dq, + dk, + dv, + None, # fake_dtype + None, # window_size + None, # d_bias, # attn_bias + None, # cu_seqlens_q_padded + None, # cu_seqlens_kv_padded + None, # page_table_k + None, # page_table_v + None, # attn_scale + None, # dropout + None, # fast_zero_fill + None, # qkv_layout + None, # attn_bias_type + None, # attn_mask_type + None, # softmax_type + None, # deterministic + None, # d_softmax_offset, # softmax_offset + ) + return output + + +fused_attn.register_autograd(fused_attn_bwd_impl, setup_context=fused_attn_bwd_setup_context) diff --git a/groot/vla/model/dreamzero/modules/flow_match_scheduler.py b/groot/vla/model/dreamzero/modules/flow_match_scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..03b53b3b26237ef82492d174c78b971a22280a0f --- /dev/null +++ b/groot/vla/model/dreamzero/modules/flow_match_scheduler.py @@ -0,0 +1,92 @@ +import torch + + + +class FlowMatchScheduler(): + + def __init__(self, num_inference_steps=100, num_train_timesteps=1000, shift=3.0, sigma_max=1.0, sigma_min=0.003/1.002, inverse_timesteps=False, extra_one_step=False, reverse_sigmas=False): + self.num_train_timesteps = num_train_timesteps + self.shift = shift + self.sigma_max = sigma_max + self.sigma_min = sigma_min + self.inverse_timesteps = inverse_timesteps + self.extra_one_step = extra_one_step + self.reverse_sigmas = reverse_sigmas + self.set_timesteps(num_inference_steps) + + + def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, training=False, shift=None): + if shift is not None: + self.shift = shift + sigma_start = self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength + if self.extra_one_step: + self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps + 1)[:-1] + else: + self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps) + if self.inverse_timesteps: + self.sigmas = torch.flip(self.sigmas, dims=[0]) + self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas) + if self.reverse_sigmas: + self.sigmas = 1 - self.sigmas + self.timesteps = self.sigmas * self.num_train_timesteps + if training: + x = self.timesteps + y = torch.exp(-2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2) + y_shifted = y - y.min() + bsmntw_weighing = y_shifted * (num_inference_steps / y_shifted.sum()) + self.linear_timesteps_weights = bsmntw_weighing + self.training = True + else: + self.training = False + + + def step(self, model_output, timestep, sample, to_final=False, **kwargs): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep_id = torch.argmin((self.timesteps - timestep).abs()) + sigma = self.sigmas[timestep_id] + if to_final or timestep_id + 1 >= len(self.timesteps): + sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0 + else: + sigma_ = self.sigmas[timestep_id + 1] + prev_sample = sample + model_output * (sigma_ - sigma) + return prev_sample + + + def return_to_timestep(self, timestep, sample, sample_stablized): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep_id = torch.argmin((self.timesteps - timestep).abs()) + sigma = self.sigmas[timestep_id] + model_output = (sample - sample_stablized) / sigma + return model_output + + + # def add_noise(self, original_samples, noise, timestep): + # if isinstance(timestep, torch.Tensor): + # timestep = timestep.cpu() + # timestep_id = torch.argmin((self.timesteps - timestep).abs()) + # sigma = self.sigmas[timestep_id] + # sample = (1 - sigma) * original_samples + sigma * noise + # return sample + + def add_noise(self, original_samples, noise, timestep): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep_id = torch.argmin((self.timesteps.unsqueeze(1) - timestep.unsqueeze(0)).abs(), dim = 0) + sigma = self.sigmas[timestep_id].to(device=original_samples.device, dtype=original_samples.dtype) + while len(sigma.shape) < len(original_samples.shape): + sigma = sigma.unsqueeze(-1) + sample = (1 - sigma) * original_samples + sigma * noise + return sample + + def training_target(self, sample, noise, timestep): + target = noise - sample + return target + + + def training_weight(self, timestep): + # timestep_id = torch.argmin((self.timesteps - timestep.to(self.timesteps.device)).abs()) + timestep_id = torch.argmin((self.timesteps.unsqueeze(1) - timestep.unsqueeze(0).to(self.timesteps.device)).abs(), dim = 0) + weights = self.linear_timesteps_weights[timestep_id] + return weights \ No newline at end of file diff --git a/groot/vla/model/dreamzero/modules/flow_unipc_multistep_scheduler.py b/groot/vla/model/dreamzero/modules/flow_unipc_multistep_scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..df7ea56e278813a242cbb1cec873edc79649f8a5 --- /dev/null +++ b/groot/vla/model/dreamzero/modules/flow_unipc_multistep_scheduler.py @@ -0,0 +1,680 @@ +# Copied from https://github.com/huggingface/diffusers/blob/v0.31.0/src/diffusers/schedulers/scheduling_unipc_multistep.py +# Convert unipc for flow matching +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. + +import math +from typing import List, Optional, Tuple, Union + +import numpy as np +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import ( + KarrasDiffusionSchedulers, + SchedulerMixin, + SchedulerOutput, +) + + +class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin): + """ + `UniPCMultistepScheduler` is a training-free framework designed for the fast sampling of diffusion models. + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. + solver_order (`int`, default `2`): + The UniPC order which can be any positive integer. The effective order of accuracy is `solver_order + 1` + due to the UniC. It is recommended to use `solver_order=2` for guided sampling, and `solver_order=3` for + unconditional sampling. + prediction_type (`str`, defaults to "flow_prediction"): + Prediction type of the scheduler function; must be `flow_prediction` for this scheduler, which predicts + the flow of the diffusion process. + thresholding (`bool`, defaults to `False`): + Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such + as Stable Diffusion. + dynamic_thresholding_ratio (`float`, defaults to 0.995): + The ratio for the dynamic thresholding method. Valid only when `thresholding=True`. + sample_max_value (`float`, defaults to 1.0): + The threshold value for dynamic thresholding. Valid only when `thresholding=True` and `predict_x0=True`. + predict_x0 (`bool`, defaults to `True`): + Whether to use the updating algorithm on the predicted x0. + solver_type (`str`, default `bh2`): + Solver type for UniPC. It is recommended to use `bh1` for unconditional sampling when steps < 10, and `bh2` + otherwise. + lower_order_final (`bool`, default `True`): + Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can + stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10. + disable_corrector (`list`, default `[]`): + Decides which step to disable the corrector to mitigate the misalignment between `epsilon_theta(x_t, c)` + and `epsilon_theta(x_t^c, c)` which can influence convergence for a large guidance scale. Corrector is + usually disabled during the first few steps. + solver_p (`SchedulerMixin`, default `None`): + Any other scheduler that if specified, the algorithm becomes `solver_p + UniC`. + use_karras_sigmas (`bool`, *optional*, defaults to `False`): + Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`, + the sigmas are determined according to a sequence of noise levels {σi}. + use_exponential_sigmas (`bool`, *optional*, defaults to `False`): + Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process. + timestep_spacing (`str`, defaults to `"linspace"`): + The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. + steps_offset (`int`, defaults to 0): + An offset added to the inference steps, as required by some model families. + final_sigmas_type (`str`, defaults to `"zero"`): + The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final + sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0. + """ + + _compatibles = [e.name for e in KarrasDiffusionSchedulers] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + solver_order: int = 2, + prediction_type: str = "flow_prediction", + shift: Optional[float] = 1.0, + use_dynamic_shifting=False, + thresholding: bool = False, + dynamic_thresholding_ratio: float = 0.995, + sample_max_value: float = 1.0, + predict_x0: bool = True, + solver_type: str = "bh2", + lower_order_final: bool = True, + disable_corrector: List[int] = [], + solver_p: SchedulerMixin = None, + timestep_spacing: str = "linspace", + steps_offset: int = 0, + final_sigmas_type: Optional[str] = "zero", # "zero", "sigma_min" + ): + + if solver_type not in ["bh1", "bh2"]: + if solver_type in ["midpoint", "heun", "logrho"]: + self.register_to_config(solver_type="bh2") + else: + raise NotImplementedError( + f"{solver_type} is not implemented for {self.__class__}") + + self.predict_x0 = predict_x0 + # setable values + self.num_inference_steps = None + alphas = np.linspace(1, 1 / num_train_timesteps, + num_train_timesteps)[::-1].copy() + sigmas = 1.0 - alphas + sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device='cuda') + + if not use_dynamic_shifting: + # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution + assert shift is not None, "shift must be not None when use_dynamic_shifting is False" + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + + self.sigmas = sigmas + self.timesteps = sigmas * num_train_timesteps + + self.model_outputs = [None] * solver_order + self.timestep_list = [None] * solver_order + self.lower_order_nums = 0 + self.disable_corrector = disable_corrector + self.solver_p = solver_p + self.last_sample = None + + self.sigma_min = self.sigmas[-1].item() + self.sigma_max = self.sigmas[0].item() + + # Modified from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.set_timesteps + def set_timesteps( + self, + num_inference_steps: int, + device: str | torch.device = None, + sigmas: np.ndarray[float] | None = None, + mu: float | None = None, + shift: float | None = None, + ): + """ + Sets the discrete timesteps used for the diffusion chain (to be run before inference). + Args: + num_inference_steps (`int`): + Total number of the spacing of the time steps. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + """ + self.num_inference_steps = num_inference_steps + + if self.config.use_dynamic_shifting and mu is None: + raise ValueError( + " you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`" + ) + + if sigmas is None: + sigmas = np.linspace( + self.sigma_max, + self.sigma_min, + num_inference_steps + 1, + ).copy()[:-1] + + if self.config.use_dynamic_shifting: + assert mu is not None + sigmas = self.time_shift(mu, 1.0, sigmas) + else: + if shift is None: + shift = self.config.shift + assert isinstance(shift, float) + assert sigmas is not None + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + + if self.config.final_sigmas_type == "sigma_min": + sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0])**0.5 + elif self.config.final_sigmas_type == "zero": + sigma_last = 0 + else: + raise ValueError( + f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" + ) + + timesteps = sigmas * self.config.num_train_timesteps + sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) + + self.sigmas = torch.from_numpy(sigmas).to(device=device) + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64) + + self.model_outputs = [None] * self.config.solver_order + self.lower_order_nums = 0 + self.last_sample = None + if self.solver_p: + self.solver_p.set_timesteps(self.num_inference_steps, device=device) + + # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample + def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: + """ + "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the + prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by + s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing + pixels from saturation at each step. We find that dynamic thresholding results in significantly better + photorealism as well as better image-text alignment, especially when using very large guidance weights." + + https://arxiv.org/abs/2205.11487 + """ + dtype = sample.dtype + batch_size, channels, *remaining_dims = sample.shape + + if dtype not in (torch.float32, torch.float64): + sample = sample.float( + ) # upcast for quantile calculation, and clamp not implemented for cpu half + + # Flatten sample for doing quantile calculation along each image + sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) + + abs_sample = sample.abs() # "a certain percentile absolute pixel value" + + s = torch.quantile( + abs_sample, self.config.dynamic_thresholding_ratio, dim=1) + s = torch.clamp( + s, min=1, max=self.config.sample_max_value + ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] + s = s.unsqueeze( + 1) # (batch_size, 1) because clamp will broadcast along dim=0 + sample = torch.clamp( + sample, -s, s + ) / s # "we threshold xt0 to the range [-s, s] and then divide by s" + + sample = sample.reshape(batch_size, channels, *remaining_dims) + sample = sample.to(dtype) + + return sample + + # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._sigma_to_t + def _sigma_to_t(self, sigma): + return sigma * self.config.num_train_timesteps + + @staticmethod + def _sigma_to_alpha_sigma_t(sigma): + return 1 - sigma, sigma + + # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.set_timesteps + def time_shift(self, mu: float, sigma: float, t: torch.Tensor): + return math.exp(mu) / (math.exp(mu) + (1 / t - 1)**sigma) + + def convert_model_output( + self, + model_output: torch.Tensor, + sample: torch.Tensor, + step_index: int, + ) -> torch.Tensor: + r""" + Convert the model output to the corresponding type the UniPC algorithm needs. + + Args: + model_output (`torch.Tensor`): + The direct output from the learned diffusion model. + timestep (`int`): + The current discrete timestep in the diffusion chain. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + step_index (`int`): + The current timestep index. + + Returns: + `torch.Tensor`: + The converted model output. + """ + if self.predict_x0: + if self.config.prediction_type == "flow_prediction": + sigma_t = self.sigmas[step_index] + x0_pred = sample - sigma_t * model_output + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`," + " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler." + ) + + if self.config.thresholding: + x0_pred = self._threshold_sample(x0_pred) + + return x0_pred + else: + if self.config.prediction_type == "flow_prediction": + sigma_t = self.sigmas[step_index] + epsilon = sample - (1 - sigma_t) * model_output + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`," + " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler." + ) + + if self.config.thresholding: + sigma_t = self.sigmas[step_index] + x0_pred = sample - sigma_t * model_output + x0_pred = self._threshold_sample(x0_pred) + epsilon = model_output + x0_pred + + return epsilon + + @torch.compile(mode="reduce-overhead", fullgraph=True, dynamic=False) + def multistep_uni_p_bh_update( + self, + model_output: torch.Tensor, + sample: torch.Tensor, + order: int, + step_index: int, + ) -> torch.Tensor: + """ + One step for the UniP (B(h) version). Alternatively, `self.solver_p` is used if is specified. + + Args: + model_output (`torch.Tensor`): + The direct output from the learned diffusion model at the current timestep. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + order (`int`): + The order of UniP at this timestep (corresponds to the *p* in UniPC-p). + step_index (`int`): + The current timestep index. + + Returns: + `torch.Tensor`: + The sample tensor at the previous timestep. + """ + model_output_list = self.model_outputs + + s0 = self.timestep_list[-1] + m0 = model_output_list[-1] + x = sample + + if self.solver_p: + x_t = self.solver_p.step(model_output, s0, x).prev_sample + return x_t + + sigma_t, sigma_s0 = self.sigmas[step_index + 1], self.sigmas[step_index] + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + + h = lambda_t - lambda_s0 + + rks = [] + D1s = [] + for i in range(1, order): + si = step_index - i + mi = model_output_list[-(i + 1)] + alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) + lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + rk = (lambda_si - lambda_s0) / h + rks.append(rk) + D1s.append((mi - m0) / rk) + + rks.append(torch.ones((), dtype=self.sigmas.dtype, device=self.sigmas.device)) + rks = torch.stack(rks, dim=0) + + R = [] + b = [] + + hh = -h if self.predict_x0 else h + h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1 + h_phi_k = h_phi_1 / hh - 1 + + factorial_i = 1 + + if self.config.solver_type == "bh1": + B_h = hh + elif self.config.solver_type == "bh2": + B_h = torch.expm1(hh) + else: + raise NotImplementedError() + + for i in range(1, order + 1): + R.append(torch.pow(rks, i - 1)) + b.append(h_phi_k * factorial_i / B_h) + factorial_i *= i + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + + R = torch.stack(R, dim=0) + b = torch.stack(b, dim=0) + + if len(D1s) > 0: + D1s = torch.stack(D1s, dim=1) # (B, K) + # for order 2, we use a simplified version + if order == 2: + rhos_p = torch.full((1,), 0.5, dtype=x.dtype, device=self.sigmas.device) + else: + rhos_p = torch.linalg.solve_ex(R[:-1, :-1], b[:-1])[0].to(x.dtype) + else: + D1s = None + rhos_p = None + + if self.predict_x0: + x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 + if D1s is not None: + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) + else: + pred_res = 0 + x_t = x_t_ - alpha_t * B_h * pred_res + else: + x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 + if D1s is not None: + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) + else: + pred_res = 0 + x_t = x_t_ - sigma_t * B_h * pred_res + + x_t = x_t.to(x.dtype) + return x_t + + @torch.compile(mode="reduce-overhead", fullgraph=True, dynamic=False) + def multistep_uni_c_bh_update( + self, + this_model_output: torch.Tensor, + last_sample: torch.Tensor, + this_sample: torch.Tensor, + order: int, + step_index: int, + ) -> torch.Tensor: + """ + One step for the UniC (B(h) version). + + Args: + this_model_output (`torch.Tensor`): + The model outputs at `x_t`. + this_timestep (`int`): + The current timestep `t`. + last_sample (`torch.Tensor`): + The generated sample before the last predictor `x_{t-1}`. + this_sample (`torch.Tensor`): + The generated sample after the last predictor `x_{t}`. + order (`int`): + The `p` of UniC-p at this step. The effective order of accuracy should be `order + 1`. + step_index (`int`): + The current timestep index. + + Returns: + `torch.Tensor`: + The corrected sample tensor at the current timestep. + """ + model_output_list = self.model_outputs + + m0 = model_output_list[-1] + x = last_sample + x_t = this_sample + model_t = this_model_output + + sigma_t, sigma_s0 = self.sigmas[step_index], self.sigmas[step_index - 1] + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + + h = lambda_t - lambda_s0 + + rks = [] + D1s = [] + for i in range(1, order): + si = step_index - (i + 1) + mi = model_output_list[-(i + 1)] + alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) + lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + rk = (lambda_si - lambda_s0) / h + rks.append(rk) + D1s.append((mi - m0) / rk) + + rks.append(torch.ones((), dtype=self.sigmas.dtype, device=self.sigmas.device)) + rks = torch.stack(rks, dim=0) + + R = [] + b = [] + + hh = -h if self.predict_x0 else h + h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1 + h_phi_k = h_phi_1 / hh - 1 + + factorial_i = 1 + + if self.config.solver_type == "bh1": + B_h = hh + elif self.config.solver_type == "bh2": + B_h = torch.expm1(hh) + else: + raise NotImplementedError() + + for i in range(1, order + 1): + R.append(torch.pow(rks, i - 1)) + b.append(h_phi_k * factorial_i / B_h) + factorial_i *= i + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + + R = torch.stack(R, dim=0) + b = torch.stack(b, dim=0) + + if len(D1s) > 0: + D1s = torch.stack(D1s, dim=1) + else: + D1s = None + + # for order 1, we use a simplified version + if order == 1: + rhos_c = torch.full((1,), 0.5, dtype=x.dtype, device=self.sigmas.device) + else: + rhos_c = torch.linalg.solve_ex(R, b)[0].to(x.dtype) + + if self.predict_x0: + x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 + if D1s is not None: + corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) + else: + corr_res = 0 + D1_t = model_t - m0 + x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t) + else: + x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 + if D1s is not None: + corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) + else: + corr_res = 0 + D1_t = model_t - m0 + x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t) + x_t = x_t.to(x.dtype) + return x_t + + def index_for_timestep(self, timestep, schedule_timesteps): + indices = (schedule_timesteps == timestep).nonzero() + + # The sigma index that is taken for the **very** first `step` + # is always the second index (or the last index if there is only 1) + # This way we can ensure we don't accidentally skip a sigma in + # case we start in the middle of the denoising schedule (e.g. for image-to-image) + pos = 1 if len(indices) > 1 else 0 + + return indices[pos].item() + + def step( + self, + model_output: torch.Tensor, + timestep: torch.Tensor, + sample: torch.Tensor, + step_index: int, + return_dict: bool = True, + ) -> SchedulerOutput | tuple: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with + the multistep UniPC. + + Args: + model_output (`torch.Tensor`): + The direct output from learned diffusion model. + timestep (`torch.Tensor`): + The current discrete timestep in the diffusion chain. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + step_index (`int`): + The current timestep index. + return_dict (`bool`): + Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`. + + Returns: + [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`: + If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a + tuple is returned where the first element is the sample tensor. + + """ + if self.num_inference_steps is None: + raise ValueError( + "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" + ) + + use_corrector = ( + step_index > 0 and + step_index - 1 not in self.disable_corrector and + self.last_sample is not None + ) + + model_output_convert = self.convert_model_output( + model_output=model_output, + sample=sample, + step_index=step_index, + ) + if use_corrector: + sample = self.multistep_uni_c_bh_update( + this_model_output=model_output_convert, + last_sample=self.last_sample, + this_sample=sample, + order=self.this_order, + step_index=step_index, + ) + # We must clone the outputs of a CUDA graph'd computation. + sample = sample.clone() + + for i in range(self.config.solver_order - 1): + self.model_outputs[i] = self.model_outputs[i + 1] + self.timestep_list[i] = self.timestep_list[i + 1] + + self.model_outputs[-1] = model_output_convert + self.timestep_list[-1] = timestep + + if self.config.lower_order_final: + this_order = min( + self.config.solver_order, + len(self.timesteps) - step_index, + ) + else: + this_order = self.config.solver_order + + # Warmup for multistep. + self.this_order = min(this_order, self.lower_order_nums + 1) + assert self.this_order > 0 + + self.last_sample = sample + # Pass the original non-converted model output, in case solver-p is used. + prev_sample = self.multistep_uni_p_bh_update( + model_output=model_output, + sample=sample, + order=self.this_order, + step_index=step_index, + ) + # We must clone the outputs of a CUDA graph'd computation. + prev_sample = prev_sample.clone() + + if self.lower_order_nums < self.config.solver_order: + self.lower_order_nums += 1 + + if not return_dict: + return (prev_sample,) + + return SchedulerOutput(prev_sample=prev_sample) + + def scale_model_input(self, sample: torch.Tensor, *args, + **kwargs) -> torch.Tensor: + """ + Ensures interchangeability with schedulers that need to scale the denoising model input depending on the + current timestep. + + Args: + sample (`torch.Tensor`): + The input sample. + + Returns: + `torch.Tensor`: + A scaled input sample. + """ + return sample + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.add_noise + def add_noise( + self, + original_samples: torch.Tensor, + noise: torch.Tensor, + timesteps: torch.IntTensor, + ) -> torch.Tensor: + # Make sure sigmas and timesteps have the same device and dtype as original_samples + sigmas = self.sigmas.to( + device=original_samples.device, dtype=original_samples.dtype) + if original_samples.device.type == "mps" and torch.is_floating_point(timesteps): + # mps does not support float64 + schedule_timesteps = self.timesteps.to( + original_samples.device, dtype=torch.float32) + timesteps = timesteps.to( + original_samples.device, dtype=torch.float32) + else: + schedule_timesteps = self.timesteps.to(original_samples.device) + timesteps = timesteps.to(original_samples.device) + + # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index + step_indices = [ + self.index_for_timestep(t, schedule_timesteps) + for t in timesteps + ] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < len(original_samples.shape): + sigma = sigma.unsqueeze(-1) + + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) + noisy_samples = alpha_t * original_samples + sigma_t * noise + return noisy_samples + + def __len__(self): + return self.config.num_train_timesteps diff --git a/groot/vla/model/dreamzero/modules/utils.py b/groot/vla/model/dreamzero/modules/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0d58e4e11dc844822d0588dac7f432e7e31b8daf --- /dev/null +++ b/groot/vla/model/dreamzero/modules/utils.py @@ -0,0 +1,182 @@ +import torch, os +from safetensors import safe_open +from contextlib import contextmanager +import hashlib + +@contextmanager +def init_weights_on_device(device = torch.device("meta"), include_buffers :bool = False): + + old_register_parameter = torch.nn.Module.register_parameter + if include_buffers: + old_register_buffer = torch.nn.Module.register_buffer + + def register_empty_parameter(module, name, param): + old_register_parameter(module, name, param) + if param is not None: + param_cls = type(module._parameters[name]) + kwargs = module._parameters[name].__dict__ + kwargs["requires_grad"] = param.requires_grad + module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs) + + def register_empty_buffer(module, name, buffer, persistent=True): + old_register_buffer(module, name, buffer, persistent=persistent) + if buffer is not None: + module._buffers[name] = module._buffers[name].to(device) + + def patch_tensor_constructor(fn): + def wrapper(*args, **kwargs): + kwargs["device"] = device + return fn(*args, **kwargs) + + return wrapper + + if include_buffers: + tensor_constructors_to_patch = { + torch_function_name: getattr(torch, torch_function_name) + for torch_function_name in ["empty", "zeros", "ones", "full"] + } + else: + tensor_constructors_to_patch = {} + + try: + torch.nn.Module.register_parameter = register_empty_parameter + if include_buffers: + torch.nn.Module.register_buffer = register_empty_buffer + for torch_function_name in tensor_constructors_to_patch.keys(): + setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name))) + yield + finally: + torch.nn.Module.register_parameter = old_register_parameter + if include_buffers: + torch.nn.Module.register_buffer = old_register_buffer + for torch_function_name, old_torch_function in tensor_constructors_to_patch.items(): + setattr(torch, torch_function_name, old_torch_function) + +def load_state_dict_from_folder(file_path, torch_dtype=None): + state_dict = {} + for file_name in os.listdir(file_path): + if "." in file_name and file_name.split(".")[-1] in [ + "safetensors", "bin", "ckpt", "pth", "pt" + ]: + state_dict.update(load_state_dict(os.path.join(file_path, file_name), torch_dtype=torch_dtype)) + return state_dict + + +def load_state_dict(file_path, torch_dtype=None, device="cpu"): + if file_path.endswith(".safetensors"): + return load_state_dict_from_safetensors(file_path, torch_dtype=torch_dtype, device=device) + else: + return load_state_dict_from_bin(file_path, torch_dtype=torch_dtype, device=device) + + +def load_state_dict_from_safetensors(file_path, torch_dtype=None, device="cpu"): + state_dict = {} + with safe_open(file_path, framework="pt", device=device) as f: + for k in f.keys(): + state_dict[k] = f.get_tensor(k) + if torch_dtype is not None: + state_dict[k] = state_dict[k].to(torch_dtype) + return state_dict + + +def load_state_dict_from_bin(file_path, torch_dtype=None, device="cpu"): + state_dict = torch.load(file_path, map_location=device, weights_only=True) + if torch_dtype is not None: + for i in state_dict: + if isinstance(state_dict[i], torch.Tensor): + state_dict[i] = state_dict[i].to(torch_dtype) + return state_dict + + +def search_for_embeddings(state_dict): + embeddings = [] + for k in state_dict: + if isinstance(state_dict[k], torch.Tensor): + embeddings.append(state_dict[k]) + elif isinstance(state_dict[k], dict): + embeddings += search_for_embeddings(state_dict[k]) + return embeddings + + +def search_parameter(param, state_dict): + for name, param_ in state_dict.items(): + if param.numel() == param_.numel(): + if param.shape == param_.shape: + if torch.dist(param, param_) < 1e-3: + return name + else: + if torch.dist(param.flatten(), param_.flatten()) < 1e-3: + return name + return None + + +def build_rename_dict(source_state_dict, target_state_dict, split_qkv=False): + matched_keys = set() + with torch.no_grad(): + for name in source_state_dict: + rename = search_parameter(source_state_dict[name], target_state_dict) + if rename is not None: + print(f'"{name}": "{rename}",') + matched_keys.add(rename) + elif split_qkv and len(source_state_dict[name].shape)>=1 and source_state_dict[name].shape[0]%3==0: + length = source_state_dict[name].shape[0] // 3 + rename = [] + for i in range(3): + rename.append(search_parameter(source_state_dict[name][i*length: i*length+length], target_state_dict)) + if None not in rename: + print(f'"{name}": {rename},') + for rename_ in rename: + matched_keys.add(rename_) + for name in target_state_dict: + if name not in matched_keys: + print("Cannot find", name, target_state_dict[name].shape) + + +def search_for_files(folder, extensions): + files = [] + if os.path.isdir(folder): + for file in sorted(os.listdir(folder)): + files += search_for_files(os.path.join(folder, file), extensions) + elif os.path.isfile(folder): + for extension in extensions: + if folder.endswith(extension): + files.append(folder) + break + return files + + +def convert_state_dict_keys_to_single_str(state_dict, with_shape=True): + keys = [] + for key, value in state_dict.items(): + if isinstance(key, str): + if isinstance(value, torch.Tensor): + if with_shape: + shape = "_".join(map(str, list(value.shape))) + keys.append(key + ":" + shape) + keys.append(key) + elif isinstance(value, dict): + keys.append(key + "|" + convert_state_dict_keys_to_single_str(value, with_shape=with_shape)) + keys.sort() + keys_str = ",".join(keys) + return keys_str + + +def split_state_dict_with_prefix(state_dict): + keys = sorted([key for key in state_dict if isinstance(key, str)]) + prefix_dict = {} + for key in keys: + prefix = key if "." not in key else key.split(".")[0] + if prefix not in prefix_dict: + prefix_dict[prefix] = [] + prefix_dict[prefix].append(key) + state_dicts = [] + for prefix, keys in prefix_dict.items(): + sub_state_dict = {key: state_dict[key] for key in keys} + state_dicts.append(sub_state_dict) + return state_dicts + + +def hash_state_dict_keys(state_dict, with_shape=True): + keys_str = convert_state_dict_keys_to_single_str(state_dict, with_shape=with_shape) + keys_str = keys_str.encode(encoding="UTF-8") + return hashlib.md5(keys_str).hexdigest() \ No newline at end of file diff --git a/groot/vla/model/dreamzero/modules/vram_management.py b/groot/vla/model/dreamzero/modules/vram_management.py new file mode 100644 index 0000000000000000000000000000000000000000..381d564caed84403b8b1062b53fed41b110668f2 --- /dev/null +++ b/groot/vla/model/dreamzero/modules/vram_management.py @@ -0,0 +1,212 @@ +import torch, copy +from groot.vla.model.dreamzero.modules.utils import init_weights_on_device + + +def cast_to(weight, dtype, device): + r = torch.empty_like(weight, dtype=dtype, device=device) + r.copy_(weight) + return r + + +class AutoTorchModule(torch.nn.Module): + def __init__(self): + super().__init__() + + def check_free_vram(self): + gpu_mem_state = torch.cuda.mem_get_info(self.computation_device) + used_memory = (gpu_mem_state[1] - gpu_mem_state[0]) / (1024 ** 3) + return used_memory < self.vram_limit + + def offload(self): + if self.state != 0: + self.to(dtype=self.offload_dtype, device=self.offload_device) + self.state = 0 + + def onload(self): + if self.state != 1: + self.to(dtype=self.onload_dtype, device=self.onload_device) + self.state = 1 + + def keep(self): + if self.state != 2: + self.to(dtype=self.computation_dtype, device=self.computation_device) + self.state = 2 + + +class AutoWrappedModule(AutoTorchModule): + def __init__(self, module: torch.nn.Module, offload_dtype, offload_device, onload_dtype, onload_device, computation_dtype, computation_device, vram_limit, **kwargs): + super().__init__() + self.module = module.to(dtype=offload_dtype, device=offload_device) + self.offload_dtype = offload_dtype + self.offload_device = offload_device + self.onload_dtype = onload_dtype + self.onload_device = onload_device + self.computation_dtype = computation_dtype + self.computation_device = computation_device + self.vram_limit = vram_limit + self.state = 0 + + def forward(self, *args, **kwargs): + if self.state == 2: + module = self.module + else: + if self.onload_dtype == self.computation_dtype and self.onload_device == self.computation_device: + module = self.module + elif self.vram_limit is not None and self.check_free_vram(): + self.keep() + module = self.module + else: + module = copy.deepcopy(self.module).to(dtype=self.computation_dtype, device=self.computation_device) + return module(*args, **kwargs) + + +class WanAutoCastLayerNorm(torch.nn.LayerNorm, AutoTorchModule): + def __init__(self, module: torch.nn.LayerNorm, offload_dtype, offload_device, onload_dtype, onload_device, computation_dtype, computation_device, vram_limit, **kwargs): + with init_weights_on_device(device=torch.device("meta")): + super().__init__(module.normalized_shape, eps=module.eps, elementwise_affine=module.elementwise_affine, bias=module.bias is not None, dtype=offload_dtype, device=offload_device) + self.weight = module.weight + self.bias = module.bias + self.offload_dtype = offload_dtype + self.offload_device = offload_device + self.onload_dtype = onload_dtype + self.onload_device = onload_device + self.computation_dtype = computation_dtype + self.computation_device = computation_device + self.vram_limit = vram_limit + self.state = 0 + + def forward(self, x, *args, **kwargs): + if self.state == 2: + weight, bias = self.weight, self.bias + else: + if self.onload_dtype == self.computation_dtype and self.onload_device == self.computation_device: + weight, bias = self.weight, self.bias + elif self.vram_limit is not None and self.check_free_vram(): + self.keep() + weight, bias = self.weight, self.bias + else: + weight = None if self.weight is None else cast_to(self.weight, self.computation_dtype, self.computation_device) + bias = None if self.bias is None else cast_to(self.bias, self.computation_dtype, self.computation_device) + with torch.amp.autocast(device_type=x.device.type): + x = torch.nn.functional.layer_norm(x.float(), self.normalized_shape, weight, bias, self.eps).type_as(x) + return x + + +class AutoWrappedLinear(torch.nn.Linear, AutoTorchModule): + def __init__(self, module: torch.nn.Linear, offload_dtype, offload_device, onload_dtype, onload_device, computation_dtype, computation_device, vram_limit, name="", **kwargs): + with init_weights_on_device(device=torch.device("meta")): + super().__init__(in_features=module.in_features, out_features=module.out_features, bias=module.bias is not None, dtype=offload_dtype, device=offload_device) + self.weight = module.weight + self.bias = module.bias + self.offload_dtype = offload_dtype + self.offload_device = offload_device + self.onload_dtype = onload_dtype + self.onload_device = onload_device + self.computation_dtype = computation_dtype + self.computation_device = computation_device + self.vram_limit = vram_limit + self.state = 0 + self.name = name + self.lora_A_weights = [] + self.lora_B_weights = [] + self.lora_merger = None + self.enable_fp8 = computation_dtype in [torch.float8_e4m3fn, torch.float8_e4m3fnuz] + + def fp8_linear( + self, + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + device = input.device + origin_dtype = input.dtype + origin_shape = input.shape + input = input.reshape(-1, origin_shape[-1]) + + x_max = torch.max(torch.abs(input), dim=-1, keepdim=True).values + fp8_max = 448.0 + # For float8_e4m3fnuz, the maximum representable value is half of that of e4m3fn. + # To avoid overflow and ensure numerical compatibility during FP8 computation, + # we scale down the input by 2.0 in advance. + # This scaling will be compensated later during the final result scaling. + if self.computation_dtype == torch.float8_e4m3fnuz: + fp8_max = fp8_max / 2.0 + scale_a = torch.clamp(x_max / fp8_max, min=1.0).float().to(device=device) + scale_b = torch.ones((weight.shape[0], 1)).to(device=device) + input = input / (scale_a + 1e-8) + input = input.to(self.computation_dtype) + weight = weight.to(self.computation_dtype) + bias = bias.to(torch.bfloat16) + + result = torch._scaled_mm( + input, + weight.T, + scale_a=scale_a, + scale_b=scale_b.T, + bias=bias, + out_dtype=origin_dtype, + ) + new_shape = origin_shape[:-1] + result.shape[-1:] + result = result.reshape(new_shape) + return result + + def forward(self, x, *args, **kwargs): + # VRAM management + if self.state == 2: + weight, bias = self.weight, self.bias + else: + if self.onload_dtype == self.computation_dtype and self.onload_device == self.computation_device: + weight, bias = self.weight, self.bias + elif self.vram_limit is not None and self.check_free_vram(): + self.keep() + weight, bias = self.weight, self.bias + else: + weight = cast_to(self.weight, self.computation_dtype, self.computation_device) + bias = None if self.bias is None else cast_to(self.bias, self.computation_dtype, self.computation_device) + + # Linear forward + if self.enable_fp8: + out = self.fp8_linear(x, weight, bias) + else: + out = torch.nn.functional.linear(x, weight, bias) + + # LoRA + if len(self.lora_A_weights) == 0: + # No LoRA + return out + elif self.lora_merger is None: + # Native LoRA inference + for lora_A, lora_B in zip(self.lora_A_weights, self.lora_B_weights): + out = out + x @ lora_A.T @ lora_B.T + else: + # LoRA fusion + lora_output = [] + for lora_A, lora_B in zip(self.lora_A_weights, self.lora_B_weights): + lora_output.append(x @ lora_A.T @ lora_B.T) + lora_output = torch.stack(lora_output) + out = self.lora_merger(out, lora_output) + return out + + +def enable_vram_management_recursively(model: torch.nn.Module, module_map: dict, module_config: dict, max_num_param=None, overflow_module_config: dict = None, total_num_param=0, vram_limit=None, name_prefix=""): + for name, module in model.named_children(): + layer_name = name if name_prefix == "" else name_prefix + "." + name + for source_module, target_module in module_map.items(): + if isinstance(module, source_module): + num_param = sum(p.numel() for p in module.parameters()) + if max_num_param is not None and total_num_param + num_param > max_num_param: + module_config_ = overflow_module_config + else: + module_config_ = module_config + module_ = target_module(module, **module_config_, vram_limit=vram_limit, name=layer_name) + setattr(model, name, module_) + total_num_param += num_param + break + else: + total_num_param = enable_vram_management_recursively(module, module_map, module_config, max_num_param, overflow_module_config, total_num_param, vram_limit=vram_limit, name_prefix=layer_name) + return total_num_param + + +def enable_vram_management(model: torch.nn.Module, module_map: dict, module_config: dict, max_num_param=None, overflow_module_config: dict = None, vram_limit=None): + enable_vram_management_recursively(model, module_map, module_config, max_num_param, overflow_module_config, total_num_param=0, vram_limit=vram_limit) + model.vram_management_enabled = True \ No newline at end of file diff --git a/groot/vla/model/dreamzero/modules/wan2_1_attention.py b/groot/vla/model/dreamzero/modules/wan2_1_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..5ffb826bf16fa222a9fdf1a35ae27aeae0dbbf03 --- /dev/null +++ b/groot/vla/model/dreamzero/modules/wan2_1_attention.py @@ -0,0 +1,348 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. + +import contextlib +import torch +from torch.profiler import profile, ProfilerActivity +import time +from typing import Optional +import os + +try: + import flash_attn_interface + + def is_hopper_gpu(): + if not torch.cuda.is_available(): + return False + device_name = torch.cuda.get_device_name(0).lower() + return "h100" in device_name or "hopper" in device_name + FLASH_ATTN_3_AVAILABLE = is_hopper_gpu() +except ModuleNotFoundError: + FLASH_ATTN_3_AVAILABLE = False + +try: + import flash_attn + FLASH_ATTN_2_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_2_AVAILABLE = False + +try: + import transformer_engine + from groot.vla.model.dreamzero.modules.cudnn_attention import DotProductAttention + TRANSFORMER_ENGINE_AVAILABLE = True +except ModuleNotFoundError: + TRANSFORMER_ENGINE_AVAILABLE = False + +import warnings + + +def _gpu_supports_flash_attention(): + """FlashAttention requires Ampere (compute capability 8.0) or newer.""" + if not (FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE): + return False + try: + if not torch.cuda.is_available(): + return False + cap = torch.cuda.get_device_capability() + return cap[0] >= 8 + except Exception: + return False + + +def _sdpa_attention_fallback( + q, k, v, + q_lens=None, + k_lens=None, + dropout_p=0., + softmax_scale=None, + q_scale=None, + causal=False, + dtype=torch.bfloat16, +): + """PyTorch SDPA fallback for GPUs that don't support FlashAttention (e.g. pre-Ampere).""" + if q_lens is not None or k_lens is not None: + warnings.warn( + 'Padding mask is disabled when using scaled_dot_product_attention on this GPU. ' + 'It can have a slight impact on quality.' + ) + q = q.transpose(1, 2).to(dtype) + k = k.transpose(1, 2).to(dtype) + v = v.transpose(1, 2).to(dtype) + if q_scale is not None: + q = q * q_scale + if softmax_scale is not None: + q = q * softmax_scale + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=None, is_causal=causal, dropout_p=dropout_p + ) + return out.transpose(1, 2).contiguous() + + +def flash_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_lens: Optional[torch.Tensor] = None, + k_lens: Optional[torch.Tensor] = None, + dropout_p: float = 0., + softmax_scale: Optional[float] = None, + q_scale: Optional[float] = None, + causal: bool = False, + window_size: Optional[tuple[int, int]] = None, + deterministic: bool = False, + dtype: torch.dtype = torch.bfloat16, + version: Optional[int] = None, +): + """ + q: [B, Lq, Nq, C1]. + k: [B, Lk, Nk, C1]. + v: [B, Lk, Nk, C2]. Nq must be divisible by Nk. + q_lens: [B]. + k_lens: [B]. + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + causal: bool. Whether to apply causal attention mask. + window_size: (left right). If not (-1, -1), apply sliding window local attention. + deterministic: bool. If True, slightly slower and uses more memory. + dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16. + version: int. 2 for flash attention 2, 3 for flash attention 3. + + Returns: + x: [B, Lq, Nq, C2]. + """ + if window_size is None: + window_size = (-1, -1) + if version is None: + version = 3 + + half_dtypes = (torch.float16, torch.bfloat16) + assert dtype in half_dtypes + assert q.device.type == 'cuda' and q.size(-1) <= 256 + + # Use PyTorch SDPA on pre-Ampere GPUs (FlashAttention requires Ampere or newer) + if not _gpu_supports_flash_attention(): + return _sdpa_attention_fallback( + q, k, v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + dtype=dtype, + ) + + # params + b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype + + def half(x): + return x if x.dtype in half_dtypes else x.to(dtype) + + # preprocess query + if q_lens is None: + q = half(q.flatten(0, 1)) + q_lens = torch.tensor([lq] * b, dtype=torch.int32, device=q.device) + else: + q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)])) + + # preprocess key, value + if k_lens is None: + k = half(k.flatten(0, 1)) + v = half(v.flatten(0, 1)) + k_lens = torch.tensor([lk] * b, dtype=torch.int32, device=k.device) + else: + k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)])) + v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)])) + + q = q.to(v.dtype) + k = k.to(v.dtype) + + if q_scale is not None: + q = q * q_scale + + if version == 3 and not FLASH_ATTN_3_AVAILABLE: + warnings.warn( + 'Flash attention 3 is not available, use flash attention 2 instead.' + ) + zeros = torch.zeros([1], dtype=torch.int32, device=q.device) + cu_seqlens_q = torch.cat([zeros, q_lens]).cumsum(0).to(torch.int32) + cu_seqlens_k = torch.cat([zeros, k_lens]).cumsum(0).to(torch.int32) + + # apply attention + if version == 3 and FLASH_ATTN_3_AVAILABLE: + # Note: dropout_p, window_size are not supported in FA3 now. + x = flash_attn_interface.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=lq, + max_seqlen_k=lk, + softmax_scale=softmax_scale, + causal=causal, + deterministic=deterministic)[0].unflatten(0, (b, lq)) + elif FLASH_ATTN_2_AVAILABLE: + x = flash_attn.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=lq, + max_seqlen_k=lk, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic).unflatten(0, (b, lq)) + else: + raise ValueError(f"Invalid version: {version}") + + # output + return x.type(out_dtype) + + +class AttentionModule(torch.nn.Module): + def __init__( + self, + num_heads: int, + head_dim: int, + dropout_p: float = 0., + softmax_scale: Optional[float] = None, + q_scale: Optional[float] = None, + causal: bool = False, + window_size: Optional[tuple[int, int]] = None, + deterministic: bool = False, + dtype: torch.dtype = torch.bfloat16, + backend: Optional[str] = None, + ): + super().__init__() + if backend is None: + backend = "torch" + + if os.getenv("ATTENTION_BACKEND") is not None: + backend = os.getenv("ATTENTION_BACKEND") + else: + backend = "FA2" + + # Check for TensorRT at runtime, not import time + if os.getenv("ENABLE_TENSORRT", "False").lower() == "true": + backend = "torch" + + # Fall back to FA backend if TE is specified but not available + if backend == "TE" and not TRANSFORMER_ENGINE_AVAILABLE: + print("Warning: Transformer Engine is not available. Falling back to FA2 backend.") + backend = "FA2" + + assert backend in ["torch", "FA2", "FA3", "TE", "torch_onnx"] + self.backend = backend + + if backend == "torch": + def _torch_impl(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + out_dtype = q.dtype + q = q.transpose(1, 2).to(dtype) + k = k.transpose(1, 2).to(dtype) + v = v.transpose(1, 2).to(dtype) + + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, + attn_mask=None, + is_causal=causal, + dropout_p=dropout_p, + scale=softmax_scale, + ) + + out = out.transpose(1, 2).contiguous() + return out.to(out_dtype) + self.attn_func = _torch_impl + + elif backend == "torch_onnx": + def _torch_onnx_impl(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + out_dtype = q.dtype + # use torch.nn.functional.scaled_dot_product_attention for tensorrt export + + # The input is (s, n, d), but sdpa needs (b, h, s, d). + # We add a batch dimension and transpose. + q = q.unsqueeze(0).transpose(1, 2).to(dtype) + k = k.unsqueeze(0).transpose(1, 2).to(dtype) + v = v.unsqueeze(0).transpose(1, 2).to(dtype) + + # Fix for ONNX export: repeat k and v to match q's batch size in cross-attention + if q.shape[0] != k.shape[0] and k.shape[0] == 1: + k = k.repeat(q.shape[0], 1, 1, 1) + v = v.repeat(q.shape[0], 1, 1, 1) + + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, + attn_mask=None, + is_causal=causal, + dropout_p=dropout_p, + scale=softmax_scale, + ) + + # Transpose back to (b, s, n, d) format. + out = out.transpose(1, 2).contiguous() + return out.to(out_dtype) + self.attn_func = _torch_onnx_impl + + elif backend == "TE" and TRANSFORMER_ENGINE_AVAILABLE: + self.attn_backend = DotProductAttention( + num_attention_heads=num_heads, + kv_channels=head_dim, + qkv_format="bshd", + attn_mask_type="causal" if causal else "no_mask", + window_size=window_size, + attention_dropout=dropout_p, + ) + + def _te_impl(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + out_dtype = q.dtype + return self.attn_backend( + query_layer=q.to(dtype), + key_layer=k.to(dtype), + value_layer=v.to(dtype), + ).to(out_dtype) + self.attn_func = _te_impl + + elif backend == "FA2" or backend == "FA3": + def _flash_attn_impl( + q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + q_lens: Optional[torch.Tensor], k_lens: Optional[torch.Tensor], + ) -> torch.Tensor: + return flash_attention( + q=q, k=k, v=v, + q_lens=q_lens, k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + dtype=dtype, + version=3 if backend == "FA3" else 2, + ) + self.attn_func = _flash_attn_impl + + else: + raise ValueError(f"Invalid backend: {backend}") + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_lens: Optional[torch.Tensor] = None, + k_lens: Optional[torch.Tensor] = None, + ): + if ( + self.backend == "torch" or + self.backend == "torch_onnx" or + (self.backend == "TE" and TRANSFORMER_ENGINE_AVAILABLE) + ): + if q_lens is not None or k_lens is not None: + warnings.warn( + 'Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance.' + ) + return self.attn_func(q, k, v) # type: ignore[call-arg] + else: + return self.attn_func(q, k, v, q_lens, k_lens) # type: ignore[call-arg] diff --git a/groot/vla/model/dreamzero/modules/wan2_1_submodule.py b/groot/vla/model/dreamzero/modules/wan2_1_submodule.py new file mode 100644 index 0000000000000000000000000000000000000000..b7bbe1187dddb4adf9ec2e9b2eb715b58c20e1e0 --- /dev/null +++ b/groot/vla/model/dreamzero/modules/wan2_1_submodule.py @@ -0,0 +1,904 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import math + +import torch +import torch.nn as nn +import os +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +from einops import repeat +from groot.vla.model.dreamzero.modules.attention import flash_attention + +__all__ = ['WanModel'] + +ENABLE_TENSORRT = os.getenv("ENABLE_TENSORRT", "False").lower() == "true" + +def sinusoidal_embedding_1d(dim: int, position: torch.Tensor) -> torch.Tensor: + # preprocess + assert dim % 2 == 0 + half = dim // 2 + position = position.type(torch.float64) + + # calculation + sinusoid = torch.outer( + position, torch.pow(10000, -torch.arange(half, dtype=position.dtype, device=position.device).div(half))) + x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) + return x + + +def rope_params(max_seq_len, dim, theta=10000): + if ENABLE_TENSORRT: + return rope_params_no_polar(max_seq_len, dim, theta) + else: + return rope_params_polar(max_seq_len, dim, theta) + + +# @amp.autocast(enabled=False) +def rope_params_polar(max_seq_len: int, dim: int, theta: float = 10000) -> torch.Tensor: + assert dim % 2 == 0 + freqs = torch.outer( + torch.arange(max_seq_len), + 1.0 / torch.pow(theta, + torch.arange(0, dim, 2).to(torch.float64).div(dim))) + freqs = torch.polar(torch.ones_like(freqs), freqs) + return freqs + +def rope_params_no_polar(max_seq_len: int, dim: int, theta: float = 10000) -> torch.Tensor: + assert dim % 2 == 0 + inv_freq = 1.0 / torch.pow( + theta, + torch.arange(0, dim, 2).to(torch.float32) / dim + ) + t = torch.arange(max_seq_len, dtype=inv_freq.dtype) + freqs = torch.outer(t, inv_freq) + emb = torch.stack((freqs.cos(), freqs.sin()), dim=-1).flatten(-2) + return emb + +def rope_apply(x, grid_sizes, freqs): + if ENABLE_TENSORRT: + return rope_apply_no_polar(x, freqs) + else: + return rope_apply_polar(x, freqs) + +# @amp.autocast(enabled=False) +def rope_apply_polar(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + B, seq_len, n, _ = x.shape + + # precompute multipliers + x = torch.view_as_complex( + x.to(torch.float64).reshape(B, seq_len, n, -1, 2) + ) + + # apply rotary embedding + freqs = freqs.unsqueeze(0) + x = torch.view_as_real(x * freqs).flatten(3) + return x + +def rope_apply_no_polar(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + B, seq_len, n, D = x.shape + + # Reshape freqs to be broadcastable: (1, seq_len, 1, D) + freqs = freqs.unsqueeze(0).unsqueeze(2) + + x0, x1 = x.chunk(2, dim=-1) + freqs_cos, freqs_sin = freqs.chunk(2, dim=-1) + + rotated_x0 = x0 * freqs_cos - x1 * freqs_sin + rotated_x1 = x1 * freqs_cos + x0 * freqs_sin + x_rotated = torch.cat((rotated_x0, rotated_x1), dim=-1) + return x_rotated + + +def rope_action_apply(x, freqs, freqs_action, freqs_state, action_register_length, num_action_per_block=32, num_state_per_block=1): + if ENABLE_TENSORRT: + return rope_action_apply_no_polar(x, freqs, freqs_action, freqs_state, action_register_length, num_action_per_block, num_state_per_block) + else: + return rope_action_apply_polar(x, freqs, freqs_action, freqs_state, action_register_length, num_action_per_block, num_state_per_block) + + +def rope_action_apply_no_polar( + x: torch.Tensor, + freqs: torch.Tensor, + freqs_action: torch.Tensor, + freqs_state: torch.Tensor, + action_register_length: int, + num_action_per_block: int = 32, + num_state_per_block: int = 1, +) -> torch.Tensor: + B, seq_len, n, D = x.shape + + if action_register_length is not None: + chunk_size = action_register_length // (num_action_per_block + num_state_per_block) + freqs_1d_action = freqs_action[:chunk_size * num_action_per_block] + freqs_1d_state = freqs_state[:chunk_size * num_state_per_block] + freqs = torch.cat([freqs, freqs_1d_action, freqs_1d_state], dim=0) + + # Reshape freqs to be broadcastable: (1, seq_len, 1, D) + freqs = freqs.unsqueeze(0).unsqueeze(2) + + x0, x1 = x.chunk(2, dim=-1) + freqs_cos, freqs_sin = freqs.chunk(2, dim=-1) + + rotated_x0 = x0 * freqs_cos - x1 * freqs_sin + rotated_x1 = x1 * freqs_cos + x0 * freqs_sin + x_rotated = torch.cat((rotated_x0, rotated_x1), dim=-1) + + return x_rotated + + +# @amp.autocast(enabled=False) +def rope_action_apply_polar( + x: torch.Tensor, + freqs: torch.Tensor, + freqs_action: torch.Tensor, + freqs_state: torch.Tensor, + action_register_length: int | None, + num_action_per_block: int | None = None, + num_state_per_block: int | None = None, +) -> torch.Tensor: + B, seq_len, n, _ = x.shape + + # precompute multipliers + x = torch.view_as_complex( + x.to(torch.float64).reshape(B, seq_len, n, -1, 2) + ) + + if action_register_length is not None: + assert num_action_per_block is not None + assert num_state_per_block is not None + + chunk_size = action_register_length // (num_action_per_block + num_state_per_block) + + freqs_1d_action = freqs_action[:chunk_size * num_action_per_block].view(chunk_size * num_action_per_block, 1, -1) + freqs_1d_state = freqs_state[:chunk_size * num_state_per_block].view(chunk_size * num_state_per_block, 1, -1) + freqs = torch.cat([freqs, freqs_1d_action, freqs_1d_state], dim=0) + + # apply rotary embedding + freqs = freqs.unsqueeze(0) + x = torch.view_as_real(x * freqs).flatten(3) + return x + + +class WanRMSNorm(nn.Module): + + def __init__(self, dim, eps=1e-5): + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + r""" + Args: + x(Tensor): Shape [B, L, C] + """ + return self._norm(x.float()).type_as(x) * self.weight + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) + + +class WanLayerNorm(nn.LayerNorm): + + def __init__(self, dim, eps=1e-6, elementwise_affine=False): + super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps) + + +class WanSelfAttention(nn.Module): + + def __init__(self, + dim, + num_heads, + window_size=(-1, -1), + qk_norm=True, + eps=1e-6): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.eps = eps + + # layers + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + + def forward(self, x, seq_lens, freqs): + r""" + Args: + x(Tensor): Shape [B, L, num_heads, C / num_heads] + seq_lens(Tensor): Shape [B] + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + + # query, key, value function + def qkv_fn(x): + q = self.norm_q(self.q(x)).view(b, s, n, d) + k = self.norm_k(self.k(x)).view(b, s, n, d) + v = self.v(x).view(b, s, n, d) + return q, k, v + + q, k, v = qkv_fn(x) + + x = flash_attention( + q=rope_apply(q, freqs), + k=rope_apply(k, freqs), + v=v, + k_lens=seq_lens, + window_size=self.window_size) + + # output + x = x.flatten(2) + x = self.o(x) + return x + + +class WanT2VCrossAttention(WanSelfAttention): + + def forward(self, x, context, context_lens, crossattn_cache=None): + r""" + Args: + x(Tensor): Shape [B, L1, C] + context(Tensor): Shape [B, L2, C] + context_lens(Tensor): Shape [B] + crossattn_cache (List[dict], *optional*): Contains the cached key and value tensors for context embedding. + """ + b, n, d = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.norm_q(self.q(x)).view(b, -1, n, d) + + if crossattn_cache is not None: + if not crossattn_cache["is_init"]: + crossattn_cache["is_init"] = True + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + crossattn_cache["k"] = k + crossattn_cache["v"] = v + else: + k = crossattn_cache["k"] + v = crossattn_cache["v"] + else: + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + + # compute attention + x = flash_attention(q, k, v, k_lens=context_lens) + + # output + x = x.flatten(2) + x = self.o(x) + return x + + +class WanGanCrossAttention(WanSelfAttention): + + def forward(self, x, context, crossattn_cache=None): + r""" + Args: + x(Tensor): Shape [B, L1, C] + context(Tensor): Shape [B, L2, C] + context_lens(Tensor): Shape [B] + crossattn_cache (List[dict], *optional*): Contains the cached key and value tensors for context embedding. + """ + b, n, d = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + qq = self.norm_q(self.q(context)).view(b, 1, -1, d) + + kk = self.norm_k(self.k(x)).view(b, -1, n, d) + vv = self.v(x).view(b, -1, n, d) + + # compute attention + x = flash_attention(qq, kk, vv) + + # output + x = x.flatten(2) + x = self.o(x) + return x + + +class WanI2VCrossAttention(WanSelfAttention): + + def __init__(self, + dim, + num_heads, + window_size=(-1, -1), + qk_norm=True, + eps=1e-6): + super().__init__(dim, num_heads, window_size, qk_norm, eps) + + self.k_img = nn.Linear(dim, dim) + self.v_img = nn.Linear(dim, dim) + # self.alpha = nn.Parameter(torch.zeros((1, ))) + self.norm_k_img = WanRMSNorm( + dim, eps=eps) if qk_norm else nn.Identity() + + def forward(self, x, context, crossattn_cache=None): + r""" + Args: + x(Tensor): Shape [B, L1, C] + context(Tensor): Shape [B, L2, C] + """ + context_img = context[:, :257] + context = context[:, 257:] + b, n, d = x.size(0), self.num_heads, self.head_dim + + q = self.norm_q(self.q(x)).view(b, -1, n, d) + + if crossattn_cache is not None: + if not crossattn_cache["is_init"]: + crossattn_cache["is_init"] = True + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + crossattn_cache["k"] = k + crossattn_cache["v"] = v + else: + k = crossattn_cache["k"] + v = crossattn_cache["v"] + else: + # compute query, key, value + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + x = flash_attention(q, k, v, k_lens=None) + + k_img = self.norm_k_img(self.k_img(context_img)).view(b, -1, n, d) + v_img = self.v_img(context_img).view(b, -1, n, d) + img_x = flash_attention(q, k_img, v_img, k_lens=None) + + # output + x = x.flatten(2) + img_x = img_x.flatten(2) + x = x + img_x + x = self.o(x) + return x + + +WAN_CROSSATTENTION_CLASSES = { + 't2v_cross_attn': WanT2VCrossAttention, + 'i2v_cross_attn': WanI2VCrossAttention, +} + + +class WanAttentionBlock(nn.Module): + + def __init__(self, + cross_attn_type, + dim, + ffn_dim, + num_heads, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=False, + eps=1e-6): + super().__init__() + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + # layers + self.norm1 = WanLayerNorm(dim, eps) + self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm, + eps) + self.norm3 = WanLayerNorm( + dim, eps, + elementwise_affine=True) if cross_attn_norm else nn.Identity() + self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim, + num_heads, + (-1, -1), + qk_norm, + eps) + self.norm2 = WanLayerNorm(dim, eps) + self.ffn = nn.Sequential( + nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'), + nn.Linear(ffn_dim, dim)) + + # modulation + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + x, + e, + seq_lens, + grid_sizes, + freqs, + context, + context_lens, + ): + r""" + Args: + x(Tensor): Shape [B, L, C] + e(Tensor): Shape [B, 6, C] + seq_lens(Tensor): Shape [B], length of each sequence in batch + grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + # assert e.dtype == torch.float32 + # with amp.autocast(dtype=torch.float32): + e = (self.modulation + e).chunk(6, dim=1) + # assert e[0].dtype == torch.float32 + + # self-attention + y = self.self_attn( + self.norm1(x) * (1 + e[1]) + e[0], seq_lens, grid_sizes, + freqs, + ) + # with amp.autocast(dtype=torch.float32): + x = x + y * e[2] + + # cross-attention & ffn function + def cross_attn_ffn(x, context, context_lens, e): + x = x + self.cross_attn(self.norm3(x), context, context_lens) + y = self.ffn(self.norm2(x) * (1 + e[4]) + e[3]) + # with amp.autocast(dtype=torch.float32): + x = x + y * e[5] + return x + + x = cross_attn_ffn(x, context, context_lens, e) + return x + + +class GanAttentionBlock(nn.Module): + + def __init__(self, + dim=1536, + ffn_dim=8192, + num_heads=12, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=True, + eps=1e-6): + super().__init__() + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + # layers + # self.norm1 = WanLayerNorm(dim, eps) + # self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm, + # eps) + self.norm3 = WanLayerNorm( + dim, eps, + elementwise_affine=True) if cross_attn_norm else nn.Identity() + + self.norm2 = WanLayerNorm(dim, eps) + self.ffn = nn.Sequential( + nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'), + nn.Linear(ffn_dim, dim)) + + self.cross_attn = WanGanCrossAttention(dim, num_heads, + (-1, -1), + qk_norm, + eps) + + # modulation + # self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + x, + context, + # seq_lens, + # grid_sizes, + # freqs, + # context, + # context_lens, + ): + r""" + Args: + x(Tensor): Shape [B, L, C] + e(Tensor): Shape [B, 6, C] + seq_lens(Tensor): Shape [B], length of each sequence in batch + grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + # assert e.dtype == torch.float32 + # with amp.autocast(dtype=torch.float32): + # e = (self.modulation + e).chunk(6, dim=1) + # assert e[0].dtype == torch.float32 + + # # self-attention + # y = self.self_attn( + # self.norm1(x) * (1 + e[1]) + e[0], seq_lens, grid_sizes, + # freqs) + # # with amp.autocast(dtype=torch.float32): + # x = x + y * e[2] + + # cross-attention & ffn function + def cross_attn_ffn(x, context): + token = context + self.cross_attn(self.norm3(x), context) + y = self.ffn(self.norm2(token)) + token # * (1 + e[4]) + e[3]) + # with amp.autocast(dtype=torch.float32): + # x = x + y * e[5] + return y + + x = cross_attn_ffn(x, context) + return x + + +class Head(nn.Module): + + def __init__(self, dim, out_dim, patch_size, eps=1e-6): + super().__init__() + self.dim = dim + self.out_dim = out_dim + self.patch_size = patch_size + self.eps = eps + + # layers + out_dim = math.prod(patch_size) * out_dim + self.norm = WanLayerNorm(dim, eps) + self.head = nn.Linear(dim, out_dim) + + # modulation + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x, e): + r""" + Args: + x(Tensor): Shape [B, L1, C] + e(Tensor): Shape [B, C] + """ + # assert e.dtype == torch.float32 + # with amp.autocast(dtype=torch.float32): + e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1) + x = (self.head(self.norm(x) * (1 + e[1]) + e[0])) + return x + + +class MLPProj(torch.nn.Module): + + def __init__(self, in_dim, out_dim): + super().__init__() + + self.proj = torch.nn.Sequential( + torch.nn.LayerNorm(in_dim), torch.nn.Linear(in_dim, in_dim), + torch.nn.GELU(), torch.nn.Linear(in_dim, out_dim), + torch.nn.LayerNorm(out_dim)) + + def forward(self, image_embeds): + clip_extra_context_tokens = self.proj(image_embeds) + return clip_extra_context_tokens + + +class RegisterTokens(nn.Module): + def __init__(self, num_registers: int, dim: int): + super().__init__() + self.register_tokens = nn.Parameter(torch.randn(num_registers, dim) * 0.02) + self.rms_norm = WanRMSNorm(dim, eps=1e-6) + + def forward(self): + return self.rms_norm(self.register_tokens) + + def reset_parameters(self): + nn.init.normal_(self.register_tokens, std=0.02) + + +class WanModel(ModelMixin, ConfigMixin): + r""" + Wan diffusion backbone supporting both text-to-video and image-to-video. + """ + + ignore_for_config = [ + 'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim', 'window_size' + ] + _no_split_modules = ['WanAttentionBlock'] + _supports_gradient_checkpointing = True + + @register_to_config + def __init__(self, + model_type='t2v', + patch_size=(1, 2, 2), + text_len=512, + in_dim=16, + dim=2048, + ffn_dim=8192, + freq_dim=256, + text_dim=4096, + out_dim=16, + num_heads=16, + num_layers=32, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=True, + eps=1e-6): + r""" + Initialize the diffusion model backbone. + + Args: + model_type (`str`, *optional*, defaults to 't2v'): + Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video) + patch_size (`tuple`, *optional*, defaults to (1, 2, 2)): + 3D patch dimensions for video embedding (t_patch, h_patch, w_patch) + text_len (`int`, *optional*, defaults to 512): + Fixed length for text embeddings + in_dim (`int`, *optional*, defaults to 16): + Input video channels (C_in) + dim (`int`, *optional*, defaults to 2048): + Hidden dimension of the transformer + ffn_dim (`int`, *optional*, defaults to 8192): + Intermediate dimension in feed-forward network + freq_dim (`int`, *optional*, defaults to 256): + Dimension for sinusoidal time embeddings + text_dim (`int`, *optional*, defaults to 4096): + Input dimension for text embeddings + out_dim (`int`, *optional*, defaults to 16): + Output video channels (C_out) + num_heads (`int`, *optional*, defaults to 16): + Number of attention heads + num_layers (`int`, *optional*, defaults to 32): + Number of transformer blocks + window_size (`tuple`, *optional*, defaults to (-1, -1)): + Window size for local attention (-1 indicates global attention) + qk_norm (`bool`, *optional*, defaults to True): + Enable query/key normalization + cross_attn_norm (`bool`, *optional*, defaults to False): + Enable cross-attention normalization + eps (`float`, *optional*, defaults to 1e-6): + Epsilon value for normalization layers + """ + + super().__init__() + + assert model_type in ['t2v', 'i2v'] + self.model_type = model_type + + self.patch_size = patch_size + self.text_len = text_len + self.in_dim = in_dim + self.dim = dim + self.ffn_dim = ffn_dim + self.freq_dim = freq_dim + self.text_dim = text_dim + self.out_dim = out_dim + self.num_heads = num_heads + self.num_layers = num_layers + self.window_size = window_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + self.local_attn_size = 21 + + # embeddings + self.patch_embedding = nn.Conv3d( + in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'), + nn.Linear(dim, dim)) + + self.time_embedding = nn.Sequential( + nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + self.time_projection = nn.Sequential( + nn.SiLU(), nn.Linear(dim, dim * 6)) + + # blocks + cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn' + self.blocks = nn.ModuleList([ + WanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads, + window_size, qk_norm, cross_attn_norm, eps) + for _ in range(num_layers) + ]) + + # head + self.head = Head(dim, out_dim, patch_size, eps) + + # buffers (don't use register_buffer otherwise dtype will be changed in to()) + assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0 + d = dim // num_heads + self.freqs = [ + rope_params(1024, d - 4 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + rope_params(1024, 2 * (d // 6)) + ] + + if model_type == 'i2v': + self.img_emb = MLPProj(1280, dim) + + # initialize weights + self.init_weights() + + self.gradient_checkpointing = False + + def _set_gradient_checkpointing(self, module, value=False): + self.gradient_checkpointing = value + + def forward( + self, + x, + t, + context, + seq_len, + classify_mode=False, + concat_time_embeddings=False, + register_tokens=None, + cls_pred_branch=None, + gan_ca_blocks=None, + clip_fea=None, + y=None, + ): + r""" + Forward pass through the diffusion model + + Args: + x (List[Tensor]): + List of input video tensors, each with shape [C_in, F, H, W] + t (Tensor): + Diffusion timesteps tensor of shape [B] + context (List[Tensor]): + List of text embeddings each with shape [L, C] + seq_len (`int`): + Maximum sequence length for positional encoding + clip_fea (Tensor, *optional*): + CLIP image features for image-to-video mode + y (List[Tensor], *optional*): + Conditional video inputs for image-to-video mode, same shape as x + + Returns: + List[Tensor]: + List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8] + """ + if self.model_type == 'i2v': + assert clip_fea is not None and y is not None + # params + device = self.patch_embedding.weight.device + if any(freqs_i.device != device for freqs_i in self.freqs): + self.freqs = [freqs_i.to(device) for freqs_i in self.freqs] + + if y is not None: + x = [torch.cat([u, v.to(dtype=u.dtype)], dim=0) for u, v in zip(x, y)] + + # embeddings + x = [self.patch_embedding(u.unsqueeze(0)) for u in x] + grid_sizes = torch.stack( + [torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) + x = [u.flatten(2).transpose(1, 2) for u in x] + seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long) + assert seq_lens.max() <= seq_len + x = torch.cat([ + torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], + dim=1) for u in x + ]) + + # time embeddings + # with amp.autocast(dtype=torch.float32): + e = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, t).type_as(x)) + e0 = self.time_projection(e).unflatten(1, (6, self.dim)) + # assert e.dtype == torch.float32 and e0.dtype == torch.float32 + + # context + context_lens = None + context = self.text_embedding( + torch.stack([ + torch.cat( + [u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) + for u in context + ])) + + if clip_fea is not None: + context_clip = self.img_emb(clip_fea) # bs x 257 x dim + context = torch.concat([context_clip, context], dim=1) + + # arguments + kwargs = dict( + e=e0, + seq_lens=seq_lens, + grid_sizes=grid_sizes, + freqs=self.freqs, + context=context, + context_lens=context_lens) + + def create_custom_forward(module): + def custom_forward(*inputs, **kwargs): + return module(*inputs, **kwargs) + return custom_forward + + # TODO: Tune the number of blocks for feature extraction + final_x = None + if classify_mode: + assert register_tokens is not None + assert gan_ca_blocks is not None + assert cls_pred_branch is not None + + final_x = [] + registers = repeat(register_tokens(), "n d -> b n d", b=x.shape[0]) + # x = torch.cat([registers, x], dim=1) + + gan_idx = 0 + for ii, block in enumerate(self.blocks): + if torch.is_grad_enabled() and self.gradient_checkpointing: + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, **kwargs, + use_reentrant=False, + ) + else: + x = block(x, **kwargs) + + if classify_mode and ii in [13, 21, 29]: + gan_token = registers[:, gan_idx: gan_idx + 1] + final_x.append(gan_ca_blocks[gan_idx](x, gan_token)) + gan_idx += 1 + + if classify_mode: + final_x = torch.cat(final_x, dim=1) + if concat_time_embeddings: + final_x = cls_pred_branch(torch.cat([final_x, 10 * e[:, None, :]], dim=1).view(final_x.shape[0], -1)) + else: + final_x = cls_pred_branch(final_x.view(final_x.shape[0], -1)) + + # head + x = self.head(x, e) + + # unpatchify + x = self.unpatchify(x, grid_sizes) + + if classify_mode: + return torch.stack(x), final_x + + return torch.stack(x) + + def unpatchify(self, x, grid_sizes, c=None): + r""" + Reconstruct video tensors from patch embeddings. + + Args: + x (List[Tensor]): + List of patchified features, each with shape [L, C_out * prod(patch_size)] + grid_sizes (Tensor): + Original spatial-temporal grid dimensions before patching, + shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches) + + Returns: + List[Tensor]: + Reconstructed video tensors with shape [C_out, F, H / 8, W / 8] + """ + + c = self.out_dim if c is None else c + out = [] + for u, v in zip(x, grid_sizes.tolist()): + u = u[:math.prod(v)].view(*v, *self.patch_size, c) + u = torch.einsum('fhwpqrc->cfphqwr', u) + u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) + out.append(u) + return out + + def init_weights(self): + r""" + Initialize model parameters using Xavier initialization. + """ + + # basic init + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + if m.bias is not None: + nn.init.zeros_(m.bias) + + # init embeddings + nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1)) + for m in self.text_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=.02) + for m in self.time_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=.02) + + # init output layer + nn.init.zeros_(self.head.head.weight) diff --git a/groot/vla/model/dreamzero/modules/wan_video_camera_controller.py b/groot/vla/model/dreamzero/modules/wan_video_camera_controller.py new file mode 100644 index 0000000000000000000000000000000000000000..8cbe658699e52c8cf3854d94ed4d0c356bb96989 --- /dev/null +++ b/groot/vla/model/dreamzero/modules/wan_video_camera_controller.py @@ -0,0 +1,202 @@ +import torch +import torch.nn as nn +import numpy as np +from einops import rearrange +import os +from typing_extensions import Literal + +class SimpleAdapter(nn.Module): + def __init__(self, in_dim, out_dim, kernel_size, stride, num_residual_blocks=1): + super(SimpleAdapter, self).__init__() + + # Pixel Unshuffle: reduce spatial dimensions by a factor of 8 + self.pixel_unshuffle = nn.PixelUnshuffle(downscale_factor=8) + + # Convolution: reduce spatial dimensions by a factor + # of 2 (without overlap) + self.conv = nn.Conv2d(in_dim * 64, out_dim, kernel_size=kernel_size, stride=stride, padding=0) + + # Residual blocks for feature extraction + self.residual_blocks = nn.Sequential( + *[ResidualBlock(out_dim) for _ in range(num_residual_blocks)] + ) + + def forward(self, x): + # Reshape to merge the frame dimension into batch + bs, c, f, h, w = x.size() + x = x.permute(0, 2, 1, 3, 4).contiguous().view(bs * f, c, h, w) + + # Pixel Unshuffle operation + x_unshuffled = self.pixel_unshuffle(x) + + # Convolution operation + x_conv = self.conv(x_unshuffled) + + # Feature extraction with residual blocks + out = self.residual_blocks(x_conv) + + # Reshape to restore original bf dimension + out = out.view(bs, f, out.size(1), out.size(2), out.size(3)) + + # Permute dimensions to reorder (if needed), e.g., swap channels and feature frames + out = out.permute(0, 2, 1, 3, 4) + + return out + + def process_camera_coordinates( + self, + direction: Literal["Left", "Right", "Up", "Down", "LeftUp", "LeftDown", "RightUp", "RightDown"], + length: int, + height: int, + width: int, + speed: float = 1/54, + origin=(0, 0.532139961, 0.946026558, 0.5, 0.5, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0) + ): + if origin is None: + origin = (0, 0.532139961, 0.946026558, 0.5, 0.5, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0) + coordinates = generate_camera_coordinates(direction, length, speed, origin) + plucker_embedding = process_pose_file(coordinates, width, height) + return plucker_embedding + + + +class ResidualBlock(nn.Module): + def __init__(self, dim): + super(ResidualBlock, self).__init__() + self.conv1 = nn.Conv2d(dim, dim, kernel_size=3, padding=1) + self.relu = nn.ReLU(inplace=True) + self.conv2 = nn.Conv2d(dim, dim, kernel_size=3, padding=1) + + def forward(self, x): + residual = x + out = self.relu(self.conv1(x)) + out = self.conv2(out) + out += residual + return out + +class Camera(object): + """Copied from https://github.com/hehao13/CameraCtrl/blob/main/inference.py + """ + def __init__(self, entry): + fx, fy, cx, cy = entry[1:5] + self.fx = fx + self.fy = fy + self.cx = cx + self.cy = cy + w2c_mat = np.array(entry[7:]).reshape(3, 4) + w2c_mat_4x4 = np.eye(4) + w2c_mat_4x4[:3, :] = w2c_mat + self.w2c_mat = w2c_mat_4x4 + self.c2w_mat = np.linalg.inv(w2c_mat_4x4) + +def get_relative_pose(cam_params): + """Copied from https://github.com/hehao13/CameraCtrl/blob/main/inference.py + """ + abs_w2cs = [cam_param.w2c_mat for cam_param in cam_params] + abs_c2ws = [cam_param.c2w_mat for cam_param in cam_params] + cam_to_origin = 0 + target_cam_c2w = np.array([ + [1, 0, 0, 0], + [0, 1, 0, -cam_to_origin], + [0, 0, 1, 0], + [0, 0, 0, 1] + ]) + abs2rel = target_cam_c2w @ abs_w2cs[0] + ret_poses = [target_cam_c2w, ] + [abs2rel @ abs_c2w for abs_c2w in abs_c2ws[1:]] + ret_poses = np.array(ret_poses, dtype=np.float32) + return ret_poses + +def custom_meshgrid(*args): + # torch>=2.0.0 only + return torch.meshgrid(*args, indexing='ij') + + +def ray_condition(K, c2w, H, W, device): + """Copied from https://github.com/hehao13/CameraCtrl/blob/main/inference.py + """ + # c2w: B, V, 4, 4 + # K: B, V, 4 + + B = K.shape[0] + + j, i = custom_meshgrid( + torch.linspace(0, H - 1, H, device=device, dtype=c2w.dtype), + torch.linspace(0, W - 1, W, device=device, dtype=c2w.dtype), + ) + i = i.reshape([1, 1, H * W]).expand([B, 1, H * W]) + 0.5 # [B, HxW] + j = j.reshape([1, 1, H * W]).expand([B, 1, H * W]) + 0.5 # [B, HxW] + + fx, fy, cx, cy = K.chunk(4, dim=-1) # B,V, 1 + + zs = torch.ones_like(i) # [B, HxW] + xs = (i - cx) / fx * zs + ys = (j - cy) / fy * zs + zs = zs.expand_as(ys) + + directions = torch.stack((xs, ys, zs), dim=-1) # B, V, HW, 3 + directions = directions / directions.norm(dim=-1, keepdim=True) # B, V, HW, 3 + + rays_d = directions @ c2w[..., :3, :3].transpose(-1, -2) # B, V, 3, HW + rays_o = c2w[..., :3, 3] # B, V, 3 + rays_o = rays_o[:, :, None].expand_as(rays_d) # B, V, 3, HW + # c2w @ dirctions + rays_dxo = torch.linalg.cross(rays_o, rays_d) + plucker = torch.cat([rays_dxo, rays_d], dim=-1) + plucker = plucker.reshape(B, c2w.shape[1], H, W, 6) # B, V, H, W, 6 + # plucker = plucker.permute(0, 1, 4, 2, 3) + return plucker + + +def process_pose_file(cam_params, width=672, height=384, original_pose_width=1280, original_pose_height=720, device='cpu', return_poses=False): + if return_poses: + return cam_params + else: + cam_params = [Camera(cam_param) for cam_param in cam_params] + + sample_wh_ratio = width / height + pose_wh_ratio = original_pose_width / original_pose_height # Assuming placeholder ratios, change as needed + + if pose_wh_ratio > sample_wh_ratio: + resized_ori_w = height * pose_wh_ratio + for cam_param in cam_params: + cam_param.fx = resized_ori_w * cam_param.fx / width + else: + resized_ori_h = width / pose_wh_ratio + for cam_param in cam_params: + cam_param.fy = resized_ori_h * cam_param.fy / height + + intrinsic = np.asarray([[cam_param.fx * width, + cam_param.fy * height, + cam_param.cx * width, + cam_param.cy * height] + for cam_param in cam_params], dtype=np.float32) + + K = torch.as_tensor(intrinsic)[None] # [1, 1, 4] + c2ws = get_relative_pose(cam_params) # Assuming this function is defined elsewhere + c2ws = torch.as_tensor(c2ws)[None] # [1, n_frame, 4, 4] + plucker_embedding = ray_condition(K, c2ws, height, width, device=device)[0].permute(0, 3, 1, 2).contiguous() # V, 6, H, W + plucker_embedding = plucker_embedding[None] + plucker_embedding = rearrange(plucker_embedding, "b f c h w -> b f h w c")[0] + return plucker_embedding + + + +def generate_camera_coordinates( + direction: Literal["Left", "Right", "Up", "Down", "LeftUp", "LeftDown", "RightUp", "RightDown"], + length: int, + speed: float = 1/54, + origin=(0, 0.532139961, 0.946026558, 0.5, 0.5, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0) +): + coordinates = [list(origin)] + while len(coordinates) < length: + coor = coordinates[-1].copy() + if "Left" in direction: + coor[9] += speed + if "Right" in direction: + coor[9] -= speed + if "Up" in direction: + coor[13] += speed + if "Down" in direction: + coor[13] -= speed + coordinates.append(coor) + return coordinates \ No newline at end of file diff --git a/groot/vla/model/dreamzero/modules/wan_video_dit.py b/groot/vla/model/dreamzero/modules/wan_video_dit.py new file mode 100644 index 0000000000000000000000000000000000000000..305d3203cba05577ee0b2881de716999c8dc9e88 --- /dev/null +++ b/groot/vla/model/dreamzero/modules/wan_video_dit.py @@ -0,0 +1,818 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import math +import os +from typing import Tuple, Optional +from einops import rearrange +from groot.vla.model.dreamzero.modules.utils import hash_state_dict_keys +from groot.vla.model.dreamzero.modules.wan_video_camera_controller import SimpleAdapter +try: + import flash_attn_interface + FLASH_ATTN_3_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_3_AVAILABLE = False + +try: + import flash_attn + FLASH_ATTN_2_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_2_AVAILABLE = False + +try: + from sageattention import sageattn + SAGE_ATTN_AVAILABLE = True +except ModuleNotFoundError: + SAGE_ATTN_AVAILABLE = False + + +def _gpu_supports_flash_attention(): + """FlashAttention requires Ampere (compute capability 8.0) or newer.""" + if not (FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE): + return False + try: + if not torch.cuda.is_available(): + return False + cap = torch.cuda.get_device_capability() + return cap[0] >= 8 + except Exception: + return False + + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin + +ENABLE_TENSORRT = os.getenv("ENABLE_TENSORRT", "False").lower() == "true" +if ENABLE_TENSORRT: + # disable torch compile and transformer engine and flash attention for onnx/tensorrt export + FLASH_ATTN_COMPATIBILITY_MODE = True + DISABLE_TORCH_COMPILE = True +else: + DISABLE_TORCH_COMPILE = False + FLASH_ATTN_COMPATIBILITY_MODE = False +def flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, num_heads: int, compatibility_mode=False): + # Use PyTorch SDPA on pre-Ampere GPUs or when compatibility_mode (FlashAttention requires Ampere or newer) + if compatibility_mode or not _gpu_supports_flash_attention(): + q = rearrange(q, "b s (n d) -> b n s d", n=num_heads) + k = rearrange(k, "b s (n d) -> b n s d", n=num_heads) + v = rearrange(v, "b s (n d) -> b n s d", n=num_heads) + x = F.scaled_dot_product_attention(q, k, v) + x = rearrange(x, "b n s d -> b s (n d)", n=num_heads) + elif FLASH_ATTN_3_AVAILABLE: + q = rearrange(q, "b s (n d) -> b s n d", n=num_heads) + k = rearrange(k, "b s (n d) -> b s n d", n=num_heads) + v = rearrange(v, "b s (n d) -> b s n d", n=num_heads) + x = flash_attn_interface.flash_attn_func(q, k, v) + if isinstance(x,tuple): + x = x[0] + x = rearrange(x, "b s n d -> b s (n d)", n=num_heads) + elif FLASH_ATTN_2_AVAILABLE: + q = rearrange(q, "b s (n d) -> b s n d", n=num_heads) + k = rearrange(k, "b s (n d) -> b s n d", n=num_heads) + v = rearrange(v, "b s (n d) -> b s n d", n=num_heads) + x = flash_attn.flash_attn_func(q, k, v) + x = rearrange(x, "b s n d -> b s (n d)", n=num_heads) + elif SAGE_ATTN_AVAILABLE: + q = rearrange(q, "b s (n d) -> b n s d", n=num_heads) + k = rearrange(k, "b s (n d) -> b n s d", n=num_heads) + v = rearrange(v, "b s (n d) -> b n s d", n=num_heads) + x = sageattn(q, k, v) + x = rearrange(x, "b n s d -> b s (n d)", n=num_heads) + else: + q = rearrange(q, "b s (n d) -> b n s d", n=num_heads) + k = rearrange(k, "b s (n d) -> b n s d", n=num_heads) + v = rearrange(v, "b s (n d) -> b n s d", n=num_heads) + x = F.scaled_dot_product_attention(q, k, v) + x = rearrange(x, "b n s d -> b s (n d)", n=num_heads) + return x + + +@torch.compile(disable=DISABLE_TORCH_COMPILE) +def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor): + return (x * (1 + scale) + shift) + + +@torch.compile(disable=DISABLE_TORCH_COMPILE) +def sinusoidal_embedding_1d(dim, position): + sinusoid = torch.outer( + position.type(torch.float64), + torch.pow(10000, -torch.arange(dim//2, dtype=torch.float64, device=position.device).div(dim//2)) + ) + x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) + return x.to(position.dtype) + +@torch.compile(disable=DISABLE_TORCH_COMPILE) +def rope_apply_no_polar_op(x, freqs_cis, num_heads): + freqs_cos, freqs_sin = freqs_cis + + x_reshaped = rearrange(x, "b s (h d) -> b s h d", h=num_heads) + x_as_complex = x_reshaped.float().reshape(*x_reshaped.shape[:-1], -1, 2) + x_real, x_imag = x_as_complex.unbind(-1) + + freqs_cos = freqs_cos.to(x.device).squeeze(1).unsqueeze(0).unsqueeze(2) + freqs_sin = freqs_sin.to(x.device).squeeze(1).unsqueeze(0).unsqueeze(2) + + x_out_real = x_real * freqs_cos - x_imag * freqs_sin + x_out_imag = x_real * freqs_sin + x_imag * freqs_cos + + x_out = torch.stack([x_out_real, x_out_imag], dim=-1).flatten(start_dim=-2) + + return rearrange(x_out, "b s h d -> b s (h d)").to(x.dtype) + + +@torch.compile(disable=DISABLE_TORCH_COMPILE) +def rope_apply_polar_op(x, freqs, num_heads): + x = rearrange(x, "b s (n d) -> b s n d", n=num_heads) + x_out = torch.view_as_complex(x.to(torch.float64).reshape( + x.shape[0], x.shape[1], x.shape[2], -1, 2)) + freqs = freqs.to(x_out.device) + x_out = torch.view_as_real(x_out * freqs).flatten(2) + return x_out.to(x.dtype) + + +class RMSNorm(nn.Module): + def __init__(self, dim, eps=1e-5): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) + + def forward(self, x): + dtype = x.dtype + return self.norm(x.float()).to(dtype) * self.weight + + +def _RMSNorm(normalized_shape, eps): + return RMSNorm(normalized_shape, eps=eps) + + +def RotaryPositionEmbedding(num_heads, head_dim): + if ENABLE_TENSORRT: + return RotaryPositionEmbeddingNoPolarOp(num_heads, head_dim) + else: + return RotaryPositionEmbeddingWithPolarOp(num_heads, head_dim) + + +def rope_apply(x, freqs, num_heads): + if ENABLE_TENSORRT: + return rope_apply_no_polar_op(x, freqs, num_heads) + else: + return rope_apply_polar_op(x, freqs, num_heads) + + +class RotaryPositionEmbeddingNoPolarOp(nn.Module): + def __init__(self, num_heads: int, head_dim: int): + super().__init__() + self.num_heads = num_heads + self.head_dim = head_dim + + self.freqs = self.precompute_freqs_cis_3d(head_dim) + + def precompute_freqs_cis_3d(self, dim: int, end: int = 1024, theta: float = 10000.0): + # 3d rope precompute + f_freqs_cis = self.precompute_freqs_cis(dim - 2 * (dim // 3), end, theta) + h_freqs_cis = self.precompute_freqs_cis(dim // 3, end, theta) + w_freqs_cis = self.precompute_freqs_cis(dim // 3, end, theta) + return {"f": f_freqs_cis, "h": h_freqs_cis, "w": w_freqs_cis} + + def precompute_freqs_cis(self, dim: int, end: int = 1024, theta: float = 10000.0): + # 1d rope precompute + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) + freqs = torch.outer(torch.arange(end, device=freqs.device), freqs) + freqs_real = torch.cos(freqs) + freqs_imag = torch.sin(freqs) + return (freqs_real, freqs_imag) + + def forward(self, f: int, h: int, w: int, a: int) -> torch.Tensor: + freqs_cos_3d = torch.cat( + [ + self.freqs["f"][0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), + self.freqs["h"][0][:h].view(1, h, 1, -1).expand(f, h, w, -1), + self.freqs["w"][0][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], + dim=-1, + ).reshape(f * h * w, 1, -1) + freqs_sin_3d = torch.cat( + [ + self.freqs["f"][1][:f].view(f, 1, 1, -1).expand(f, h, w, -1), + self.freqs["h"][1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + self.freqs["w"][1][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], + dim=-1, + ).reshape(f * h * w, 1, -1) + + + return torch.cat([freqs_cos_3d, freqs_sin_3d], dim=0) + + def post_initialize(self): + self.freqs = { + key: (value[0].to("cuda"), value[1].to("cuda")) for key, value in self.freqs.items() + } + + +class RotaryPositionEmbeddingWithPolarOp(nn.Module): + def __init__(self, num_heads: int, head_dim: int): + super().__init__() + self.num_heads = num_heads + self.head_dim = head_dim + + self.freqs = self.precompute_freqs_cis_3d(head_dim) + + def precompute_freqs_cis_3d(self, dim: int, end: int = 1024, theta: float = 10000.0): + # 3d rope precompute + f_freqs_cis = self.precompute_freqs_cis(dim - 2 * (dim // 3), end, theta) + h_freqs_cis = self.precompute_freqs_cis(dim // 3, end, theta) + w_freqs_cis = self.precompute_freqs_cis(dim // 3, end, theta) + return {"f": f_freqs_cis, "h": h_freqs_cis, "w": w_freqs_cis} + + def precompute_freqs_cis(self, dim: int, end: int = 1024, theta: float = 10000.0): + # 1d rope precompute + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].double() / dim)) + freqs = torch.outer(torch.arange(end), freqs) + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 + return freqs_cis + + def forward(self, f: int, h: int, w: int, a: int) -> torch.Tensor: + freqs = torch.cat( + [ + self.freqs["f"][:f].view(f, 1, 1, -1).expand(f, h, w, -1), + self.freqs["h"][:h].view(1, h, 1, -1).expand(f, h, w, -1), + self.freqs["w"][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], + dim=-1, + ).reshape(f * h * w, 1, -1) + return freqs + + def post_initialize(self): + self.freqs = {key: value.to(device="cuda") for key, value in self.freqs.items()} + + +class AttentionModule(nn.Module): + def __init__(self, num_heads): + super().__init__() + self.num_heads = num_heads + + def forward(self, q, k, v): + x = flash_attention(q=q, k=k, v=v, num_heads=self.num_heads, compatibility_mode=FLASH_ATTN_COMPATIBILITY_MODE) + return x + + +class SelfAttention(nn.Module): + def __init__(self, dim: int, num_heads: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = _RMSNorm(dim, eps=eps) + self.norm_k = _RMSNorm(dim, eps=eps) + + self.attn = AttentionModule(self.num_heads) + + def forward(self, x, freqs): + q = self.norm_q(self.q(x)) + k = self.norm_k(self.k(x)) + v = self.v(x) + q = rope_apply(q, freqs, self.num_heads) + k = rope_apply(k, freqs, self.num_heads) + x = self.attn(q, k, v) + return self.o(x) + + +class CrossAttention(nn.Module): + def __init__(self, dim: int, num_heads: int, eps: float = 1e-6, has_image_input: bool = False): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = _RMSNorm(dim, eps=eps) + self.norm_k = _RMSNorm(dim, eps=eps) + self.has_image_input = has_image_input + if has_image_input: + self.k_img = nn.Linear(dim, dim) + self.v_img = nn.Linear(dim, dim) + self.norm_k_img = _RMSNorm(dim, eps=eps) + + self.attn = AttentionModule(self.num_heads) + + def forward(self, x: torch.Tensor, y: torch.Tensor): + if self.has_image_input: + img = y[:, :257] + ctx = y[:, 257:] + else: + ctx = y + q = self.norm_q(self.q(x)) + k = self.norm_k(self.k(ctx)) + v = self.v(ctx) + x = self.attn(q, k, v) + if self.has_image_input: + k_img = self.norm_k_img(self.k_img(img)) + v_img = self.v_img(img) + y = flash_attention(q, k_img, v_img, num_heads=self.num_heads, compatibility_mode=FLASH_ATTN_COMPATIBILITY_MODE) + x = x + y + return self.o(x) + + +class GateModule(nn.Module): + def __init__(self,): + super().__init__() + + @torch.compile(disable=DISABLE_TORCH_COMPILE) + def forward(self, x, gate, residual): + return x + gate * residual + +class DiTBlock(nn.Module): + def __init__(self, has_image_input: bool, dim: int, num_heads: int, ffn_dim: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.ffn_dim = ffn_dim + + self.self_attn = SelfAttention(dim, num_heads, eps) + self.cross_attn = CrossAttention( + dim, num_heads, eps, has_image_input=has_image_input) + self.norm1 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) + self.norm2 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) + self.norm3 = nn.LayerNorm(dim, eps=eps) + self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU( + approximate='tanh'), nn.Linear(ffn_dim, dim)) + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + self.gate = GateModule() + + def forward(self, x, context, t_mod, freqs): + # msa: multi-head self-attention mlp: multi-layer perceptron + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(6, dim=1) + input_x = modulate(self.norm1(x), shift_msa, scale_msa).to(dtype=t_mod.dtype) + x = self.gate(x, gate_msa, self.self_attn(input_x, freqs)) + x = x + self.cross_attn(self.norm3(x).to(dtype=t_mod.dtype), context.to(dtype=t_mod.dtype)) + input_x = modulate(self.norm2(x), shift_mlp, scale_mlp) + x = self.gate(x, gate_mlp, self.ffn(input_x)) + return x + + +class MLP(torch.nn.Module): + def __init__(self, in_dim, out_dim, has_pos_emb=False): + super().__init__() + self.proj = torch.nn.Sequential( + nn.LayerNorm(in_dim), + nn.Linear(in_dim, in_dim), + nn.GELU(), + nn.Linear(in_dim, out_dim), + nn.LayerNorm(out_dim) + ) + self.has_pos_emb = has_pos_emb + if has_pos_emb: + self.emb_pos = torch.nn.Parameter(torch.zeros((1, 514, 1280))) + + def forward(self, x): + if self.has_pos_emb: + x = x + self.emb_pos.to(dtype=x.dtype, device=x.device) + return self.proj(x) + + +class Head(nn.Module): + def __init__(self, dim: int, out_dim: int, patch_size: Tuple[int, int, int], eps: float): + super().__init__() + self.dim = dim + self.patch_size = patch_size + self.norm = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) + self.head = nn.Linear(dim, out_dim * math.prod(patch_size)) + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x, t_mod): + batch_size = x.shape[0] + shift, scale = (self.modulation.repeat(batch_size,1,1).to(dtype=t_mod.dtype, device=t_mod.device) + t_mod.unsqueeze(1)).chunk(2, dim=1) + x = (self.head(self.norm(x) * (1 + scale) + shift)) + return x + + +class WanModel(ModelMixin, ConfigMixin): + @register_to_config + def __init__( + self, + dim: int, + in_dim: int, + ffn_dim: int, + out_dim: int, + freq_dim: int, + eps: float, + num_heads: int, + num_layers: int, + text_dim: int = 4096, + patch_size: Tuple[int, int, int] = [1, 2, 2], + has_image_input: bool = True, + has_image_pos_emb: bool = False, + has_ref_conv: bool = False, + add_control_adapter: bool = False, + in_dim_control_adapter: int = 24, + diffusion_model_pretrained_path: str = None, + ): + super().__init__() + self.dim = dim + self.freq_dim = freq_dim + self.has_image_input = has_image_input + self.patch_size = patch_size + self.diffusion_model_pretrained_path = diffusion_model_pretrained_path + + + self.patch_embedding = nn.Conv3d( + in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, dim), + nn.GELU(approximate='tanh'), + nn.Linear(dim, dim) + ) + self.time_embedding = nn.Sequential( + nn.Linear(freq_dim, dim), + nn.SiLU(), + nn.Linear(dim, dim) + ) + self.time_projection = nn.Sequential( + nn.SiLU(), nn.Linear(dim, dim * 6)) + self.blocks = nn.ModuleList([ + DiTBlock(has_image_input, dim, num_heads, ffn_dim, eps) + for _ in range(num_layers) + ]) + self.head = Head(dim, out_dim, patch_size, eps) + head_dim = dim // num_heads + + self.rope = RotaryPositionEmbedding(num_heads=num_heads, head_dim=head_dim) + + if has_image_input: + self.img_emb = MLP(1280, dim, has_pos_emb=has_image_pos_emb) # clip_feature_dim = 1280 + if has_ref_conv: + self.ref_conv = nn.Conv2d(16, dim, kernel_size=(2, 2), stride=(2, 2)) + self.has_image_pos_emb = has_image_pos_emb + self.has_ref_conv = has_ref_conv + if add_control_adapter: + self.control_adapter = SimpleAdapter(in_dim_control_adapter, dim, kernel_size=patch_size[1:], stride=patch_size[1:]) + else: + self.control_adapter = None + + self.use_gradient_checkpointing = False + self.use_gradient_checkpointing_offload = False + + def patchify(self, x: torch.Tensor,control_camera_latents_input: torch.Tensor = None): + x = self.patch_embedding(x) + if self.control_adapter is not None and control_camera_latents_input is not None: + y_camera = self.control_adapter(control_camera_latents_input) + x = [u + v for u, v in zip(x, y_camera)] + x = x[0].unsqueeze(0) + grid_size = x.shape[2:] + x = rearrange(x, 'b c f h w -> b (f h w) c').contiguous() + return x, grid_size # x, grid_size: (f, h, w) + + def unpatchify(self, x: torch.Tensor, grid_size: torch.Tensor): + return rearrange( + x, 'b (f h w) (x y z c) -> b c (f x) (h y) (w z)', + f=grid_size[0], h=grid_size[1], w=grid_size[2], + x=self.patch_size[0], y=self.patch_size[1], z=self.patch_size[2] + ) + + def forward(self, + x: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + clip_feature: Optional[torch.Tensor] = None, + y: Optional[torch.Tensor] = None, + use_gradient_checkpointing: bool = False, + use_gradient_checkpointing_offload: bool = False, + **kwargs, + ): + t = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, timestep)) + t_mod = self.time_projection(t).unflatten(1, (6, self.dim)) + context = self.text_embedding(context) + + if self.has_image_input: + x = torch.cat([x, y], dim=1) # (b, c_x + c_y, f, h, w) + clip_embdding = self.img_emb(clip_feature) + context = torch.cat([clip_embdding, context], dim=1) + # print("clip embedding shape", clip_embdding.shape) + # print("x before patchify", x.shape) + # clip embedding shape of B * 257 * dim (5120) + # context shape of B * 769 * dim (5120) - 257 + 512 + # x before patchify B * 36 * l_t * l_h * l_w + x, (f, h, w) = self.patchify(x) + # x after patchify B * 512 * 5120 + # f = 2, h = 16, w = 16 + # print("x and context shape", x.shape, context.shape, f,h,w) + + freqs = self.rope(f=f, h=h, w=w, a=x.shape[1]) + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + return custom_forward + + for block in self.blocks: + if self.training and self.use_gradient_checkpointing: + if self.use_gradient_checkpointing_offload: + with torch.autograd.graph.save_on_cpu(): + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, context, t_mod, freqs, + use_reentrant=False, + ) + else: + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, context, t_mod, freqs, + use_reentrant=False, + ) + else: + x = block(x, context, t_mod, freqs) + + # print("x before head", x.shape) + x = self.head(x, t) + # print("x before unpatchify", x.shape) + x = self.unpatchify(x, (f, h, w)) + # print("x after patchify", x.shape) + return x + def post_initialize(self): + self.rope.post_initialize() + @staticmethod + def state_dict_converter(): + return WanModelStateDictConverter() + + +class WanModelStateDictConverter: + def __init__(self): + pass + + def from_diffusers(self, state_dict): + rename_dict = { + "blocks.0.attn1.norm_k.weight": "blocks.0.self_attn.norm_k.weight", + "blocks.0.attn1.norm_q.weight": "blocks.0.self_attn.norm_q.weight", + "blocks.0.attn1.to_k.bias": "blocks.0.self_attn.k.bias", + "blocks.0.attn1.to_k.weight": "blocks.0.self_attn.k.weight", + "blocks.0.attn1.to_out.0.bias": "blocks.0.self_attn.o.bias", + "blocks.0.attn1.to_out.0.weight": "blocks.0.self_attn.o.weight", + "blocks.0.attn1.to_q.bias": "blocks.0.self_attn.q.bias", + "blocks.0.attn1.to_q.weight": "blocks.0.self_attn.q.weight", + "blocks.0.attn1.to_v.bias": "blocks.0.self_attn.v.bias", + "blocks.0.attn1.to_v.weight": "blocks.0.self_attn.v.weight", + "blocks.0.attn2.norm_k.weight": "blocks.0.cross_attn.norm_k.weight", + "blocks.0.attn2.norm_q.weight": "blocks.0.cross_attn.norm_q.weight", + "blocks.0.attn2.to_k.bias": "blocks.0.cross_attn.k.bias", + "blocks.0.attn2.to_k.weight": "blocks.0.cross_attn.k.weight", + "blocks.0.attn2.to_out.0.bias": "blocks.0.cross_attn.o.bias", + "blocks.0.attn2.to_out.0.weight": "blocks.0.cross_attn.o.weight", + "blocks.0.attn2.to_q.bias": "blocks.0.cross_attn.q.bias", + "blocks.0.attn2.to_q.weight": "blocks.0.cross_attn.q.weight", + "blocks.0.attn2.to_v.bias": "blocks.0.cross_attn.v.bias", + "blocks.0.attn2.to_v.weight": "blocks.0.cross_attn.v.weight", + "blocks.0.ffn.net.0.proj.bias": "blocks.0.ffn.0.bias", + "blocks.0.ffn.net.0.proj.weight": "blocks.0.ffn.0.weight", + "blocks.0.ffn.net.2.bias": "blocks.0.ffn.2.bias", + "blocks.0.ffn.net.2.weight": "blocks.0.ffn.2.weight", + "blocks.0.norm2.bias": "blocks.0.norm3.bias", + "blocks.0.norm2.weight": "blocks.0.norm3.weight", + "blocks.0.scale_shift_table": "blocks.0.modulation", + "condition_embedder.text_embedder.linear_1.bias": "text_embedding.0.bias", + "condition_embedder.text_embedder.linear_1.weight": "text_embedding.0.weight", + "condition_embedder.text_embedder.linear_2.bias": "text_embedding.2.bias", + "condition_embedder.text_embedder.linear_2.weight": "text_embedding.2.weight", + "condition_embedder.time_embedder.linear_1.bias": "time_embedding.0.bias", + "condition_embedder.time_embedder.linear_1.weight": "time_embedding.0.weight", + "condition_embedder.time_embedder.linear_2.bias": "time_embedding.2.bias", + "condition_embedder.time_embedder.linear_2.weight": "time_embedding.2.weight", + "condition_embedder.time_proj.bias": "time_projection.1.bias", + "condition_embedder.time_proj.weight": "time_projection.1.weight", + "patch_embedding.bias": "patch_embedding.bias", + "patch_embedding.weight": "patch_embedding.weight", + "scale_shift_table": "head.modulation", + "proj_out.bias": "head.head.bias", + "proj_out.weight": "head.head.weight", + } + state_dict_ = {} + for name, param in state_dict.items(): + if name in rename_dict: + state_dict_[rename_dict[name]] = param + else: + name_ = ".".join(name.split(".")[:1] + ["0"] + name.split(".")[2:]) + if name_ in rename_dict: + name_ = rename_dict[name_] + name_ = ".".join(name_.split(".")[:1] + [name.split(".")[1]] + name_.split(".")[2:]) + state_dict_[name_] = param + if hash_state_dict_keys(state_dict) == "cb104773c6c2cb6df4f9529ad5c60d0b": + config = { + "model_type": "t2v", + "patch_size": (1, 2, 2), + "text_len": 512, + "in_dim": 16, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "window_size": (-1, -1), + "qk_norm": True, + "cross_attn_norm": True, + "eps": 1e-6, + } + else: + config = {} + return state_dict_, config + + def from_civitai(self, state_dict): + state_dict = {name: param for name, param in state_dict.items() if not name.startswith("vace")} + if hash_state_dict_keys(state_dict) == "9269f8db9040a9d860eaca435be61814": + config = { + "has_image_input": False, + "patch_size": [1, 2, 2], + "in_dim": 16, + "dim": 1536, + "ffn_dim": 8960, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 12, + "num_layers": 30, + "eps": 1e-6 + } + elif hash_state_dict_keys(state_dict) == "aafcfd9672c3a2456dc46e1cb6e52c70": + config = { + "has_image_input": False, + "patch_size": [1, 2, 2], + "in_dim": 16, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "eps": 1e-6 + } + elif hash_state_dict_keys(state_dict) == "6bfcfb3b342cb286ce886889d519a77e": + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 36, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "eps": 1e-6 + } + elif hash_state_dict_keys(state_dict) == "6d6ccde6845b95ad9114ab993d917893": + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 36, + "dim": 1536, + "ffn_dim": 8960, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 12, + "num_layers": 30, + "eps": 1e-6 + } + elif hash_state_dict_keys(state_dict) == "6bfcfb3b342cb286ce886889d519a77e": + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 36, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "eps": 1e-6 + } + elif hash_state_dict_keys(state_dict) == "349723183fc063b2bfc10bb2835cf677": + # 1.3B PAI control + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 48, + "dim": 1536, + "ffn_dim": 8960, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 12, + "num_layers": 30, + "eps": 1e-6 + } + elif hash_state_dict_keys(state_dict) == "efa44cddf936c70abd0ea28b6cbe946c": + # 14B PAI control + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 48, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "eps": 1e-6 + } + elif hash_state_dict_keys(state_dict) == "3ef3b1f8e1dab83d5b71fd7b617f859f": + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 36, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "eps": 1e-6, + "has_image_pos_emb": True + } + elif hash_state_dict_keys(state_dict) == "70ddad9d3a133785da5ea371aae09504": + # 1.3B PAI control v1.1 + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 48, + "dim": 1536, + "ffn_dim": 8960, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 12, + "num_layers": 30, + "eps": 1e-6, + "has_ref_conv": True + } + elif hash_state_dict_keys(state_dict) == "26bde73488a92e64cc20b0a7485b9e5b": + # 14B PAI control v1.1 + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 48, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "eps": 1e-6, + "has_ref_conv": True + } + elif hash_state_dict_keys(state_dict) == "ac6a5aa74f4a0aab6f64eb9a72f19901": + # 1.3B PAI control-camera v1.1 + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 32, + "dim": 1536, + "ffn_dim": 8960, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 12, + "num_layers": 30, + "eps": 1e-6, + "has_ref_conv": False, + "add_control_adapter": True, + "in_dim_control_adapter": 24, + } + elif hash_state_dict_keys(state_dict) == "b61c605c2adbd23124d152ed28e049ae": + # 14B PAI control-camera v1.1 + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 32, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "eps": 1e-6, + "has_ref_conv": False, + "add_control_adapter": True, + "in_dim_control_adapter": 24, + } + else: + config = {} + return state_dict, config diff --git a/groot/vla/model/dreamzero/modules/wan_video_dit_action_casual_chunk.py b/groot/vla/model/dreamzero/modules/wan_video_dit_action_casual_chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..c5e2528401c0d7a1c69daafb29eb4ad50c06f30b --- /dev/null +++ b/groot/vla/model/dreamzero/modules/wan_video_dit_action_casual_chunk.py @@ -0,0 +1,2245 @@ +from typing import Any, TypeAlias + +from groot.vla.model.dreamzero.modules.wan2_1_attention import AttentionModule +from groot.vla.model.n1_5.modules.action_encoder import ( + SinusoidalPositionalEncoding, + swish, +) +from groot.vla.model.dreamzero.modules.wan2_1_submodule import ( + WanRMSNorm, + rope_action_apply, + WanLayerNorm, + WAN_CROSSATTENTION_CLASSES, + rope_params, + MLPProj, + sinusoidal_embedding_1d +) +from torch.nn.attention.flex_attention import create_block_mask, create_mask +from torch.nn.attention.flex_attention import BlockMask +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +import torch.nn as nn +import torch.nn.functional as F +import torch +import math +import torch.distributed as dist +import os + +ENABLE_TENSORRT = os.getenv("ENABLE_TENSORRT", "False").lower() == "true" + + +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] + 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 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 + + # 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 + + +def causal_rope_action_apply(x, freqs, freqs_action, freqs_state, action_register_length, num_action_per_block, num_state_per_block, action_state_index): + if ENABLE_TENSORRT: + return causal_rope_action_apply_no_polar(x, freqs, freqs_action, freqs_state, action_register_length, num_action_per_block, num_state_per_block, action_state_index) + else: + return causal_rope_action_apply_polar(x, freqs, freqs_action, freqs_state, action_register_length, num_action_per_block, num_state_per_block, action_state_index) + + +def causal_rope_action_apply_no_polar( + x: torch.Tensor, + freqs: torch.Tensor, + freqs_action: torch.Tensor, + freqs_state: torch.Tensor, + action_register_length: int | None, + num_action_per_block: int, + num_state_per_block: int, + action_state_index: int, +): + B, seq_len, n, d = x.shape + + # (B, seq_len, n, d) -> (B, seq_len, n, d/2, 2) + x = x.reshape(B, seq_len, n, -1, 2) + x_real = x[..., 0] + x_imag = x[..., 1] + + # Split freqs into cos and sin components + freqs = freqs.unsqueeze(0).view(1, freqs.shape[0], 1, -1, 2) + freqs_cos = freqs[..., 0] # Shape: (1, seq_len', 1, d/2) + freqs_sin = freqs[..., 1] # Shape: (1, seq_len', 1, d/2) + + # Handle the Action/State Register Frequencies + if action_register_length is not None: + assert action_register_length == (num_action_per_block + num_state_per_block) + + freqs_action_slice = freqs_action[ + action_state_index * num_action_per_block:(action_state_index + 1) * num_action_per_block + ] + freqs_state_slice = freqs_state[ + action_state_index * num_state_per_block:(action_state_index + 1) * num_state_per_block + ] + + # Combine the action/state tokens for this frame + freqs_1d = torch.cat([freqs_action_slice, freqs_state_slice], dim=0).view( + action_register_length, 1, -1, 2 + ) + + # Split the new action/state frequencies + freqs_cos_1d = freqs_1d[..., 0] + freqs_sin_1d = freqs_1d[..., 1] + + # Append the action/state register sin/cos to the main sequence sin/cos + freqs_cos = torch.cat([freqs_cos[0], freqs_cos_1d], dim=0).unsqueeze(0) + freqs_sin = torch.cat([freqs_sin[0], freqs_sin_1d], dim=0).unsqueeze(0) + + x_real_rotated = x_real * freqs_cos - x_imag * freqs_sin + x_imag_rotated = x_real * freqs_sin + x_imag * freqs_cos + + x_rotated = torch.stack((x_real_rotated, x_imag_rotated), dim=-1) + + return x_rotated.flatten(3) + +def causal_rope_action_apply_polar( + x: torch.Tensor, + freqs: torch.Tensor, + freqs_action: torch.Tensor, + freqs_state: torch.Tensor, + action_register_length: int | None, + num_action_per_block: int, + num_state_per_block: int, + action_state_index: int, +): + B, seq_len, n, _ = x.shape + + # precompute multipliers + x = torch.view_as_complex( + x.to(torch.float64).reshape(B, seq_len, n, -1, 2) + ) + + if action_register_length is not None: + assert action_register_length == (num_action_per_block + num_state_per_block) + freqs_action = freqs_action[ + action_state_index * num_action_per_block:(action_state_index + 1) * num_action_per_block + ] + freqs_state = freqs_state[ + action_state_index * num_state_per_block:(action_state_index + 1) * num_state_per_block + ] + freqs_1d = torch.cat([freqs_action, freqs_state], dim=0).view(action_register_length, 1, -1) + freqs = torch.cat([freqs, freqs_1d], dim=0) + + # apply rotary embedding + freqs = freqs.unsqueeze(0) + x = torch.view_as_real(x * freqs).flatten(3) + + return x + + +class CausalWanSelfAttention(nn.Module): + + def __init__(self, + dim, + num_heads, + frame_seqlen, + local_attn_size=-1, + sink_size=0, + num_frame_per_block=1, + qk_norm=True, + eps=1e-6, + num_action_per_block=32, + num_state_per_block=1): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.local_attn_size = local_attn_size + self.sink_size = sink_size + self.num_frame_per_block = num_frame_per_block + self.qk_norm = qk_norm + self.eps = eps + self.max_attention_size = 21 * frame_seqlen if local_attn_size == -1 else local_attn_size * frame_seqlen + self.frame_seqlen = frame_seqlen + self.num_action_per_block = num_action_per_block + self.num_state_per_block = num_state_per_block + # layers + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.attn = AttentionModule(num_heads=self.num_heads, head_dim=self.head_dim) + self.causal_attn = AttentionModule(num_heads=self.num_heads, head_dim=self.head_dim, causal=True) + + def _visualize_attention_mask(self, total_len, first_image_len, image_blocks_len, + action_len, state_len, num_image_blocks, + num_action_blocks, num_state_blocks, + num_frame_per_block, frame_seqlen, + num_action_per_block, num_state_per_block): + """ + Create and print a visualization of the attention mask pattern. + Returns a binary mask [total_len, total_len] where 1 = can attend, 0 = cannot attend. + """ + # Token ranges + first_image_start = 0 + first_image_end = first_image_len + image_blocks_start = first_image_end + image_blocks_end = image_blocks_start + image_blocks_len + action_start = image_blocks_end + action_end = action_start + action_len + state_start = action_end + state_end = state_start + state_len + + # Create mask tensor + mask = torch.zeros(total_len, total_len, dtype=torch.bool) + + # First image: self-attention only + mask[first_image_start:first_image_end, first_image_start:first_image_end] = True + + # Image blocks + for block_idx in range(num_image_blocks): + block_start = image_blocks_start + block_idx * num_frame_per_block * frame_seqlen + block_end = image_blocks_start + (block_idx + 1) * num_frame_per_block * frame_seqlen + + # Attend to first image + mask[block_start:block_end, first_image_start:first_image_end] = True + + # Attend to previous and current image blocks + if self.local_attn_size != -1: + image_kv_start = max(image_blocks_start, block_end - self.local_attn_size * frame_seqlen) + else: + image_kv_start = image_blocks_start + mask[block_start:block_end, image_kv_start:block_end] = True + + # Attend to current action block + action_block_start = action_start + block_idx * num_action_per_block + action_block_end = action_start + (block_idx + 1) * num_action_per_block + mask[block_start:block_end, action_block_start:action_block_end] = True + + # Attend to current state block + state_block_start = state_start + block_idx * num_state_per_block + state_block_end = state_start + (block_idx + 1) * num_state_per_block + mask[block_start:block_end, state_block_start:state_block_end] = True + + # Action blocks + for block_idx in range(num_action_blocks): + action_block_start = action_start + block_idx * num_action_per_block + action_block_end = action_start + (block_idx + 1) * num_action_per_block + + # Attend to first image + mask[action_block_start:action_block_end, first_image_start:first_image_end] = True + + # Attend to previous and current image blocks + image_block_end = image_blocks_start + (block_idx + 1) * num_frame_per_block * frame_seqlen + if self.local_attn_size != -1: + image_kv_start = max(image_blocks_start, image_block_end - self.local_attn_size * frame_seqlen) + else: + image_kv_start = image_blocks_start + mask[action_block_start:action_block_end, image_kv_start:image_block_end] = True + + # Self-attention + mask[action_block_start:action_block_end, action_block_start:action_block_end] = True + + # Attend to current state block + state_block_start = state_start + block_idx * num_state_per_block + state_block_end = state_start + (block_idx + 1) * num_state_per_block + mask[action_block_start:action_block_end, state_block_start:state_block_end] = True + + # State blocks: self-attention only + for block_idx in range(num_state_blocks): + state_block_start = state_start + block_idx * num_state_per_block + state_block_end = state_start + (block_idx + 1) * num_state_per_block + mask[state_block_start:state_block_end, state_block_start:state_block_end] = True + + return mask + + def _blockwise_causal_flash_attn(self, q, k, v, frame_seqlen, num_frame_per_block=1, + action_horizon=None, state_horizon=None, + num_action_per_block=None, num_state_per_block=None, + visualize_mask=False): + """ + Implement blockwise causal attention using flash_attention. + Matches the pattern from _prepare_blockwise_causal_attn_mask: + + Structure: + - First image: conditioning only, cannot attend to anything + - Image blocks: can attend to first image + previous image blocks + current action block + current state block + - Action blocks: can attend to previous image blocks + current image block + current state block + first image + - State blocks: conditioning only, cannot attend to anything + + Args: + q, k, v: Query, key, value tensors [B, L, num_heads, head_dim] + frame_seqlen: Number of tokens per frame + num_frame_per_block: Number of frames per attention block + action_horizon: Total number of action tokens (if None, no action/state tokens) + state_horizon: Total number of state tokens (if None, no action/state tokens) + num_action_per_block: Number of action tokens per block + num_state_per_block: Number of state tokens per block + visualize_mask: If True, print the attention mask pattern + + Returns: + Attention output [B, L, num_heads, head_dim] + """ + b, total_len, n, d = q.shape + + # Check if we have action/state tokens + has_action_state = (action_horizon is not None and state_horizon is not None) + + if not has_action_state: + # OPTIMIZED: Simple blockwise causal attention (without action/state tokens) + num_frames = total_len // frame_seqlen + block_size = frame_seqlen * num_frame_per_block + num_blocks = (num_frames - 1) // num_frame_per_block + + # Handle edge case when sequence is too short (no blocks to process) + if num_blocks <= 0: + # Process entire sequence as a single block + return self.attn(q, k, v) + + # OPTIMIZATION: For global attention, process all blocks in one call with causal masking + if self.local_attn_size == -1: + # Single flash_attention call with causal=True for all blocks at once + # This is much faster than looping! + return self.causal_attn(q, k, v) + + # With local attention, still need loop but optimize it + # Pre-allocate output tensor + output = torch.empty_like(q) + + # Pre-compute block boundaries + block_starts = [frame_seqlen + i * block_size for i in range(num_blocks)] + block_ends = [min(start + block_size, total_len) for start in block_starts] + kv_starts = [max(0, end - self.local_attn_size * frame_seqlen) for end in block_ends] + + for block_idx in range(num_blocks): + block_start = block_starts[block_idx] + block_end = block_ends[block_idx] + kv_start = kv_starts[block_idx] + + output[:, block_start:block_end] = self.attn( + q[:, block_start:block_end], + k[:, kv_start:block_end], + v[:, kv_start:block_end] + ) + + return output + + assert action_horizon is not None and state_horizon is not None + assert num_action_per_block is not None and num_state_per_block is not None + + # Multi-modal structure: [first image] [image blocks] [action blocks] [state blocks] + # Calculate block structure + first_image_len = frame_seqlen + action_len = action_horizon + state_len = state_horizon + image_blocks_len = total_len - first_image_len - action_len - state_len + + num_image_blocks = image_blocks_len // (num_frame_per_block * frame_seqlen) + num_action_blocks = action_horizon // num_action_per_block + num_state_blocks = state_horizon // num_state_per_block + + assert num_image_blocks == num_action_blocks == num_state_blocks + + # Token ranges + first_image_start = 0 + first_image_end = first_image_len + image_blocks_start = first_image_end + image_blocks_end = image_blocks_start + image_blocks_len + action_start = image_blocks_end + action_end = action_start + action_len + state_start = action_end + state_end = state_start + state_len + + # Visualize attention mask if requested + if visualize_mask: + mask = self._visualize_attention_mask( + total_len, first_image_len, image_blocks_len, + action_len, state_len, num_image_blocks, + num_action_blocks, num_state_blocks, + num_frame_per_block, frame_seqlen, + num_action_per_block, num_state_per_block + ) + + print("\n" + "="*80) + print("ATTENTION MASK VISUALIZATION") + print("="*80) + print(f"Total length: {total_len}") + print(f"First image: [{first_image_start}:{first_image_end}] (len={first_image_len})") + print(f"Image blocks: [{image_blocks_start}:{image_blocks_end}] (len={image_blocks_len}, num_blocks={num_image_blocks})") + print(f"Action tokens: [{action_start}:{action_end}] (len={action_len}, num_blocks={num_action_blocks})") + print(f"State tokens: [{state_start}:{state_end}] (len={state_len}, num_blocks={num_state_blocks})") + print(f"Local attention size: {self.local_attn_size}") + print("-"*80) + + # Print a downsampled version of the mask if it's too large + if total_len <= 100: + # Print full mask for small sequences + print("Attention mask (1=can attend, 0=cannot attend):") + print("Rows=Query tokens, Cols=Key tokens") + for i in range(total_len): + row = "".join(["1" if mask[i, j] else "." for j in range(total_len)]) + print(f"{i:4d}: {row}") + else: + # Print downsampled version for large sequences + downsample = max(1, total_len // 100) + print(f"Attention mask (downsampled by {downsample}x):") + print("Rows=Query tokens, Cols=Key tokens (1=can attend, .=cannot attend)") + for i in range(0, total_len, downsample): + row = "".join(["1" if mask[i, j] else "." for j in range(0, total_len, downsample)]) + print(f"{i:4d}: {row}") + + # Save mask as image + try: + import cv2 + import numpy as np + mask_np = mask.cpu().float().numpy() + # Resize for visualization if needed + if total_len > 1000: + mask_np = cv2.resize(mask_np, (1000, 1000), interpolation=cv2.INTER_NEAREST) + mask_img = (mask_np * 255).astype(np.uint8) + cv2.imwrite("attention_mask_blockwise_flash.png", mask_img) + print(f"\nMask saved to: attention_mask_blockwise_flash.png") + except Exception as e: + print(f"Could not save mask image: {e}") + + print("="*80 + "\n") + + # OPTIMIZED: Pre-allocate output tensor and pre-compute all indices + output = torch.empty_like(q) + + # Process first image (conditioning, can only self-attend) + output[:, first_image_start:first_image_end] = self.attn( + q[:, first_image_start:first_image_end], + k[:, first_image_start:first_image_end], + v[:, first_image_start:first_image_end] + ) + + # Pre-compute all block indices for image blocks + image_block_starts = [image_blocks_start + i * num_frame_per_block * frame_seqlen for i in range(num_image_blocks)] + image_block_ends = [image_blocks_start + (i + 1) * num_frame_per_block * frame_seqlen for i in range(num_image_blocks)] + if self.local_attn_size != -1: + image_kv_starts = [max(image_blocks_start, end - self.local_attn_size * frame_seqlen) for end in image_block_ends] + else: + image_kv_starts = [image_blocks_start] * num_image_blocks + + # Pre-compute action and state block indices + action_block_starts = [action_start + i * num_action_per_block for i in range(num_action_blocks)] + action_block_ends = [action_start + (i + 1) * num_action_per_block for i in range(num_action_blocks)] + state_block_starts = [state_start + i * num_state_per_block for i in range(num_state_blocks)] + state_block_ends = [state_start + (i + 1) * num_state_per_block for i in range(num_state_blocks)] + + # Process each image block + for block_idx in range(num_image_blocks): + block_start = image_block_starts[block_idx] + block_end = image_block_ends[block_idx] + image_kv_start = image_kv_starts[block_idx] + action_block_start = action_block_starts[block_idx] + action_block_end = action_block_ends[block_idx] + state_block_start = state_block_starts[block_idx] + state_block_end = state_block_ends[block_idx] + + # Build context: first image + relevant image blocks + current action + current state + k_context = torch.cat([ + k[:, first_image_start:first_image_end], # First image + k[:, image_kv_start:block_end], # Image blocks + k[:, action_block_start:action_block_end], # Current action block + k[:, state_block_start:state_block_end] # Current state block + ], dim=1) + v_context = torch.cat([ + v[:, first_image_start:first_image_end], + v[:, image_kv_start:block_end], + v[:, action_block_start:action_block_end], + v[:, state_block_start:state_block_end] + ], dim=1) + + output[:, block_start:block_end] = self.attn( + q[:, block_start:block_end], k_context, v_context + ) + + # Process each action block + for block_idx in range(num_action_blocks): + action_block_start = action_block_starts[block_idx] + action_block_end = action_block_ends[block_idx] + image_block_end = image_block_ends[block_idx] + state_block_start = state_block_starts[block_idx] + state_block_end = state_block_ends[block_idx] + + # Determine image context range + if self.local_attn_size != -1: + image_kv_start = max(image_blocks_start, image_block_end - self.local_attn_size * frame_seqlen) + else: + image_kv_start = image_blocks_start + + # Build context + k_context = torch.cat([ + k[:, first_image_start:first_image_end], # First image + k[:, image_kv_start:image_block_end], # Image blocks + k[:, action_block_start:action_block_end], # Current action block + k[:, state_block_start:state_block_end] # Current state block + ], dim=1) + v_context = torch.cat([ + v[:, first_image_start:first_image_end], + v[:, image_kv_start:image_block_end], + v[:, action_block_start:action_block_end], + v[:, state_block_start:state_block_end] + ], dim=1) + + output[:, action_block_start:action_block_end] = self.attn( + q[:, action_block_start:action_block_end], k_context, v_context + ) + + # Process state blocks (conditioning, can only self-attend) + for block_idx in range(num_state_blocks): + state_block_start = state_block_starts[block_idx] + state_block_end = state_block_ends[block_idx] + + output[:, state_block_start:state_block_end] = self.attn( + q[:, state_block_start:state_block_end], + k[:, state_block_start:state_block_end], + v[:, state_block_start:state_block_end] + ) + + return output + + def _process_clean_image_only(self, clean_image_q, clean_image_k, clean_image_v, clean_frames): + """Process clean image blocks with causal attention pattern - OPTIMIZED + + First frame: conditioning, cannot attend to anything (self-attention only) + Block i: attends to first frame + previous blocks (0 to i-1) + current block + + OPTIMIZATION: Instead of looping through blocks, we batch process them together + by using a single flash_attention call with properly structured KV cache. + """ + block_size = self.frame_seqlen * self.num_frame_per_block + num_blocks = (clean_frames - 1) // self.num_frame_per_block + + if num_blocks == 0: + # Only first frame - single attention call + return self.attn( + clean_image_q[:, :self.frame_seqlen], + clean_image_k[:, :self.frame_seqlen], + clean_image_v[:, :self.frame_seqlen] + ) + + # Pre-allocate output tensor (avoids list append + cat overhead) + b, total_len, n, d = clean_image_q.shape + output = torch.empty_like(clean_image_q) + + # First frame: conditioning, self-attention only + output[:, :self.frame_seqlen] = self.attn( + clean_image_q[:, :self.frame_seqlen], + clean_image_k[:, :self.frame_seqlen], + clean_image_v[:, :self.frame_seqlen] + ) + + # OPTIMIZATION: Process all blocks together with causal masking + # For global attention (no local_attn_size), we can process all blocks in one call + if self.local_attn_size == -1: + # Single attention call for all blocks! + # Each position can attend to first_frame + everything up to itself + blocks_q = clean_image_q[:, self.frame_seqlen:] + blocks_k = clean_image_k # Can attend to everything including first frame + blocks_v = clean_image_v + + # Use causal masking: each block token can see first frame + all previous tokens + output[:, self.frame_seqlen:] = self.causal_attn( + blocks_q, blocks_k, blocks_v + ) + else: + # With local attention, we still need to loop but with optimizations + # Pre-compute all block boundaries to reduce overhead + block_starts = [self.frame_seqlen + i * block_size for i in range(num_blocks)] + block_ends = [min(start + block_size, total_len) for start in block_starts] + + for block_idx in range(num_blocks): + block_start = block_starts[block_idx] + block_end = block_ends[block_idx] + + q_block = clean_image_q[:, block_start:block_end] + + # Context: first frame + recent blocks within local_attn_size + image_kv_start = max(self.frame_seqlen, block_end - self.local_attn_size * self.frame_seqlen) + k_context = torch.cat([ + clean_image_k[:, :self.frame_seqlen], # First frame + clean_image_k[:, image_kv_start:block_end] # Recent blocks + current + ], dim=1) + v_context = torch.cat([ + clean_image_v[:, :self.frame_seqlen], + clean_image_v[:, image_kv_start:block_end] + ], dim=1) + + output[:, block_start:block_end] = self.attn(q_block, k_context, v_context) + + return output + + def _process_state_blocks(self, state_q, state_k, state_v, state_horizon): + """Process state blocks: self-attention only - OPTIMIZED + + OPTIMIZATION: State blocks only do self-attention within each block. + Instead of looping, we can process all blocks in a single call with block-diagonal masking, + or even simpler: just one attention call since they're independent. + """ + num_blocks = state_horizon // self.num_state_per_block + + if num_blocks == 1: + # Single block - one attention call + return self.attn(state_q, state_k, state_v) + + # OPTIMIZATION: Since each state block only attends to itself (no cross-block attention), + # we can process all blocks in a single batched call. Flash attention will handle this + # efficiently. The blocks are independent, so this is safe. + # Alternative: reshape and process as separate batch items + + # Pre-allocate output + output = torch.empty_like(state_q) + + # Process all blocks (keeping loop for now due to block-diagonal pattern) + # This could be further optimized with custom masking + for block_idx in range(num_blocks): + state_block_start = block_idx * self.num_state_per_block + state_block_end = state_block_start + self.num_state_per_block + + output[:, state_block_start:state_block_end] = self.attn( + state_q[:, state_block_start:state_block_end], + state_k[:, state_block_start:state_block_end], + state_v[:, state_block_start:state_block_end] + ) + + return output + + def _process_noisy_image_blocks(self, noisy_image_q, noisy_image_k, noisy_image_v, + clean_image_k, clean_image_v, + noisy_action_k, noisy_action_v, noisy_state_k, noisy_state_v, + half_frames, action_horizon, state_horizon): + """Process noisy image blocks with teacher forcing pattern - OPTIMIZED + + First frame: conditioning, cannot attend to anything (self-attention only) + Block i: attends to action[i] + state[i] + first_clean_frame + clean_blocks[0:i] + current_noisy_block + + OPTIMIZATION: Pre-allocate output, pre-compute indices, reduce memory allocations + """ + block_size = self.frame_seqlen * self.num_frame_per_block + num_blocks = (half_frames - 1) // self.num_frame_per_block + + # Pre-allocate output tensor + output = torch.empty_like(noisy_image_q) + + # First noisy frame: conditioning, self-attention only + output[:, :self.frame_seqlen] = self.attn( + noisy_image_q[:, :self.frame_seqlen], + noisy_image_k[:, :self.frame_seqlen], + noisy_image_v[:, :self.frame_seqlen] + ) + + if num_blocks == 0: + return output + + # Pre-compute all block indices to reduce loop overhead + noisy_block_starts = [self.frame_seqlen + i * block_size for i in range(num_blocks)] + noisy_block_ends = [min(start + block_size, noisy_image_q.shape[1]) for start in noisy_block_starts] + clean_context_ends = [self.frame_seqlen + i * block_size for i in range(num_blocks)] + action_block_starts = [i * self.num_action_per_block for i in range(num_blocks)] + action_block_ends = [start + self.num_action_per_block for start in action_block_starts] + state_block_starts = [i * self.num_state_per_block for i in range(num_blocks)] + state_block_ends = [start + self.num_state_per_block for start in state_block_starts] + + # Process noisy image blocks + for block_idx in range(num_blocks): + noisy_start = noisy_block_starts[block_idx] + noisy_end = noisy_block_ends[block_idx] + clean_end = clean_context_ends[block_idx] + action_start = action_block_starts[block_idx] + action_end = action_block_ends[block_idx] + state_start = state_block_starts[block_idx] + state_end = state_block_ends[block_idx] + + q_block = noisy_image_q[:, noisy_start:noisy_end] + + # Build context: first_clean_frame + clean_blocks[0:i] + current_noisy_block + action[i] + state[i] + k_context = torch.cat([ + clean_image_k[:, :clean_end], + noisy_image_k[:, noisy_start:noisy_end], + noisy_action_k[:, action_start:action_end], + noisy_state_k[:, state_start:state_end] + ], dim=1) + v_context = torch.cat([ + clean_image_v[:, :clean_end], + noisy_image_v[:, noisy_start:noisy_end], + noisy_action_v[:, action_start:action_end], + noisy_state_v[:, state_start:state_end] + ], dim=1) + + output[:, noisy_start:noisy_end] = self.attn(q_block, k_context, v_context) + + return output + + def _process_noisy_action_blocks(self, noisy_action_q, noisy_action_k, noisy_action_v, + clean_image_k, clean_image_v, + noisy_image_k, noisy_image_v, + noisy_state_k, noisy_state_v, + half_frames, action_horizon, state_horizon): + """Process noisy action blocks with teacher forcing pattern - OPTIMIZED + + First action (for first frame): cannot attend to anything (self-attention only) + Action block i: attends to first_clean_frame + clean_blocks[0:i] + noisy_image[i] + action[i] + state[i] + + OPTIMIZATION: Pre-allocate output, pre-compute indices, reduce memory allocations + """ + num_blocks = (half_frames - 1) // self.num_frame_per_block + + if num_blocks == 0: + return torch.empty_like(noisy_action_q) + + # Pre-allocate output tensor + output = torch.empty_like(noisy_action_q) + + # Pre-compute all block indices + action_block_starts = [i * self.num_action_per_block for i in range(num_blocks)] + action_block_ends = [start + self.num_action_per_block for start in action_block_starts] + clean_context_ends = [self.frame_seqlen + i * self.frame_seqlen * self.num_frame_per_block for i in range(num_blocks)] + noisy_image_block_starts = [self.frame_seqlen + i * self.frame_seqlen * self.num_frame_per_block for i in range(num_blocks)] + noisy_image_block_ends = [start + self.frame_seqlen * self.num_frame_per_block for start in noisy_image_block_starts] + state_block_starts = [i * self.num_state_per_block for i in range(num_blocks)] + state_block_ends = [start + self.num_state_per_block for start in state_block_starts] + + # Process noisy action blocks + for block_idx in range(num_blocks): + action_start = action_block_starts[block_idx] + action_end = action_block_ends[block_idx] + clean_end = clean_context_ends[block_idx] + noisy_img_start = noisy_image_block_starts[block_idx] + noisy_img_end = noisy_image_block_ends[block_idx] + state_start = state_block_starts[block_idx] + state_end = state_block_ends[block_idx] + + q_block = noisy_action_q[:, action_start:action_end] + + # Build context: first_clean_frame + clean_blocks[0:i] + noisy_image[i] + action[i] + state[i] + k_context = torch.cat([ + clean_image_k[:, :clean_end], + noisy_image_k[:, noisy_img_start:noisy_img_end], + noisy_action_k[:, action_start:action_end], + noisy_state_k[:, state_start:state_end] + ], dim=1) + v_context = torch.cat([ + clean_image_v[:, :clean_end], + noisy_image_v[:, noisy_img_start:noisy_img_end], + noisy_action_v[:, action_start:action_end], + noisy_state_v[:, state_start:state_end] + ], dim=1) + + output[:, action_start:action_end] = self.attn(q_block, k_context, v_context) + + return output + + def forward( + self, + x: torch.Tensor, + freqs: torch.Tensor, + freqs_action: torch.Tensor, + freqs_state: torch.Tensor, + action_register_length: int | None, + kv_cache: torch.Tensor | None = None, + current_start_frame: int = 0, + is_tf: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + r""" + Args: + x(Tensor): Shape [B, L, num_heads, C / num_heads] + seq_lens(Tensor): Shape [B] + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + + # query, key, value function + def qkv_fn(x): + q = self.norm_q(self.q(x)).view(b, s, n, d) + k = self.norm_k(self.k(x)).view(b, s, n, d) + v = self.v(x).view(b, s, n, d) + return q, k, v + + q, k, v = qkv_fn(x) + + updated_kv_cache: torch.Tensor | None = None + + if kv_cache is None: + if is_tf: + # Teacher forcing training. + if action_register_length is not None: + q_context = q[:, :(s-action_register_length)//2] + k_context = k[:, :(s-action_register_length)//2] + q_noisy = q[:, (s-action_register_length)//2:] + k_noisy = k[:, (s-action_register_length)//2:] + else: + q_context = q[:, :s//2] + k_context = k[:, :s//2] + q_noisy = q[:, s//2:] + k_noisy = k[:, s//2:] + roped_query = [] + roped_key = [] + + # rope should be same for clean and noisy parts + rq_context = rope_action_apply( + x=q_context, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=None, + ).type_as(v) + rk_context = rope_action_apply( + x=k_context, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=None, + ).type_as(v) + + rq_noisy = rope_action_apply( + x=q_noisy, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=action_register_length, + num_action_per_block=self.num_action_per_block, + num_state_per_block=self.num_state_per_block, + ).type_as(v) + rk_noisy = rope_action_apply( + x=k_noisy, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=action_register_length, + num_action_per_block=self.num_action_per_block, + num_state_per_block=self.num_state_per_block, + ).type_as(v) + + roped_query.append(rq_context) + roped_key.append(rk_context) + roped_query.append(rq_noisy) + roped_key.append(rk_noisy) + + roped_query = torch.cat(roped_query, dim=1) + roped_key = torch.cat(roped_key, dim=1) + # Calculate sequence dimensions + half_seq_len = (s - (action_register_length if action_register_length is not None else 0)) // 2 + + if action_register_length is not None: + # Teacher forcing structure: + # Clean half: [image tokens only] + # Noisy half: [image tokens][action tokens][state tokens] + # Causality only applies to image blocks! + + # Clean half contains ONLY image tokens + clean_image_seq_len = half_seq_len + clean_frames = clean_image_seq_len // self.frame_seqlen + + # Noisy half contains image + action + state tokens + noisy_image_seq_len = half_seq_len + noisy_frames = noisy_image_seq_len // self.frame_seqlen + num_image_blocks = (noisy_frames - 1) // self.num_frame_per_block + action_horizon = num_image_blocks * self.num_action_per_block + state_horizon = num_image_blocks * self.num_state_per_block + + # Block layout must match actual register length. For 5B use 320x176 so latent frame_seqlen=55. + if roped_query.shape[1] != half_seq_len + noisy_image_seq_len + action_horizon + state_horizon: + raise ValueError( + "Sequence length does not match block layout. " + "For 5B use 320x176 (e.g. data=dreamzero/droid_relative_wan22 or image_resolution_width=320, image_resolution_height=176). " + f"Got noisy_frames={noisy_frames}, num_image_blocks={num_image_blocks}, " + f"action_register_length={action_register_length}. " + "Ensure (noisy_frames - 1) // num_frame_per_block >= 1 and register length equals " + "num_blocks * (num_action_per_block + num_state_per_block)." + ) + + # Split clean and noisy parts + # Clean: [image tokens only] + clean_image_q = roped_query[:, :clean_image_seq_len] + clean_image_k = roped_key[:, :clean_image_seq_len] + clean_image_v = v[:, :clean_image_seq_len] + + # Noisy: [image tokens][action tokens][state tokens] + noisy_image_q = roped_query[:, half_seq_len:half_seq_len + noisy_image_seq_len] + noisy_action_q = roped_query[:, half_seq_len + noisy_image_seq_len:half_seq_len + noisy_image_seq_len + action_horizon] + noisy_state_q = roped_query[:, half_seq_len + noisy_image_seq_len + action_horizon:] + + noisy_image_k = roped_key[:, half_seq_len:half_seq_len + noisy_image_seq_len] + noisy_action_k = roped_key[:, half_seq_len + noisy_image_seq_len:half_seq_len + noisy_image_seq_len + action_horizon] + noisy_state_k = roped_key[:, half_seq_len + noisy_image_seq_len + action_horizon:] + + noisy_image_v = v[:, half_seq_len:half_seq_len + noisy_image_seq_len] + noisy_action_v = v[:, half_seq_len + noisy_image_seq_len:half_seq_len + noisy_image_seq_len + action_horizon] + noisy_state_v = v[:, half_seq_len + noisy_image_seq_len + action_horizon:] + + # ========== Process CLEAN (context) image tokens ========== + # Clean images: simple blockwise causal attention (no action/state) + clean_image_outputs = self._process_clean_image_only( + clean_image_q, clean_image_k, clean_image_v, clean_frames) + + # ========== Process NOISY tokens ========== + # Noisy image blocks: attend to previous clean image blocks + current noisy image + current noisy action + current noisy state + noisy_image_outputs = self._process_noisy_image_blocks( + noisy_image_q, noisy_image_k, noisy_image_v, + clean_image_k, clean_image_v, + noisy_action_k, noisy_action_v, noisy_state_k, noisy_state_v, + noisy_frames, action_horizon, state_horizon) + + # Noisy action blocks: attend to previous clean image blocks (including first) + current noisy image + current noisy action + same state + noisy_action_outputs = self._process_noisy_action_blocks( + noisy_action_q, noisy_action_k, noisy_action_v, + clean_image_k, clean_image_v, + noisy_image_k, noisy_image_v, + noisy_state_k, noisy_state_v, + noisy_frames, action_horizon, state_horizon) + + # Noisy state blocks: self-attention only + noisy_state_outputs = self._process_state_blocks( + noisy_state_q, noisy_state_k, noisy_state_v, state_horizon) + + # Concatenate all outputs in order: clean_img, noisy_img, noisy_act, noisy_state + x = torch.cat([ + clean_image_outputs, + noisy_image_outputs, noisy_action_outputs, noisy_state_outputs + ], dim=1) + else: + # No action/state tokens, fall back to simple image-only teacher forcing + half_frames = half_seq_len // self.frame_seqlen + clean_q = roped_query[:, :half_seq_len] + clean_k = roped_key[:, :half_seq_len] + clean_v = v[:, :half_seq_len] + noisy_q = roped_query[:, half_seq_len:] + noisy_k = roped_key[:, half_seq_len:] + noisy_v = v[:, half_seq_len:] + + # Process clean frames with blockwise causal attention + x_clean = self._blockwise_causal_flash_attn( + clean_q, clean_k, clean_v, self.frame_seqlen, self.num_frame_per_block, + action_horizon=None, state_horizon=None, + num_action_per_block=None, num_state_per_block=None, + visualize_mask=False) + + # Process noisy frames: attend to all clean frames + themselves + full_k = torch.cat([clean_k, noisy_k], dim=1) + full_v = torch.cat([clean_v, noisy_v], dim=1) + x_noisy = self.attn(noisy_q, full_k, full_v) + + x = torch.cat([x_clean, x_noisy], dim=1) + + else: + roped_query = rope_action_apply( + x=q, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=action_register_length, + num_action_per_block=self.num_action_per_block, + num_state_per_block=self.num_state_per_block, + ).type_as(v) + roped_key = rope_action_apply( + x=k, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=action_register_length, + num_action_per_block=self.num_action_per_block, + num_state_per_block=self.num_state_per_block, + ).type_as(v) + + # Calculate dynamic action and state horizons + if action_register_length is not None: + chunk_size = action_register_length // (self.num_action_per_block + self.num_state_per_block) + action_horizon = chunk_size * self.num_action_per_block + state_horizon = chunk_size * self.num_state_per_block + else: + action_horizon = None + state_horizon = None + + # Use blockwise causal flash attention without massive padding + visualize = False + x = self._blockwise_causal_flash_attn( + roped_query, roped_key, v, self.frame_seqlen, self.num_frame_per_block, + action_horizon=action_horizon, + state_horizon=state_horizon, + num_action_per_block=self.num_action_per_block if action_register_length else None, + num_state_per_block=self.num_state_per_block if action_register_length else None, + visualize_mask=visualize) + + else: + action_state_index = (current_start_frame - 1) // self.num_frame_per_block + + roped_query = causal_rope_action_apply( + x=q, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=action_register_length, + num_action_per_block=self.num_action_per_block, + num_state_per_block=self.num_state_per_block, + action_state_index=action_state_index, + ).type_as(v) + roped_key = causal_rope_action_apply( + x=k, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=action_register_length, + num_action_per_block=self.num_action_per_block, + num_state_per_block=self.num_state_per_block, + action_state_index=action_state_index, + ).type_as(v) + + # split roped_query and roped_action_query (the last action_register_length tokens) + roped_action_query: torch.Tensor | None = None + roped_action_key: torch.Tensor | None = None + action_v: torch.Tensor | None = None + + if action_register_length is not None: + roped_action_query = roped_query[:, -action_register_length:] + roped_query = roped_query[:, :-action_register_length] + roped_action_key = roped_key[:, -action_register_length:] + roped_key = roped_key[:, :-action_register_length] + action_v = v[:, -action_register_length:] + v = v[:, :-action_register_length] + assert roped_action_query is not None + assert roped_action_key is not None + assert action_v is not None + + num_new_tokens = roped_query.shape[1] + assert roped_key.shape[1] == num_new_tokens + assert v.shape[1] == num_new_tokens + + # If we are using local attention and the current KV cache size is larger + # than the local attention size, we need to truncate the KV cache + + updated_kv_cache = kv_cache + updated_k = updated_kv_cache[0] + updated_v = updated_kv_cache[1] + # Assign new keys/values directly up to current_end + new_k = torch.cat([updated_k, roped_key], dim=1) + new_v = torch.cat([updated_v, v], dim=1) + + # We may need to truncate the KV cache if it's size is larger than the max attention size. + new_k = new_k[:, -self.max_attention_size:] + new_v = new_v[:, -self.max_attention_size:] + + if action_register_length is not None: + x = self.attn( + torch.cat([roped_query, roped_action_query], dim=1), + torch.cat([new_k, roped_action_key], dim=1), + torch.cat([new_v, action_v], dim=1), + ) + else: + x = self.attn( + roped_query, + new_k, + new_v, + ) + updated_kv_cache = torch.stack([new_k, new_v], dim=0) + + + # output + x = x.flatten(2) + x = self.o(x) + return x, updated_kv_cache + + +class CausalWanAttentionBlock(nn.Module): + + def __init__(self, + cross_attn_type, + dim, + ffn_dim, + num_heads, + frame_seqlen, + local_attn_size=-1, + sink_size=0, + num_frame_per_block=1, + qk_norm=True, + cross_attn_norm=False, + eps=1e-6, + num_action_per_block=32, + num_state_per_block=1): + super().__init__() + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.local_attn_size = local_attn_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + # layers + self.norm1 = WanLayerNorm(dim, eps) + self.self_attn = CausalWanSelfAttention( + dim=dim, + num_heads=num_heads, + frame_seqlen=frame_seqlen, + local_attn_size=local_attn_size, + sink_size=sink_size, + num_frame_per_block=num_frame_per_block, + qk_norm=qk_norm, + eps=eps, + num_action_per_block=num_action_per_block, + num_state_per_block=num_state_per_block, + ) + self.norm3 = WanLayerNorm( + dim, eps, + elementwise_affine=True) if cross_attn_norm else nn.Identity() + self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim, + num_heads, + (-1, -1), + qk_norm, + eps) + self.norm2 = WanLayerNorm(dim, eps) + self.ffn = nn.Sequential( + nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'), + nn.Linear(ffn_dim, dim)) + + # modulation + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + x: torch.Tensor, + e: torch.Tensor, + freqs: torch.Tensor, + freqs_action: torch.Tensor, + freqs_state: torch.Tensor, + action_register_length: int | None, + context: torch.Tensor, + kv_cache: torch.Tensor | None = None, + crossattn_cache: torch.Tensor | None = None, + current_start_frame: int = 0, + is_tf: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + r""" + Args: + x(Tensor): Shape [B, L, C] + e(Tensor): Shape [B, F, 6, C] + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + e = (self.modulation.unsqueeze(1) + e).chunk(6, dim=2) + + # Align modulation sequence length to x so mul/add broadcast (e.g. when F != L under compile) + L = x.shape[1] + aligned = [] + for part in e: + L_e = part.shape[1] + if L_e == L: + aligned.append(part) + elif L_e >= L: + aligned.append(part[:, :L]) + else: + repeat = (L + L_e - 1) // L_e + aligned.append(part.repeat_interleave(repeat, dim=1)[:, :L]) + e = tuple(aligned) + + # self-attention + y, updated_kv_cache = self.self_attn( + x=(self.norm1(x) * (1 + e[1].squeeze(2)) + e[0].squeeze(2)), + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=action_register_length, + kv_cache=kv_cache, + is_tf=is_tf, + current_start_frame=current_start_frame, + ) + x = x + (y * e[2].squeeze(2)) + + # cross-attention & ffn function + def cross_attn_ffn(x, context, e): + x = x + self.cross_attn(self.norm3(x), context) + y = self.ffn( + (self.norm2(x) * (1 + e[4].squeeze(2)) + e[3].squeeze(2)) + ) + x = x + (y * e[5].squeeze(2)) + return x + + x = cross_attn_ffn(x, context, e) + return x, updated_kv_cache + + +class CausalHead(nn.Module): + + def __init__(self, dim, out_dim, patch_size, eps=1e-6): + super().__init__() + self.dim = dim + self.out_dim = out_dim + self.patch_size = patch_size + self.eps = eps + + # layers + out_dim = math.prod(patch_size) * out_dim + self.norm = WanLayerNorm(dim, eps) + self.head = nn.Linear(dim, out_dim) + + # modulation + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x, e): + r""" + Args: + x(Tensor): Shape [B, L1, C] + e(Tensor): Shape [B, F, 1, C] + """ + e = (self.modulation.unsqueeze(1) + e).chunk(2, dim=2) + # Align modulation sequence length to x (e.g. when F != L1 under compile) + L = x.shape[1] + aligned = [] + for part in e: + L_e = part.shape[1] + if L_e == L: + aligned.append(part) + elif L_e >= L: + aligned.append(part[:, :L]) + else: + repeat = (L + L_e - 1) // L_e + aligned.append(part.repeat_interleave(repeat, dim=1)[:, :L]) + e = tuple(aligned) + x = (self.head(self.norm(x) * (1 + e[1].squeeze(2)) + e[0].squeeze(2))) + return x + + +class CausalWanModel(ModelMixin, ConfigMixin): + r""" + Wan diffusion backbone supporting both text-to-video and image-to-video. + """ + + ignore_for_config = [ + 'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim' + ] + _no_split_modules = ['WanAttentionBlock'] + _supports_gradient_checkpointing = True + + @register_to_config + def __init__(self, + model_type='t2v', + patch_size=(1, 2, 2), + frame_seqlen=220, + text_len=512, + in_dim=16, + dim=2048, + ffn_dim=8192, + freq_dim=256, + text_dim=4096, + out_dim=16, + num_heads=16, + num_layers=32, + max_chunk_size=-1, + sink_size=0, + qk_norm=True, + cross_attn_norm=True, + eps=1e-6, + num_frame_per_block=1, + action_dim=32, + num_registers=8, + max_state_dim=64, + max_num_embodiments=32, + hidden_size=1024, + diffusion_model_pretrained_path=None, + num_action_per_block=32, + num_state_per_block=1, + concat_first_frame_latent=True): + r""" + Initialize the diffusion model backbone. + + Args: + concat_first_frame_latent (`bool`, *optional*, defaults to True): + If True, concat [x; y] before patch_embedding (14B I2V style). If False, latent only (5B pretrained style; first-frame via CLIP). + model_type (`str`, *optional*, defaults to 't2v'): + Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video) + patch_size (`tuple`, *optional*, defaults to (1, 2, 2)): + 3D patch dimensions for video embedding (t_patch, h_patch, w_patch) + text_len (`int`, *optional*, defaults to 512): + Fixed length for text embeddings + in_dim (`int`, *optional*, defaults to 16): + Input video channels (C_in) + dim (`int`, *optional*, defaults to 2048): + Hidden dimension of the transformer + ffn_dim (`int`, *optional*, defaults to 8192): + Intermediate dimension in feed-forward network + freq_dim (`int`, *optional*, defaults to 256): + Dimension for sinusoidal time embeddings + text_dim (`int`, *optional*, defaults to 4096): + Input dimension for text embeddings + out_dim (`int`, *optional*, defaults to 16): + Output video channels (C_out) + num_heads (`int`, *optional*, defaults to 16): + Number of attention heads + num_layers (`int`, *optional*, defaults to 32): + Number of transformer blocks + local_attn_size (`int`, *optional*, defaults to -1): + Window size for temporal local attention (-1 indicates global attention) + sink_size (`int`, *optional*, defaults to 0): + Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache + qk_norm (`bool`, *optional*, defaults to True): + Enable query/key normalization + cross_attn_norm (`bool`, *optional*, defaults to False): + Enable cross-attention normalization + eps (`float`, *optional*, defaults to 1e-6): + Epsilon value for normalization layers + """ + + super().__init__() + + assert model_type in ['t2v', 'i2v', 'ti2v'] + self.model_type = model_type + + self.patch_size = patch_size + self.frame_seqlen = frame_seqlen + self.text_len = text_len + self.in_dim = in_dim + self.dim = dim + self.ffn_dim = ffn_dim + self.freq_dim = freq_dim + self.text_dim = text_dim + self.out_dim = out_dim + self.num_heads = num_heads + self.num_layers = num_layers + self.local_attn_size = max_chunk_size * num_frame_per_block + 1 if max_chunk_size != -1 else -1 + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + self.num_frame_per_block = num_frame_per_block + self.diffusion_model_pretrained_path = diffusion_model_pretrained_path + self.action_dim = action_dim + self.num_registers = num_registers + self.max_state_dim = max_state_dim + self.max_num_embodiments = max_num_embodiments + self.hidden_size = hidden_size + self.num_action_per_block = num_action_per_block + self.num_state_per_block = num_state_per_block + self.concat_first_frame_latent = concat_first_frame_latent + + max_num_embodiments = 1 + + self.state_encoder = CategorySpecificMLP( + num_categories=max_num_embodiments, + input_dim=max_state_dim, + hidden_dim=self.hidden_size, + output_dim=self.dim, + ) + self.action_encoder = MultiEmbodimentActionEncoder( + action_dim=action_dim, + hidden_size=self.dim, + num_embodiments=max_num_embodiments, + ) + self.action_decoder = CategorySpecificMLP( + num_categories=max_num_embodiments, + input_dim=dim, + hidden_dim=self.hidden_size, + output_dim=action_dim, + ) + + # embeddings + self.patch_embedding = nn.Conv3d( + in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'), + nn.Linear(dim, dim)) + + self.time_embedding = nn.Sequential( + nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + self.time_projection = nn.Sequential( + nn.SiLU(), nn.Linear(dim, dim * 6)) + + # blocks + cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn' + self.blocks = nn.ModuleList([ + CausalWanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads, frame_seqlen, + self.local_attn_size, sink_size, num_frame_per_block, qk_norm, cross_attn_norm, eps, + num_action_per_block, num_state_per_block) + for _ in range(num_layers) + ]) + + # head + self.head = CausalHead(dim, out_dim, patch_size, eps) + + # buffers (don't use register_buffer otherwise dtype will be changed in to()) + assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0 + d = dim // num_heads + + self.freqs_action = rope_params(1024*10, d) + self.freqs_state = rope_params(1024, d) + self.freqs = [ + rope_params(1024, d - 4 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + ] + if model_type in ('i2v', 'ti2v'): + self.img_emb = MLPProj(1280, dim) + + # initialize weights + self.init_weights() + + self.gradient_checkpointing = True + self.independent_first_frame = False if self.num_frame_per_block == 1 else True + + + def _set_gradient_checkpointing(self, module, value=False): + self.gradient_checkpointing = value + + @staticmethod + def _prepare_blockwise_causal_attn_mask( + device: torch.device | str, num_frames: int = 21, + frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1, action_horizon=1, state_horizon=1, num_action_per_block=30, num_state_per_block=1 + ) -> BlockMask: + """ + We will divide the token sequence into the following format: + [first image (conditioning)] [image blocks] [action blocks] [state blocks] + + Structure: + - First image: conditioning only, cannot attend to anything + - Image blocks: can attend to first image + previous image block + current action block + current state block + - Action blocks: can attend to previous image block + current image block + current state block + - State blocks: conditioning only, cannot attend to anything + + Block alignment: + - num_image_blocks = (num_frames - 1) // num_frame_per_block + - num_action_blocks = action_horizon // num_action_per_block + - num_state_blocks = state_horizon // num_state_per_block + - num_image_blocks = num_action_blocks + 1 = num_state_blocks + 1 + """ + # Calculate block structure + num_image_blocks = (num_frames - 1) // num_frame_per_block + num_action_blocks = action_horizon // num_action_per_block + num_state_blocks = state_horizon // num_state_per_block + + # Verify the relationship: num_image_blocks = num_action_blocks + 1 = num_state_blocks + 1 + assert num_image_blocks == num_action_blocks, \ + f"image_blocks mismatch: {num_image_blocks} != {num_action_blocks}" + assert num_image_blocks == num_state_blocks, \ + f"image_blocks mismatch: {num_image_blocks} != {num_state_blocks}" + + # Token ranges + first_image_len = frame_seqlen # First image (conditioning) + image_blocks_len = num_image_blocks * num_frame_per_block * frame_seqlen + action_len = action_horizon + state_len = state_horizon + total_length = first_image_len + image_blocks_len + action_len + state_len + + # print("total_length", total_length, first_image_len, image_blocks_len, action_len, state_len) + # Padding to multiple of 128 + # padded_length = math.ceil(total_length / 128) * 128 - total_length + padded_length = math.ceil((local_attn_size * frame_seqlen + (local_attn_size - 1) + 32 * (local_attn_size - 1))/128) * 128 - total_length + total_padded_length = total_length + padded_length + # print("total_padded_length", total_padded_length, total_length, padded_length) + + # Define token ranges for each modality + first_image_start = 0 + first_image_end = first_image_len + image_blocks_start = first_image_end + image_blocks_end = image_blocks_start + image_blocks_len + action_start = image_blocks_end + action_end = action_start + action_len + state_start = action_end + state_end = state_start + state_len + + # Precompute block indices for each token + block_indices = torch.zeros(total_padded_length, device=device, dtype=torch.long) + + # First image gets special block index -1 (conditioning, cannot attend to anything) + block_indices[first_image_start:first_image_end] = -1 + + # Assign block indices for image blocks (0 to num_image_blocks-1) + for block_idx in range(num_image_blocks): + start_idx = image_blocks_start + block_idx * num_frame_per_block * frame_seqlen + end_idx = image_blocks_start + (block_idx + 1) * num_frame_per_block * frame_seqlen + block_indices[start_idx:end_idx] = block_idx + + # Assign block indices for action tokens (0 to num_action_blocks-1) + for block_idx in range(num_action_blocks): + start_idx = action_start + block_idx * num_action_per_block + end_idx = action_start + (block_idx + 1) * num_action_per_block + block_indices[start_idx:end_idx] = block_idx + + # Assign block indices for state tokens (0 to num_state_blocks-1) + for block_idx in range(num_state_blocks): + start_idx = state_start + block_idx * num_state_per_block + end_idx = state_start + (block_idx + 1) * num_state_per_block + block_indices[start_idx:end_idx] = block_idx + + # Padding tokens get block index of last block + 1 (won't attend to anything) + block_indices[total_length:] = num_image_blocks + + def attention_mask(b, h, q_idx, kv_idx): + # Self-attention + self_attn = (q_idx == kv_idx) + + # Determine which modality q and kv belong to + q_is_first_image = (q_idx >= first_image_start) & (q_idx < first_image_end) + q_is_image_block = (q_idx >= image_blocks_start) & (q_idx < image_blocks_end) + q_is_action = (q_idx >= action_start) & (q_idx < action_end) + q_is_state = (q_idx >= state_start) & (q_idx < state_end) + + kv_is_first_image = (kv_idx >= first_image_start) & (kv_idx < first_image_end) + kv_is_image_block = (kv_idx >= image_blocks_start) & (kv_idx < image_blocks_end) + kv_is_action = (kv_idx >= action_start) & (kv_idx < action_end) + kv_is_state = (kv_idx >= state_start) & (kv_idx < state_end) + + q_block = block_indices[q_idx] + kv_block = block_indices[kv_idx] + + # First image query (conditioning) - cannot attend to anything + first_image_mask = q_is_first_image & False + + # Image block query + image_to_first = q_is_image_block & kv_is_first_image # Image block to first image: always allowed + image_to_image = q_is_image_block & kv_is_image_block & (kv_block <= q_block) # Image block to image block: can attend to current and previous image blocks + image_to_action = q_is_image_block & kv_is_action & (kv_block == q_block) # Image block to action: can attend to current action block + image_to_state = q_is_image_block & kv_is_state & (kv_block == q_block) # Image block to state: can attend to current state block + + image_block_mask = image_to_first | image_to_image | image_to_action | image_to_state + + # Action query + action_to_image = q_is_action & kv_is_image_block & (kv_block <= q_block) # Action to image block: can attend to current and all previous image blocks + action_to_action = q_is_action & kv_is_action & (kv_block == q_block) # Action to action: only same block + action_to_state = q_is_action & kv_is_state & (kv_block == q_block) # Action to state: only same block + action_to_first = q_is_action & kv_is_first_image # Action to first image: always allowed + + action_mask = action_to_image | action_to_action | action_to_state | action_to_first + + # State query (conditioning) - cannot attend to anything + state_mask = q_is_state & False + + # Combine all masks + return self_attn | first_image_mask | image_block_mask | action_mask | state_mask + + block_mask = create_block_mask( + attention_mask, B=None, H=None, + Q_LEN=total_padded_length, + KV_LEN=total_padded_length, + _compile=False, device=device + ) + + if not dist.is_initialized() or dist.get_rank() == 0: + print(f"Created blockwise causal attention mask:") + print(f" first_image_tokens={first_image_len} (conditioning)") + print(f" num_image_blocks={num_image_blocks} (blocks of {num_frame_per_block * frame_seqlen})") + print(f" num_action_blocks={num_action_blocks} (blocks of {num_action_per_block})") + print(f" num_state_blocks={num_state_blocks} (blocks of {num_state_per_block})") + print(f" total_length={total_length}, padded_length={padded_length}") + print(block_mask) + + # Debug: materialize a small slice of the mask into 0/1 strings + try: + dense_mask = create_mask( + attention_mask, + B=None, + H=None, + Q_LEN=total_padded_length, + KV_LEN=total_padded_length, + device=device, + )[0, 0] # [Q, K] + preview_q = min(979, dense_mask.shape[0]) + preview_k = min(979, dense_mask.shape[1]) + print("Block mask (preview):") + for qi in range(preview_q): + row = dense_mask[qi, :preview_k].to(torch.int8).tolist() + print(" ".join(str(int(v)) for v in row)) + except Exception as err: + print("[warn] Failed to materialize block mask preview:", err) + + return block_mask + + @staticmethod + def _prepare_teacher_forcing_mask( + device: torch.device | str, num_frames: int = 21, + frame_seqlen: int = 1560, num_frame_per_block=1 + ) -> BlockMask: + """ + we will divide the token sequence into the following format + [1 latent frame] [1 latent frame] ... [1 latent frame] + We use flexattention to construct the attention mask + """ + total_length = num_frames * frame_seqlen * 2 + + # we do right padding to get to a multiple of 128 + padded_length = math.ceil(self.local_attn_size * frame_seqlen/128) * 128 - total_length + # padded_length = math.ceil(total_length / 128) * 128 - total_length + + clean_ends = num_frames * frame_seqlen + # for clean context frames, we can construct their flex attention mask based on a [start, end] interval + context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + # for noisy frames, we need two intervals to construct the flex attention mask [context_start, context_end] [noisy_start, noisy_end] + noise_context_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + noise_context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + noise_noise_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + noise_noise_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + + # Block-wise causal mask will attend to all elements that are before the end of the current chunk + attention_block_size = frame_seqlen * num_frame_per_block + frame_indices = torch.arange( + start=0, + end=num_frames * frame_seqlen, + step=attention_block_size, + device=device, dtype=torch.long + ) + + # attention for clean context frames + for start in frame_indices: + context_ends[start:start + attention_block_size] = start + attention_block_size + + noisy_image_start_list = torch.arange( + num_frames * frame_seqlen, total_length, + step=attention_block_size, + device=device, dtype=torch.long + ) + noisy_image_end_list = noisy_image_start_list + attention_block_size + + # attention for noisy frames + for block_index, (start, end) in enumerate(zip(noisy_image_start_list, noisy_image_end_list)): + # attend to noisy tokens within the same block + noise_noise_starts[start:end] = start + noise_noise_ends[start:end] = end + # attend to context tokens in previous blocks + # noise_context_starts[start:end] = 0 + noise_context_ends[start:end] = block_index * attention_block_size + + def attention_mask(b, h, q_idx, kv_idx): + # first design the mask for clean frames + clean_mask = (q_idx < clean_ends) & (kv_idx < context_ends[q_idx]) + # then design the mask for noisy frames + # noisy frames will attend to all clean preceeding clean frames + itself + C1 = (kv_idx < noise_noise_ends[q_idx]) & (kv_idx >= noise_noise_starts[q_idx]) + C2 = (kv_idx < noise_context_ends[q_idx]) & (kv_idx >= noise_context_starts[q_idx]) + noise_mask = (q_idx >= clean_ends) & (C1 | C2) + + eye_mask = q_idx == kv_idx + return eye_mask | clean_mask | noise_mask + + block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length, + KV_LEN=total_length + padded_length, _compile=False, device=device) + + if DEBUG: + print(block_mask) + import imageio + import numpy as np + from torch.nn.attention.flex_attention import create_mask + + mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length + + padded_length, KV_LEN=total_length + padded_length, device=device) + import cv2 + mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024)) + imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask)) + + return block_mask + + @staticmethod + def _prepare_blockwise_causal_attn_mask_i2v( + device: torch.device | str, num_frames: int = 21, + frame_seqlen: int = 1560, num_frame_per_block=4, local_attn_size=-1 + ) -> BlockMask: + """ + we will divide the token sequence into the following format + [1 latent frame] [N latent frame] ... [N latent frame] + The first frame is separated out to support I2V generation + We use flexattention to construct the attention mask + """ + total_length = num_frames * frame_seqlen + + # we do right padding to get to a multiple of 128 + padded_length = math.ceil(local_attn_size * frame_seqlen/128) * 128 - total_length + # padded_length = math.ceil(total_length / 128) * 128 - total_length + + ends = torch.zeros(total_length + padded_length, + device=device, dtype=torch.long) + + # special handling for the first frame + ends[:frame_seqlen] = frame_seqlen + + # Block-wise causal mask will attend to all elements that are before the end of the current chunk + frame_indices = torch.arange( + start=frame_seqlen, + end=total_length, + step=frame_seqlen * num_frame_per_block, + device=device + ) + + for idx, tmp in enumerate(frame_indices): + ends[tmp:tmp + frame_seqlen * num_frame_per_block] = tmp + \ + frame_seqlen * num_frame_per_block + + def attention_mask(b, h, q_idx, kv_idx): + if local_attn_size == -1: + return (kv_idx < ends[q_idx]) | (q_idx == kv_idx) + else: + return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | \ + (q_idx == kv_idx) + + block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length, + KV_LEN=total_length + padded_length, _compile=False, device=device) + + if not dist.is_initialized() or dist.get_rank() == 0: + print( + f" cache a block wise causal mask with block size of {num_frame_per_block} frames") + print(block_mask) + + return block_mask + + def _forward_blocks( + self, + x: torch.Tensor, + seq_len: int, + freqs: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + clip_feature: torch.Tensor | None, + embodiment_id: torch.Tensor | None, + action: torch.Tensor | None, + timestep_action: torch.Tensor | None, + state: torch.Tensor | None, + kv_cache: list[torch.Tensor], + current_start_frame: int, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]: + r""" + Forward pass through the diffusion model blocks. + """ + x = x.flatten(start_dim=2).transpose(1, 2) + + B = x.shape[0] + F = timestep.shape[1] + + if action is not None: + embodiment_id = torch.tensor([0], device=x.device).repeat(x.shape[0]) + action_features = self.action_encoder(action, timestep_action, embodiment_id) + state_features = self.state_encoder(state, embodiment_id) + action_register = torch.cat([action_features, state_features], dim=1) + action_length = action_features.shape[1] + action_register_length = action_register.shape[1] + x = torch.cat([x, action_register], dim=1) + else: + action_features = None + state_features = None + action_length = 0 + action_register_length = None + + # time embeddings: expand to exactly seq_len so e matches x (5B: frame_seqlen=50, 1 frame -> 50 tokens) + if F <= seq_len: + repeat = (seq_len + F - 1) // F + timestep = timestep.repeat_interleave(repeat, dim=1)[:, :seq_len] + else: + indices = torch.linspace(0, F - 1, seq_len, device=timestep.device, dtype=torch.long) + timestep = timestep[:, indices] + + if action is not None: + assert timestep_action is not None + assert state_features is not None + stride = timestep_action.shape[1] // state_features.shape[1] + timestep_state = timestep_action[:, ::stride] + timestep = torch.cat([timestep, timestep_action, timestep_state], dim=1) + + e = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, timestep.flatten()).type_as(x)) + e = e.unflatten(dim=0, sizes=(B, -1)) + e0 = self.time_projection(e) + e0 = e0.unflatten(dim=2, sizes=(6, self.dim)) + + # context + context = self.text_embedding(context) + + if clip_feature is not None: + clip_embedding = self.img_emb(clip_feature) + context = torch.cat([clip_embedding, context], dim=1) + + updated_kv_caches: list[torch.Tensor] = [] + for block_index, block in enumerate(self.blocks): + x, updated_kv_cache = block( + x=x, + e=e0, + freqs=freqs, + freqs_action=self.freqs_action, + freqs_state=self.freqs_state, + context=context, + action_register_length=action_register_length, + kv_cache=kv_cache[block_index], + current_start_frame=current_start_frame, + ) + updated_kv_caches.append(updated_kv_cache) + + if action is not None: + action_noise_pred = x[:, seq_len: seq_len + action_length] + action_noise_pred = self.action_decoder(action_noise_pred, embodiment_id) + else: + action_noise_pred = None + + # Build a tensor that contains only video tokens per sample with length = max(video_lens) + x_video = x[:, :seq_len] + e_video = e[:, :seq_len] + + # Unpatchify video-only tokens + x_video = self.head(x_video, e_video.unsqueeze(2)) + + return x_video, action_noise_pred, updated_kv_caches + + + def _forward_inference_trt( + self, + x, + timestep, + context, + kv_cache_packed: torch.Tensor, + y, + clip_feature, + action, + timestep_action, + state, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + + + frame_seqlen = 880 + seq_len = 2*frame_seqlen + kv_cache_seq_len = kv_cache_packed.shape[3] + current_start_frame = kv_cache_seq_len // frame_seqlen + + kv_cache_list = [] + for block_index in range(len(self.blocks)): + kv_cache_list.append(kv_cache_packed[block_index]) + + x_video, action_noise_pred, _ = self._forward_inference( + x=x, + timestep=timestep, + context=context, + seq_len=int(seq_len), + kv_cache=kv_cache_list, + crossattn_cache=None, + y=y, + clip_feature=clip_feature, + action=action, + timestep_action=timestep_action, + state=state, + current_start_frame = current_start_frame, + ) + + return x_video, action_noise_pred + + def _forward_inference_trt_droid( + self, + x, + timestep, + context, + kv_cache_packed: torch.Tensor, + y, + clip_feature, + action, + timestep_action, + state, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + + + frame_seqlen = 880 + seq_len = 2*frame_seqlen + kv_cache_seq_len = kv_cache_packed.shape[3] + current_start_frame = kv_cache_seq_len // frame_seqlen + + kv_cache_list = [] + for block_index in range(len(self.blocks)): + kv_cache_list.append(kv_cache_packed[block_index]) + + x_video, action_noise_pred, _ = self._forward_inference( + x=x, + timestep=timestep, + context=context, + seq_len=int(seq_len), + kv_cache=kv_cache_list, + crossattn_cache=None, + y=y, + clip_feature=clip_feature, + action=action, + timestep_action=timestep_action, + state=state, + current_start_frame = current_start_frame, + ) + + return x_video, action_noise_pred + + + def _forward_inference( + self, + x, + timestep, + context, + seq_len, + kv_cache: list[torch.Tensor], + crossattn_cache: list[torch.Tensor], + current_start_frame: int, + y=None, + clip_feature=None, + action=None, + timestep_action=None, + state=None, + embodiment_id=None, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]: + r""" + Run the diffusion model with kv caching. + See Algorithm 2 of CausVid paper https://arxiv.org/abs/2412.07772 for details. + This function will be run for num_frame times. + Process the latent frames one by one (1560 tokens each) + + Args: + x (List[Tensor]): + List of input video tensors, each with shape [C_in, F, H, W] + timestep (Tensor): + Diffusion timesteps tensor of shape [B] + context (List[Tensor]): + List of text embeddings each with shape [L, C] + seq_len (`int`): + Maximum sequence length for positional encoding + action (Tensor, *optional*): + Action tensor of shape [B, H, D] + state (Tensor, *optional*): + State tensor of shape [B, H, D] + embodiment_id (Tensor, *optional*): + Embodiment ID tensor of shape [B] + y (List[Tensor], *optional*): + Conditional video inputs for image-to-video mode, same shape as x + clip_feature (Tensor, *optional*): + CLIP image features for image-to-video mode + timestep_action (Tensor, *optional*): + Action timestep tensor of shape [B] + Returns: + List[Tensor]: + List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8] + """ + if self.model_type == 'i2v': + assert clip_feature is not None and y is not None + assert context.shape[1] == self.text_len + + # Concat [x; y] only when pretrained that way (14B). 5B uses latent only, first-frame via CLIP. + if y is not None and self.concat_first_frame_latent: + x = torch.cat([x, y.to(dtype=x.dtype)], dim=1) + + # embeddings + x = self.patch_embedding(x) + grid_size = torch.tensor(x.shape[2:], dtype=torch.long) + + freqs = self._create_freqs( + grid_size=grid_size, + start_frame=current_start_frame, + ) + + x_video, action_noise_pred, updated_kv_caches = self._forward_blocks( + x=x, + seq_len=seq_len, + freqs=freqs, + timestep=timestep, + context=context, + clip_feature=clip_feature, + embodiment_id=embodiment_id, + action=action, + timestep_action=timestep_action, + state=state, + kv_cache=kv_cache, + current_start_frame=current_start_frame, + ) + + # Copy the updated KV caches back to the original KV cache. + x_video = x_video.clone() + if action_noise_pred is not None: + action_noise_pred = action_noise_pred.clone() + #for block_index, updated_kv_cache in enumerate(updated_kv_caches): + # kv_cache[block_index] = updated_kv_cache.clone() + + video_noise_pred = self.unpatchify(x_video, grid_size) + + return video_noise_pred, action_noise_pred, updated_kv_caches + + def _forward_train( + self, + x, + timestep, + timestep_action, + context, + seq_len, + clean_x=None, + aug_t=None, + y=None, + clip_feature=None, + action=None, + state=None, + embodiment_id=None, + ): + r""" + Forward pass through the diffusion model + + Args: + x (List[Tensor]): + List of input video tensors, each with shape [C_in, F, H, W] + t (Tensor): + Diffusion timesteps tensor of shape [B] + context (List[Tensor]): + List of text embeddings each with shape [L, C] + seq_len (`int`): + Maximum sequence length for positional encoding + clip_fea (Tensor, *optional*): + CLIP image features for image-to-video mode + y (List[Tensor], *optional*): + Conditional video inputs for image-to-video mode, same shape as x + + Returns: + List[Tensor]: + List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8] + """ + if self.model_type == 'i2v': + assert clip_feature is not None and y is not None + + # Concat [x; y] only when pretrained that way (14B). 5B uses latent only, first-frame via CLIP. + if y is not None and self.concat_first_frame_latent: + x = torch.cat([x, y.to(dtype=x.dtype)], dim=1) + + # embeddings + x = self.patch_embedding(x) + + grid_size = torch.tensor(x.shape[2:], dtype=torch.long) + freqs = self._create_freqs( + grid_size=grid_size, + start_frame=0, + ) + + x = x.flatten(start_dim=2).transpose(1, 2) + assert x.shape[1] == seq_len + + B = x.shape[0] + F = timestep.shape[1] + + # time embeddings + if action is not None: + embodiment_id = torch.tensor([0]).repeat(x.shape[0]).to(device=embodiment_id.device) + action_features = self.action_encoder(action, timestep_action, embodiment_id) + action_length = action_features.shape[1] + state_features = self.state_encoder(state, embodiment_id) + action_register = torch.cat([action_features, state_features], dim=1) + action_register_length = action_register.shape[1] + x = torch.cat([x, action_register], dim=1) + else: + action_features = None + action_length = None + state_features = None + action_register = None + action_register_length = None + + # time embeddings + timestep = timestep.unsqueeze(-1).expand(B, F, seq_len // F).reshape(B, -1) + timestep_original = timestep.clone() + + if action is not None: + assert timestep_action is not None + assert state_features is not None + stride = timestep_action.shape[1] // state_features.shape[1] + timestep_state = timestep_action[:, ::stride] + timestep = torch.cat([timestep, timestep_action, timestep_state], dim=1) + + e = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, timestep.flatten()).type_as(x)) + e = e.unflatten(dim=0, sizes=(B, -1)) + e0 = self.time_projection(e) + e0 = e0.unflatten(dim=2, sizes=(6, self.dim)) + + # context + assert context.shape[1] == self.text_len + context = self.text_embedding(context) + + if clip_feature is not None: + clip_embedding = self.img_emb(clip_feature) + context = torch.cat([clip_embedding, context], dim=1) + + if clean_x is not None: + if y is not None and self.concat_first_frame_latent: + clean_x = torch.cat([clean_x, y.to(dtype=clean_x.dtype)], dim=1) + clean_x = self.patch_embedding(clean_x) + clean_x = clean_x.flatten(start_dim=2).transpose(1, 2) + assert clean_x.shape[1] == seq_len + + x = torch.cat([clean_x, x], dim=1) + + if aug_t is None: + aug_t = torch.zeros_like(timestep_original) + assert aug_t is not None + + e_clean = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, aug_t.flatten()).type_as(x)) + e_clean = e_clean.unflatten(dim=0, sizes=timestep_original.shape) + e0_clean = self.time_projection(e_clean) + e0_clean = e0_clean.unflatten(dim=2, sizes=(6, self.dim)) + e0 = torch.cat([e0_clean, e0], dim=1) + + # arguments + kwargs = dict( + e=e0, + freqs=freqs, + freqs_action=self.freqs_action, + freqs_state=self.freqs_state, + action_register_length=action_register_length, + context=context, + is_tf=clean_x is not None, + ) + + def create_custom_forward(module): + def custom_forward(*inputs, **kwargs): + outputs, updated_kv_cache = module(*inputs, **kwargs) + assert updated_kv_cache is None + return outputs + return custom_forward + + for block in self.blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, **kwargs, + use_reentrant=False, + ) + else: + x = block(x, **kwargs) + + if clean_x is not None: + x = x[:, clean_x.shape[1]:] + + if action is not None: + action_noise_pred = x[:, seq_len: seq_len + action_length] + action_noise_pred = self.action_decoder(action_noise_pred, embodiment_id) + else: + action_noise_pred = None + + # Build a tensor that contains only video tokens per sample with length = max(video_lens) + x_video = x[:, :seq_len] + e_video = e[:, :seq_len] + + # Unpatchify video-only tokens + x_video = self.head(x_video, e_video.unsqueeze(2)) + video_noise_pred = self.unpatchify(x_video, grid_size) + + return video_noise_pred, action_noise_pred + + def forward( + self, + *args, + **kwargs + ): + if kwargs.get('kv_cache', None) is not None: + return self._forward_inference(*args, **kwargs) + else: + return self._forward_train(*args, **kwargs) + + def unpatchify(self, x, grid_size): + r""" + Reconstruct video tensors from patch embeddings. + + Args: + x (Tensor): + Patchified features, with shape [B, L, C_out * prod(patch_size)]. + grid_size (Tensor): + Spatial-temporal grid dimensions before patching, with shape [3] + (3 dimensions correspond to F_patches, H_patches, W_patches) + + Returns: + Tensor: + Reconstructed video tensors with shape [B, C_out, F, H / 8, W / 8] + """ + B = x.shape[0] + c = self.out_dim + grid_size = grid_size.tolist() + assert x.shape[1] == math.prod(grid_size) + x = x.view(B, *grid_size, *self.patch_size, c) + x = torch.einsum('bfhwpqrc->bcfphqwr', x) + x = x.reshape(B, c, *[i * j for i, j in zip(grid_size, self.patch_size)]) + return x + + def _create_freqs( + self, + grid_size: torch.Tensor, + start_frame: int, + ): + device = self.patch_embedding.weight.device + if any(freq.device != device for freq in self.freqs): + self.freqs = [freq.to(device) for freq in self.freqs] + if self.freqs_action.device != device: + self.freqs_action = self.freqs_action.to(device) + if self.freqs_state.device != device: + self.freqs_state = self.freqs_state.to(device) + + f, h, w = grid_size.tolist() + freqs = torch.cat( + [ + self.freqs[0][start_frame:start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1), + self.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + self.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], + dim=-1 + ).reshape(f * h * w, 1, -1) + + return freqs + + def init_weights(self): + r""" + Initialize model parameters using Xavier initialization. + """ + + # basic init + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + if m.bias is not None: + nn.init.zeros_(m.bias) + + # init embeddings + nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1)) + for m in self.text_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=.02) + for m in self.time_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=.02) + + # init output layer + nn.init.zeros_(self.head.head.weight) diff --git a/groot/vla/model/dreamzero/modules/wan_video_image_encoder.py b/groot/vla/model/dreamzero/modules/wan_video_image_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..ec0523226f3019b78d2d9d4dc6b17476e010384e --- /dev/null +++ b/groot/vla/model/dreamzero/modules/wan_video_image_encoder.py @@ -0,0 +1,908 @@ +""" +Concise re-implementation of +``https://github.com/openai/CLIP'' and +``https://github.com/mlfoundations/open_clip''. +""" +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchvision.transforms as T +from groot.vla.model.dreamzero.modules.wan_video_dit import flash_attention + + +class SelfAttention(nn.Module): + + def __init__(self, dim, num_heads, dropout=0.1, eps=1e-5): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.eps = eps + + # layers + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.dropout = nn.Dropout(dropout) + + def forward(self, x, mask): + """ + x: [B, L, C]. + """ + b, s, c, n, d = *x.size(), self.num_heads, self.head_dim + + # compute query, key, value + q = self.q(x).reshape(b, s, n, d).permute(0, 2, 1, 3) + k = self.k(x).reshape(b, s, n, d).permute(0, 2, 1, 3) + v = self.v(x).reshape(b, s, n, d).permute(0, 2, 1, 3) + + # compute attention + p = self.dropout.p if self.training else 0.0 + x = F.scaled_dot_product_attention(q, k, v, mask, p) + x = x.permute(0, 2, 1, 3).reshape(b, s, c) + + # output + x = self.o(x) + x = self.dropout(x) + return x + + +class AttentionBlock(nn.Module): + + def __init__(self, dim, num_heads, post_norm, dropout=0.1, eps=1e-5): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.post_norm = post_norm + self.eps = eps + + # layers + self.attn = SelfAttention(dim, num_heads, dropout, eps) + self.norm1 = nn.LayerNorm(dim, eps=eps) + self.ffn = nn.Sequential( + nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim), + nn.Dropout(dropout)) + self.norm2 = nn.LayerNorm(dim, eps=eps) + + def forward(self, x, mask): + if self.post_norm: + x = self.norm1(x + self.attn(x, mask)) + x = self.norm2(x + self.ffn(x)) + else: + x = x + self.attn(self.norm1(x), mask) + x = x + self.ffn(self.norm2(x)) + return x + + +class XLMRoberta(nn.Module): + """ + XLMRobertaModel with no pooler and no LM head. + """ + + def __init__(self, + vocab_size=250002, + max_seq_len=514, + type_size=1, + pad_id=1, + dim=1024, + num_heads=16, + num_layers=24, + post_norm=True, + dropout=0.1, + eps=1e-5): + super().__init__() + self.vocab_size = vocab_size + self.max_seq_len = max_seq_len + self.type_size = type_size + self.pad_id = pad_id + self.dim = dim + self.num_heads = num_heads + self.num_layers = num_layers + self.post_norm = post_norm + self.eps = eps + + # embeddings + self.token_embedding = nn.Embedding(vocab_size, dim, padding_idx=pad_id) + self.type_embedding = nn.Embedding(type_size, dim) + self.pos_embedding = nn.Embedding(max_seq_len, dim, padding_idx=pad_id) + self.dropout = nn.Dropout(dropout) + + # blocks + self.blocks = nn.ModuleList([ + AttentionBlock(dim, num_heads, post_norm, dropout, eps) + for _ in range(num_layers) + ]) + + # norm layer + self.norm = nn.LayerNorm(dim, eps=eps) + + def forward(self, ids): + """ + ids: [B, L] of torch.LongTensor. + """ + b, s = ids.shape + mask = ids.ne(self.pad_id).long() + + # embeddings + x = self.token_embedding(ids) + \ + self.type_embedding(torch.zeros_like(ids)) + \ + self.pos_embedding(self.pad_id + torch.cumsum(mask, dim=1) * mask) + if self.post_norm: + x = self.norm(x) + x = self.dropout(x) + + # blocks + mask = torch.where( + mask.view(b, 1, 1, s).gt(0), 0.0, + torch.finfo(x.dtype).min) + for block in self.blocks: + x = block(x, mask) + + # output + if not self.post_norm: + x = self.norm(x) + return x + + +def xlm_roberta_large(pretrained=False, + return_tokenizer=False, + device='cpu', + **kwargs): + """ + XLMRobertaLarge adapted from Huggingface. + """ + # params + cfg = dict( + vocab_size=250002, + max_seq_len=514, + type_size=1, + pad_id=1, + dim=1024, + num_heads=16, + num_layers=24, + post_norm=True, + dropout=0.1, + eps=1e-5) + cfg.update(**kwargs) + + # init model + if pretrained: + from sora import DOWNLOAD_TO_CACHE + + # init a meta model + with torch.device('meta'): + model = XLMRoberta(**cfg) + + # load checkpoint + model.load_state_dict( + torch.load( + DOWNLOAD_TO_CACHE('models/xlm_roberta/xlm_roberta_large.pth'), + map_location=device), + assign=True) + else: + # init a model on device + with torch.device(device): + model = XLMRoberta(**cfg) + + # init tokenizer + if return_tokenizer: + from sora.data import HuggingfaceTokenizer + tokenizer = HuggingfaceTokenizer( + name='xlm-roberta-large', + seq_len=model.text_len, + clean='whitespace') + return model, tokenizer + else: + return model + + + +def pos_interpolate(pos, seq_len): + if pos.size(1) == seq_len: + return pos + else: + src_grid = int(math.sqrt(pos.size(1))) + tar_grid = int(math.sqrt(seq_len)) + n = pos.size(1) - src_grid * src_grid + return torch.cat([ + pos[:, :n], + F.interpolate( + pos[:, n:].float().reshape(1, src_grid, src_grid, -1).permute( + 0, 3, 1, 2), + size=(tar_grid, tar_grid), + mode='bicubic', + align_corners=False).flatten(2).transpose(1, 2) + ], + dim=1) + + +class QuickGELU(nn.Module): + + def forward(self, x): + return x * torch.sigmoid(1.702 * x) + + +class LayerNorm(nn.LayerNorm): + + def forward(self, x): + return super().forward(x).type_as(x) + + +class SelfAttention(nn.Module): + + def __init__(self, + dim, + num_heads, + causal=False, + attn_dropout=0.0, + proj_dropout=0.0): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.causal = causal + self.attn_dropout = attn_dropout + self.proj_dropout = proj_dropout + + # layers + self.to_qkv = nn.Linear(dim, dim * 3) + self.proj = nn.Linear(dim, dim) + + def forward(self, x): + """ + x: [B, L, C]. + """ + # compute query, key, value + q, k, v = self.to_qkv(x).chunk(3, dim=-1) + + # compute attention + x = flash_attention(q, k, v, num_heads=self.num_heads, compatibility_mode=True) + + # output + x = self.proj(x) + x = F.dropout(x, self.proj_dropout, self.training) + return x + + +class SwiGLU(nn.Module): + + def __init__(self, dim, mid_dim): + super().__init__() + self.dim = dim + self.mid_dim = mid_dim + + # layers + self.fc1 = nn.Linear(dim, mid_dim) + self.fc2 = nn.Linear(dim, mid_dim) + self.fc3 = nn.Linear(mid_dim, dim) + + def forward(self, x): + x = F.silu(self.fc1(x)) * self.fc2(x) + x = self.fc3(x) + return x + + +class AttentionBlock(nn.Module): + + def __init__(self, + dim, + mlp_ratio, + num_heads, + post_norm=False, + causal=False, + activation='quick_gelu', + attn_dropout=0.0, + proj_dropout=0.0, + norm_eps=1e-5): + assert activation in ['quick_gelu', 'gelu', 'swi_glu'] + super().__init__() + self.dim = dim + self.mlp_ratio = mlp_ratio + self.num_heads = num_heads + self.post_norm = post_norm + self.causal = causal + self.norm_eps = norm_eps + + # layers + self.norm1 = LayerNorm(dim, eps=norm_eps) + self.attn = SelfAttention(dim, num_heads, causal, attn_dropout, + proj_dropout) + self.norm2 = LayerNorm(dim, eps=norm_eps) + if activation == 'swi_glu': + self.mlp = SwiGLU(dim, int(dim * mlp_ratio)) + else: + self.mlp = nn.Sequential( + nn.Linear(dim, int(dim * mlp_ratio)), + QuickGELU() if activation == 'quick_gelu' else nn.GELU(), + nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(proj_dropout)) + + def forward(self, x): + if self.post_norm: + x = x + self.norm1(self.attn(x)) + x = x + self.norm2(self.mlp(x)) + else: + x = x + self.attn(self.norm1(x)) + x = x + self.mlp(self.norm2(x)) + return x + + +class AttentionPool(nn.Module): + + def __init__(self, + dim, + mlp_ratio, + num_heads, + activation='gelu', + proj_dropout=0.0, + norm_eps=1e-5): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.mlp_ratio = mlp_ratio + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.proj_dropout = proj_dropout + self.norm_eps = norm_eps + + # layers + gain = 1.0 / math.sqrt(dim) + self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim)) + self.to_q = nn.Linear(dim, dim) + self.to_kv = nn.Linear(dim, dim * 2) + self.proj = nn.Linear(dim, dim) + self.norm = LayerNorm(dim, eps=norm_eps) + self.mlp = nn.Sequential( + nn.Linear(dim, int(dim * mlp_ratio)), + QuickGELU() if activation == 'quick_gelu' else nn.GELU(), + nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(proj_dropout)) + + def forward(self, x): + """ + x: [B, L, C]. + """ + b, s, c, n, d = *x.size(), self.num_heads, self.head_dim + + # compute query, key, value + q = self.to_q(self.cls_embedding).view(1, 1, n*d).expand(b, -1, -1) + k, v = self.to_kv(x).chunk(2, dim=-1) + + # compute attention + x = flash_attention(q, k, v, num_heads=self.num_heads, compatibility_mode=True) + x = x.reshape(b, 1, c) + + # output + x = self.proj(x) + x = F.dropout(x, self.proj_dropout, self.training) + + # mlp + x = x + self.mlp(self.norm(x)) + return x[:, 0] + + +class VisionTransformer(nn.Module): + + def __init__(self, + image_size=224, + patch_size=16, + dim=768, + mlp_ratio=4, + out_dim=512, + num_heads=12, + num_layers=12, + pool_type='token', + pre_norm=True, + post_norm=False, + activation='quick_gelu', + attn_dropout=0.0, + proj_dropout=0.0, + embedding_dropout=0.0, + norm_eps=1e-5): + if image_size % patch_size != 0: + print( + '[WARNING] image_size is not divisible by patch_size', + flush=True) + assert pool_type in ('token', 'token_fc', 'attn_pool') + out_dim = out_dim or dim + super().__init__() + self.image_size = image_size + self.patch_size = patch_size + self.num_patches = (image_size // patch_size)**2 + self.dim = dim + self.mlp_ratio = mlp_ratio + self.out_dim = out_dim + self.num_heads = num_heads + self.num_layers = num_layers + self.pool_type = pool_type + self.post_norm = post_norm + self.norm_eps = norm_eps + + # embeddings + gain = 1.0 / math.sqrt(dim) + self.patch_embedding = nn.Conv2d( + 3, + dim, + kernel_size=patch_size, + stride=patch_size, + bias=not pre_norm) + if pool_type in ('token', 'token_fc'): + self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim)) + self.pos_embedding = nn.Parameter(gain * torch.randn( + 1, self.num_patches + + (1 if pool_type in ('token', 'token_fc') else 0), dim)) + self.dropout = nn.Dropout(embedding_dropout) + + # transformer + self.pre_norm = LayerNorm(dim, eps=norm_eps) if pre_norm else None + self.transformer = nn.Sequential(*[ + AttentionBlock(dim, mlp_ratio, num_heads, post_norm, False, + activation, attn_dropout, proj_dropout, norm_eps) + for _ in range(num_layers) + ]) + self.post_norm = LayerNorm(dim, eps=norm_eps) + + # head + if pool_type == 'token': + self.head = nn.Parameter(gain * torch.randn(dim, out_dim)) + elif pool_type == 'token_fc': + self.head = nn.Linear(dim, out_dim) + elif pool_type == 'attn_pool': + self.head = AttentionPool(dim, mlp_ratio, num_heads, activation, + proj_dropout, norm_eps) + + def forward(self, x, interpolation=False, use_31_block=False): + b = x.size(0) + + # embeddings + x = self.patch_embedding(x).flatten(2).permute(0, 2, 1) + # print("x in forward: ", x[0,0,100:105], x.shape) + # print("patch_embedding: ", self.patch_embedding.module.weight[0:10, 0, 0, 3], self.patch_embedding.module.weight.shape) + if self.pool_type in ('token', 'token_fc'): + x = torch.cat([self.cls_embedding.expand(b, -1, -1).to(dtype=x.dtype, device=x.device), x], dim=1) + if interpolation: + e = pos_interpolate(self.pos_embedding, x.size(1)) + else: + e = self.pos_embedding + e = e.to(dtype=x.dtype, device=x.device) + x = self.dropout(x + e) + if self.pre_norm is not None: + x = self.pre_norm(x) + + # transformer + if use_31_block: + # print("x before transformer: ", x[0,0,100:105], x.shape) + x = self.transformer[:-1](x) + return x + else: + # print("x before transformer: ", x[0,0,100:105], x.shape) + x = self.transformer(x) + return x + + +class CLIP(nn.Module): + + def __init__(self, + embed_dim=512, + image_size=224, + patch_size=16, + vision_dim=768, + vision_mlp_ratio=4, + vision_heads=12, + vision_layers=12, + vision_pool='token', + vision_pre_norm=True, + vision_post_norm=False, + vocab_size=49408, + text_len=77, + text_dim=512, + text_mlp_ratio=4, + text_heads=8, + text_layers=12, + text_causal=True, + text_pool='argmax', + text_head_bias=False, + logit_bias=None, + activation='quick_gelu', + attn_dropout=0.0, + proj_dropout=0.0, + embedding_dropout=0.0, + norm_eps=1e-5): + super().__init__() + self.embed_dim = embed_dim + self.image_size = image_size + self.patch_size = patch_size + self.vision_dim = vision_dim + self.vision_mlp_ratio = vision_mlp_ratio + self.vision_heads = vision_heads + self.vision_layers = vision_layers + self.vision_pool = vision_pool + self.vision_pre_norm = vision_pre_norm + self.vision_post_norm = vision_post_norm + self.vocab_size = vocab_size + self.text_len = text_len + self.text_dim = text_dim + self.text_mlp_ratio = text_mlp_ratio + self.text_heads = text_heads + self.text_layers = text_layers + self.text_causal = text_causal + self.text_pool = text_pool + self.text_head_bias = text_head_bias + self.norm_eps = norm_eps + + # models + self.visual = VisionTransformer( + image_size=image_size, + patch_size=patch_size, + dim=vision_dim, + mlp_ratio=vision_mlp_ratio, + out_dim=embed_dim, + num_heads=vision_heads, + num_layers=vision_layers, + pool_type=vision_pool, + pre_norm=vision_pre_norm, + post_norm=vision_post_norm, + activation=activation, + attn_dropout=attn_dropout, + proj_dropout=proj_dropout, + embedding_dropout=embedding_dropout, + norm_eps=norm_eps) + self.textual = TextTransformer( + vocab_size=vocab_size, + text_len=text_len, + dim=text_dim, + mlp_ratio=text_mlp_ratio, + out_dim=embed_dim, + num_heads=text_heads, + num_layers=text_layers, + causal=text_causal, + pool_type=text_pool, + head_bias=text_head_bias, + activation=activation, + attn_dropout=attn_dropout, + proj_dropout=proj_dropout, + embedding_dropout=embedding_dropout, + norm_eps=norm_eps) + self.log_scale = nn.Parameter(math.log(1 / 0.07) * torch.ones([])) + if logit_bias is not None: + self.logit_bias = nn.Parameter(logit_bias * torch.ones([])) + + # initialize weights + self.init_weights() + + def forward(self, imgs, txt_ids): + """ + imgs: [B, 3, H, W] of torch.float32. + - mean: [0.48145466, 0.4578275, 0.40821073] + - std: [0.26862954, 0.26130258, 0.27577711] + txt_ids: [B, L] of torch.long. Encoded by data.CLIPTokenizer. + """ + xi = self.visual(imgs) + xt = self.textual(txt_ids) + return xi, xt + + def init_weights(self): + # embeddings + nn.init.normal_(self.textual.token_embedding.weight, std=0.02) + nn.init.normal_(self.visual.patch_embedding.weight, std=0.1) + + # attentions + for modality in ['visual', 'textual']: + dim = self.vision_dim if modality == 'visual' else self.text_dim + transformer = getattr(self, modality).transformer + proj_gain = (1.0 / math.sqrt(dim)) * ( + 1.0 / math.sqrt(2 * len(transformer))) + attn_gain = 1.0 / math.sqrt(dim) + mlp_gain = 1.0 / math.sqrt(2.0 * dim) + for block in transformer: + nn.init.normal_(block.attn.to_qkv.weight, std=attn_gain) + nn.init.normal_(block.attn.proj.weight, std=proj_gain) + nn.init.normal_(block.mlp[0].weight, std=mlp_gain) + nn.init.normal_(block.mlp[2].weight, std=proj_gain) + + def param_groups(self): + groups = [{ + 'params': [ + p for n, p in self.named_parameters() + if 'norm' in n or n.endswith('bias') + ], + 'weight_decay': 0.0 + }, { + 'params': [ + p for n, p in self.named_parameters() + if not ('norm' in n or n.endswith('bias')) + ] + }] + return groups + + +class XLMRobertaWithHead(XLMRoberta): + + def __init__(self, **kwargs): + self.out_dim = kwargs.pop('out_dim') + super().__init__(**kwargs) + + # head + mid_dim = (self.dim + self.out_dim) // 2 + self.head = nn.Sequential( + nn.Linear(self.dim, mid_dim, bias=False), nn.GELU(), + nn.Linear(mid_dim, self.out_dim, bias=False)) + + def forward(self, ids): + # xlm-roberta + x = super().forward(ids) + + # average pooling + mask = ids.ne(self.pad_id).unsqueeze(-1).to(x) + x = (x * mask).sum(dim=1) / mask.sum(dim=1) + + # head + x = self.head(x) + return x + + +class XLMRobertaCLIP(nn.Module): + + def __init__(self, + embed_dim=1024, + image_size=224, + patch_size=14, + vision_dim=1280, + vision_mlp_ratio=4, + vision_heads=16, + vision_layers=32, + vision_pool='token', + vision_pre_norm=True, + vision_post_norm=False, + activation='gelu', + vocab_size=250002, + max_text_len=514, + type_size=1, + pad_id=1, + text_dim=1024, + text_heads=16, + text_layers=24, + text_post_norm=True, + text_dropout=0.1, + attn_dropout=0.0, + proj_dropout=0.0, + embedding_dropout=0.0, + norm_eps=1e-5): + super().__init__() + self.embed_dim = embed_dim + self.image_size = image_size + self.patch_size = patch_size + self.vision_dim = vision_dim + self.vision_mlp_ratio = vision_mlp_ratio + self.vision_heads = vision_heads + self.vision_layers = vision_layers + self.vision_pre_norm = vision_pre_norm + self.vision_post_norm = vision_post_norm + self.activation = activation + self.vocab_size = vocab_size + self.max_text_len = max_text_len + self.type_size = type_size + self.pad_id = pad_id + self.text_dim = text_dim + self.text_heads = text_heads + self.text_layers = text_layers + self.text_post_norm = text_post_norm + self.norm_eps = norm_eps + + # models + self.visual = VisionTransformer( + image_size=image_size, + patch_size=patch_size, + dim=vision_dim, + mlp_ratio=vision_mlp_ratio, + out_dim=embed_dim, + num_heads=vision_heads, + num_layers=vision_layers, + pool_type=vision_pool, + pre_norm=vision_pre_norm, + post_norm=vision_post_norm, + activation=activation, + attn_dropout=attn_dropout, + proj_dropout=proj_dropout, + embedding_dropout=embedding_dropout, + norm_eps=norm_eps) + self.textual = None + self.log_scale = nn.Parameter(math.log(1 / 0.07) * torch.ones([])) + + def forward(self, imgs, txt_ids): + """ + imgs: [B, 3, H, W] of torch.float32. + - mean: [0.48145466, 0.4578275, 0.40821073] + - std: [0.26862954, 0.26130258, 0.27577711] + txt_ids: [B, L] of torch.long. + Encoded by data.CLIPTokenizer. + """ + xi = self.visual(imgs) + xt = self.textual(txt_ids) + return xi, xt + + def param_groups(self): + groups = [{ + 'params': [ + p for n, p in self.named_parameters() + if 'norm' in n or n.endswith('bias') + ], + 'weight_decay': 0.0 + }, { + 'params': [ + p for n, p in self.named_parameters() + if not ('norm' in n or n.endswith('bias')) + ] + }] + return groups + + +def _clip(pretrained=False, + pretrained_name=None, + model_cls=CLIP, + return_transforms=False, + return_tokenizer=False, + tokenizer_padding='eos', + dtype=torch.float32, + device='cpu', + **kwargs): + # init model + if pretrained and pretrained_name: + from sora import BUCKET, DOWNLOAD_TO_CACHE + + # init a meta model + with torch.device('meta'): + model = model_cls(**kwargs) + + # checkpoint path + checkpoint = f'models/clip/{pretrained_name}' + if dtype in (torch.float16, torch.bfloat16): + suffix = '-' + { + torch.float16: 'fp16', + torch.bfloat16: 'bf16' + }[dtype] + if object_exists(BUCKET, f'{checkpoint}{suffix}.pth'): + checkpoint = f'{checkpoint}{suffix}' + checkpoint += '.pth' + + # load + model.load_state_dict( + torch.load(DOWNLOAD_TO_CACHE(checkpoint), map_location=device), + assign=True, + strict=False) + else: + # init a model on device + with torch.device(device): + model = model_cls(**kwargs) + + # set device + output = (model,) + + # init transforms + if return_transforms: + # mean and std + if 'siglip' in pretrained_name.lower(): + mean, std = [0.5, 0.5, 0.5], [0.5, 0.5, 0.5] + else: + mean = [0.48145466, 0.4578275, 0.40821073] + std = [0.26862954, 0.26130258, 0.27577711] + + # transforms + transforms = T.Compose([ + T.Resize((model.image_size, model.image_size), + interpolation=T.InterpolationMode.BICUBIC), + T.ToTensor(), + T.Normalize(mean=mean, std=std) + ]) + output += (transforms,) + + # init tokenizer + if return_tokenizer: + from sora import data + if 'siglip' in pretrained_name.lower(): + tokenizer = data.HuggingfaceTokenizer( + name=f'timm/{pretrained_name}', + seq_len=model.text_len, + clean='canonicalize') + elif 'xlm' in pretrained_name.lower(): + tokenizer = data.HuggingfaceTokenizer( + name='xlm-roberta-large', + seq_len=model.max_text_len - 2, + clean='whitespace') + elif 'mba' in pretrained_name.lower(): + tokenizer = data.HuggingfaceTokenizer( + name='facebook/xlm-roberta-xl', + seq_len=model.max_text_len - 2, + clean='whitespace') + else: + tokenizer = data.CLIPTokenizer( + seq_len=model.text_len, padding=tokenizer_padding) + output += (tokenizer,) + return output[0] if len(output) == 1 else output + + +def clip_xlm_roberta_vit_h_14( + pretrained=False, + pretrained_name='open-clip-xlm-roberta-large-vit-huge-14', + **kwargs): + cfg = dict( + embed_dim=1024, + image_size=224, + patch_size=14, + vision_dim=1280, + vision_mlp_ratio=4, + vision_heads=16, + vision_layers=32, + vision_pool='token', + activation='gelu', + vocab_size=250002, + max_text_len=514, + type_size=1, + pad_id=1, + text_dim=1024, + text_heads=16, + text_layers=24, + text_post_norm=True, + text_dropout=0.1, + attn_dropout=0.0, + proj_dropout=0.0, + embedding_dropout=0.0) + cfg.update(**kwargs) + return _clip(pretrained, pretrained_name, XLMRobertaCLIP, **cfg) + + +class WanImageEncoder(torch.nn.Module): + + def __init__(self, image_encoder_pretrained_path: str=None): + super().__init__() + # init model + self.model, self.transforms = clip_xlm_roberta_vit_h_14( + pretrained=False, + return_transforms=True, + return_tokenizer=False, + dtype=torch.float32, + device="cpu") + self.image_encoder_pretrained_path = image_encoder_pretrained_path + + def encode_image(self, videos): + # preprocess + size = (self.model.image_size,) * 2 + videos = torch.cat([ + F.interpolate( + u, + size=size, + mode='bicubic', + align_corners=False) for u in videos + ]) + videos = self.transforms.transforms[-1](videos.mul_(0.5).add_(0.5)) + + # forward + dtype = next(iter(self.model.visual.parameters())).dtype + videos = videos.to(dtype) + out = self.model.visual(videos, use_31_block=True) + # The outputs of torch compile always need to be cloned before being used. + out = out.clone() + return out + + @staticmethod + def state_dict_converter(): + return WanImageEncoderStateDictConverter() + + +class WanImageEncoderStateDictConverter: + def __init__(self): + pass + + def from_diffusers(self, state_dict): + return state_dict + + def from_civitai(self, state_dict): + state_dict_ = {} + for name, param in state_dict.items(): + if name.startswith("textual."): + continue + name = "model." + name + state_dict_[name] = param + return state_dict_ diff --git a/groot/vla/model/dreamzero/modules/wan_video_text_encoder.py b/groot/vla/model/dreamzero/modules/wan_video_text_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..abeb018333f6f61f1af12aa39ab064105274f8d1 --- /dev/null +++ b/groot/vla/model/dreamzero/modules/wan_video_text_encoder.py @@ -0,0 +1,280 @@ +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def fp16_clamp(x): + if x.dtype == torch.float16 and torch.isinf(x).any(): + clamp = torch.finfo(x.dtype).max - 1000 + x = torch.clamp(x, min=-clamp, max=clamp) + return x + + +class GELU(nn.Module): + + def forward(self, x): + return 0.5 * x * (1.0 + torch.tanh( + math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0)))) + + +class T5LayerNorm(nn.Module): + + def __init__(self, dim, eps=1e-6): + super(T5LayerNorm, self).__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + x = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + + self.eps) + if self.weight.dtype in [torch.float16, torch.bfloat16]: + x = x.type_as(self.weight) + return self.weight * x + + +class T5Attention(nn.Module): + + def __init__(self, dim, dim_attn, num_heads, dropout=0.1): + assert dim_attn % num_heads == 0 + super(T5Attention, self).__init__() + self.dim = dim + self.dim_attn = dim_attn + self.num_heads = num_heads + self.head_dim = dim_attn // num_heads + + # layers + self.q = nn.Linear(dim, dim_attn, bias=False) + self.k = nn.Linear(dim, dim_attn, bias=False) + self.v = nn.Linear(dim, dim_attn, bias=False) + self.o = nn.Linear(dim_attn, dim, bias=False) + self.dropout = nn.Dropout(dropout) + + def forward(self, x, context=None, mask=None, pos_bias=None): + """ + x: [B, L1, C]. + context: [B, L2, C] or None. + mask: [B, L2] or [B, L1, L2] or None. + """ + # check inputs + context = x if context is None else context + b, n, c = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.q(x).view(b, -1, n, c) + k = self.k(context).view(b, -1, n, c) + v = self.v(context).view(b, -1, n, c) + + # attention bias + attn_bias = x.new_zeros(b, n, q.size(1), k.size(1)) + if pos_bias is not None: + attn_bias += pos_bias + if mask is not None: + assert mask.ndim in [2, 3] + mask = mask.view(b, 1, 1, + -1) if mask.ndim == 2 else mask.unsqueeze(1) + attn_bias.masked_fill_(mask == 0, torch.finfo(x.dtype).min) + + # compute attention (T5 does not use scaling) + attn = torch.einsum('binc,bjnc->bnij', q, k) + attn_bias + attn = F.softmax(attn.float(), dim=-1).type_as(attn) + x = torch.einsum('bnij,bjnc->binc', attn, v) + + # output + x = x.reshape(b, -1, n * c) + x = self.o(x) + x = self.dropout(x) + return x + + +class T5FeedForward(nn.Module): + + def __init__(self, dim, dim_ffn, dropout=0.1): + super(T5FeedForward, self).__init__() + self.dim = dim + self.dim_ffn = dim_ffn + + # layers + self.gate = nn.Sequential(nn.Linear(dim, dim_ffn, bias=False), GELU()) + self.fc1 = nn.Linear(dim, dim_ffn, bias=False) + self.fc2 = nn.Linear(dim_ffn, dim, bias=False) + self.dropout = nn.Dropout(dropout) + + def forward(self, x): + x = self.fc1(x) * self.gate(x) + x = self.dropout(x) + x = self.fc2(x) + x = self.dropout(x) + return x + + +class T5SelfAttention(nn.Module): + + def __init__(self, + dim, + dim_attn, + dim_ffn, + num_heads, + num_buckets, + shared_pos=True, + dropout=0.1): + super(T5SelfAttention, self).__init__() + self.dim = dim + self.dim_attn = dim_attn + self.dim_ffn = dim_ffn + self.num_heads = num_heads + self.num_buckets = num_buckets + self.shared_pos = shared_pos + + # layers + self.norm1 = T5LayerNorm(dim) + self.attn = T5Attention(dim, dim_attn, num_heads, dropout) + self.norm2 = T5LayerNorm(dim) + self.ffn = T5FeedForward(dim, dim_ffn, dropout) + self.pos_embedding = None if shared_pos else T5RelativeEmbedding( + num_buckets, num_heads, bidirectional=True) + + def forward(self, x, mask=None, pos_bias=None): + e = pos_bias if self.shared_pos else self.pos_embedding( + x.size(1), x.size(1)) + x = fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e)) + x = fp16_clamp(x + self.ffn(self.norm2(x))) + # print("x after attn: ", x[0, 0:10], x.shape) + return x + + +class T5RelativeEmbedding(nn.Module): + + def __init__(self, num_buckets, num_heads, bidirectional, max_dist=128): + super(T5RelativeEmbedding, self).__init__() + self.num_buckets = num_buckets + self.num_heads = num_heads + self.bidirectional = bidirectional + self.max_dist = max_dist + + # layers + self.embedding = nn.Embedding(num_buckets, num_heads) + + def forward(self, lq, lk): + device = self.embedding.weight.device + # rel_pos = torch.arange(lk).unsqueeze(0).to(device) - \ + # torch.arange(lq).unsqueeze(1).to(device) + rel_pos = torch.arange(lk, device=device).unsqueeze(0) - \ + torch.arange(lq, device=device).unsqueeze(1) + rel_pos = self._relative_position_bucket(rel_pos) + rel_pos_embeds = self.embedding(rel_pos) + rel_pos_embeds = rel_pos_embeds.permute(2, 0, 1).unsqueeze( + 0) # [1, N, Lq, Lk] + return rel_pos_embeds.contiguous() + + def _relative_position_bucket(self, rel_pos): + # preprocess + if self.bidirectional: + num_buckets = self.num_buckets // 2 + rel_buckets = (rel_pos > 0).long() * num_buckets + rel_pos = torch.abs(rel_pos) + else: + num_buckets = self.num_buckets + rel_buckets = 0 + rel_pos = -torch.min(rel_pos, torch.zeros_like(rel_pos)) + + # embeddings for small and large positions + max_exact = num_buckets // 2 + rel_pos_large = max_exact + (torch.log(rel_pos.float() / max_exact) / + math.log(self.max_dist / max_exact) * + (num_buckets - max_exact)).long() + rel_pos_large = torch.min( + rel_pos_large, torch.full_like(rel_pos_large, num_buckets - 1)) + rel_buckets += torch.where(rel_pos < max_exact, rel_pos, rel_pos_large) + return rel_buckets + +def init_weights(m): + if isinstance(m, T5LayerNorm): + nn.init.ones_(m.weight) + elif isinstance(m, T5FeedForward): + nn.init.normal_(m.gate[0].weight, std=m.dim**-0.5) + nn.init.normal_(m.fc1.weight, std=m.dim**-0.5) + nn.init.normal_(m.fc2.weight, std=m.dim_ffn**-0.5) + elif isinstance(m, T5Attention): + nn.init.normal_(m.q.weight, std=(m.dim * m.dim_attn)**-0.5) + nn.init.normal_(m.k.weight, std=m.dim**-0.5) + nn.init.normal_(m.v.weight, std=m.dim**-0.5) + nn.init.normal_(m.o.weight, std=(m.num_heads * m.dim_attn)**-0.5) + elif isinstance(m, T5RelativeEmbedding): + nn.init.normal_( + m.embedding.weight, std=(2 * m.num_buckets * m.num_heads)**-0.5) + + +class WanTextEncoder(torch.nn.Module): + + def __init__(self, + vocab: int | nn.Embedding = 256384, + dim=4096, + dim_attn=4096, + dim_ffn=10240, + num_heads=64, + num_layers=24, + num_buckets=32, + shared_pos=False, + dropout=0.1, + text_encoder_pretrained_path: str=None): + super(WanTextEncoder, self).__init__() + self.dim = dim + self.dim_attn = dim_attn + self.dim_ffn = dim_ffn + self.num_heads = num_heads + self.num_layers = num_layers + self.num_buckets = num_buckets + self.shared_pos = shared_pos + self.text_encoder_pretrained_path = text_encoder_pretrained_path + + # layers + if isinstance(vocab, int): + self.token_embedding = nn.Embedding(vocab, dim) + else: + self.token_embedding = vocab + if shared_pos: + self.pos_embedding = T5RelativeEmbedding( + num_buckets, num_heads, bidirectional=True) + else: + self.pos_embedding = None + self.dropout = nn.Dropout(dropout) + self.blocks = nn.ModuleList([ + T5SelfAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets, + shared_pos, dropout) for _ in range(num_layers) + ]) + self.norm = T5LayerNorm(dim) + + # initialize weights + self.apply(init_weights) + + def forward(self, ids, mask=None): + x = self.token_embedding(ids) + x = self.dropout(x) + if self.shared_pos: + assert self.pos_embedding is not None + e = self.pos_embedding(x.size(1), x.size(1)) + else: + e = None + for block in self.blocks: + x = block(x, mask, pos_bias=e) + x = self.norm(x) + x = self.dropout(x) + return x + + @staticmethod + def state_dict_converter(): + return WanTextEncoderStateDictConverter() + + +class WanTextEncoderStateDictConverter: + def __init__(self): + pass + + def from_diffusers(self, state_dict): + return state_dict + + def from_civitai(self, state_dict): + return state_dict \ No newline at end of file diff --git a/groot/vla/model/dreamzero/modules/wan_video_vae.py b/groot/vla/model/dreamzero/modules/wan_video_vae.py new file mode 100644 index 0000000000000000000000000000000000000000..4e2ff44f0db1e043b3ba9914962040f5ca0d36fa --- /dev/null +++ b/groot/vla/model/dreamzero/modules/wan_video_vae.py @@ -0,0 +1,1373 @@ +from einops import rearrange, repeat + +import torch +import torch.nn as nn +import torch.nn.functional as F +from tqdm import tqdm + +CACHE_T = 2 + + +def check_is_instance(model, module_class): + if isinstance(model, module_class): + return True + if hasattr(model, "module") and isinstance(model.module, module_class): + return True + return False + + +def block_causal_mask(x, block_size): + # params + b, n, s, _ = x.shape + assert s % block_size == 0 + num_blocks = s // block_size + + # build mask + mask = torch.zeros(b, n, s, s, dtype=torch.bool, device=x.device) + for i in range(num_blocks): + mask[:, :, + i * block_size:(i + 1) * block_size, :(i + 1) * block_size] = 1 + return mask + + +class CausalConv3d(nn.Conv3d): + """ + Causal 3d convolusion. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._padding = (self.padding[2], self.padding[2], self.padding[1], + self.padding[1], 2 * self.padding[0], 0) + self.padding = (0, 0, 0) + + def forward(self, x, cache_x=None): + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding) + + return super().forward(x) + + +class RMS_norm(nn.Module): + + def __init__(self, dim, channel_first=True, images=True, bias=False): + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0. + + def forward(self, x): + return F.normalize( + x, dim=(1 if self.channel_first else + -1)) * self.scale * self.gamma + self.bias + + +class Upsample(nn.Upsample): + + def forward(self, x): + """ + Fix bfloat16 support for nearest neighbor interpolation. + """ + return super().forward(x.float()).type_as(x) + + +class Resample(nn.Module): + + def __init__(self, dim, mode): + assert mode in ('none', 'upsample2d', 'upsample3d', 'downsample2d', + 'downsample3d') + super().__init__() + self.dim = dim + self.mode = mode + + # layers + if mode == 'upsample2d': + self.resample = nn.Sequential( + Upsample(scale_factor=(2., 2.), mode='nearest-exact'), + nn.Conv2d(dim, dim // 2, 3, padding=1)) + elif mode == 'upsample3d': + self.resample = nn.Sequential( + Upsample(scale_factor=(2., 2.), mode='nearest-exact'), + nn.Conv2d(dim, dim // 2, 3, padding=1)) + self.time_conv = CausalConv3d(dim, + dim * 2, (3, 1, 1), + padding=(1, 0, 0)) + + elif mode == 'downsample2d': + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + elif mode == 'downsample3d': + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.time_conv = CausalConv3d(dim, + dim, (3, 1, 1), + stride=(2, 1, 1), + padding=(0, 0, 0)) + + else: + self.resample = nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + b, c, t, h, w = x.size() + if self.mode == 'upsample3d': + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = 'Rep' + feat_idx[0] += 1 + else: + + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[ + idx] is not None and feat_cache[idx] != 'Rep': + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], dim=2) + if cache_x.shape[2] < 2 and feat_cache[ + idx] is not None and feat_cache[idx] == 'Rep': + cache_x = torch.cat([ + torch.zeros_like(cache_x), + cache_x, + ], dim=2) + if feat_cache[idx] == 'Rep': + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), + 3) + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + x = rearrange(x, 'b c t h w -> (b t) c h w') + x = self.resample(x) + x = rearrange(x, '(b t) c h w -> b c t h w', t=t) + + if self.mode == 'downsample3d': + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv( + torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return x + + def init_weight(self, conv): + conv_weight = conv.weight + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + one_matrix = torch.eye(c1, c2) + init_matrix = one_matrix + nn.init.zeros_(conv_weight) + conv_weight.data[:, :, 1, 0, 0] = init_matrix + conv.weight.data.copy_(conv_weight) + nn.init.zeros_(conv.bias.data) + + def init_weight2(self, conv): + conv_weight = conv.weight.data + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + init_matrix = torch.eye(c1 // 2, c2) + conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix + conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix + conv.weight.data.copy_(conv_weight) + nn.init.zeros_(conv.bias.data) + + + +def patchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 4: + x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size) + elif x.dim() == 5: + x = rearrange(x, + "b c f (h q) (w r) -> b (c r q) f h w", + q=patch_size, + r=patch_size) + else: + raise ValueError(f"Invalid input shape: {x.shape}") + return x + + +def unpatchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 4: + x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size) + elif x.dim() == 5: + x = rearrange(x, + "b (c r q) f h w -> b c f (h q) (w r)", + q=patch_size, + r=patch_size) + return x + + +class Resample38(Resample): + + def __init__(self, dim, mode): + assert mode in ( + "none", + "upsample2d", + "upsample3d", + "downsample2d", + "downsample3d", + ) + super(Resample, self).__init__() + self.dim = dim + self.mode = mode + + # layers + if mode == "upsample2d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim, 3, padding=1), + ) + elif mode == "upsample3d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim, 3, padding=1), + ) + self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) + elif mode == "downsample2d": + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)) + ) + elif mode == "downsample3d": + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)) + ) + self.time_conv = CausalConv3d( + dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0) + ) + else: + self.resample = nn.Identity() + +class ResidualBlock(nn.Module): + + def __init__(self, in_dim, out_dim, dropout=0.0): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + # layers + self.residual = nn.Sequential( + RMS_norm(in_dim, images=False), nn.SiLU(), + CausalConv3d(in_dim, out_dim, 3, padding=1), + RMS_norm(out_dim, images=False), nn.SiLU(), nn.Dropout(dropout), + CausalConv3d(out_dim, out_dim, 3, padding=1)) + self.shortcut = CausalConv3d(in_dim, out_dim, 1) \ + if in_dim != out_dim else nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + h = self.shortcut(x) + for layer in self.residual: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + h + + +class AttentionBlock(nn.Module): + """ + Causal self-attention with a single head. + """ + + def __init__(self, dim): + super().__init__() + self.dim = dim + + # layers + self.norm = RMS_norm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + + # zero out the last layer params + nn.init.zeros_(self.proj.weight) + + def forward(self, x): + identity = x + b, c, t, h, w = x.size() + x = rearrange(x, 'b c t h w -> (b t) c h w') + x = self.norm(x) + # compute query, key, value + q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute( + 0, 1, 3, 2).contiguous().chunk(3, dim=-1) + + # apply attention + x = F.scaled_dot_product_attention( + q, + k, + v, + #attn_mask=block_causal_mask(q, block_size=h * w) + ) + x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) + + # output + x = self.proj(x) + x = rearrange(x, '(b t) c h w-> b c t h w', t=t) + return x + identity + + +class AvgDown3D(nn.Module): + def __init__( + self, + in_channels, + out_channels, + factor_t, + factor_s=1, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + + assert in_channels * self.factor % out_channels == 0 + self.group_size = in_channels * self.factor // out_channels + + def forward(self, x: torch.Tensor) -> torch.Tensor: + pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t + pad = (0, 0, 0, 0, pad_t, 0) + x = F.pad(x, pad) + B, C, T, H, W = x.shape + x = x.view( + B, + C, + T // self.factor_t, + self.factor_t, + H // self.factor_s, + self.factor_s, + W // self.factor_s, + self.factor_s, + ) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view( + B, + C * self.factor, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.view( + B, + self.out_channels, + self.group_size, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.mean(dim=2) + return x + + +class DupUp3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + factor_t, + factor_s=1, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + + assert out_channels * self.factor % in_channels == 0 + self.repeats = out_channels * self.factor // in_channels + + def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor: + x = x.repeat_interleave(self.repeats, dim=1) + x = x.view( + x.size(0), + self.out_channels, + self.factor_t, + self.factor_s, + self.factor_s, + x.size(2), + x.size(3), + x.size(4), + ) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + x = x.view( + x.size(0), + self.out_channels, + x.size(2) * self.factor_t, + x.size(4) * self.factor_s, + x.size(6) * self.factor_s, + ) + if first_chunk: + x = x[:, :, self.factor_t - 1 :, :, :] + return x + + +class Down_ResidualBlock(nn.Module): + def __init__( + self, in_dim, out_dim, dropout, mult, temperal_downsample=False, down_flag=False + ): + super().__init__() + + # Shortcut path with downsample + self.avg_shortcut = AvgDown3D( + in_dim, + out_dim, + factor_t=2 if temperal_downsample else 1, + factor_s=2 if down_flag else 1, + ) + + # Main path with residual blocks and downsample + downsamples = [] + for _ in range(mult): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + + # Add the final downsample block + if down_flag: + mode = "downsample3d" if temperal_downsample else "downsample2d" + downsamples.append(Resample38(out_dim, mode=mode)) + + self.downsamples = nn.Sequential(*downsamples) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + x_copy = x.clone() + for module in self.downsamples: + x = module(x, feat_cache, feat_idx) + + return x + self.avg_shortcut(x_copy) + + +class Up_ResidualBlock(nn.Module): + def __init__( + self, in_dim, out_dim, dropout, mult, temperal_upsample=False, up_flag=False + ): + super().__init__() + # Shortcut path with upsample + if up_flag: + self.avg_shortcut = DupUp3D( + in_dim, + out_dim, + factor_t=2 if temperal_upsample else 1, + factor_s=2 if up_flag else 1, + ) + else: + self.avg_shortcut = None + + # Main path with residual blocks and upsample + upsamples = [] + for _ in range(mult): + upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + + # Add the final upsample block + if up_flag: + mode = "upsample3d" if temperal_upsample else "upsample2d" + upsamples.append(Resample38(out_dim, mode=mode)) + + self.upsamples = nn.Sequential(*upsamples) + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + x_main = x.clone() + for module in self.upsamples: + x_main = module(x_main, feat_cache, feat_idx) + if self.avg_shortcut is not None: + x_shortcut = self.avg_shortcut(x, first_chunk) + return x_main + x_shortcut + else: + return x_main + + +class Encoder3d(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + + # dimensions + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + # init block + self.conv1 = CausalConv3d(3, dims[0], 3, padding=1) + + # downsample blocks + downsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + for _ in range(num_res_blocks): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + downsamples.append(AttentionBlock(out_dim)) + in_dim = out_dim + + # downsample block + if i != len(dim_mult) - 1: + mode = 'downsample3d' if temperal_downsample[ + i] else 'downsample2d' + downsamples.append(Resample(out_dim, mode=mode)) + scale /= 2.0 + self.downsamples = nn.Sequential(*downsamples) + + # middle blocks + self.middle = nn.Sequential(ResidualBlock(out_dim, out_dim, dropout), + AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout)) + + # output blocks + self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), + CausalConv3d(out_dim, z_dim, 3, padding=1)) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], dim=2) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## downsamples + for layer in self.downsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## middle + for layer in self.middle: + if check_is_instance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +class Encoder3d_38(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + + # dimensions + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + # init block + self.conv1 = CausalConv3d(12, dims[0], 3, padding=1) + + # downsample blocks + downsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + t_down_flag = ( + temperal_downsample[i] if i < len(temperal_downsample) else False + ) + downsamples.append( + Down_ResidualBlock( + in_dim=in_dim, + out_dim=out_dim, + dropout=dropout, + mult=num_res_blocks, + temperal_downsample=t_down_flag, + down_flag=i != len(dim_mult) - 1, + ) + ) + scale /= 2.0 + self.downsamples = nn.Sequential(*downsamples) + + # middle blocks + self.middle = nn.Sequential( + ResidualBlock(out_dim, out_dim, dropout), + AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout), + ) + + # # output blocks + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), + nn.SiLU(), + CausalConv3d(out_dim, z_dim, 3, padding=1), + ) + + + def forward(self, x, feat_cache=None, feat_idx=[0]): + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], + dim=2, + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## downsamples + for layer in self.downsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## middle + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], + dim=2, + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + + return x + + +class Decoder3d(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_upsample = temperal_upsample + + # dimensions + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + scale = 1.0 / 2**(len(dim_mult) - 2) + + # init block + self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) + + # middle blocks + self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout), + AttentionBlock(dims[0]), + ResidualBlock(dims[0], dims[0], dropout)) + + # upsample blocks + upsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + if i == 1 or i == 2 or i == 3: + in_dim = in_dim // 2 + for _ in range(num_res_blocks + 1): + upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + upsamples.append(AttentionBlock(out_dim)) + in_dim = out_dim + + # upsample block + if i != len(dim_mult) - 1: + mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d' + upsamples.append(Resample(out_dim, mode=mode)) + scale *= 2.0 + self.upsamples = nn.Sequential(*upsamples) + + # output blocks + self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), + CausalConv3d(out_dim, 3, 3, padding=1)) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + ## conv1 + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], dim=2) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## middle + for layer in self.middle: + if check_is_instance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## upsamples + for layer in self.upsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + + +class Decoder3d_38(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_upsample = temperal_upsample + + # dimensions + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + scale = 1.0 / 2 ** (len(dim_mult) - 2) + # init block + self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) + + # middle blocks + self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout), + AttentionBlock(dims[0]), + ResidualBlock(dims[0], dims[0], dropout)) + + # upsample blocks + upsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + t_up_flag = temperal_upsample[i] if i < len(temperal_upsample) else False + upsamples.append( + Up_ResidualBlock(in_dim=in_dim, + out_dim=out_dim, + dropout=dropout, + mult=num_res_blocks + 1, + temperal_upsample=t_up_flag, + up_flag=i != len(dim_mult) - 1)) + self.upsamples = nn.Sequential(*upsamples) + + # output blocks + self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), + CausalConv3d(out_dim, 12, 3, padding=1)) + + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], + dim=2, + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + for layer in self.middle: + if check_is_instance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## upsamples + for layer in self.upsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx, first_chunk) + else: + x = layer(x) + + ## head + for layer in self.head: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], + dim=2, + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +def count_conv3d(model): + count = 0 + for m in model.modules(): + if isinstance(m, CausalConv3d): + count += 1 + return count + + +class VideoVAE_(nn.Module): + + def __init__(self, + dim=96, + z_dim=16, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + # modules + self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks, + attn_scales, self.temperal_downsample, dropout) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks, + attn_scales, self.temperal_upsample, dropout) + + self._enc_conv_num = count_conv3d(self.encoder) + self._dec_conv_num = count_conv3d(self.decoder) + + def encode(self, x, scale): + feat_map = [None] * self._enc_conv_num + + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + + out = self.encoder( + x[:, :, :1, :, :], + feat_cache=feat_map, + feat_idx=[0], + ) + + for i in range(1, iter_): + out_ = self.encoder( + x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], + feat_cache=feat_map, + feat_idx=[0], + ) + out = torch.cat([out, out_], dim=2) + mu, _ = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=mu.dtype, device=mu.device) for s in scale] + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=mu.dtype, device=mu.device) + mu = (mu - scale[0]) * scale[1] + return mu + + def decode(self, z, scale): + feat_map = [None] * self._dec_conv_num + + # z: [b,c,t,h,w] + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=z.dtype, device=z.device) for s in scale] + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=z.dtype, device=z.device) + z = z / scale[1] + scale[0] + iter_ = z.shape[2] + x = self.conv2(z) + + out = self.decoder( + x[:, :, 0:1, :, :], + feat_cache=feat_map, + feat_idx=[0], + ) + + for i in range(1, iter_): + out_ = self.decoder( + x[:, :, i:i + 1, :, :], + feat_cache=feat_map, + feat_idx=[0], + ) + out = torch.cat([out, out_], dim=2) + return out + + def reparameterize(self, mu, log_var): + std = torch.exp(0.5 * log_var) + eps = torch.randn_like(std) + return eps * std + mu + + def sample(self, imgs, deterministic=False): + mu, log_var = self.encode(imgs) + if deterministic: + return mu + std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0)) + return mu + std * torch.randn_like(std) + + +class WanVideoVAE(nn.Module): + + def __init__(self, z_dim=16, vae_pretrained_path: str | None = None): + super().__init__() + + mean = [ + -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, + 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921 + ] + std = [ + 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, + 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160 + ] + self.mean = torch.tensor(mean, device='cuda') + self.std = torch.tensor(std, device='cuda') + self.scale = [self.mean, 1.0 / self.std] + + # init model + self.model = VideoVAE_(z_dim=z_dim).eval().requires_grad_(False) + self.upsampling_factor = 8 + self.z_dim = z_dim + self.vae_pretrained_path = vae_pretrained_path + + + def build_1d_mask(self, length, left_bound, right_bound, border_width, device): + x = torch.ones((length,), device=device) + border = (torch.arange(border_width, device=device) + 1) + if not left_bound: + x[:border_width] = border / border_width + if not right_bound: + x[-border_width:] = torch.flip(border / border_width, dims=(0,)) + return x + + + def build_mask(self, data, is_bound, border_width): + _, _, _, H, W = data.shape + h = self.build_1d_mask(H, is_bound[0], is_bound[1], border_width[0], device=data.device) + w = self.build_1d_mask(W, is_bound[2], is_bound[3], border_width[1], device=data.device) + + h = repeat(h, "H -> H W", H=H, W=W) + w = repeat(w, "W -> H W", H=H, W=W) + + mask = torch.stack([h, w]).min(dim=0).values + mask = rearrange(mask, "H W -> 1 1 1 H W") + return mask + + + def tiled_decode(self, hidden_states, tile_size, tile_stride): + _, _, T, H, W = hidden_states.shape + size_h, size_w = tile_size + stride_h, stride_w = tile_stride + + # Split tasks + tasks = [] + for h in range(0, H, stride_h): + if (h-stride_h >= 0 and h-stride_h+size_h >= H): continue + for w in range(0, W, stride_w): + if (w-stride_w >= 0 and w-stride_w+size_w >= W): continue + h_, w_ = h + size_h, w + size_w + tasks.append((h, h_, w, w_)) + + out_T = T * 4 - 3 + weight = torch.zeros( + (1, 1, out_T, H * self.upsampling_factor, W * self.upsampling_factor), + dtype=hidden_states.dtype, + device=hidden_states.device + ) + values = torch.zeros( + (1, 3, out_T, H * self.upsampling_factor, W * self.upsampling_factor), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + + for h, h_, w, w_ in tqdm(tasks, desc="VAE decoding"): + hidden_states_batch = hidden_states[:, :, :, h:h_, w:w_] + hidden_states_batch = self.model.decode(hidden_states_batch, self.scale) + + mask = self.build_mask( + hidden_states_batch, + is_bound=(h==0, h_>=H, w==0, w_>=W), + border_width=((size_h - stride_h) * self.upsampling_factor, (size_w - stride_w) * self.upsampling_factor) + ).to(dtype=hidden_states.dtype) + + target_h = h * self.upsampling_factor + target_w = w * self.upsampling_factor + values[ + :, + :, + :, + target_h:target_h + hidden_states_batch.shape[3], + target_w:target_w + hidden_states_batch.shape[4], + ] += hidden_states_batch * mask + weight[ + :, + :, + :, + target_h: target_h + hidden_states_batch.shape[3], + target_w: target_w + hidden_states_batch.shape[4], + ] += mask + values = values / weight + values = values.clamp_(-1, 1) + return values + + def tiled_encode(self, video, tile_size, tile_stride): + _, _, T, H, W = video.shape + size_h, size_w = tile_size + stride_h, stride_w = tile_stride + + # Split tasks + tasks = [] + for h in range(0, H, stride_h): + if (h-stride_h >= 0 and h-stride_h+size_h >= H): continue + for w in range(0, W, stride_w): + if (w-stride_w >= 0 and w-stride_w+size_w >= W): continue + h_, w_ = h + size_h, w + size_w + tasks.append((h, h_, w, w_)) + + out_T = (T + 3) // 4 + weight = torch.zeros( + (1, 1, out_T, H // self.upsampling_factor, W // self.upsampling_factor), + dtype=video.dtype, + device=video.device, + ) + values = torch.zeros( + (1, self.z_dim, out_T, H // self.upsampling_factor, W // self.upsampling_factor), + dtype=video.dtype, + device=video.device, + ) + + for h, h_, w, w_ in tqdm(tasks, desc="VAE encoding"): + hidden_states_batch = video[:, :, :, h:h_, w:w_] + hidden_states_batch = self.model.encode(hidden_states_batch, self.scale) + + mask = self.build_mask( + hidden_states_batch, + is_bound=(h==0, h_>=H, w==0, w_>=W), + border_width=((size_h - stride_h) // self.upsampling_factor, (size_w - stride_w) // self.upsampling_factor) + ).to(dtype=video.dtype) + + target_h = h // self.upsampling_factor + target_w = w // self.upsampling_factor + values[ + :, + :, + :, + target_h:target_h + hidden_states_batch.shape[3], + target_w:target_w + hidden_states_batch.shape[4], + ] += hidden_states_batch * mask + weight[ + :, + :, + :, + target_h: target_h + hidden_states_batch.shape[3], + target_w: target_w + hidden_states_batch.shape[4], + ] += mask + values = values / weight + return values + + def single_encode(self, video): + x = self.model.encode(video, self.scale) + # The outputs of torch compile always need to be cloned before being used. + x = x.clone() + return x + + def single_decode(self, hidden_state): + video = self.model.decode(hidden_state, self.scale) + return video.clamp_(-1, 1) + + def encode(self, videos, tiled=False, tile_size=(34, 34), tile_stride=(18, 16)): + hidden_states = [] + for video in videos: + video = video.unsqueeze(0) + if tiled: + tile_size = (tile_size[0] * self.upsampling_factor, tile_size[1] * self.upsampling_factor) + tile_stride = (tile_stride[0] * self.upsampling_factor, tile_stride[1] * self.upsampling_factor) + hidden_state = self.tiled_encode(video, tile_size, tile_stride) + else: + hidden_state = self.single_encode(video) + hidden_state = hidden_state.squeeze(0) + hidden_states.append(hidden_state) + hidden_states = torch.stack(hidden_states) + return hidden_states + + def decode(self, hidden_states, tiled=False, tile_size=(34, 34), tile_stride=(18, 16)): + if tiled: + video = self.tiled_decode(hidden_states, tile_size, tile_stride) + else: + video = self.single_decode(hidden_states) + return video + + + @staticmethod + def state_dict_converter(): + return WanVideoVAEStateDictConverter() + + +class WanVideoVAEStateDictConverter: + + def __init__(self): + pass + + def from_civitai(self, state_dict): + state_dict_ = {} + if 'model_state' in state_dict: + state_dict = state_dict['model_state'] + for name in state_dict: + state_dict_['model.' + name] = state_dict[name] + return state_dict_ + + +class VideoVAE38_(VideoVAE_): + + def __init__(self, + dim=160, + z_dim=48, + dec_dim=256, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0): + super(VideoVAE_, self).__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + # modules + self.encoder = Encoder3d_38(dim, z_dim * 2, dim_mult, num_res_blocks, + attn_scales, self.temperal_downsample, dropout) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d_38(dec_dim, z_dim, dim_mult, num_res_blocks, + attn_scales, self.temperal_upsample, dropout) + + self._enc_conv_num = count_conv3d(self.encoder) + self._dec_conv_num = count_conv3d(self.decoder) + + def encode(self, x, scale): + feat_map = [None] * self._enc_conv_num + + x = patchify(x, patch_size=2) + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + + out = self.encoder( + x[:, :, :1, :, :], + feat_cache=feat_map, + feat_idx=[0], + ) + + for i in range(1, iter_): + out_ = self.encoder( + x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], + feat_cache=feat_map, + feat_idx=[0], + ) + out = torch.cat([out, out_], dim=2) + + mu, _ = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=mu.dtype, device=mu.device) for s in scale] + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=mu.dtype, device=mu.device) + mu = (mu - scale[0]) * scale[1] + return mu + + def decode(self, z, scale): + feat_map = [None] * self._dec_conv_num + + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=z.dtype, device=z.device) for s in scale] + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=z.dtype, device=z.device) + z = z / scale[1] + scale[0] + iter_ = z.shape[2] + x = self.conv2(z) + + out = self.decoder( + x[:, :, 0:1, :, :], + feat_cache=feat_map, + feat_idx=[0], + first_chunk=True, + ) + + for i in range(1, iter_): + out_ = self.decoder( + x[:, :, i:i + 1, :, :], + feat_cache=feat_map, + feat_idx=[0], + ) + out = torch.cat([out, out_], 2) + + out = unpatchify(out, patch_size=2) + return out + + +class WanVideoVAE38(WanVideoVAE): + + def __init__(self, z_dim=48, dim=160, vae_pretrained_path: str | None = None): + super(WanVideoVAE, self).__init__() + + mean = [ + -0.2289, -0.0052, -0.1323, -0.2339, -0.2799, 0.0174, 0.1838, 0.1557, + -0.1382, 0.0542, 0.2813, 0.0891, 0.1570, -0.0098, 0.0375, -0.1825, + -0.2246, -0.1207, -0.0698, 0.5109, 0.2665, -0.2108, -0.2158, 0.2502, + -0.2055, -0.0322, 0.1109, 0.1567, -0.0729, 0.0899, -0.2799, -0.1230, + -0.0313, -0.1649, 0.0117, 0.0723, -0.2839, -0.2083, -0.0520, 0.3748, + 0.0152, 0.1957, 0.1433, -0.2944, 0.3573, -0.0548, -0.1681, -0.0667 + ] + std = [ + 0.4765, 1.0364, 0.4514, 1.1677, 0.5313, 0.4990, 0.4818, 0.5013, + 0.8158, 1.0344, 0.5894, 1.0901, 0.6885, 0.6165, 0.8454, 0.4978, + 0.5759, 0.3523, 0.7135, 0.6804, 0.5833, 1.4146, 0.8986, 0.5659, + 0.7069, 0.5338, 0.4889, 0.4917, 0.4069, 0.4999, 0.6866, 0.4093, + 0.5709, 0.6065, 0.6415, 0.4944, 0.5726, 1.2042, 0.5458, 1.6887, + 0.3971, 1.0600, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744 + ] + self.mean = torch.tensor(mean, device='cuda') + self.std = torch.tensor(std, device='cuda') + self.scale = [self.mean, 1.0 / self.std] + + # init model + self.model = VideoVAE38_(z_dim=z_dim, dim=dim).eval().requires_grad_(False) + self.upsampling_factor = 16 + self.z_dim = z_dim + self.vae_pretrained_path = vae_pretrained_path \ No newline at end of file diff --git a/groot/vla/model/dreamzero/transform/__init__.py b/groot/vla/model/dreamzero/transform/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/model/dreamzero/transform/common.py b/groot/vla/model/dreamzero/transform/common.py new file mode 100644 index 0000000000000000000000000000000000000000..913add82986c026fb17f0c0ca6224102dd2a8cc1 --- /dev/null +++ b/groot/vla/model/dreamzero/transform/common.py @@ -0,0 +1,15 @@ +""" +Maybe there's a better place for the functions. +""" + +import re + + +def formalize_language(language: str) -> str: + """ + 1. Force lowercase + 2. Remove all punctuations + """ + language = language.lower() + language = re.sub(r"[^\w\s]", "", language) + return language diff --git a/groot/vla/model/dreamzero/transform/dreamzero_cotrain.py b/groot/vla/model/dreamzero/transform/dreamzero_cotrain.py new file mode 100644 index 0000000000000000000000000000000000000000..155b1d648c5740d868e9721edb76556319343d19 --- /dev/null +++ b/groot/vla/model/dreamzero/transform/dreamzero_cotrain.py @@ -0,0 +1,634 @@ +import os +import random +from typing import Any, Dict, List, Optional + +from einops import rearrange +import numpy as np +from pydantic import Field, PrivateAttr +import torch +from transformers import AutoProcessor, ProcessorMixin, AutoTokenizer +from transformers.data.data_collator import DataCollatorMixin +from transformers.feature_extraction_utils import BatchFeature +import tree +import re +import ftfy +import html +import regex as re +import ast + +from groot.vla.data.schema import ( + EmbodimentTag, + DatasetMetadata, +) +from groot.vla.data.transform.base import InvertibleModalityTransform +from groot.vla.model.dreamzero.transform.common import formalize_language + + +def basic_clean(text): + text = ftfy.fix_text(text) + text = html.unescape(html.unescape(text)) + return text.strip() + +def whitespace_clean(text): + text = re.sub(r'\s+', ' ', text) + text = text.strip() + return text + + +class HuggingfaceTokenizer: + + def __init__(self, name, seq_len=None, clean=None, **kwargs): + assert clean in (None, 'whitespace') + self.name = name + self.seq_len = seq_len + self.clean = clean + + # When loading from a local checkpoint path (e.g. from training runs), pass + # local_files_only=True to avoid HFValidationError from validate_repo_id. + load_kwargs = dict(kwargs) + if os.path.isdir(name): + load_kwargs.setdefault("local_files_only", True) + # init tokenizer + self.tokenizer = AutoTokenizer.from_pretrained(name, **load_kwargs) + self.vocab_size = self.tokenizer.vocab_size + + def __call__(self, sequence, **kwargs): + return_mask = kwargs.pop('return_mask', False) + + # arguments + _kwargs = {'return_tensors': 'pt'} + if self.seq_len is not None: + _kwargs.update({ + 'padding': 'max_length', + 'truncation': True, + 'max_length': self.seq_len + }) + _kwargs.update(**kwargs) + + + # tokenization + if isinstance(sequence, str): + sequence = [sequence] + if self.clean: + sequence = [self._clean(u) for u in sequence] + ids = self.tokenizer(sequence, **_kwargs) + + # output + if return_mask: + return ids.input_ids, ids.attention_mask + else: + return ids.input_ids + + def _clean(self, text): + if self.clean == 'whitespace': + text = whitespace_clean(basic_clean(text)) + # elif self.clean == 'lower': + # text = whitespace_clean(basic_clean(text)).lower() + # elif self.clean == 'canonicalize': + # text = canonicalize(basic_clean(text)) + return text + + +def collate(features: List[dict], tokenizer: AutoTokenizer, num_views=3, embodiment_tag_mapping=None) -> dict: + batch = {} + keys = features[0].keys() + + for key in keys: + if key == "text": + output_values = [] + for elem in features: + item = elem[key] + try: + parsed_item = ast.literal_eval(item) + # Handle different return types from ast.literal_eval + if isinstance(parsed_item, (list, tuple)): + processed_item = str(parsed_item[0]) + else: + # If it's already a scalar (string, float, int, etc.), convert to string + processed_item = str(parsed_item) + + if num_views > 1 and elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.AGIBOT.value]: + processed_item = "A multi-view video shows that a robot " + processed_item.lower() + " The video is split into four views: The top-left view shows the camera view from the robot's head, the top-right view shows the camera view from the right hand, the bottom-left view shows the camera view from the left hand, and the bottom-right view is a black screen (inactive view). The robot " + processed_item.lower() + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.OXE_DROID.value]: + processed_item = ( + "A multi-view video shows that a robot " + + processed_item.lower() + + " The video is split into three views: The top view shows the camera view from the robot's wrist, the bottom-left view shows the camera view from the left exterior camera, and the bottom-right view shows the camera view from the right exterior camera. During training, one of the two bottom exterior views may be a black screen (dropped view). The robot " + + processed_item.lower() + ) + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.GR1_UNIFIED.value]: + processed_item = "A single view video shows that a human " + processed_item.lower() + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.MECKA_HANDS.value]: + processed_item = "A single view video shows that a human " + processed_item.lower() + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.XDOF.value]: + processed_item = "A multi-view video shows that a robot " + processed_item.lower() + " The video is split into four views: The top-left view shows the camera view from the robot's head, the top-right view shows the camera view from the right hand, the bottom-left view shows the camera view from the left hand, and the bottom-right view is a black screen (inactive view). The robot " + processed_item.lower() + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.YAM.value]: + processed_item = "A multi-view video shows that a robot " + processed_item.lower() + " The video is split into four views: The top-left view shows the top camera, the top-right view shows the right camera, the bottom-left view shows the left camera, and the bottom-right view is a black screen. The robot " + processed_item.lower() + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.TREX.value]: + processed_item = "A multi-view video shows that a bimanual robot with dexterous hands " + processed_item.lower() + " The video is split into four views: The top-left view shows the camera view from the robot's head, the top-right view shows the camera view from the right wrist, the bottom-left view shows the camera view from the left wrist, and the bottom-right view is a black screen. The robot " + processed_item.lower() + else: + raise ValueError(f"Embodiment ID {elem['embodiment_id']} not supported.") + output_values.append(processed_item) + except (ValueError, SyntaxError, TypeError): + # If parsing fails or item is already a string, use it directly + if num_views > 1 and elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.AGIBOT.value]: + item = "A multi-view video shows that a robot " + str(item).lower() + " The video is split into four views: The top-left view shows the camera view from the robot's head, the top-right view shows the camera view from the right hand, the bottom-left view shows the camera view from the left hand, and the bottom-right view is a black screen (inactive view). The robot " + str(item).lower() + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.OXE_DROID.value]: + item = ( + "A multi-view video shows that a robot " + + str(item).lower() + + " The video is split into three views: The top view shows the camera view from the robot's wrist, the bottom-left view shows the camera view from the left exterior camera, and the bottom-right view shows the camera view from the right exterior camera. During training, one of the two bottom exterior views may be a black screen (dropped view). The robot " + + str(item).lower() + ) + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.GR1_UNIFIED.value]: + item = "A single view video shows that a human " + str(item).lower() + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.MECKA_HANDS.value]: + item = "A single view video shows that a human " + str(item).lower() + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.XDOF.value]: + item = "A multi-view video shows that a robot " + str(item).lower() + " The video is split into four views: The top-left view shows the camera view from the robot's head, the top-right view shows the camera view from the right hand, the bottom-left view shows the camera view from the left hand, and the bottom-right view is a black screen (inactive view). The robot " + str(item).lower() + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.YAM.value]: + item = "A multi-view video shows that a robot " + str(item).lower() + " The video is split into four views: The top-left view shows the top camera, the top-right view shows the right camera, the bottom-left view shows the left camera, and the bottom-right view is a black screen. The robot " + str(item).lower() + elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.TREX.value]: + item = "A multi-view video shows that a bimanual robot with dexterous hands " + str(item).lower() + " The video is split into four views: The top-left view shows the camera view from the robot's head, the top-right view shows the camera view from the right wrist, the bottom-left view shows the camera view from the left wrist, and the bottom-right view is a black screen. The robot " + str(item).lower() + else: + raise ValueError(f"Embodiment ID {elem['embodiment_id']} not supported.") + output_values.append(item) + # print("output_values", output_values) + ids, mask = tokenizer(output_values, return_mask=True, add_special_tokens=True) + batch[key] = ids + batch['text_attention_mask'] = mask + elif key == "text_negative": + values = [elem[key] for elem in features] + ids, mask = tokenizer(values, return_mask=True, add_special_tokens=True) + batch[key] = ids + batch['text_attention_mask_negative'] = mask + else: + values = [elem[key] for elem in features] + batch[key] = torch.from_numpy(np.stack(values)) + return batch + + + +class DefaultDataCollator(DataCollatorMixin): + def __init__(self, tokenizer_path: str="google/umt5-xxl", max_length: int=512, num_views: int=1, embodiment_tag_mapping=None): + super().__init__() + self.tokenizer = HuggingfaceTokenizer(name=tokenizer_path, seq_len=max_length, clean='whitespace') + self.num_views = num_views + self.embodiment_tag_mapping = embodiment_tag_mapping + + def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]: + return collate(features, self.tokenizer, self.num_views, self.embodiment_tag_mapping) + + +class DreamTransform(InvertibleModalityTransform): + + # -- 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." + ) + training: bool = Field( + default=True, description="Whether to apply the transform in training mode." + ) + + formalize_language: bool = Field(default=False, description="Formalize language if True.") + + embodiment_tag_mapping: dict[str, int] = Field( + default_factory=dict, + description="The projector index of each embodiment tag.", + ) + + language_dropout_prob: float = Field( + default=0.0, + description="Dropout probability for language.", + ) + always_use_default_instruction: bool = Field( + default=False, + description="Whether to always use the default instruction. For studying how much the language helps.", + ) + + # Private attributes to keep track of shapes/dimensions across apply/unapply + _language_key: Optional[str] = PrivateAttr(default=None) + _language_keys: Optional[list[str]] = PrivateAttr(default=None) + + # XEmbDiT arguments + default_instruction: str + max_state_dim: int + max_action_dim: int + max_length: int = 512 + embodiment_tag: EmbodimentTag | None = None + state_horizon: int + action_horizon: int + num_views: int = 3 + + # Add tokenizer attribute + tokenizer_path: str = Field( + default="google/umt5-xxl", + description="Path to the tokenizer." + ) + _tokenizer: Optional[HuggingfaceTokenizer] = PrivateAttr(default=None) + + def __init__(self, **kwargs): + super().__init__(**kwargs) + # Initialize the tokenizer + self._tokenizer = HuggingfaceTokenizer( + name=self.tokenizer_path, + seq_len=self.max_length, + clean='whitespace' + ) + + @property + def tokenizer(self): + return self._tokenizer + + def set_metadata( + self, dataset_metadata: DatasetMetadata + ): + self.embodiment_tag = dataset_metadata.embodiment_tag + + def get_embodiment_tag(self) -> int: + """Get the embodiment tag from the data.""" + assert ( + self.embodiment_tag is not None + ), "Embodiment tag not set. Please call set_metadata first." + return self.embodiment_tag_mapping[self.embodiment_tag.value] + + def check_keys_and_batch_size(self, data): + grouped_keys = {} + for key in data.keys(): + try: + modality, _ = key.split(".") + if "annotation" in key: + modality = "language" + except: # noqa: E722 + ### Handle language annotation special case + if "annotation" in key: + modality = "language" + else: + modality = "others" # will contain the video, state, and action + if modality not in grouped_keys: + grouped_keys[modality] = [] + grouped_keys[modality].append(key) + # Use video key to determine batch size. + video_ndim = data["video"].ndim + if video_ndim == 5: # Interpret as [T, V, H, W, C] + is_batched = False + batch_size = 1 + elif video_ndim == 6: # Interpret as [B, T, V, H, W, C] + is_batched = True + batch_size = data["video"].shape[0] + else: + raise ValueError(f"Unsupported video number of dimensions: {video_ndim}") + + # Handle language + if "language" in grouped_keys: + language_keys = grouped_keys["language"] + self._language_keys = language_keys # Store all keys for random selection + if len(language_keys) == 1: + self._language_key = language_keys[0] + else: + self._language_key = None # Will be selected randomly in _prepare_language + return is_batched, batch_size + + def _apply_vlm_processing(self, batch: dict) -> BatchFeature: + """ + Args: + batch: + video: [V, T, C, H, W] + Returns: required input with the format `BatchFeature` + """ + images = batch["images"] # [V, T, C, H, W] + + np_images = rearrange(images, "v t c h w -> (t v) h w c") + if "language" in batch: + lang = batch["language"] + if isinstance(lang, list) or isinstance(lang, np.ndarray): + lang = lang[0] + + inputs = {} + inputs["images"] = np_images + inputs["text"] = lang + + return inputs + + def _prepare_video(self, data: dict): + """Process, stack, and pad images from data['video'].""" + images = rearrange( + data["video"], + "t v h w c -> v t c h w", + ) + if images.shape[0] > 1: + v, t, c, h, w = images.shape + + # For DROID embodiment: 2x2 grid where the wrist view spans the full top row, + # and the two exterior views occupy the bottom row. + # + # View indices (expected): + # - View 0: left exterior + # - View 1: right exterior + # - View 2: wrist + # + # Layout: + # [wrist, wrist] (wrist duplicated to have 2x width) + # [left_ext | right_ext] + # + # Training-time augmentation: + # - Randomly drop (black out) either left_ext or right_ext. + if self.embodiment_tag == EmbodimentTag.OXE_DROID and v >= 3: + left_exterior = images[0] # (t, c, h, w) + right_exterior = images[1] # (t, c, h, w) + wrist_image = images[2] # (t, c, h, w) + + concat_images = np.zeros((1, t, c, 2 * h, 2 * w), dtype=images.dtype) + + # Top row: a SINGLE wrist view, resized to be 2x wider (same height). + # We use nearest-neighbor upscaling by repeating pixels along width. + wrist_wide = np.repeat(wrist_image, 2, axis=-1) # (t, c, h, 2w) + concat_images[0, :, :, :h, :] = wrist_wide + + # # Bottom row: left/right exteriors. + # drop_exterior_idx = None + # if self.training: + # # Always drop exactly one exterior view during training. + # drop_exterior_idx = random.choice([0, 1]) # 0=left, 1=right + + # if drop_exterior_idx != 0: + concat_images[0, :, :, h:, :w] = left_exterior + # if drop_exterior_idx != 1: + concat_images[0, :, :, h:, w:] = right_exterior + + return concat_images + + # For other embodiments: use 2x2 grid layout + # Layout: [head, right] + # [left, black] + + # Create output tensor with doubled height and width + concat_images = np.zeros((1, t, c, 2*h, 2*w), dtype=images.dtype) + + # Place images in the 2x2 grid + # Left upper: head image (view 0) + if v > 0: + concat_images[0, :, :, :h, :w] = images[0] + + # Left bottom: left image (view 1) + if v > 1: + concat_images[0, :, :, h:, :w] = images[1] + + # Right top: right image (view 2) + if v > 2: + concat_images[0, :, :, :h, w:] = images[2] + + # Right bottom: black pixels (already zeros from initialization) + + return concat_images + + return images + + def _prepare_language(self, data: dict): + """Tokenize data['language'] (or default_instruction if missing).""" + # Determine which language key to use + selected_key = self._language_key + + # For DROID embodiment during training, randomly select from available language keys + if (self._language_keys is not None and + len(self._language_keys) > 1 and + self.training and + self.embodiment_tag == EmbodimentTag.OXE_DROID): + selected_key = random.choice(self._language_keys) + elif self._language_keys is not None and len(self._language_keys) > 0 and selected_key is None: + selected_key = self._language_keys[0] + + if selected_key is not None: + raw_language = data[selected_key] + if isinstance(raw_language, np.ndarray): + raw_language = raw_language.item() if raw_language.size == 1 else raw_language[0] + if isinstance(raw_language, list): + raw_language = raw_language[0] + + # Language dropout + # WARNING: this is not compatible with LAPA and DREAM + if self.training and self.language_dropout_prob > 1e-9: + if random.random() < self.language_dropout_prob: + raw_language = self.default_instruction + else: + raw_language = self.default_instruction + + if "" in raw_language: + raw_language = raw_language.replace("", "") + is_lapa_instance = True + else: + is_lapa_instance = False + + if "" in raw_language: + raw_language = raw_language.replace("", "") + is_dream_instance = True + else: + is_dream_instance = False + + if "" in raw_language: + raw_language = raw_language.replace("", "") + is_cotrain_instance = True + else: + is_cotrain_instance = False + + if self.always_use_default_instruction: + raw_language = self.default_instruction + + # print("raw_language", raw_language) + + # Formalize language + if self.formalize_language: + formalized_language = formalize_language(raw_language) + return formalized_language, is_lapa_instance, is_dream_instance, is_cotrain_instance + else: + return raw_language, is_lapa_instance, is_dream_instance, is_cotrain_instance + + def _prepare_state(self, data: dict): + """ + Gathers final state from data['state'], then pads to max_state_dim. + Return (state, state_mask, n_state_tokens). + """ + + if "state" not in data: + state = np.zeros((self.state_horizon, self.max_state_dim)) + state_mask = np.zeros((self.state_horizon, self.max_state_dim), dtype=bool) + n_state_tokens = self.state_horizon + return state, state_mask, n_state_tokens + + state = data["state"] + assert state.shape[0] % self.state_horizon == 0, f"{state.shape=}, {self.state_horizon=}" + + n_state_dims = state.shape[-1] + + # Instead of asserting, just take the first max_state_dim dimensions if needed + if n_state_dims > self.max_state_dim: + state = state[:, : self.max_state_dim] + n_state_dims = self.max_state_dim + else: + # Pad up to max_state_dim if smaller + state = np.pad(state, ((0, 0), (0, self.max_state_dim - n_state_dims)), "constant") + + # Create mask for real state dims + state_mask = np.zeros_like(state).astype(bool) + state_mask[:, :n_state_dims] = True + + # We only have 1 "proprio" token to represent the entire state + n_state_tokens = state.shape[0] + return state, state_mask, n_state_tokens + + def _prepare_action(self, data: dict): + """ + Pad to max_action_dim, return masks. + """ + if "action" not in data: + actions = np.zeros((self.action_horizon, self.max_action_dim)) + actions_mask = np.zeros((self.action_horizon, self.max_action_dim), dtype=bool) + n_action_tokens = self.action_horizon + return actions, actions_mask, n_action_tokens + + actions = data["action"] + assert actions.shape[0] % self.action_horizon == 0, f"{actions.shape=}, {self.action_horizon=}" + + n_action_tokens = actions.shape[0] # T + n_action_dims = actions.shape[1] + + assert ( + n_action_dims <= self.max_action_dim + ), f"Action dim {n_action_dims} exceeds max allowed {self.max_action_dim}." + + # Pad the channel dimension + actions = np.pad(actions, ((0, 0), (0, self.max_action_dim - n_action_dims)), "constant") + + # Create mask: [T, max_action_dim] + actions_mask = np.zeros((n_action_tokens, self.max_action_dim), dtype=bool) + actions_mask[:, :n_action_dims] = True + + return actions, actions_mask, n_action_tokens + + def apply_single(self, data: dict) -> dict: + transformed_data = {} + + # 1) Prepare video and language with vlm processing. + images = self._prepare_video(data) + images = images.astype(np.uint8) + language, is_lapa_instance, is_dream_instance, is_cotrain_instance = self._prepare_language(data) + batch_data = {"images": images, "language": language} + vlm_outputs = self._apply_vlm_processing(batch_data) + + # 2) Prepare state + state, state_mask, _ = self._prepare_state(data) + transformed_data["state"] = state + transformed_data["state_mask"] = state_mask + + if self.training: + # 3) Prepare actions + is_detection_instance = self.embodiment_tag == EmbodimentTag.GR1_UNIFIED_SEGMENTATION + if is_detection_instance: + transformed_data["segmentation_target"] = data["action"][0, -3:-1] + transformed_data["segmentation_target_mask"] = data["action"][0, -1:] + transformed_data["has_real_action"] = np.zeros((), dtype=bool) + else: + transformed_data["segmentation_target"] = np.zeros((2,)) + transformed_data["segmentation_target_mask"] = np.zeros((1,)) + transformed_data["has_real_action"] = np.ones((), dtype=bool) + actions, actions_mask, _ = self._prepare_action(data) + transformed_data["action"] = actions + transformed_data["action_mask"] = actions_mask + + # default for lapa instance + transformed_data["lapa_action"] = np.zeros_like(transformed_data["action"]) + transformed_data["lapa_action_mask"] = np.zeros_like(transformed_data["action_mask"]) + # else: + transformed_data["text_negative"] = "Vibrant colors, overexposed, static, blurry details, text, subtitles, style, artwork, painting, image, still, grayscale, dull, worst quality, low quality, JPEG artifacts, ugly, mutilated, extra fingers, bad hands, bad face, deformed, disfigured, mutated limbs, fused fingers, stagnant image, cluttered background, three legs, many people in the background, walking backwards." + + for k, v in vlm_outputs.items(): + assert k not in transformed_data, f"Key {k} already exists in transformed_data." + transformed_data[k] = v + + transformed_data["embodiment_id"] = self.get_embodiment_tag() + + if self.embodiment_tag == EmbodimentTag.MECKA_HANDS: + is_cotrain_instance = True + else: + is_cotrain_instance = False + + transformed_data["has_lapa_action"] = np.zeros((), dtype=bool) + # print("dreamzero_fixed", is_cotrain_instance) + if is_cotrain_instance: + transformed_data["is_cotrain_instance"] = np.ones((), dtype=bool) + else: + transformed_data["is_cotrain_instance"] = np.zeros((), dtype=bool) + + if is_dream_instance: + assert "dream_actions" in data + transformed_data["embodiment_id"] = self.embodiment_tag_mapping["dream"] + transformed_data["state"] = np.zeros_like(transformed_data["state"]) + actions_shape = transformed_data["action"].shape + + # Treat the "dream" IDM action as a real action so that flow matching loss will be applied. + transformed_data["has_real_action"] = np.ones((), dtype=bool) + transformed_data["has_lapa_action"] = np.zeros((), dtype=bool) + + dream_actions = data["dream_actions"] + assert ( + dream_actions.size == actions_shape[0] * actions_shape[1] + ), f"dream_actions size {dream_actions.size} does not match action shape {actions_shape}" + transformed_data["action"] = dream_actions.reshape(actions_shape) + + if is_lapa_instance: + assert "lapa_action" in data + transformed_data["has_real_action"] = np.ones((), dtype=bool) + transformed_data["has_lapa_action"] = np.zeros((), dtype=bool) + transformed_data["embodiment_id"] = self.embodiment_tag_mapping["lapa"] + transformed_data["state"] = np.zeros_like(transformed_data["state"]) + actions_shape = transformed_data["action"].shape + lapa_actions = data["lapa_action"] + # Ensure total elements match before reshaping + assert ( + lapa_actions.size == actions_shape[0] * actions_shape[1] + ), f"Cannot reshape lapa_actions of size {lapa_actions.size} to {actions_shape}" + # Reshape the lapa_actions to match the expected shape + reshaped_lapa_actions = lapa_actions.reshape(actions_shape) + # lapa_action should be between -1 and 1 + assert np.all(reshaped_lapa_actions >= -1) and np.all( + reshaped_lapa_actions <= 1 + ), "LAPA action values should be between -1 and 1" + transformed_data["action"] = reshaped_lapa_actions + transformed_data["action_mask"] = np.ones(actions_shape, dtype=bool) + + if self.training: + action_and_mask_keys = ["action", "action_mask", "lapa_action", "lapa_action_mask"] + assert all( + transformed_data[key].shape == transformed_data["action"].shape + for key in action_and_mask_keys + ), f"Shape mismatch: {[(key, transformed_data[key].shape) for key in action_and_mask_keys]}" + + return transformed_data + + def apply_batch(self, data: dict, batch_size: int) -> dict: + # Split on batch dimension. + # delete lapa_action and lapa_action_mask + data.pop("lapa_action", None) + # data.pop("lapa_action_mask", None) + data.pop("dream_actions", None) + data_split = [tree.map_structure(lambda x: x[i], data) for i in range(batch_size)] + # Process each element. + data_split_processed = [self.apply_single(elem) for elem in data_split] + return collate(data_split_processed, self.tokenizer, self.num_views, self.embodiment_tag_mapping) + + def apply(self, data: dict) -> dict: + if not self.training and data["video"].ndim == 5: + data["video"] = data["video"][None, ...] + is_batched, batch_size = self.check_keys_and_batch_size(data) + if is_batched: + return self.apply_batch(data, batch_size) + else: + return self.apply_single(data) + + def unapply(self, data: dict) -> dict: + # Leave as is so that ConcatTransform can split the values + return data + + def __call__(self, data: dict) -> dict: + return self.apply(data) + diff --git a/groot/vla/model/n1_5/__init__.py b/groot/vla/model/n1_5/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/model/n1_5/action_head/__init__.py b/groot/vla/model/n1_5/action_head/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/model/n1_5/action_head/base_action_head.py b/groot/vla/model/n1_5/action_head/base_action_head.py new file mode 100644 index 0000000000000000000000000000000000000000..e8398622d4da362796b0fdb8a807a1592b55270e --- /dev/null +++ b/groot/vla/model/n1_5/action_head/base_action_head.py @@ -0,0 +1,32 @@ +from abc import ABC, abstractmethod + +from torch import nn +from transformers.feature_extraction_utils import BatchFeature + + +class ActionHead(ABC, nn.Module): + def __init__(self): + super(ActionHead, self).__init__() + + @abstractmethod + def forward(self, backbone_output: BatchFeature, action_input: BatchFeature) -> BatchFeature: + pass + + def get_action( + self, + backbone_output: BatchFeature, + action_input: BatchFeature, + num_action_samples: int = 1, + inference_batch_size: int = 32, + ) -> BatchFeature: + # Used for predicting actions during inference + # By default, the action head does the same thing as a normal forward pass + return self.forward(backbone_output, action_input) + + def prepare_input(self, batch: dict) -> BatchFeature: + pass + + def set_override_kwargs(self, **kwargs): + for key, value in kwargs.items(): + setattr(self.config, key, value) + setattr(self, key, value) diff --git a/groot/vla/model/n1_5/modules/__init__.py b/groot/vla/model/n1_5/modules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/model/n1_5/modules/action_encoder.py b/groot/vla/model/n1_5/modules/action_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..55a39303c2c69913be01c7a86fc5cea7af6921bb --- /dev/null +++ b/groot/vla/model/n1_5/modules/action_encoder.py @@ -0,0 +1,87 @@ +### Implemented based on Pi0 Action Encoding +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/groot/vla/model/n1_5/sim_policy.py b/groot/vla/model/n1_5/sim_policy.py new file mode 100644 index 0000000000000000000000000000000000000000..f1d01b510928deaeb69ad454c55388df97e6ef61 --- /dev/null +++ b/groot/vla/model/n1_5/sim_policy.py @@ -0,0 +1,1197 @@ +import importlib +import json +from pathlib import Path +from typing import Any, Callable + +import gymnasium as gym +from hydra.utils import instantiate +import numpy as np +from omegaconf import OmegaConf +from tianshou.data import Batch +from tianshou.policy import BasePolicy as BaseTianshouPolicy +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh +import tree +import time + +from groot.vla.data.schema import DatasetMetadata, EmbodimentTag +from groot.vla.data.transform import ComposedModalityTransform + + +class ModelManager: + """ + Manages model loading/offloading to handle memory efficiently when using multiple models. + Modified to keep VLM always loaded and only manage text_encoder/vae components. + """ + def __init__(self): + self.active_components = None # Track which components are active for action_head + self.models = {} + self.vlm_policy = None # Keep VLM always loaded + self.action_head_policy = None # Action head policy for component management + + def register_model(self, name: str, policy_instance): + """Register a policy instance with the manager.""" + self.models[name] = policy_instance + + if name == "vlm": + self.vlm_policy = policy_instance + # Load VLM to GPU and keep it there + self.load_vlm_model() + elif name == "action_head": + self.action_head_policy = policy_instance + # Ensure vram management is enabled for action head + self.enable_action_head_vram_management() + # Initially offload action head components to save memory + self.offload_action_head_components() + + def load_vlm_model(self): + """Load VLM model to GPU and keep it there.""" + if self.vlm_policy is None: + return + + policy = self.vlm_policy + print(f"Loading VLM model to GPU (keeping it loaded)...") + + # Move VLM model to GPU + policy.trained_model.to(device=policy.device) + + # Apply bf16 if needed + if policy.eval_bf16: + policy.trained_model = policy.trained_model.to(dtype=torch.bfloat16) + + torch.cuda.empty_cache() # Clear cache after loading + + def activate_model(self, name: str): + """Activate a model - for VLM this is a no-op, for action_head this manages components.""" + if name not in self.models: + raise ValueError(f"Model {name} not registered") + + if name == "vlm": + # VLM is always loaded, just ensure it's ready + if self.vlm_policy: + self.vlm_policy.ensure_model_on_gpu() + return + elif name == "action_head": + # For action head, activate the text_encoder and vae components + self.activate_action_head_components(['text_encoder', 'vae']) + self.active_components = ['text_encoder', 'vae'] + + def activate_action_head_components(self, component_names): + """Activate specific components of the action head model.""" + if self.action_head_policy is None: + return + + policy = self.action_head_policy + print(f"Loading action head components {component_names} to GPU...") + + # Check if the action head has the load_models_to_device method + if (hasattr(policy.trained_model, "action_head") and + hasattr(policy.trained_model.action_head, "load_models_to_device")): + # Use the selective loading method + policy.trained_model.action_head.load_models_to_device(component_names) + else: + # Fallback to loading the entire action head + if hasattr(policy.trained_model, "action_head") and hasattr(policy.trained_model.action_head, "image_encoder"): + policy.trained_model.action_head.enable_vram_management() + else: + policy.trained_model.to(device=policy.device) + + # Apply bf16 if needed + if policy.eval_bf16: + policy.trained_model = policy.trained_model.to(dtype=torch.bfloat16) + + torch.cuda.empty_cache() # Clear cache after loading + + def offload_action_head_components(self): + """Offload action head components to CPU.""" + if self.action_head_policy is None: + return + + policy = self.action_head_policy + print(f"Offloading action head components to CPU...") + + # Check if the action head has the load_models_to_device method + if (hasattr(policy.trained_model, "action_head") and + hasattr(policy.trained_model.action_head, "load_models_to_device")): + # Use the selective offloading method - pass empty list to offload all + policy.trained_model.action_head.load_models_to_device([]) + else: + # Fallback to offloading the entire action head + if hasattr(policy.trained_model, "action_head") and hasattr(policy.trained_model.action_head, "image_encoder"): + if hasattr(policy.trained_model.action_head, 'disable_vram_management'): + policy.trained_model.action_head.disable_vram_management() + policy.trained_model.to(device='cpu') + else: + policy.trained_model.to(device='cpu') + + torch.cuda.empty_cache() # Clear cache after offloading + self.active_components = None + + def load_model(self, name: str): + """Load a model to GPU. (Deprecated - use activate_model instead)""" + print(f"Warning: load_model is deprecated. Use activate_model instead.") + self.activate_model(name) + + def offload_model(self, name: str): + """Offload a model to CPU. (Deprecated - components are managed automatically)""" + if name == "action_head": + self.offload_action_head_components() + # VLM is never offloaded in this new approach + + def get_status(self): + """Get current status of model manager for debugging.""" + status = { + "vlm_loaded": self.vlm_policy is not None, + "action_head_available": self.action_head_policy is not None, + "active_components": self.active_components + } + return status + + def enable_action_head_vram_management(self): + """Enable vram management for the action head model.""" + if self.action_head_policy is None: + return + + policy = self.action_head_policy + + # Check if the action head has vram management capability + if (hasattr(policy.trained_model, "action_head") and + hasattr(policy.trained_model.action_head, "enable_vram_management")): + print("Enabling vram management for action head...") + policy.trained_model.action_head.enable_vram_management() + else: + print("Warning: Action head does not support vram management") + + +class BaseGrootSimPolicy(BaseTianshouPolicy): + def __init__(self, embodiment_tag: EmbodimentTag, model_path: str, device: int | str): + super().__init__() + self.embodiment_tag = embodiment_tag + self.model_path = model_path + self.device = device + + def forward(self, batch, state=None, **kwargs): + raise NotImplementedError + + @property + def video_delta_indices(self) -> np.ndarray: + return np.array([0]) + + @property + def state_delta_indices(self) -> np.ndarray: + return np.array([0]) + + @property + def raw_data_image_transform(self) -> Callable: + return lambda x: x + + def on_env_init(self, env: gym.Env): + pass + + def learn(self, batch: Batch, **kwargs) -> dict: + """Dummy learn method for BasePolicy.learn as this is an inference-only wrapper. + + Args: + batch: Input batch of experiences + **kwargs: Additional arguments + + Returns: + Empty dict as no learning occurs + """ + return {} + + +def _update_tokenizer_path_in_config(cfg, new_path: str) -> None: + """Update tokenizer_path in the transforms subtree only (avoids trainer.model etc.).""" + from omegaconf import DictConfig, ListConfig + if isinstance(cfg, DictConfig): + if "tokenizer_path" in cfg: + cfg.tokenizer_path = new_path + # Only recurse via "transforms" to avoid triggering resolution of trainer, etc. + if "transforms" not in cfg: + return + sub = cfg.transforms + if isinstance(sub, DictConfig): + for v in sub.values(): + _update_tokenizer_path_in_config(v, new_path) + elif isinstance(sub, ListConfig): + for v in sub: + _update_tokenizer_path_in_config(v, new_path) + elif isinstance(cfg, ListConfig): + for v in cfg: + _update_tokenizer_path_in_config(v, new_path) + + +class GrootSimPolicy(BaseGrootSimPolicy): + def __init__( + self, + embodiment_tag: EmbodimentTag, + model_path: str, + device: int | str, + model_config_overrides: list[str] | None = [], + tokenizer_path_override: str | None = None, + skip_assert_delta_indices: bool = False, + skip_img_transform: bool = False, + lazy_load: bool = False, + device_mesh: DeviceMesh | None = None, + ): + """ + Initialize the GrootSimPolicy. + + Args: + env_name (str): The name of the environment. + model_path (str): Path to the model checkpoint. + device (int | str): Device to run the model on. + lazy_load (bool): If True, don't load model to GPU immediately. + device_mesh (DeviceMesh | None): Device mesh to parallelize the model across. + """ + super().__init__(embodiment_tag=embodiment_tag, model_path=model_path, device=device) + model_dir = Path(model_path) + self.rank = dist.get_rank() + + exp_cfg_dir = model_dir / "experiment_cfg" + train_cfg_path = exp_cfg_dir / "conf.yaml" + train_cfg = OmegaConf.load(train_cfg_path) + if tokenizer_path_override is not None: + _update_tokenizer_path_in_config(train_cfg, tokenizer_path_override) + self.train_cfg = train_cfg + self.lazy_load = lazy_load + + # Store model loading parameters for lazy loading + self.model_config_overrides = model_config_overrides + self.model_dir = model_dir + + # 1. Load the model + if ( + train_cfg.model._target_.endswith(".from_pretrained") + or train_cfg.model._target_.endswith(".from_pretrained_for_tuning") + or train_cfg.model._target_.endswith(".from_pretrained_with_wrapped_action_head") + ): + # Compatibility with the finetuned model + model_target = train_cfg.model._target_.rsplit(".", 1)[0] + else: + model_target = train_cfg.model._target_ + + self.model_target = model_target + + if model_config_overrides is not None and len(model_config_overrides) != 0: + print(f"Applying model config overrides: {model_config_overrides}") + # Only apply new logic if there are config overrides + + # Import the model class from the target module path + module_path, class_name = model_target.rsplit(".", 1) + if "lora" in class_name: + module_path, class_name = module_path.rsplit(".", 1) + module = importlib.import_module(module_path) + model_class = getattr(module, class_name) + model_config_class = model_class.config_class + + # Load the model config from the model directory + model_config = json.load(open(model_dir / "config.json", "r")) + model_config = OmegaConf.create(model_config) + model_config.merge_with_dotlist(list(model_config_overrides)) + model_config = OmegaConf.to_container( + model_config, resolve=True + ) # need dict for later model instantiation + model_config = model_config_class.from_dict(model_config) + + # Instantiate the model + if hasattr(train_cfg, "save_lora_only") and train_cfg.save_lora_only is True: + print(f"Loading LoRA weights from pretrained") + model = model_class.load_lora(model_path) + else: + print(f"Loading model from pretrained directly") + model = model_class.from_pretrained(model_path, config=model_config) + else: + print(f"No model config overrides provided") + # Otherwise, just call from_pretrained directly + cls_module, cls_name = model_target.rsplit(".", 1) + if 'lora' in cls_name: + cls_module, cls_name = cls_module.rsplit(".", 1) + if hasattr(train_cfg, "save_lora_only") and train_cfg.save_lora_only is True: + print(f"Loading LoRA weights from pretrained") + cls = getattr(importlib.import_module(cls_module), cls_name) + model = cls.load_lora(model_path) + else: + print(f"Loading model from pretrained directly") + cls = getattr(importlib.import_module(cls_module), cls_name) + from_pretrained = getattr(cls, "from_pretrained") + model = from_pretrained(model_path) + + model.eval() + model.requires_grad_(False) + if model.action_head.train_architecture == "lora": + print(f"Merging LoRA weights into main model weights") + # Merge the LoRA weights into the main model weights, and delete the LoRA + # weights to save memory. Note that the WanModel is in model.action_head.model. + model.action_head.model = model.action_head.model.merge_and_unload() + else: + print(f"Skipping merging LoRA weights into main model weights") + + self.eval_bf16 = self.train_cfg.get("eval_bf16", False) + if self.eval_bf16 and not lazy_load: + model = model.to(dtype=torch.bfloat16) + + # Store model initially on CPU if lazy loading + if lazy_load: + model.to(device='cpu') + else: + model.to(device=device) + + # Post initialize, move RoPE freqs to cuda. + model.post_initialize() + + # Parallelize the model across devices. + try: + model.parallelize(device_mesh=device_mesh) + except Exception as e: + print("Skipping parallelization") + + torch.cuda.empty_cache() + + self.trained_model = model + + # 2. Load the action, video, and state transforms + # 2.1. Load the metadata for normalization stats + # We have an assumption: one policy is only for rolling out one type of env, i.e., one embodiment_tag + # metadata_versions = train_cfg.metadata_versions + # metadata = get_metadata(self.embodiment_tag, metadata_versions[self.embodiment_tag.value]) + metadata_path = exp_cfg_dir / "metadata.json" + with open(metadata_path, "r") as f: + metadatas = json.load(f) + if "gr1_unified_offline_rl" in metadatas and self.embodiment_tag.value == "gr1_unified": + self.embodiment_tag = EmbodimentTag.GR1_UNIFIED_OFFLINE_RL + metadata = DatasetMetadata.model_validate(metadatas[self.embodiment_tag.value]) + + # If the model's action head has target_video_height/width (e.g. DreamZero Wan 5B), use that + # as the expected video resolution so the transform matches the model. metadata.json can + # otherwise contain a different resolution (e.g. 180x320) from dataset config. + if hasattr(self.trained_model, "action_head") and hasattr( + self.trained_model.action_head, "config" + ): + cfg = self.trained_model.action_head.config + target_h = getattr(cfg, "target_video_height", None) + target_w = getattr(cfg, "target_video_width", None) + if target_h is not None and target_w is not None and metadata.modalities.video: + for key in metadata.modalities.video.keys(): + metadata.modalities.video[key].resolution = (int(target_w), int(target_h)) + + # 2.2. Get the eval transforms + assert ( + self.embodiment_tag.value in train_cfg.transforms + ), f"{self.embodiment_tag.value=}, {train_cfg.transforms.keys()=}" + eval_transform_cfg = train_cfg.transforms[self.embodiment_tag.value] + if skip_img_transform: + for t in eval_transform_cfg.transforms: + if t._target_ == "groot.vla.data.transform.VideoCrop": + scale = t.scale + for key in metadata.modalities.video.keys(): + metadata.modalities.video[key].resolution = ( + int(metadata.modalities.video[key].resolution[0] * scale), + int(metadata.modalities.video[key].resolution[1] * scale), + ) + elif t._target_ == "groot.vla.data.transform.VideoResize": + height, width = t.height, t.width + for key in metadata.modalities.video.keys(): + metadata.modalities.video[key].resolution = (height, width) + + # assume it's always ComposedModalityTransform + assert ( + eval_transform_cfg._target_ == "groot.vla.data.transform.ComposedModalityTransform" + ) + skipped_transforms = [ + "groot.vla.data.transform.VideoCrop", + "groot.vla.data.transform.VideoResize", + "groot.vla.data.transform.VideoColorJitter", + ] + eval_transform_cfg.transforms = [ + t for t in eval_transform_cfg.transforms if t._target_ not in skipped_transforms + ] + + eval_transform = instantiate(train_cfg.transforms[self.embodiment_tag.value]) + assert isinstance(eval_transform, ComposedModalityTransform), f"{eval_transform=}" + eval_transform.set_metadata(metadata) + + # Set per-horizon statistics for PerHorizonActionTransform if using relative_action_per_horizon + relative_action_per_horizon = self.train_cfg.get('relative_action_per_horizon', False) + print(f"DEBUG: relative_action_per_horizon = {relative_action_per_horizon}") + if relative_action_per_horizon: + # Extract per-horizon statistics from metadata + # The metadata has format: {embodiment: {statistics: {action: {key: {stat: [[h0], [h1], ...]}}}} + action_stats = metadata.statistics.action + print(f"DEBUG: action_stats keys = {list(action_stats.keys())}") + per_horizon_stats = {} + for action_key in action_stats: + stats_dict = action_stats[action_key].model_dump() + print(f"DEBUG: action_key={action_key}, stats_dict keys={list(stats_dict.keys())}") + # Check if stats are per-horizon (2D lists) by examining q01/q99 + if 'q01' in stats_dict: + q01_val = stats_dict['q01'] + print(f"DEBUG: q01 type={type(q01_val)}, value sample={q01_val[:2] if hasattr(q01_val, '__getitem__') else q01_val}") + # Handle both list and numpy array + is_2d = False + if isinstance(q01_val, (list, np.ndarray)) and len(q01_val) > 0: + first_elem = q01_val[0] + print(f"DEBUG: q01[0] type={type(first_elem)}") + if isinstance(first_elem, (list, np.ndarray)): + is_2d = True + + if is_2d: + # This is per-horizon stats (2D array) - convert to list if numpy + if isinstance(q01_val, np.ndarray): + for k in stats_dict: + if isinstance(stats_dict[k], np.ndarray): + stats_dict[k] = stats_dict[k].tolist() + per_horizon_stats[action_key] = stats_dict + print(f"DEBUG: Added {action_key} to per_horizon_stats") + + if per_horizon_stats: + print(f"Setting per-horizon statistics for keys: {list(per_horizon_stats.keys())}") + eval_transform.set_per_horizon_statistics(per_horizon_stats) + else: + print(f"WARNING: No per-horizon statistics found despite relative_action_per_horizon=True") + + eval_transform.eval() + self.eval_transform = eval_transform + + # 3. Load horizons needed + if self.embodiment_tag.value in train_cfg.modality_configs: + self.modality_configs = instantiate( + train_cfg.modality_configs[self.embodiment_tag.value] + ) + else: + self.modality_configs = instantiate(train_cfg.modality_configs) + + self._video_delta_indices = np.array(self.modality_configs.video.eval_delta_indices) + # self._video_delta_indices = np.array([0]) + # self.assert_delta_indices(self._video_delta_indices) + self._video_horizon = len(self._video_delta_indices) + + # We might not use `state`, which will be a vision-only policy + if "state" in self.modality_configs: + self._state_delta_indices = np.array(self.modality_configs.state.eval_delta_indices) + if not skip_assert_delta_indices: + self.assert_delta_indices(self._state_delta_indices) + self._state_horizon = len(self._state_delta_indices) + else: + self._state_horizon = None + self._state_delta_indices = None + self._raw_data_image_transform = None + + def offload_to_cpu(self): + """Offload the model to CPU to free GPU memory.""" + if hasattr(self.trained_model, "action_head") and hasattr(self.trained_model.action_head, "image_encoder"): + # For models with vram management, disable it and move to CPU + if hasattr(self.trained_model.action_head, 'disable_vram_management'): + self.trained_model.action_head.disable_vram_management() + + self.trained_model.to(device='cpu') + torch.cuda.empty_cache() + print(f"Model offloaded to CPU") + + def load_to_gpu(self): + """Load the model to GPU for inference.""" + print(f"Loading model to GPU...") + + # Move model to GPU + if hasattr(self.trained_model, "action_head") and hasattr(self.trained_model.action_head, "image_encoder"): + self.trained_model.action_head.enable_vram_management() + else: + self.trained_model.to(device=self.device) + + # Apply bf16 if needed + if self.eval_bf16: + self.trained_model = self.trained_model.to(dtype=torch.bfloat16) + + torch.cuda.empty_cache() + + def ensure_model_on_gpu(self): + """Ensure the model is loaded on GPU before inference.""" + # Check if model is on CPU + model_device = next(self.trained_model.parameters()).device + if model_device.type == 'cpu': + self.load_to_gpu() + + def assert_delta_indices(self, delta_indices: np.ndarray): + # All delta indices should be non-positive because there's no way to get the future observations + assert np.all(delta_indices <= 0), f"{delta_indices=}" + # The last delta index should be 0 because it doesn't make sense to not use the latest observation + assert delta_indices[-1] == 0, f"{delta_indices=}" + if len(delta_indices) > 1: + # The step is consistent + assert np.all( + np.diff(delta_indices) == delta_indices[1] - delta_indices[0] + ), f"{delta_indices=}" + # And the step is positive + assert (delta_indices[1] - delta_indices[0]) > 0, f"{delta_indices=}" + + def apply(self, batch: Batch, **kwargs) -> Batch: + """Normalize inputs""" + obs = batch.obs + + normalized_input = self.eval_transform(obs) + batch.normalized_obs = normalized_input + return batch + + def unapply(self, batch: Batch, obs: dict = None, **kwargs): + """Unnormalize actions and convert relative actions to absolute if needed""" + unnormalized_action = self.eval_transform.unapply( + dict(action=batch.normalized_action.cpu()) + ) + + # Check if relative_action is enabled and convert relative to absolute + relative_action = self.train_cfg.get('relative_action', False) + relative_action_per_horizon = self.train_cfg.get('relative_action_per_horizon', False) + relative_action_keys = self.train_cfg.get('relative_action_keys', []) + print("relative_action_per_horizon", relative_action_per_horizon) + if (relative_action or relative_action_per_horizon) and relative_action_keys and obs is not None: + for key in relative_action_keys: + action_key = f"action.{key}" + state_key = f"state.{key}" + + if action_key not in unnormalized_action: + continue + + # Try to find the state data - check multiple possible key formats + last_state = None + + + if last_state is None and state_key in obs: + # Format 1: Direct key like "state.joint_position" + last_state = obs[state_key] + elif last_state is None: + # Format 2: Search for keys containing both "state" and the key name + for obs_key in obs.keys(): + if 'state' in obs_key and key in obs_key: + last_state = obs[obs_key] + break + + # Format 3: If key is "joint_position" and obs has "state" key directly + # This handles cases where the observation uses modality-level keys + if last_state is None and 'state' in obs: + state_data = obs['state'] + # Check if the state data shape matches the action shape + action_dim = unnormalized_action[action_key].shape[-1] + if torch.is_tensor(state_data): + state_dim = state_data.shape[-1] + elif isinstance(state_data, np.ndarray): + state_dim = state_data.shape[-1] + else: + state_dim = None + + if state_dim == action_dim: + last_state = state_data + + if last_state is None: + continue + + if torch.is_tensor(last_state): + last_state = last_state.cpu().numpy() + + # Shape is (B, T, D) or (T, D), we want the last timestep + # After indexing: (B, D) or (D,) + if len(last_state.shape) >= 2: + last_state = last_state[..., -1, :] # Get the last timestep + + # Action shape is (horizon, D) or (B, horizon, D) + # Expand dims to broadcast: (D,) -> (1, D) or (B, D) -> (B, 1, D) + if len(unnormalized_action[action_key].shape) > len(last_state.shape): + last_state = np.expand_dims(last_state, axis=-2) # Add horizon dimension + + # Add state to relative action to get absolute action + print("last_state", last_state.shape, "unnormalized_action[action_key]", unnormalized_action[action_key].shape) + unnormalized_action[action_key] = unnormalized_action[action_key] + last_state + + batch.act = unnormalized_action + return batch + + def forward(self, batch, state=None, **kwargs): + + # 1. Check if input is batched and add batch dimension if needed + is_batched = self._check_state_is_batched(batch.obs) + if not is_batched: + batch.obs = unsqueeze_dict_values(batch.obs) + + # 2. Apply transforms/normalization + batch = self.apply(batch) + normalized_input = batch.normalized_obs + for k, v in normalized_input.items(): + if torch.is_tensor(v) and v.dtype == torch.float32 and self.eval_bf16: + normalized_input[k] = v.to(dtype=torch.bfloat16) + + # 3. Model inference + with torch.inference_mode(): + # with maybe_autocast: + model_pred = self.trained_model.get_action(normalized_input) + normalized_action = model_pred["action_pred"].float() + + # 4. Unnormalize actions (pass obs for relative action conversion) + original_obs = batch.obs + batch = self.unapply(Batch(normalized_action=normalized_action), obs=original_obs) + + # 5. Remove batch dimension if we added it + if not is_batched: + batch.act = squeeze_dict_values(batch.act) + return batch + + def joint_forward(self, batch, video=None, state=None, **kwargs): + # 0. Ensure model is on GPU + + # 1. Check if input is batched and add batch dimension if needed + is_batched = self._check_state_is_batched(batch.obs) + if not is_batched: + batch.obs = unsqueeze_dict_values(batch.obs) + + # 2. Apply transforms/normalization + batch = self.apply(batch) + normalized_input = batch.normalized_obs + + if isinstance(normalized_input, Batch): + normalized_input = normalized_input.__getstate__() + + if video is not None: + for key in normalized_input: + if 'images' in key: + print("key", key, normalized_input[key].shape) + normalized_input[key] = video + + for k, v in normalized_input.items(): + if torch.is_tensor(v) and v.dtype == torch.float32 and self.eval_bf16: + normalized_input[k] = v.to(dtype=torch.bfloat16) + + # 3. Model inference + with torch.inference_mode(): + # with maybe_autocast: + model_pred = self.trained_model.joint_video_action(normalized_input) + normalized_action = model_pred["action_pred"].float() + video_pred = model_pred["video_pred"] + + # 4. Unnormalize actions (pass obs for relative action conversion) + original_obs = batch.obs + batch = self.unapply(Batch(normalized_action=normalized_action), obs=original_obs) + + # 5. Remove batch dimension if we added it + if not is_batched: + batch.act = squeeze_dict_values(batch.act) + return batch, video_pred + + def lazy_joint_forward_causal(self, batch, video=None, latent_video=None, state=None, video_only=False, **kwargs): + + transform_start_time = time.perf_counter() + + # Save original observation before any modification (for relative action conversion) + original_obs_for_relative = {k: v.copy() if isinstance(v, np.ndarray) else v.clone() if torch.is_tensor(v) else v + for k, v in batch.obs.items()} + + # 1. Check if input is batched and add batch dimension if needed + is_batched = self._check_state_is_batched(batch.obs) + if not is_batched: + batch.obs = unsqueeze_dict_values(batch.obs) + # Also unsqueeze the saved original obs + original_obs_for_relative = unsqueeze_dict_values(original_obs_for_relative) + + # 2. Apply transforms/normalization + batch = self.apply(batch) + normalized_input = batch.normalized_obs + + transform_time = time.perf_counter() - transform_start_time + + if isinstance(normalized_input, Batch): + normalized_input = normalized_input.__getstate__() + + if video is not None: + for key in normalized_input: + if 'images' in key: + print("key", key, normalized_input[key].shape) + normalized_input[key] = video + + for k, v in normalized_input.items(): + if torch.is_tensor(v) and v.dtype == torch.float32 and self.eval_bf16: + normalized_input[k] = v.to(dtype=torch.bfloat16) + + model_start_time = time.perf_counter() + + # 3. Model inference + with torch.inference_mode(): + # with maybe_autocast: + model_pred = self.trained_model.lazy_joint_video_action_causal(normalized_input, latent_video=latent_video) + normalized_action = model_pred["action_pred"].float() + video_pred = model_pred["video_pred"] + + model_time = time.perf_counter() - model_start_time + + untransform_start_time = time.perf_counter() + + # 4. Unnormalize actions (pass obs for relative action conversion) + if not video_only: + batch = self.unapply(Batch(normalized_action=normalized_action), obs=original_obs_for_relative) + else: + batch = Batch(normalized_action=normalized_action) + + # 5. Remove batch dimension if we added it + if not is_batched: + batch.act = squeeze_dict_values(batch.act) + + untransform_time = time.perf_counter() - untransform_start_time + total_time = transform_time + model_time + untransform_time + + if self.rank == 0: + print(f"Inference Time: Total {total_time:.3f} seconds, " + f"Transform: {transform_time:.3f} seconds, " + f"Model: {model_time:.3f} seconds, " + f"Untransform: {untransform_time:.3f} seconds") + + return batch, video_pred + + def lazy_joint_forward_causal_gt_cond(self, batch, video=None, latent_video=None, state=None, **kwargs): + + # 1. Check if input is batched and add batch dimension if needed + is_batched = self._check_state_is_batched(batch.obs) + if not is_batched: + batch.obs = unsqueeze_dict_values(batch.obs) + + # 2. Apply transforms/normalization + batch = self.apply(batch) + normalized_input = batch.normalized_obs + + if isinstance(normalized_input, Batch): + normalized_input = normalized_input.__getstate__() + + if video is not None: + for key in normalized_input: + if 'images' in key: + print("key", key, normalized_input[key].shape) + normalized_input[key] = video + + for k, v in normalized_input.items(): + if torch.is_tensor(v) and v.dtype == torch.float32 and self.eval_bf16: + normalized_input[k] = v.to(dtype=torch.bfloat16) + + # 3. Model inference + with torch.inference_mode(): + # with maybe_autocast: + model_pred = self.trained_model.lazy_joint_video_action_causal_gt_cond(normalized_input, latent_video=latent_video) + normalized_action = model_pred["action_pred"].float() + video_pred = model_pred["video_pred"] + + # 4. Unnormalize actions (pass obs for relative action conversion) + original_obs = batch.obs + batch = self.unapply(Batch(normalized_action=normalized_action), obs=original_obs) + + # 5. Remove batch dimension if we added it + if not is_batched: + batch.act = squeeze_dict_values(batch.act) + return batch, video_pred + + def lazy_joint_forward(self, batch, video=None, state=None, **kwargs): + + # 1. Check if input is batched and add batch dimension if needed + is_batched = self._check_state_is_batched(batch.obs) + if not is_batched: + batch.obs = unsqueeze_dict_values(batch.obs) + + # 2. Apply transforms/normalization + batch = self.apply(batch) + normalized_input = batch.normalized_obs + + if isinstance(normalized_input, Batch): + normalized_input = normalized_input.__getstate__() + + if video is not None: + for key in normalized_input: + if 'images' in key: + print("key", key, normalized_input[key].shape) + normalized_input[key] = video + + for k, v in normalized_input.items(): + if torch.is_tensor(v) and v.dtype == torch.float32 and self.eval_bf16: + normalized_input[k] = v.to(dtype=torch.bfloat16) + + # 3. Model inference + with torch.inference_mode(): + # with maybe_autocast: + model_pred = self.trained_model.lazy_joint_video_action(normalized_input) + normalized_action = model_pred["action_pred"].float() + video_pred = model_pred["video_pred"] + + # 4. Unnormalize actions (pass obs for relative action conversion) + original_obs = batch.obs + batch = self.unapply(Batch(normalized_action=normalized_action), obs=original_obs) + + # 5. Remove batch dimension if we added it + if not is_batched: + batch.act = squeeze_dict_values(batch.act) + return batch, video_pred + + def lazy_joint_forward_efficient(self, batch, video=None, state=None, prompt_embs=None, prompt_emb_nega=None, **kwargs): + + # 1. Check if input is batched and add batch dimension if needed + is_batched = self._check_state_is_batched(batch.obs) + if not is_batched: + batch.obs = unsqueeze_dict_values(batch.obs) + + # 2. Apply transforms/normalization + batch = self.apply(batch) + normalized_input = batch.normalized_obs + + if isinstance(normalized_input, Batch): + normalized_input = normalized_input.__getstate__() + + if video is not None: + for key in normalized_input: + if 'images' in key: + print("key", key, normalized_input[key].shape) + normalized_input[key] = video + + for k, v in normalized_input.items(): + if torch.is_tensor(v) and v.dtype == torch.float32 and self.eval_bf16: + normalized_input[k] = v.to(dtype=torch.bfloat16) + + # 3. Model inference + with torch.inference_mode(): + # with maybe_autocast: + model_pred = self.trained_model.lazy_joint_video_action_efficient(normalized_input, prompt_embs=prompt_embs, prompt_emb_nega=prompt_emb_nega) + normalized_action = model_pred["action_pred"].float() + video_pred = model_pred["video_pred"] + + # 4. Unnormalize actions (pass obs for relative action conversion) + original_obs = batch.obs + batch = self.unapply(Batch(normalized_action=normalized_action), obs=original_obs) + + # 5. Remove batch dimension if we added it + if not is_batched: + batch.act = squeeze_dict_values(batch.act) + return batch, video_pred + + def language_forward(self, batch, state=None, **kwargs): + + # 1. Check if input is batched and add batch dimension if needed + is_batched = self._check_state_is_batched(batch.obs) + if not is_batched: + batch.obs = unsqueeze_dict_values(batch.obs) + + # 2. Apply transforms/normalization + batch = self.apply(batch) + normalized_input = batch.normalized_obs + + if isinstance(normalized_input, Batch): + normalized_input = normalized_input.__getstate__() + + for k, v in normalized_input.items(): + if torch.is_tensor(v) and v.dtype == torch.float32 and self.eval_bf16: + normalized_input[k] = v.to(dtype=torch.bfloat16) + + # 3. Model inference + with torch.inference_mode(): + # with maybe_autocast: + model_pred = self.trained_model.get_language(normalized_input) + # normalized_action = model_pred["action_pred"].float() + output = model_pred["output"] + # 4. Unnormalize actions + # batch = self.unapply(Batch(normalized_action=normalized_action)) + + # # 5. Remove batch dimension if we added it + # if not is_batched: + # batch.act = squeeze_dict_values(batch.act) + return output + # return batch + + + def gt_video_action_pred(self, batch, video=None, state=None, **kwargs): + + # 1. Check if input is batched and add batch dimension if needed + is_batched = self._check_state_is_batched(batch.obs) + if not is_batched: + batch.obs = unsqueeze_dict_values(batch.obs) + + # 2. Apply transforms/normalization + batch = self.apply(batch) + normalized_input = batch.normalized_obs + + if isinstance(normalized_input, Batch): + normalized_input = normalized_input.__getstate__() + + if video is not None: + for key in normalized_input: + if 'images' in key: + print("key", key, normalized_input[key].shape) + normalized_input[key] = video + + for k, v in normalized_input.items(): + if torch.is_tensor(v) and v.dtype == torch.float32 and self.eval_bf16: + normalized_input[k] = v.to(dtype=torch.bfloat16) + + # 3. Model inference + with torch.inference_mode(): + # with maybe_autocast: + model_pred = self.trained_model.gt_video_action_pred(normalized_input) + normalized_action = model_pred["action_pred"].float() + video_pred = model_pred["video_pred"] + + # 4. Unnormalize actions (pass obs for relative action conversion) + original_obs = batch.obs + batch = self.unapply(Batch(normalized_action=normalized_action), obs=original_obs) + + # 5. Remove batch dimension if we added it + if not is_batched: + batch.act = squeeze_dict_values(batch.act) + return batch, video_pred + + def _check_state_is_batched(self, obs: dict[str, Any]) -> bool: + for k, v in obs.items(): + if "state" in k and len(v.shape) < 3: # (B, Time, Dim) + return False + return True + + @property + def raw_data_image_transform(self) -> Callable: + """ + Get the raw data image transform for the policy + """ + return lambda x: x + + @property + def state_delta_indices(self) -> np.ndarray: + return self._state_delta_indices + + @property + def video_delta_indices(self) -> np.ndarray: + return self._video_delta_indices + + +def unsqueeze_dict_values(data: dict[str, Any]) -> dict[str, Any]: + """ + Unsqueeze the values of a dictionary. + This converts the data to be batched of size 1. + """ + unsqueezed_data = {} + for k, v in data.items(): + if isinstance(v, np.ndarray): + unsqueezed_data[k] = np.expand_dims(v, axis=0) + elif isinstance(v, list): + unsqueezed_data[k] = np.array(v) + elif isinstance(v, torch.Tensor): + unsqueezed_data[k] = v.unsqueeze(0) + elif isinstance(v, str): + unsqueezed_data[k] = np.array([v]) + else: + unsqueezed_data[k] = v + return unsqueezed_data + + +def squeeze_dict_values(data: dict[str, Any]) -> dict[str, Any]: + """ + Squeeze the values of a dictionary. This removes the batch dimension. + """ + squeezed_data = {} + for k, v in data.items(): + if isinstance(v, np.ndarray): + squeezed_data[k] = np.squeeze(v) + elif isinstance(v, torch.Tensor): + squeezed_data[k] = v.squeeze() + else: + squeezed_data[k] = v + return squeezed_data + + +def tree_get_leading_dim(data: dict[str, Any]) -> int: + """ + Get the batch size (leading dimension) of all tensors in the structure. + + Args: + data: Nested structure containing tensors + strict: If True, raises error if batch sizes don't match. + If False, returns None if inconsistent. + + Returns: + The common batch size, or None if inconsistent (when strict=False) + + Raises: + ValueError: If batch sizes are inconsistent (when strict=True) + ValueError: If no tensors found + """ + batch_sizes = [] + tensor_paths = [] + + def collect_batch_sizes(path, x): + if hasattr(x, "shape") and len(x.shape) > 0: + batch_sizes.append(x.shape[0]) + tensor_paths.append(".".join(map(str, path))) + + # Collect all batch sizes with their paths + tree.map_structure_with_path(collect_batch_sizes, data) + + if not batch_sizes: + raise ValueError("No tensors found in the data structure") + + # Check if all batch sizes are the same + first_batch_size = batch_sizes[0] + + if not all(bs == first_batch_size for bs in batch_sizes): + inconsistent_info = [f"{path}: {bs}" for path, bs in zip(tensor_paths, batch_sizes)] + error_msg = "Inconsistent batch sizes found:\n" + "\n".join(inconsistent_info) + + raise ValueError(error_msg) + + return first_batch_size + + +class GrootSimRLPolicy(BaseGrootSimPolicy): + """ + A class for Offline RL policies, + """ + + def __init__( + self, + embodiment_tag: EmbodimentTag, + model_path: str, + device: int | str, + q_fn_path: str | None = None, + n_action_samples: int = 1, + n_denoising_steps: int = 4, + gaussian_std: float = 0.0, + inference_batch_size: int = 1, + model_config_overrides: list[str] | None = [], + ): + """ + Implicit Policy Extraction (IPE) GrootSimPolicy. + """ + super().__init__(embodiment_tag=embodiment_tag, model_path=model_path, device=device) + + self.base_policy = GrootSimPolicy( + embodiment_tag=embodiment_tag, + model_path=str(model_path), + device=device, + ) + self.base_policy.trained_model.action_head.num_inference_timesteps = n_denoising_steps + + if q_fn_path is None: + if n_action_samples > 1: + raise ValueError("If n_action_samples > 1, q_fn_path should be provided.") + self.qfn = None + else: + assert n_action_samples > 1, "If q_fn_path is provided, n_action_samples should be > 1." + self.qfn = GrootSimPolicy( + embodiment_tag=EmbodimentTag.GR1_UNIFIED_OFFLINE_RL, + model_path=q_fn_path, + device=device, + skip_assert_delta_indices=True, + ) + + self.n_action_samples = n_action_samples + self.gaussian_std = gaussian_std + self.inference_batch_size = inference_batch_size + + def batch_obs(self, obs: dict[str, Any]) -> tuple[dict[str, Any], bool]: + """ + Check if input is batched and add batch dimension if needed + """ + is_batched = self.base_policy._check_state_is_batched(obs) + if not is_batched: + obs = unsqueeze_dict_values(obs) + + return obs, is_batched + + def add_action_noise(self, action: torch.Tensor, std: float) -> torch.Tensor: + if std <= 0.0: + return action + noise = torch.randn_like(action) * std + noisy_action = action + noise + return noisy_action + + def sample_actions_from_base_policy( + self, + batch: Batch, + n_action_samples: int, # number of actions to sample PER item in the batch + inference_batch_size: int = 1, + ) -> torch.Tensor: + # a) Apply transforms/normalization + batch = self.base_policy.apply(batch) + normalized_input = batch.normalized_obs + for k, v in normalized_input.items(): + if torch.is_tensor(v) and v.dtype == torch.float32 and self.base_policy.eval_bf16: + normalized_input[k] = v.to(dtype=torch.bfloat16) + + with torch.inference_mode(): + model_pred = self.base_policy.trained_model.get_action( + normalized_input, + num_action_samples=n_action_samples, + inference_batch_size=inference_batch_size, + validate=False, + ) + normalized_action = model_pred["action_pred"] + normalized_action = self.add_action_noise(normalized_action, self.gaussian_std) + return normalized_action # [num_envs, n_action_samples, horizon, action_dim] + + def get_action_values( + self, batch: Batch, actions: torch.Tensor, inference_batch_size: int = 1 + ) -> torch.Tensor: + assert self.qfn is not None, "Q-function must be provided to evaluate actions." + + # a) Apply transforms/normalization + batch = self.qfn.apply(batch) + normalized_input = batch.normalized_obs + for k, v in normalized_input.items(): + if torch.is_tensor(v) and v.dtype == torch.float32 and self.qfn.eval_bf16: + normalized_input[k] = v.to(dtype=torch.bfloat16) + + with torch.inference_mode(): + # [num_envs, ...] + backbone_inputs, action_inputs = self.qfn.trained_model.prepare_input(normalized_input) + backbone_outputs = self.qfn.trained_model.backbone(backbone_inputs) + q_values = self.qfn.trained_model.action_head.get_q_pred( + backbone_outputs, + action_inputs, + actions, + inference_batch_size=inference_batch_size, + ) + + return q_values + + def forward(self, batch, state=None, **kwargs): + orig_obs = batch.obs.copy() + orig_obs, is_batched = self.batch_obs(orig_obs) + + # Get number of environments + num_envs = 1 if not is_batched else tree_get_leading_dim(orig_obs) + + # 1/ Sample normalized actions from the base policy + actions = self.sample_actions_from_base_policy( + Batch(obs=orig_obs), + n_action_samples=self.n_action_samples, + inference_batch_size=self.inference_batch_size, + ) # [num_envs, n_action_samples, horizon, action_dim] + + # 2/ Evaluate the sampled actions with the Q-function + if self.qfn is not None: + q_values = self.get_action_values( + Batch(obs=orig_obs), + actions=actions, + inference_batch_size=self.inference_batch_size, + ) + + best_indices = torch.argmax(q_values, dim=1) # [num_envs] + normalized_action = actions[ + torch.arange(num_envs), best_indices + ].float() # [num_envs, horizon, action_dim] + else: + normalized_action = actions.float() + + batch = self.base_policy.unapply(Batch(normalized_action=normalized_action), obs=orig_obs) + + if not is_batched: + batch.act = squeeze_dict_values(batch.act) + return batch + + @property + def video_delta_indices(self) -> np.ndarray: + return self.base_policy.video_delta_indices + + @property + def state_delta_indices(self) -> np.ndarray: + return self.base_policy.state_delta_indices + + @property + def raw_data_image_transform(self) -> Callable: + return lambda x: x diff --git a/groot/vla/model/trex_track_force/__init__.py b/groot/vla/model/trex_track_force/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8929830f8dde02bcc872cfc3d3fe84e08b96fa9d --- /dev/null +++ b/groot/vla/model/trex_track_force/__init__.py @@ -0,0 +1,72 @@ +"""Independent T-Rex track/force model package.""" + +from .attention import ( + BlockTokenLayout, + CausalWanSelfAttention, + TokenLayout, + TokenType, + build_dense_allow_matrix, +) +from .blocks import CausalWanAttentionBlock, CausalWanTrackForceBlock +from .force import ( + FORCE_OFFSETS, + ForceOnlyOutput, + ForceOnlyTransformer, + ForceTactileTransformer, + PerFingerForceVQ, + TwoStageFlowSchedule, + pad_action_62_to_64, +) +from .model import CausalWanModel, CausalWanTrackForceModel +from .policy import ( + WANPolicyHead, + WANPolicyHeadConfig, + WANTrackForcePolicyHead, + WANTrackForcePolicyHeadConfig, +) +from .track import ( + NUM_TRACK_POINTS, + TRACK_SEGMENTS, + TrackDecoder, + TrackEmbedding, + TrackEncoder, + track_motion_weights, + visibility_masked_track_loss, + visibility_masked_track_smoothness_loss, + zero_invisible_track_xy, +) +from .vla import TrexTrackForceVLA, TrexTrackForceVLAConfig + +__all__ = [ + "BlockTokenLayout", + "CausalWanAttentionBlock", + "CausalWanModel", + "CausalWanSelfAttention", + "CausalWanTrackForceBlock", + "CausalWanTrackForceModel", + "FORCE_OFFSETS", + "ForceOnlyOutput", + "ForceOnlyTransformer", + "ForceTactileTransformer", + "NUM_TRACK_POINTS", + "PerFingerForceVQ", + "TRACK_SEGMENTS", + "TokenLayout", + "TokenType", + "TrackDecoder", + "TrackEmbedding", + "TrackEncoder", + "TrexTrackForceVLA", + "TrexTrackForceVLAConfig", + "TwoStageFlowSchedule", + "WANPolicyHead", + "WANPolicyHeadConfig", + "WANTrackForcePolicyHead", + "WANTrackForcePolicyHeadConfig", + "build_dense_allow_matrix", + "pad_action_62_to_64", + "track_motion_weights", + "visibility_masked_track_loss", + "visibility_masked_track_smoothness_loss", + "zero_invisible_track_xy", +] diff --git a/groot/vla/model/trex_track_force/attention.py b/groot/vla/model/trex_track_force/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..ea7ae7e1449fdac250c841ed2c8f025938f3e794 --- /dev/null +++ b/groot/vla/model/trex_track_force/attention.py @@ -0,0 +1,402 @@ +"""Block-causal attention for the T-Rex track/force model. + +The public :func:`build_dense_allow_matrix` intentionally materializes the +mask. The production attention path can later replace it with FlexAttention, +but keeping a dense reference implementation makes the information-flow +contract exact and easy to unit test. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import IntEnum +from typing import Iterable + +import torch +from torch import nn +from torch.nn import functional as F + + +class TokenType(IntEnum): + """Token roles in the packed self-attention sequence.""" + + CONDITIONING_OBS = 0 + OBS = 1 + ACTION = 2 + STATE = 3 + TRACK_PAST = 4 + TRACK_FUTURE = 5 + + +@dataclass(frozen=True) +class BlockTokenLayout: + """Description of the packed sequence consumed by every causal WAN block. + + Sequence order is:: + + [conditioning_obs] + [obs_0, action_0, state_0, track_past_0, track_future_0] + ... + [obs_n, action_n, state_n, track_past_n, track_future_n] + + Interleaving by block keeps the cross-block causal rule explicit while + preserving bidirectional attention inside the current obs/action block. + """ + + num_blocks: int + obs_tokens_per_block: int + action_tokens_per_block: int = 16 + state_tokens_per_block: int = 1 + track_points: int = 250 + conditioning_obs_tokens: int = 0 + + def __post_init__(self) -> None: + values = { + "num_blocks": self.num_blocks, + "obs_tokens_per_block": self.obs_tokens_per_block, + "action_tokens_per_block": self.action_tokens_per_block, + "state_tokens_per_block": self.state_tokens_per_block, + "track_points": self.track_points, + "conditioning_obs_tokens": self.conditioning_obs_tokens, + } + if self.num_blocks <= 0: + raise ValueError("num_blocks must be positive") + if self.obs_tokens_per_block <= 0: + raise ValueError("obs_tokens_per_block must be positive") + if any(value < 0 for value in values.values()): + raise ValueError(f"layout sizes must be non-negative: {values}") + + @property + def tokens_per_block(self) -> int: + return ( + self.obs_tokens_per_block + + self.action_tokens_per_block + + self.state_tokens_per_block + + 2 * self.track_points + ) + + @property + def sequence_length(self) -> int: + return self.conditioning_obs_tokens + self.num_blocks * self.tokens_per_block + + @property + def conditioning_slice(self) -> slice: + return slice(0, self.conditioning_obs_tokens) + + def block_slice(self, block_index: int) -> slice: + self._validate_block(block_index) + start = self.conditioning_obs_tokens + block_index * self.tokens_per_block + return slice(start, start + self.tokens_per_block) + + def slices(self, block_index: int) -> dict[TokenType, slice]: + """Return the five modality slices for one block.""" + + block = self.block_slice(block_index) + cursor = block.start + result: dict[TokenType, slice] = {} + sizes = ( + (TokenType.OBS, self.obs_tokens_per_block), + (TokenType.ACTION, self.action_tokens_per_block), + (TokenType.STATE, self.state_tokens_per_block), + (TokenType.TRACK_PAST, self.track_points), + (TokenType.TRACK_FUTURE, self.track_points), + ) + for token_type, size in sizes: + result[token_type] = slice(cursor, cursor + size) + cursor += size + return result + + def indices( + self, + token_types: TokenType | Iterable[TokenType], + *, + device: torch.device | str | None = None, + ) -> torch.Tensor: + """Return packed indices for one or more token types.""" + + if isinstance(token_types, TokenType): + token_types = (token_types,) + requested = set(token_types) + indices: list[torch.Tensor] = [] + if TokenType.CONDITIONING_OBS in requested and self.conditioning_obs_tokens: + indices.append(torch.arange(self.conditioning_obs_tokens, device=device)) + for block_index in range(self.num_blocks): + block_slices = self.slices(block_index) + for token_type in requested: + if token_type is TokenType.CONDITIONING_OBS: + continue + token_slice = block_slices[token_type] + indices.append(torch.arange(token_slice.start, token_slice.stop, device=device)) + if not indices: + return torch.empty(0, dtype=torch.long, device=device) + return torch.cat(indices).long() + + def token_metadata( + self, *, device: torch.device | str | None = None + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return ``(token_type, block_id)`` vectors for the packed sequence. + + Conditioning observations use block id ``-1``. + """ + + token_type = torch.empty(self.sequence_length, dtype=torch.long, device=device) + block_id = torch.full( + (self.sequence_length,), -1, dtype=torch.long, device=device + ) + if self.conditioning_obs_tokens: + token_type[self.conditioning_slice] = int(TokenType.CONDITIONING_OBS) + for block_index in range(self.num_blocks): + for kind, token_slice in self.slices(block_index).items(): + token_type[token_slice] = int(kind) + block_id[token_slice] = block_index + return token_type, block_id + + def _validate_block(self, block_index: int) -> None: + if not 0 <= block_index < self.num_blocks: + raise IndexError( + f"block index {block_index} outside [0, {self.num_blocks})" + ) + + +# Backwards-friendly short name used by tests and downstream integrations. +TokenLayout = BlockTokenLayout + + +def build_dense_allow_matrix( + layout: BlockTokenLayout | None = None, + *, + num_blocks: int | None = None, + obs_tokens_per_block: int | None = None, + action_tokens_per_block: int = 16, + state_tokens_per_block: int = 1, + track_points: int = 250, + conditioning_obs_tokens: int = 0, + max_previous_blocks: int | None = None, + device: torch.device | str | None = None, +) -> torch.Tensor: + """Build the exact boolean self-attention allow matrix. + + ``True`` means that the row query may read the column key. + + Contract: + + * conditioning observations only read conditioning observations; + * obs/action are bidirectional within a block, read same-block state and + track, and read obs/action/track from earlier blocks; + * state queries only read state tokens in their own block; + * past-track queries read track history from earlier blocks and current + past-track tokens, but not the current future-track target; + * future-track queries read only past/future track tokens from current and + earlier blocks; + * no query can read a later block; ``max_previous_blocks`` optionally + applies DreamZero's finite autoregressive memory window. + + In particular, state and obs/action information is never exposed to a + track query. + """ + + if layout is None: + if num_blocks is None or obs_tokens_per_block is None: + raise TypeError( + "pass layout or both num_blocks and obs_tokens_per_block" + ) + layout = BlockTokenLayout( + num_blocks=num_blocks, + obs_tokens_per_block=obs_tokens_per_block, + action_tokens_per_block=action_tokens_per_block, + state_tokens_per_block=state_tokens_per_block, + track_points=track_points, + conditioning_obs_tokens=conditioning_obs_tokens, + ) + + if max_previous_blocks is not None and max_previous_blocks < 0: + raise ValueError("max_previous_blocks must be non-negative or None") + + allow = torch.zeros( + layout.sequence_length, + layout.sequence_length, + dtype=torch.bool, + device=device, + ) + + # The clean conditioning image is isolated as a query, but is readable by + # every obs/action block below. + if layout.conditioning_obs_tokens: + allow[layout.conditioning_slice, layout.conditioning_slice] = True + + for query_block in range(layout.num_blocks): + query_slices = layout.slices(query_block) + first_key_block = ( + 0 + if max_previous_blocks is None + else max(0, query_block - max_previous_blocks) + ) + + obs_action_queries = torch.cat( + ( + torch.arange( + query_slices[TokenType.OBS].start, + query_slices[TokenType.OBS].stop, + device=device, + ), + torch.arange( + query_slices[TokenType.ACTION].start, + query_slices[TokenType.ACTION].stop, + device=device, + ), + ) + ).long() + + if layout.conditioning_obs_tokens: + allow[obs_action_queries, layout.conditioning_slice] = True + + for key_block in range(first_key_block, query_block + 1): + key_slices = layout.slices(key_block) + # Obs/action can read all causal obs/action and track context. + for key_type in ( + TokenType.OBS, + TokenType.ACTION, + TokenType.TRACK_PAST, + TokenType.TRACK_FUTURE, + ): + allow[obs_action_queries, key_slices[key_type]] = True + + # State is a same-block condition, never a cross-block memory. + if key_block == query_block: + allow[obs_action_queries, key_slices[TokenType.STATE]] = True + + # State is a condition token, not a carrier of multimodal information. + state_slice = query_slices[TokenType.STATE] + allow[state_slice, state_slice] = True + + past_query_slice = query_slices[TokenType.TRACK_PAST] + future_query_slice = query_slices[TokenType.TRACK_FUTURE] + for key_block in range(first_key_block, query_block + 1): + key_slices = layout.slices(key_block) + # Prior block tracks are causal memory for both track roles. + if key_block < query_block: + for key_type in (TokenType.TRACK_PAST, TokenType.TRACK_FUTURE): + allow[past_query_slice, key_slices[key_type]] = True + allow[future_query_slice, key_slices[key_type]] = True + else: + # Clean past cannot consume its own noisy future target. + allow[past_query_slice, key_slices[TokenType.TRACK_PAST]] = True + allow[future_query_slice, key_slices[TokenType.TRACK_PAST]] = True + allow[future_query_slice, key_slices[TokenType.TRACK_FUTURE]] = True + + if not bool(allow.any(dim=-1).all()): + raise RuntimeError("attention layout produced a query with no allowed key") + return allow + + +def allow_matrix_to_additive_mask( + allow: torch.Tensor, + *, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Convert a boolean allow matrix to an additive ``0/-inf`` mask.""" + + if allow.dtype is not torch.bool: + raise TypeError(f"allow matrix must be bool, got {allow.dtype}") + mask = torch.zeros(allow.shape, dtype=dtype, device=allow.device) + return mask.masked_fill(~allow, torch.finfo(dtype).min) + + +class CausalWanSelfAttention(nn.Module): + """Dense reference implementation of the new CausalWan self-attention.""" + + def __init__( + self, + dim: int, + num_heads: int, + *, + dropout: float = 0.0, + qk_norm: bool = True, + ) -> None: + super().__init__() + if dim % num_heads: + raise ValueError(f"dim={dim} must be divisible by num_heads={num_heads}") + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.dropout = dropout + + # Keep WAN/PEFT-friendly projection names. + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = nn.RMSNorm(self.head_dim) if qk_norm else nn.Identity() + self.norm_k = nn.RMSNorm(self.head_dim) if qk_norm else nn.Identity() + + def forward( + self, + x: torch.Tensor, + layout: BlockTokenLayout | None = None, + *, + allow_matrix: torch.Tensor | None = None, + return_attention: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if x.ndim != 3: + raise ValueError(f"x must have shape [B,L,D], got {tuple(x.shape)}") + batch_size, seq_len, _ = x.shape + if allow_matrix is None: + if layout is None: + raise TypeError("layout is required when allow_matrix is omitted") + allow_matrix = build_dense_allow_matrix(layout, device=x.device) + if allow_matrix.shape != (seq_len, seq_len): + raise ValueError( + f"allow matrix {tuple(allow_matrix.shape)} does not match L={seq_len}" + ) + allow_matrix = allow_matrix.to(device=x.device, dtype=torch.bool) + + def project(layer: nn.Linear, values: torch.Tensor) -> torch.Tensor: + return ( + layer(values) + .view(batch_size, seq_len, self.num_heads, self.head_dim) + .transpose(1, 2) + ) + + query = self.norm_q(project(self.q, x)) + key = self.norm_k(project(self.k, x)) + value = project(self.v, x) + expanded_allow = allow_matrix.view(1, 1, seq_len, seq_len) + + if return_attention: + scale = self.head_dim**-0.5 + scores = torch.matmul(query, key.transpose(-1, -2)) * scale + scores = scores.masked_fill(~expanded_allow, torch.finfo(scores.dtype).min) + attention = scores.softmax(dim=-1) + if self.training and self.dropout: + attention = F.dropout(attention, p=self.dropout) + output = torch.matmul(attention, value) + else: + output = F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=expanded_allow, + dropout_p=self.dropout if self.training else 0.0, + ) + attention = None + + output = output.transpose(1, 2).reshape(batch_size, seq_len, self.dim) + output = self.o(output) + if return_attention: + assert attention is not None + return output, attention + return output + + +TrackForceCausalSelfAttention = CausalWanSelfAttention + + +__all__ = [ + "BlockTokenLayout", + "CausalWanSelfAttention", + "TokenLayout", + "TokenType", + "TrackForceCausalSelfAttention", + "allow_matrix_to_additive_mask", + "build_dense_allow_matrix", +] diff --git a/groot/vla/model/trex_track_force/blocks.py b/groot/vla/model/trex_track_force/blocks.py new file mode 100644 index 0000000000000000000000000000000000000000..e021ed90a8374130c51b23a1eafa348722c399f0 --- /dev/null +++ b/groot/vla/model/trex_track_force/blocks.py @@ -0,0 +1,214 @@ +"""Wan2.2-compatible blocks with track-isolated block-causal attention.""" + +from __future__ import annotations + +import torch +from torch.nn import functional as F + +from groot.vla.model.dreamzero.modules.wan_video_dit_action_casual_chunk import ( + CausalWanAttentionBlock as WanCausalAttentionBlock, + CausalWanSelfAttention as WanCausalSelfAttention, +) +from groot.vla.model.dreamzero.modules.wan2_1_submodule import ( + WanT2VCrossAttention, +) + +from .attention import BlockTokenLayout, TokenType, build_dense_allow_matrix + + +def apply_multimodal_rope( + values: torch.Tensor, frequencies: torch.Tensor +) -> torch.Tensor: + """Apply pre-packed video/action/state/track RoPE to ``[B,L,H,D]``.""" + + if values.ndim != 4 or frequencies.shape[0] != values.shape[1]: + raise ValueError( + f"RoPE values/frequencies mismatch: {tuple(values.shape)}, " + f"{tuple(frequencies.shape)}" + ) + batch, length, heads, head_dim = values.shape + if frequencies.is_complex(): + complex_values = torch.view_as_complex( + values.to(torch.float64).reshape(batch, length, heads, -1, 2) + ) + rotated = complex_values * frequencies.view(1, length, 1, -1) + return torch.view_as_real(rotated).flatten(3).type_as(values) + + # TensorRT/non-polar representation stores [cos, sin] pairs. + pairs = values.reshape(batch, length, heads, -1, 2) + frequency_pairs = frequencies.reshape(1, length, 1, -1, 2).to(values.dtype) + real, imag = pairs[..., 0], pairs[..., 1] + cosine, sine = frequency_pairs[..., 0], frequency_pairs[..., 1] + rotated = torch.stack( + (real * cosine - imag * sine, real * sine + imag * cosine), dim=-1 + ) + if rotated.shape[-2] * 2 != head_dim: + raise RuntimeError("non-polar RoPE width does not match attention head") + return rotated.flatten(3) + + +class TrackForceWanSelfAttention(WanCausalSelfAttention): + """Original WAN q/k/v/o and norms with the new exact dense topology.""" + + def forward( + self, + x: torch.Tensor, + *, + layout: BlockTokenLayout, + rope_frequencies: torch.Tensor, + allow_matrix: torch.Tensor | None = None, + **_: object, + ) -> tuple[torch.Tensor, None]: + batch, length = x.shape[:2] + heads, head_dim = self.num_heads, self.head_dim + query = self.norm_q(self.q(x)).view(batch, length, heads, head_dim) + key = self.norm_k(self.k(x)).view(batch, length, heads, head_dim) + value = self.v(x).view(batch, length, heads, head_dim) + query = apply_multimodal_rope(query, rope_frequencies).type_as(value) + key = apply_multimodal_rope(key, rope_frequencies).type_as(value) + + if allow_matrix is None: + allow_matrix = build_dense_allow_matrix(layout, device=x.device) + if allow_matrix.shape != (length, length): + raise ValueError("attention allow matrix does not match packed sequence") + output = F.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + attn_mask=allow_matrix.view(1, 1, length, length), + ) + output = output.transpose(1, 2).reshape(batch, length, self.dim) + return self.o(output), None + + +class CausalWanTrackForceBlock(WanCausalAttentionBlock): + """Parameter-key-compatible WAN block with track-safe query routing.""" + + def __init__( + self, + cross_attn_type: str, + dim: int, + ffn_dim: int, + num_heads: int, + frame_seqlen: int, + local_attn_size: int = -1, + sink_size: int = 0, + num_frame_per_block: int = 1, + qk_norm: bool = True, + cross_attn_norm: bool = False, + eps: float = 1e-6, + num_action_per_block: int = 32, + num_state_per_block: int = 1, + ) -> None: + super().__init__( + cross_attn_type, + dim, + ffn_dim, + num_heads, + frame_seqlen, + local_attn_size, + sink_size, + num_frame_per_block, + qk_norm, + cross_attn_norm, + eps, + num_action_per_block, + num_state_per_block, + ) + original_attention = self.self_attn + self.self_attn = TrackForceWanSelfAttention( + dim=dim, + num_heads=num_heads, + frame_seqlen=frame_seqlen, + local_attn_size=local_attn_size, + sink_size=sink_size, + num_frame_per_block=num_frame_per_block, + qk_norm=qk_norm, + eps=eps, + num_action_per_block=num_action_per_block, + num_state_per_block=num_state_per_block, + ) + # This is strict by construction: q/k/v/o, RMSNorms, and parameter + # shapes are inherited unchanged from the original WAN attention. + self.self_attn.load_state_dict(original_attention.state_dict(), strict=True) + + def forward( + self, + x: torch.Tensor, + *, + e: torch.Tensor, + layout: BlockTokenLayout, + rope_frequencies: torch.Tensor, + context: torch.Tensor | None, + context_mask: torch.Tensor | None = None, + allow_matrix: torch.Tensor | None = None, + ) -> torch.Tensor: + modulation = (self.modulation.unsqueeze(1) + e).chunk(6, dim=2) + shift_attn, scale_attn, gate_attn, shift_ffn, scale_ffn, gate_ffn = ( + part.squeeze(2) for part in modulation + ) + + attention_input = self.norm1(x) * (1 + scale_attn) + shift_attn + attended, _ = self.self_attn( + attention_input, + layout=layout, + rope_frequencies=rope_frequencies, + allow_matrix=allow_matrix, + ) + x = x + attended * gate_attn + + # Track queries never enter WAN cross-attention. FFN below is strictly + # point-wise over the token axis, so it cannot introduce cross-token + # leakage either. + if context is not None: + if context_mask is not None: + if context_mask.shape != context.shape[:2]: + raise ValueError("context_mask must be [B,L_context]") + context = context * context_mask.to( + device=context.device, dtype=context.dtype + ).unsqueeze(-1) + query_indices = layout.indices( + ( + TokenType.CONDITIONING_OBS, + TokenType.OBS, + TokenType.ACTION, + ), + device=x.device, + ) + selected = x.index_select(1, query_indices) + normalized = self.norm3(selected) + if isinstance(self.cross_attn, WanT2VCrossAttention): + context_lengths = ( + context_mask.sum(dim=1) + if context_mask is not None + else torch.full( + (context.shape[0],), + context.shape[1], + dtype=torch.long, + device=context.device, + ) + ) + cross_output = self.cross_attn( + normalized, context, context_lengths + ) + else: + cross_output = self.cross_attn(normalized, context) + selected = selected + cross_output + x = x.clone() + x[:, query_indices] = selected + + feed_forward = self.ffn( + self.norm2(x) * (1 + scale_ffn) + shift_ffn + ) + return x + feed_forward * gate_ffn + + +CausalWanAttentionBlock = CausalWanTrackForceBlock + + +__all__ = [ + "CausalWanAttentionBlock", + "CausalWanTrackForceBlock", + "TrackForceWanSelfAttention", + "apply_multimodal_rope", +] diff --git a/groot/vla/model/trex_track_force/dataset.py b/groot/vla/model/trex_track_force/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..b366ed6e15a0e5e98551f7c2ea924900caebb969 --- /dev/null +++ b/groot/vla/model/trex_track_force/dataset.py @@ -0,0 +1,520 @@ +"""Dedicated LeRobot-v2 loader for 20 Hz T-Rex Track-Force training. + +The source dataset stays at its original frame rate. This loader creates a +deterministic 20 Hz logical control grid from timestamps, then gathers: + +* four autoregressive blocks with 16 delta-base EEF actions each; +* 16 past and 16 target 250-point tracks per block, with target step zero + equal to the current anchor frame; +* one conditioning RGB frame plus 8 future frames per block at 10 Hz; +* four 5 Hz force refresh slots per action block. + +Only the three RGB camera keys configured for this model are decoded. Tactile +deformation videos are never requested. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable + +import numpy as np +import pandas as pd + +from groot.vla.data.dataset.lerobot_sharded import ( + ShardedLeRobotSubLangSingleActionChunkDatasetDROID, +) + +from .force import ( + ACTION_HORIZON, + FORCE_HISTORY_FRAMES, + FORCE_OFFSETS, + FORCE_SENSOR_COUNT, + FORCE_SENSOR_DIM, +) +from .track import NUM_TRACK_POINTS, TRACK_HORIZON + + +ACTION_RATE_HZ = 20.0 +TACTILE_RATE_HZ = 5.0 +VIDEO_RATE_HZ = 10.0 +AR_BLOCKS = 4 +VIDEO_FRAMES_PER_BLOCK = 8 +VIDEO_HISTORY_FRAMES = 1 +VIDEO_FUTURE_FRAMES = AR_BLOCKS * VIDEO_FRAMES_PER_BLOCK + +TRACK_XY_COLUMN = "observation.track_xy" +TRACK_VISIBILITY_COLUMN = "observation.track_visibility" +FORCE_COLUMN = "observation.tactile_force" + + +@dataclass(frozen=True) +class TimestampSelection: + indices: np.ndarray + target_timestamps: np.ndarray + padding_mask: np.ndarray + alignment_error: np.ndarray + + +def nearest_timestamp_indices( + source_timestamps: np.ndarray, + target_timestamps: np.ndarray, + *, + tie_tolerance: float = 1e-9, +) -> TimestampSelection: + """Map target times to source rows, choosing the earlier row on a tie.""" + + source = np.asarray(source_timestamps, dtype=np.float64) + target = np.asarray(target_timestamps, dtype=np.float64) + if source.ndim != 1 or source.size < 2: + raise ValueError("source timestamps must be one-dimensional with at least two rows") + if target.ndim != 1 or target.size == 0: + raise ValueError("target timestamps must be a non-empty one-dimensional array") + if not np.isfinite(source).all() or not np.isfinite(target).all(): + raise ValueError("timestamps must be finite") + if not np.all(np.diff(source) > 0): + raise ValueError("source timestamps must be strictly increasing") + if not np.all(np.diff(target) > 0): + raise ValueError("target timestamps must be strictly increasing") + + insertion = np.searchsorted(source, target, side="left") + lower = np.clip(insertion - 1, 0, source.size - 1) + upper = np.clip(insertion, 0, source.size - 1) + lower_error = np.abs(target - source[lower]) + upper_error = np.abs(source[upper] - target) + choose_upper = upper_error < lower_error - tie_tolerance + indices = np.where(choose_upper, upper, lower).astype(np.int64) + padding = (target < source[0]) | (target > source[-1]) + indices[target < source[0]] = 0 + indices[target > source[-1]] = source.size - 1 + return TimestampSelection( + indices=indices, + target_timestamps=target, + padding_mask=padding, + alignment_error=np.abs(source[indices] - target), + ) + + +def uniform_target_times( + anchor_timestamp: float, + offsets: Iterable[int], + rate_hz: float, +) -> np.ndarray: + offsets_array = np.asarray(tuple(offsets), dtype=np.float64) + if offsets_array.ndim != 1 or offsets_array.size == 0: + raise ValueError("offsets must be a non-empty sequence") + if not np.all(np.diff(offsets_array) > 0): + raise ValueError("offsets must be strictly increasing") + if not np.isfinite(rate_hz) or rate_hz <= 0: + raise ValueError("rate_hz must be positive") + return float(anchor_timestamp) + offsets_array / float(rate_hz) + + +def rotation_6d_to_matrix(rotation_6d: np.ndarray) -> np.ndarray: + """Convert matrix-column 6D rotations to proper rotation matrices.""" + + values = np.asarray(rotation_6d, dtype=np.float64) + if values.shape[-1] != 6: + raise ValueError("rotation_6d must end in six values") + first = values[..., :3] + first /= np.linalg.norm(first, axis=-1, keepdims=True).clip(min=1e-8) + second = values[..., 3:6] + second = second - np.sum(first * second, axis=-1, keepdims=True) * first + second /= np.linalg.norm(second, axis=-1, keepdims=True).clip(min=1e-8) + third = np.cross(first, second) + return np.stack((first, second, third), axis=-1) + + +def eef62_delta_base(reference_state: np.ndarray, absolute_targets: np.ndarray) -> np.ndarray: + """Apply T-Rex delta-base pose math; hand targets remain absolute.""" + + reference = np.asarray(reference_state, dtype=np.float64) + targets = np.asarray(absolute_targets, dtype=np.float64) + if reference.shape != (62,) or targets.shape[-1] != 62: + raise ValueError("expected reference [62] and targets [...,62]") + + output = np.empty_like(targets, dtype=np.float64) + for pose_slice, hand_slice in ((slice(0, 9), slice(9, 31)), (slice(31, 40), slice(40, 62))): + reference_pose = reference[pose_slice] + target_pose = targets[..., pose_slice] + reference_rotation = rotation_6d_to_matrix(reference_pose[3:9]) + target_rotation = rotation_6d_to_matrix(target_pose[..., 3:9]) + translation_delta = np.einsum( + "ji,...j->...i", + reference_rotation, + target_pose[..., :3] - reference_pose[:3], + ) + rotation_delta = np.einsum( + "ji,...jk->...ik", reference_rotation, target_rotation + ) + output[..., pose_slice] = np.concatenate( + ( + translation_delta, + rotation_delta[..., :, 0], + rotation_delta[..., :, 1], + ), + axis=-1, + ) + output[..., hand_slice] = targets[..., hand_slice] + return output.astype(np.float32) + + +class TrexTrackForceShardedDataset( + ShardedLeRobotSubLangSingleActionChunkDatasetDROID +): + """Sharded dataset with an isolated timestamp-driven T-Rex sampling path.""" + + def __init__( + self, + *args, + action_rate_hz: float = ACTION_RATE_HZ, + tactile_rate_hz: float = TACTILE_RATE_HZ, + video_rate_hz: float = VIDEO_RATE_HZ, + max_training_anchors: int | None = None, + pin_anchor_rank: int | None = None, + **kwargs, + ) -> None: + self.action_rate_hz = float(action_rate_hz) + self.tactile_rate_hz = float(tactile_rate_hz) + self.video_rate_hz = float(video_rate_hz) + if max_training_anchors is not None and int(max_training_anchors) <= 0: + raise ValueError("max_training_anchors must be positive when set") + if pin_anchor_rank is not None and int(pin_anchor_rank) < 0: + raise ValueError("pin_anchor_rank must be non-negative when set") + self.max_training_anchors = ( + None if max_training_anchors is None else int(max_training_anchors) + ) + self.pin_anchor_rank = ( + None if pin_anchor_rank is None else int(pin_anchor_rank) + ) + configured_blocks = kwargs.get("max_chunk_size", AR_BLOCKS) + self.num_ar_blocks = ( + AR_BLOCKS if configured_blocks is None else int(configured_blocks) + ) + if not np.isclose(self.action_rate_hz, 20.0): + raise ValueError("T-Rex Track-Force actions are fixed at 20 Hz") + if not np.isclose(self.tactile_rate_hz, 5.0): + raise ValueError("T-Rex Track-Force tactile updates are fixed at 5 Hz") + if self.num_ar_blocks != AR_BLOCKS: + raise ValueError( + "T-Rex Track-Force keeps DreamZero's four autoregressive blocks" + ) + kwargs["max_chunk_size"] = self.num_ar_blocks + self._logical_anchor_times: dict[int, dict[int, float]] = {} + super().__init__(*args, **kwargs) + + force_stats = self.lerobot_stats_meta.get(FORCE_COLUMN) + if force_stats is None: + raise ValueError(f"dataset stats are missing {FORCE_COLUMN}") + self._force_low = np.asarray(force_stats.q01, dtype=np.float32).reshape( + FORCE_SENSOR_COUNT, FORCE_SENSOR_DIM + ) + self._force_high = np.asarray(force_stats.q99, dtype=np.float32).reshape( + FORCE_SENSOR_COUNT, FORCE_SENSOR_DIM + ) + + def _get_max_delta_index(self) -> int: + # Future validity is enforced from timestamp coverage in _get_step_filter. + return 0 + + def _get_step_filter(self) -> dict[int, np.ndarray]: + source_filter = super()._get_step_filter() + filtered: dict[int, np.ndarray] = {} + future_seconds = max( + self.num_ar_blocks * ACTION_HORIZON / self.action_rate_hz, + self.num_ar_blocks * TRACK_HORIZON / self.action_rate_hz, + self.num_ar_blocks * VIDEO_FRAMES_PER_BLOCK / self.video_rate_hz, + ) + for trajectory_id_value in self.trajectory_ids: + trajectory_id = int(trajectory_id_value) + parquet = self.get_parquet_path(trajectory_id) + timestamps = pd.read_parquet(parquet, columns=["timestamp"])[ + "timestamp" + ].to_numpy(dtype=np.float64) + logical_times = np.arange( + timestamps[0], + timestamps[-1] + 1e-9, + 1.0 / self.action_rate_hz, + dtype=np.float64, + ) + selection = nearest_timestamp_indices(timestamps, logical_times) + valid = ( + ~selection.padding_mask + & (selection.target_timestamps + future_seconds <= timestamps[-1]) + ) + allowed_source = set(np.asarray(source_filter[trajectory_id], dtype=np.int64)) + anchors: list[int] = [] + anchor_times: dict[int, float] = {} + for source_index, target_time in zip( + selection.indices[valid], selection.target_timestamps[valid] + ): + index = int(source_index) + if index in allowed_source: + anchors.append(index) + anchor_times[index] = float(target_time) + if not anchors: + raise ValueError( + f"episode {trajectory_id} has no complete " + f"{self.num_ar_blocks}x16-step 20 Hz autoregressive samples" + ) + selected = list(anchors) + if self.pin_anchor_rank is not None: + rank = min(self.pin_anchor_rank, len(selected) - 1) + selected = [selected[rank]] + if self.max_training_anchors is not None: + selected = selected[: self.max_training_anchors] + filtered[trajectory_id] = np.asarray(selected, dtype=np.int64) + self._logical_anchor_times[trajectory_id] = { + index: anchor_times[index] for index in selected + } + return filtered + + @staticmethod + def _column_array(frame: pd.DataFrame, column: str) -> np.ndarray: + if column not in frame.columns: + raise KeyError(f"missing required T-Rex column {column!r}") + values = np.stack(frame[column].to_numpy()) + # PyArrow fixed-size-list> columns are exposed by + # pandas as an object array whose individual elements are numeric + # ndarrays (not as one dense rank-3 ndarray). Collapse one nesting + # level at a time so track_xy becomes [T, P, 2] on real parquet data. + while values.dtype == object: + try: + values = np.stack([np.stack(row) for row in values]) + except (TypeError, ValueError) as exc: + raise ValueError( + f"column {column!r} cannot be converted to a dense array" + ) from exc + if not np.isfinite(values).all(): + raise ValueError(f"column {column!r} contains NaN/Inf") + return values + + def _sample( + self, + timestamps: np.ndarray, + anchor_timestamp: float, + offsets: Iterable[int], + rate_hz: float, + ) -> TimestampSelection: + return nearest_timestamp_indices( + timestamps, + uniform_target_times(anchor_timestamp, offsets, rate_hz), + ) + + def _normalize_force(self, force: np.ndarray) -> np.ndarray: + scale = self._force_high - self._force_low + safe_scale = np.where(scale == 0, 1.0, scale) + normalized = np.where( + scale != 0, + 2.0 * (force - self._force_low) / safe_scale - 1.0, + force, + ) + return np.clip(normalized, -1.0, 1.0).astype(np.float32) + + def get_step_data( + self, trajectory_id: int, indices: dict[str, np.ndarray] + ) -> dict: + if not indices: + raise ValueError("at least one configured modality index is required") + base_index = int(next(iter(indices.values()))[0]) + trajectory_id = int(trajectory_id) + try: + anchor_timestamp = self._logical_anchor_times[trajectory_id][base_index] + except KeyError as exc: + raise ValueError( + f"source row {base_index} is not a valid 20 Hz anchor for episode " + f"{trajectory_id}" + ) from exc + + frame = self.get_trajectory_data(trajectory_id) + self.curr_traj_data = frame + timestamps = frame["timestamp"].to_numpy(dtype=np.float64) + + action_selection = self._sample( + timestamps, + anchor_timestamp, + range(self.num_ar_blocks * ACTION_HORIZON), + self.action_rate_hz, + ) + block_anchor_timestamps = [ + anchor_timestamp + + block_index * ACTION_HORIZON / self.action_rate_hz + for block_index in range(self.num_ar_blocks) + ] + track_past_selections = [ + self._sample( + timestamps, + block_anchor, + range(-(FORCE_HISTORY_FRAMES - 1), 1), + self.action_rate_hz, + ) + for block_anchor in block_anchor_timestamps + ] + track_future_selections = [ + self._sample( + timestamps, + block_anchor, + range(TRACK_HORIZON), + self.action_rate_hz, + ) + for block_anchor in block_anchor_timestamps + ] + video_history_selection = self._sample( + timestamps, + anchor_timestamp, + range(VIDEO_HISTORY_FRAMES), + self.video_rate_hz, + ) + video_future_selection = self._sample( + timestamps, + anchor_timestamp, + range( + 1, + self.num_ar_blocks * VIDEO_FRAMES_PER_BLOCK + 1, + ), + self.video_rate_hz, + ) + force_history_selections = [ + [ + self._sample( + timestamps, + block_anchor + action_offset / self.action_rate_hz, + range(-(FORCE_HISTORY_FRAMES - 1), 1), + self.tactile_rate_hz, + ) + for action_offset in FORCE_OFFSETS + ] + for block_anchor in block_anchor_timestamps + ] + if ( + action_selection.padding_mask.any() + or any( + selection.padding_mask.any() + for selection in track_future_selections + ) + or video_future_selection.padding_mask.any() + ): + raise RuntimeError("a filtered training anchor unexpectedly needs future padding") + + state_meta = self.lerobot_modality_meta.state["eef62"] + state_source = self._column_array(frame, state_meta.original_key) + state_selection = self._sample( + timestamps, + anchor_timestamp, + range( + 0, + self.num_ar_blocks * ACTION_HORIZON, + ACTION_HORIZON, + ), + self.action_rate_hz, + ) + reference_states = state_source[ + state_selection.indices, state_meta.start : state_meta.end + ] + + action_meta = self.lerobot_modality_meta.action["eef62"] + absolute_action = self._column_array(frame, action_meta.original_key)[ + action_selection.indices, action_meta.start : action_meta.end + ].reshape(self.num_ar_blocks, ACTION_HORIZON, -1) + action_delta = np.stack( + [ + eef62_delta_base(reference_states[index], absolute_action[index]) + for index in range(self.num_ar_blocks) + ] + ).reshape(self.num_ar_blocks * ACTION_HORIZON, -1) + + track_xy = self._column_array(frame, TRACK_XY_COLUMN).astype(np.float32) + track_visibility = self._column_array( + frame, TRACK_VISIBILITY_COLUMN + ).astype(np.float32) + if track_xy.shape[1:] != (NUM_TRACK_POINTS, 2): + raise ValueError(f"track XY has invalid shape {track_xy.shape}") + + force = self._column_array(frame, FORCE_COLUMN).astype(np.float32).reshape( + -1, FORCE_SENSOR_COUNT, FORCE_SENSOR_DIM + ) + + def normalized_force_history(selection: TimestampSelection) -> np.ndarray: + history = self._normalize_force(force[selection.indices]) + history[selection.padding_mask] = 0.0 + return history + + force_history = np.stack( + [ + [ + normalized_force_history(selection) + for selection in block_selections + ] + for block_selections in force_history_selections + ] + ) + current_force = force_history[:, :, -1] + + data: dict[str, object] = { + "state.eef62": reference_states.astype(np.float32), + "action.eef62": action_delta, + "track_past_xy": np.stack( + [track_xy[selection.indices] for selection in track_past_selections] + ), + "track_past_visibility": np.stack( + [ + track_visibility[selection.indices] + * (~selection.padding_mask[:, None]) + for selection in track_past_selections + ] + ), + "track_future_xy": np.stack( + [track_xy[selection.indices] for selection in track_future_selections] + ), + "track_future_visibility": np.stack( + [ + track_visibility[selection.indices] + for selection in track_future_selections + ] + ), + "current_force": current_force, + "force_history": force_history, + "force_history_padding_mask": np.stack( + [ + [ + selection.padding_mask + for selection in block_selections + ] + for block_selections in force_history_selections + ] + ), + "sample_anchor_timestamp": np.asarray(anchor_timestamp, dtype=np.float32), + } + + video_indices = np.concatenate( + (video_history_selection.indices, video_future_selection.indices) + ) + if self.cached_shard is None or self.shard_start_indices is None: + raise RuntimeError("video shard must be cached before sampling") + shard_indices = self.shard_start_indices[trajectory_id] + video_indices + for key in self.modality_keys["video"]: + data[key] = self.cached_shard[key][shard_indices] + + for key in self.modality_keys.get("language", []): + data[key] = self.get_language( + trajectory_id, key, np.asarray([base_index], dtype=np.int64) + ) + return data + + +__all__ = [ + "AR_BLOCKS", + "ACTION_RATE_HZ", + "TACTILE_RATE_HZ", + "VIDEO_FUTURE_FRAMES", + "VIDEO_FRAMES_PER_BLOCK", + "VIDEO_HISTORY_FRAMES", + "VIDEO_RATE_HZ", + "TimestampSelection", + "TrexTrackForceShardedDataset", + "eef62_delta_base", + "nearest_timestamp_indices", + "rotation_6d_to_matrix", + "uniform_target_times", +] diff --git a/groot/vla/model/trex_track_force/force.py b/groot/vla/model/trex_track_force/force.py new file mode 100644 index 0000000000000000000000000000000000000000..9645235779edee44b46598222c494f84b497abbd --- /dev/null +++ b/groot/vla/model/trex_track_force/force.py @@ -0,0 +1,495 @@ +"""Force-only transformer, per-finger tactile VQ, and two-stage flow schedule.""" + +from __future__ import annotations + +from dataclasses import dataclass +import math +from typing import Callable, Iterator, Literal + +import torch +from torch import nn +from torch.nn import functional as F + +from .tactile_vq import ( + PerFingerForceVQVAE, + PerFingerVQVAEOutput, +) + +PHYSICAL_ACTION_DIM = 62 +PADDED_ACTION_DIM = 64 +ACTION_HORIZON = 16 +FORCE_SENSOR_COUNT = 10 +FORCE_SENSOR_DIM = 6 +FORCE_HISTORY_FRAMES = 16 +FORCE_CODE_TOKENS = 10 +FORCE_STRIDE = 4 +FORCE_OFFSETS = (0, 4, 8, 12) + + +def pad_action_62_to_64(action: torch.Tensor) -> torch.Tensor: + if action.shape[-1] == PADDED_ACTION_DIM: + return action + if action.shape[-1] != PHYSICAL_ACTION_DIM: + raise ValueError( + f"action must end in {PHYSICAL_ACTION_DIM} or {PADDED_ACTION_DIM}, " + f"got {tuple(action.shape)}" + ) + return F.pad(action, (0, PADDED_ACTION_DIM - PHYSICAL_ACTION_DIM)) + + +def unpad_action_64_to_62(action: torch.Tensor) -> torch.Tensor: + if action.shape[-1] != PADDED_ACTION_DIM: + raise ValueError(f"padded action must end in {PADDED_ACTION_DIM}") + return action[..., :PHYSICAL_ACTION_DIM] + + +def force_offsets( + action_horizon: int = ACTION_HORIZON, stride: int = FORCE_STRIDE +) -> tuple[int, ...]: + if action_horizon <= 0 or stride <= 0: + raise ValueError("action_horizon and stride must be positive") + return tuple(range(0, action_horizon, stride)) + + +@dataclass(frozen=True) +class FlowStageStep: + stage: Literal["coarse", "force"] + index: int + tau: float + tau_next: float + + @property + def delta_tau(self) -> float: + return self.tau_next - self.tau + + +@dataclass(frozen=True) +class TwoStageFlowSchedule: + """T-Rex: six slow updates to 0.4, then four tactile updates to zero.""" + + coarse_steps: int = 6 + force_steps: int = 4 + tau_start: float = 1.0 + tau_split: float = 0.4 + tau_end: float = 0.0 + action_horizon: int = ACTION_HORIZON + force_stride: int = FORCE_STRIDE + + def __post_init__(self) -> None: + if self.coarse_steps <= 0 or self.force_steps <= 0: + raise ValueError("both schedule stages need at least one update") + if not self.tau_start > self.tau_split > self.tau_end: + raise ValueError("require tau_start > tau_split > tau_end") + if ( + self.action_horizon == ACTION_HORIZON + and self.force_stride == FORCE_STRIDE + and force_offsets(self.action_horizon, self.force_stride) != FORCE_OFFSETS + ): + raise RuntimeError("canonical force offsets changed unexpectedly") + + @property + def num_steps(self) -> int: + return self.coarse_steps + self.force_steps + + @property + def offsets(self) -> tuple[int, ...]: + return force_offsets(self.action_horizon, self.force_stride) + + def boundaries( + self, + stage: Literal["coarse", "force"], + *, + device: torch.device | str | None = None, + dtype: torch.dtype = torch.float32, + ) -> torch.Tensor: + if stage == "coarse": + start, end, steps = self.tau_start, self.tau_split, self.coarse_steps + elif stage == "force": + start, end, steps = self.tau_split, self.tau_end, self.force_steps + else: + raise ValueError(f"unknown stage: {stage}") + return torch.linspace(start, end, steps + 1, device=device, dtype=dtype) + + def iter_steps( + self, stage: Literal["coarse", "force"] | None = None + ) -> Iterator[FlowStageStep]: + stages = ("coarse", "force") if stage is None else (stage,) + for stage_name in stages: + boundaries = self.boundaries(stage_name).tolist() + for index, (tau, tau_next) in enumerate( + zip(boundaries[:-1], boundaries[1:]) + ): + yield FlowStageStep(stage_name, index, tau, tau_next) + + +def sample_trex_timestep( + shape: tuple[int, ...], + *, + device: torch.device | str, + dtype: torch.dtype, + upper: float = 1.0, +) -> torch.Tensor: + """Sample T-Rex Beta(1.5, 1.0) timesteps in ``(0, upper]``.""" + + if not 0.0 < upper <= 1.0: + raise ValueError("upper must be in (0,1]") + distribution = torch.distributions.Beta( + torch.tensor(1.5, device=device, dtype=torch.float32), + torch.tensor(1.0, device=device, dtype=torch.float32), + ) + sample = distribution.sample(shape) * 0.999 + 0.001 + return (sample * upper).to(dtype=dtype) + + +def euler_flow_step( + sample: torch.Tensor, + flow: torch.Tensor, + tau: float | torch.Tensor, + tau_next: float | torch.Tensor, +) -> torch.Tensor: + if sample.shape != flow.shape: + raise ValueError("sample and flow shapes must match") + delta = torch.as_tensor( + tau_next, device=sample.device, dtype=sample.dtype + ) - torch.as_tensor(tau, device=sample.device, dtype=sample.dtype) + return sample + flow * delta + + +def sinusoidal_tau_embedding(tau: torch.Tensor, dim: int) -> torch.Tensor: + if dim < 2: + raise ValueError("time embedding dim must be at least 2") + tau = tau.float().reshape(-1) + half = dim // 2 + exponent = -math.log(10_000.0) * torch.arange( + half, device=tau.device, dtype=tau.dtype + ) / max(half - 1, 1) + phase = tau[:, None] * exponent.exp()[None] * 1000.0 + result = torch.cat((phase.sin(), phase.cos()), dim=-1) + return F.pad(result, (0, max(dim - result.shape[-1], 0))) + + +PerFingerForceVQ = PerFingerForceVQVAE +PerFingerVQOutput = PerFingerVQVAEOutput + + +@dataclass +class ForceOnlyOutput: + flow: torch.Tensor + tactile_codes: torch.Tensor + vq_loss: torch.Tensor + commitment_loss: torch.Tensor + tactile_reconstruction: torch.Tensor | None = None + tactile_perplexity: torch.Tensor | None = None + tactile_active_codes: torch.Tensor | None = None + + +class ForceOnlyTransformer(nn.Module): + """Action-only denoiser with 10 fresh and 10 historical tactile tokens.""" + + def __init__( + self, + *, + dim: int = 512, + num_heads: int = 8, + num_layers: int = 6, + ffn_dim: int | None = None, + action_horizon: int = ACTION_HORIZON, + action_dim: int = PADDED_ACTION_DIM, + force_sensor_count: int = FORCE_SENSOR_COUNT, + force_sensor_dim: int = FORCE_SENSOR_DIM, + history_frames: int = FORCE_HISTORY_FRAMES, + vq_codebook_size: int = 64, + vq_code_dim: int = 256, + coarse_memory_dim: int | None = None, + force_stride: int = FORCE_STRIDE, + dropout: float = 0.0, + ) -> None: + super().__init__() + if action_dim != PADDED_ACTION_DIM: + raise ValueError("force model action width must be padded 64-D") + if dim % num_heads: + raise ValueError("dim must be divisible by num_heads") + self.dim = dim + self.action_horizon = action_horizon + self.action_dim = action_dim + self.force_sensor_count = force_sensor_count + self.force_sensor_dim = force_sensor_dim + self.history_frames = history_frames + self.force_stride = force_stride + self.offsets = force_offsets(action_horizon, force_stride) + + self.action_encoder = nn.Linear(action_dim, dim) + self.action_decoder = nn.Sequential(nn.LayerNorm(dim), nn.Linear(dim, action_dim)) + self.force_encoder = nn.Sequential( + nn.Linear(force_sensor_dim, dim), nn.SiLU(), nn.Linear(dim, dim) + ) + self.tactile_vq = PerFingerForceVQ( + history_frames=history_frames, + num_fingers=force_sensor_count, + force_dim=force_sensor_dim, + codebook_size=vq_codebook_size, + code_dim=vq_code_dim, + ) + self.history_projection = nn.Linear(vq_code_dim, dim) + if coarse_memory_dim is None: + self.coarse_projection: nn.Module | None = None + elif coarse_memory_dim == dim: + self.coarse_projection = nn.Identity() + else: + self.coarse_projection = nn.Linear(coarse_memory_dim, dim) + + self.action_position = nn.Embedding(action_horizon, dim) + self.force_sensor_position = nn.Embedding(force_sensor_count, dim) + self.force_slot_embedding = nn.Embedding(len(self.offsets), dim) + self.force_update_embedding = nn.Embedding(len(self.offsets), dim) + self.token_type_embedding = nn.Embedding(4, dim) + self.time_mlp = nn.Sequential( + nn.Linear(dim, dim), nn.SiLU(), nn.Linear(dim, dim) + ) + layer = nn.TransformerEncoderLayer( + d_model=dim, + nhead=num_heads, + dim_feedforward=ffn_dim or 4 * dim, + dropout=dropout, + activation="gelu", + batch_first=True, + norm_first=True, + ) + self.transformer = nn.TransformerEncoder(layer, num_layers=num_layers) + self.final_norm = nn.LayerNorm(dim) + + @staticmethod + def _normalize_tau( + tau: float | torch.Tensor, batch: int, device: torch.device + ) -> torch.Tensor: + value = torch.as_tensor(tau, device=device, dtype=torch.float32) + if value.ndim == 0: + value = value.expand(batch) + elif value.shape == (batch, 1): + value = value[:, 0] + if value.shape != (batch,): + raise ValueError("force tau must be scalar or [B]") + return value + + def _normalize_update_offset( + self, + update_offset: int | torch.Tensor, + batch: int, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor]: + offsets = torch.as_tensor(update_offset, device=device, dtype=torch.long) + if offsets.ndim == 0: + offsets = offsets.expand(batch) + if offsets.shape != (batch,): + raise ValueError("force update_offset must be scalar or [B]") + canonical = torch.as_tensor(self.offsets, device=device, dtype=torch.long) + matches = offsets[:, None] == canonical[None] + if not bool(matches.any(dim=1).all()): + raise ValueError( + f"force update_offset must be one of {self.offsets}, " + f"got {offsets.tolist()}" + ) + return offsets, matches.long().argmax(dim=1) + + def forward( + self, + action: torch.Tensor, + tau: float | torch.Tensor, + current_force: torch.Tensor, + tactile_history_or_codes: torch.Tensor, + *, + coarse_memory: torch.Tensor | None = None, + coarse_memory_mask: torch.Tensor | None = None, + update_offset: int | torch.Tensor = 0, + detach_coarse_memory: bool = True, + tactile_keep_mask: torch.Tensor | None = None, + tactile_history_valid_mask: torch.Tensor | None = None, + return_details: bool = False, + ) -> torch.Tensor | ForceOnlyOutput: + action = pad_action_62_to_64(action) + batch = action.shape[0] + if action.shape != (batch, self.action_horizon, self.action_dim): + raise ValueError( + f"action must be [B,{self.action_horizon},{self.action_dim}]" + ) + if current_force.shape != ( + batch, + self.force_sensor_count, + self.force_sensor_dim, + ): + raise ValueError("current_force must be [B,10,6]") + if tactile_history_or_codes.shape[0] != batch: + raise ValueError("tactile history batch does not match action") + + device, dtype = action.device, self.action_encoder.weight.dtype + action = action.to(dtype=dtype) + current_force = current_force.to(device=device, dtype=dtype) + history = tactile_history_or_codes.to(device=device) + history_valid = ( + None + if tactile_history_valid_mask is None + else tactile_history_valid_mask.to(device=device, dtype=torch.bool) + ) + vq = self.tactile_vq(history, valid_history_mask=history_valid) + _, update_slot = self._normalize_update_offset( + update_offset, batch, device + ) + update_embedding = self.force_update_embedding(update_slot).unsqueeze(1) + time = self.time_mlp( + sinusoidal_tau_embedding( + self._normalize_tau(tau, batch, device), self.dim + ).to(dtype) + ).unsqueeze(1) + action_positions = torch.arange(self.action_horizon, device=device) + slot_ids = torch.div( + action_positions, self.force_stride, rounding_mode="floor" + ).clamp_max(len(self.offsets) - 1) + action_tokens = ( + self.action_encoder(action) + + self.action_position(action_positions).unsqueeze(0) + + self.force_slot_embedding(slot_ids).unsqueeze(0) + + self.token_type_embedding.weight[0].view(1, 1, self.dim) + + update_embedding + + time + ) + finger_positions = torch.arange(self.force_sensor_count, device=device) + force_tokens = ( + self.force_encoder(current_force) + + self.force_sensor_position(finger_positions).unsqueeze(0) + + self.token_type_embedding.weight[1].view(1, 1, self.dim) + + update_embedding + ) + history_tokens = ( + self.history_projection(vq.tokens.to(dtype)) + + self.force_sensor_position(finger_positions).unsqueeze(0) + + self.token_type_embedding.weight[2].view(1, 1, self.dim) + + update_embedding + ) + if tactile_keep_mask is not None: + keep = tactile_keep_mask.to(device=device, dtype=dtype).view(batch, 1, 1) + force_tokens = force_tokens * keep + history_tokens = history_tokens * keep + + token_groups = [action_tokens, force_tokens, history_tokens] + padding_groups = [ + torch.zeros( + batch, + self.action_horizon + 2 * self.force_sensor_count, + dtype=torch.bool, + device=device, + ) + ] + if coarse_memory is not None: + if self.coarse_projection is None: + raise ValueError("coarse_memory_dim is not configured") + memory = coarse_memory.detach() if detach_coarse_memory else coarse_memory + coarse_tokens = self.coarse_projection(memory.to(device, dtype)) + token_groups.append( + coarse_tokens + + self.token_type_embedding.weight[3].view(1, 1, self.dim) + ) + if coarse_memory_mask is None: + coarse_padding = torch.zeros( + batch, coarse_tokens.shape[1], dtype=torch.bool, device=device + ) + else: + if coarse_memory_mask.shape != coarse_tokens.shape[:2]: + raise ValueError("coarse_memory_mask must be [B,L]") + coarse_padding = ~coarse_memory_mask.to(device, torch.bool) + padding_groups.append(coarse_padding) + + hidden = self.transformer( + torch.cat(token_groups, dim=1), + src_key_padding_mask=torch.cat(padding_groups, dim=1), + ) + flow = self.action_decoder( + self.final_norm(hidden[:, : self.action_horizon]) + ) + if return_details: + return ForceOnlyOutput( + flow=flow, + tactile_codes=vq.codes, + vq_loss=vq.reconstruction_loss, + commitment_loss=vq.commitment_loss, + tactile_reconstruction=vq.reconstruction, + tactile_perplexity=vq.perplexity, + tactile_active_codes=vq.active_codes, + ) + return flow + + @torch.no_grad() + def refine_action( + self, + action_at_split: torch.Tensor, + current_force: torch.Tensor, + tactile_history_or_codes: torch.Tensor, + *, + coarse_memory: torch.Tensor | None = None, + schedule: TwoStageFlowSchedule | None = None, + update_offset: int = 0, + tactile_history_valid_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + schedule = schedule or TwoStageFlowSchedule( + action_horizon=self.action_horizon, + force_stride=self.force_stride, + ) + action = pad_action_62_to_64(action_at_split).clone() + for step in schedule.iter_steps("force"): + flow = self( + action, + step.tau, + current_force, + tactile_history_or_codes, + coarse_memory=coarse_memory, + update_offset=update_offset, + tactile_history_valid_mask=tactile_history_valid_mask, + ) + assert isinstance(flow, torch.Tensor) + updated = euler_flow_step(action, flow, step.tau, step.tau_next) + action[:, update_offset:] = updated[:, update_offset:] + return action + + +def run_force_stage( + action_at_split: torch.Tensor, + velocity_fn: Callable[[torch.Tensor, float], torch.Tensor], + *, + schedule: TwoStageFlowSchedule | None = None, +) -> torch.Tensor: + schedule = schedule or TwoStageFlowSchedule() + action = pad_action_62_to_64(action_at_split) + for step in schedule.iter_steps("force"): + action = euler_flow_step( + action, velocity_fn(action, step.tau), step.tau, step.tau_next + ) + return action + + +ForceTactileTransformer = ForceOnlyTransformer + + +__all__ = [ + "ACTION_HORIZON", + "FORCE_CODE_TOKENS", + "FORCE_HISTORY_FRAMES", + "FORCE_OFFSETS", + "FORCE_SENSOR_COUNT", + "FORCE_SENSOR_DIM", + "FORCE_STRIDE", + "PADDED_ACTION_DIM", + "PHYSICAL_ACTION_DIM", + "FlowStageStep", + "ForceOnlyOutput", + "ForceOnlyTransformer", + "ForceTactileTransformer", + "PerFingerForceVQ", + "PerFingerVQOutput", + "TwoStageFlowSchedule", + "euler_flow_step", + "force_offsets", + "pad_action_62_to_64", + "run_force_stage", + "sample_trex_timestep", + "sinusoidal_tau_embedding", + "unpad_action_64_to_62", +] diff --git a/groot/vla/model/trex_track_force/lora.py b/groot/vla/model/trex_track_force/lora.py new file mode 100644 index 0000000000000000000000000000000000000000..f8958e6a264d564ae50896ca688cf0a818a5342d --- /dev/null +++ b/groot/vla/model/trex_track_force/lora.py @@ -0,0 +1,143 @@ +"""Small dependency-free LoRA adapter used by the isolated model.""" + +from __future__ import annotations + +from collections.abc import Iterable + +import torch +from torch import nn + + +class LoRALinear(nn.Module): + """Frozen linear layer plus trainable low-rank residual.""" + + def __init__( + self, + base_layer: nn.Linear, + *, + rank: int = 4, + alpha: float = 4.0, + ) -> None: + super().__init__() + if rank <= 0: + raise ValueError("LoRA rank must be positive") + self.base_layer = base_layer + self.rank = rank + self.alpha = alpha + self.scaling = alpha / rank + self.lora_A = nn.Linear(base_layer.in_features, rank, bias=False) + self.lora_B = nn.Linear(rank, base_layer.out_features, bias=False) + nn.init.kaiming_uniform_(self.lora_A.weight, a=5**0.5) + nn.init.zeros_(self.lora_B.weight) + self.base_layer.requires_grad_(False) + + @property + def in_features(self) -> int: + return self.base_layer.in_features + + @property + def out_features(self) -> int: + return self.base_layer.out_features + + def forward(self, values: torch.Tensor) -> torch.Tensor: + base = self.base_layer(values) + residual = self.lora_B(self.lora_A(values)) * self.scaling + return base + residual + + +def inject_lora( + module: nn.Module, + *, + target_modules: str | Iterable[str] = ("q", "k", "v", "o", "ffn.0", "ffn.2"), + rank: int = 4, + alpha: float = 4.0, +) -> list[str]: + """Replace matching descendant ``nn.Linear`` modules with LoRA wrappers.""" + + if isinstance(target_modules, str): + targets = tuple(item.strip() for item in target_modules.split(",") if item.strip()) + else: + targets = tuple(target_modules) + replaced: list[str] = [] + + def visit(parent: nn.Module, prefix: str) -> None: + for child_name, child in list(parent.named_children()): + full_name = f"{prefix}.{child_name}" if prefix else child_name + if isinstance(child, LoRALinear): + continue + matches = any( + full_name == target + or full_name.endswith(f".{target}") + or child_name == target + for target in targets + ) + if matches and isinstance(child, nn.Linear): + setattr( + parent, + child_name, + LoRALinear(child, rank=rank, alpha=alpha), + ) + replaced.append(full_name) + else: + visit(child, full_name) + + visit(module, "") + return replaced + + +def set_lora_and_new_modules_trainable( + root: nn.Module, + *, + new_modules: Iterable[nn.Module], +) -> None: + """Freeze the base, then enable LoRA adapters and all newly-added modules.""" + + root.requires_grad_(False) + for candidate in root.modules(): + if isinstance(candidate, LoRALinear): + candidate.lora_A.requires_grad_(True) + candidate.lora_B.requires_grad_(True) + for new_module in new_modules: + new_module.requires_grad_(True) + + +def merge_lora_layers(module: nn.Module) -> list[str]: + """Merge custom LoRA residuals into ordinary linear weights in place.""" + + merged: list[str] = [] + + def visit(parent: nn.Module, prefix: str) -> None: + for child_name, child in list(parent.named_children()): + full_name = f"{prefix}.{child_name}" if prefix else child_name + if isinstance(child, LoRALinear): + base = child.base_layer + replacement = nn.Linear( + base.in_features, + base.out_features, + bias=base.bias is not None, + device=base.weight.device, + dtype=base.weight.dtype, + ) + with torch.no_grad(): + residual = ( + child.lora_B.weight.to(base.weight.dtype) + @ child.lora_A.weight.to(base.weight.dtype) + ) * child.scaling + replacement.weight.copy_(base.weight + residual) + if base.bias is not None: + replacement.bias.copy_(base.bias) + setattr(parent, child_name, replacement) + merged.append(full_name) + else: + visit(child, full_name) + + visit(module, "") + return merged + + +__all__ = [ + "LoRALinear", + "inject_lora", + "merge_lora_layers", + "set_lora_and_new_modules_trainable", +] diff --git a/groot/vla/model/trex_track_force/model.py b/groot/vla/model/trex_track_force/model.py new file mode 100644 index 0000000000000000000000000000000000000000..20eaf90700cc70d1679ce1396adce159107e7be0 --- /dev/null +++ b/groot/vla/model/trex_track_force/model.py @@ -0,0 +1,892 @@ +"""Wan2.2-compatible CausalWan model with track tokens and exact topology.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Mapping + +import torch +from torch import nn + +from groot.vla.model.dreamzero.modules.wan2_1_submodule import ( + rope_params, + sinusoidal_embedding_1d, +) +from groot.vla.model.dreamzero.modules.wan_video_dit_action_casual_chunk import ( + CausalWanModel as DreamZeroCausalWanModel, +) + +from .attention import BlockTokenLayout, TokenType, build_dense_allow_matrix +from .blocks import CausalWanTrackForceBlock +from .force import ACTION_HORIZON, PADDED_ACTION_DIM, pad_action_62_to_64 +from .lora import merge_lora_layers +from .track import ( + NUM_TRACK_POINTS, + TRACK_HORIZON, + TrackDecoder, + TrackEncoder, +) + + +WAN_SHARED_PREFIXES = ( + "patch_embedding.", + "text_embedding.", + "time_embedding.", + "time_projection.", + "blocks.", + "head.", + "img_emb.", + "state_encoder.", + "action_encoder.", + "action_decoder.", +) + + +@dataclass +class CausalWanTrackForceOutput: + action_flow: torch.Tensor + track_flow: torch.Tensor + obs_flow_tokens: torch.Tensor + coarse_memory: torch.Tensor + coarse_memory_by_block: torch.Tensor + coarse_memory_mask_by_block: torch.Tensor + layout: BlockTokenLayout + video_flow: torch.Tensor | None = None + hidden_states: torch.Tensor | None = None + + def to_dict(self) -> dict[str, Any]: + return { + "action_flow": self.action_flow, + "track_flow": self.track_flow, + "obs_flow_tokens": self.obs_flow_tokens, + "coarse_memory": self.coarse_memory, + "coarse_memory_by_block": self.coarse_memory_by_block, + "coarse_memory_mask_by_block": self.coarse_memory_mask_by_block, + "layout": self.layout, + "video_flow": self.video_flow, + "hidden_states": self.hidden_states, + } + + def __getitem__(self, key: str) -> Any: + return self.to_dict()[key] + + +def _pad_feature_width(values: torch.Tensor, width: int, name: str) -> torch.Tensor: + if values.shape[-1] > width: + raise ValueError(f"{name} width {values.shape[-1]} exceeds {width}") + if values.shape[-1] == width: + return values + return torch.nn.functional.pad(values, (0, width - values.shape[-1])) + + +class CausalWanTrackForceModel(DreamZeroCausalWanModel): + """True subclass of DreamZero's CausalWanModel. + + All pretrained WAN modules retain their original names and shapes. Only + each block's self-attention implementation is replaced; its q/k/v/o and + normalization parameters remain strict state-dict compatible. + """ + + def __init__( + self, + model_type: str = "ti2v", + patch_size: tuple[int, int, int] = (1, 2, 2), + frame_seqlen: int = 50, + text_len: int = 512, + in_dim: int = 48, + dim: int = 512, + ffn_dim: int = 2048, + freq_dim: int = 256, + text_dim: int = 4096, + out_dim: int = 48, + num_heads: int = 8, + num_layers: int = 8, + max_chunk_size: int = -1, + sink_size: int = 0, + qk_norm: bool = True, + cross_attn_norm: bool = True, + eps: float = 1e-6, + num_frame_per_block: int = 2, + action_dim: int = PADDED_ACTION_DIM, + action_horizon: int = ACTION_HORIZON, + max_state_dim: int = PADDED_ACTION_DIM, + num_state_per_block: int = 1, + hidden_size: int = 1024, + track_points: int = NUM_TRACK_POINTS, + track_horizon: int = TRACK_HORIZON, + max_track_history: int = 32, + max_blocks: int = 32, + obs_token_dim: int | None = None, + context_dim: int | None = None, + gradient_checkpointing: bool = False, + diffusion_model_pretrained_path: str | None = None, + concat_first_frame_latent: bool = False, + **_: Any, + ) -> None: + if action_dim != PADDED_ACTION_DIM: + raise ValueError("physical 62-D action must use WAN width 64") + if action_horizon != TRACK_HORIZON: + raise ValueError("action and track future horizons must both be 16") + if track_points != NUM_TRACK_POINTS: + raise ValueError("T-Rex track layout is fixed at 250 points") + + super().__init__( + model_type=model_type, + patch_size=patch_size, + frame_seqlen=frame_seqlen, + text_len=text_len, + in_dim=in_dim, + dim=dim, + ffn_dim=ffn_dim, + freq_dim=freq_dim, + text_dim=text_dim, + out_dim=out_dim, + num_heads=num_heads, + num_layers=num_layers, + max_chunk_size=max_chunk_size, + sink_size=sink_size, + qk_norm=qk_norm, + cross_attn_norm=cross_attn_norm, + eps=eps, + num_frame_per_block=num_frame_per_block, + action_dim=action_dim, + max_state_dim=max_state_dim, + hidden_size=hidden_size, + diffusion_model_pretrained_path=diffusion_model_pretrained_path, + num_action_per_block=action_horizon, + num_state_per_block=num_state_per_block, + concat_first_frame_latent=concat_first_frame_latent, + ) + self.action_horizon = action_horizon + self.track_points = track_points + self.track_horizon = track_horizon + self.max_track_history = max_track_history + self.max_blocks = max_blocks + self.max_chunk_size = max_chunk_size + self.obs_token_dim = obs_token_dim or dim + self.gradient_checkpointing = gradient_checkpointing + + cross_attn_type = "t2v_cross_attn" if model_type == "t2v" else "i2v_cross_attn" + compatible_blocks = nn.ModuleList() + for original_block in self.blocks: + block = CausalWanTrackForceBlock( + cross_attn_type, + dim, + ffn_dim, + num_heads, + frame_seqlen, + self.local_attn_size, + sink_size, + num_frame_per_block, + qk_norm, + cross_attn_norm, + eps, + action_horizon, + num_state_per_block, + ) + block.load_state_dict(original_block.state_dict(), strict=True) + compatible_blocks.append(block) + self.blocks = compatible_blocks + + self.track_encoder = TrackEncoder( + dim, + num_points=track_points, + future_horizon=track_horizon, + max_history_frames=max_track_history, + max_blocks=max_blocks, + ) + self.track_decoder = TrackDecoder( + dim, num_points=track_points, future_horizon=track_horizon + ) + self.modality_embedding = nn.Embedding(len(TokenType), dim) + self.obs_input_projection = ( + nn.Identity() + if self.obs_token_dim == dim + else nn.Linear(self.obs_token_dim, dim) + ) + self.obs_token_decoder = nn.Linear(dim, self.obs_token_dim) + self.context_embedding: nn.Module = ( + nn.Identity() + if context_dim is None or context_dim == dim + else nn.Linear(context_dim, dim) + ) + + head_dim = dim // num_heads + self.freqs_track = rope_params(max_blocks * 2 * track_points, head_dim) + self.freqs_obs_1d = rope_params( + max_blocks * max(frame_seqlen * num_frame_per_block, 1) + frame_seqlen, + head_dim, + ) + + @staticmethod + def _normalize_checkpoint_key(key: str) -> str: + for prefix in ("module.", "model.", "action_head.model."): + if key.startswith(prefix): + return key[len(prefix) :] + return key + + def pretrained_compatibility_report( + self, state_dict: Mapping[str, torch.Tensor] + ) -> dict[str, Any]: + """Report WAN-key and parameter-count coverage before loading.""" + + current = self.state_dict() + eligible: dict[str, torch.Tensor] = {} + for source_key, value in state_dict.items(): + key = self._normalize_checkpoint_key(source_key) + if key.startswith(WAN_SHARED_PREFIXES): + eligible[key] = value + + matched: list[str] = [] + missing: list[str] = [] + shape_mismatches: dict[str, dict[str, tuple[int, ...]]] = {} + matched_numel = 0 + total_numel = 0 + for key, value in eligible.items(): + total_numel += value.numel() + if key not in current: + missing.append(key) + elif current[key].shape != value.shape: + shape_mismatches[key] = { + "checkpoint": tuple(value.shape), + "model": tuple(current[key].shape), + } + else: + matched.append(key) + matched_numel += value.numel() + total_keys = len(eligible) + return { + "wan_checkpoint_keys": total_keys, + "matched_keys": len(matched), + "key_coverage": len(matched) / max(total_keys, 1), + "matched_numel": matched_numel, + "wan_checkpoint_numel": total_numel, + "numel_coverage": matched_numel / max(total_numel, 1), + "missing_keys": sorted(missing), + "shape_mismatches": shape_mismatches, + } + + def load_pretrained_wan_state_dict( + self, + state_dict: Mapping[str, torch.Tensor], + *, + minimum_numel_coverage: float = 0.95, + ) -> dict[str, Any]: + """Validated WAN load; never silently accepts broad incompatibility.""" + + report = self.pretrained_compatibility_report(state_dict) + if report["wan_checkpoint_keys"] == 0: + raise ValueError("checkpoint contains no recognized WAN parameter keys") + if report["shape_mismatches"]: + raise ValueError( + f"WAN checkpoint shape mismatches: {report['shape_mismatches']}" + ) + if report["numel_coverage"] < minimum_numel_coverage: + raise ValueError( + "WAN checkpoint coverage " + f"{report['numel_coverage']:.3f} < {minimum_numel_coverage:.3f}; " + f"missing={report['missing_keys'][:20]}" + ) + current = self.state_dict() + compatible = { + self._normalize_checkpoint_key(key): value + for key, value in state_dict.items() + if self._normalize_checkpoint_key(key) in current + and current[self._normalize_checkpoint_key(key)].shape == value.shape + } + self.load_state_dict(compatible, strict=False) + return report + + def merge_and_unload(self) -> "CausalWanTrackForceModel": + """Merge custom LoRA layers using the interface expected by rollout.""" + + merge_lora_layers(self.blocks) + return self + + @staticmethod + def _normalize_tau( + tau: float | torch.Tensor, + batch_size: int, + num_blocks: int, + device: torch.device, + ) -> torch.Tensor: + value = torch.as_tensor(tau, device=device, dtype=torch.float32) + if value.ndim == 0: + return value.expand(batch_size, num_blocks) + if value.shape == (batch_size,): + return value[:, None].expand(-1, num_blocks) + if value.shape == (batch_size, 1): + return value.expand(-1, num_blocks) + if value.shape != (batch_size, num_blocks): + raise ValueError("tau must be scalar, [B], or [B,blocks]") + return value + + def _canonicalize_action( + self, action: torch.Tensor, num_blocks: int + ) -> torch.Tensor: + action = pad_action_62_to_64(action) + if action.ndim == 3: + action = action.view( + action.shape[0], num_blocks, self.action_horizon, self.action_dim + ) + if action.shape[1:] != ( + num_blocks, + self.action_horizon, + self.action_dim, + ): + raise ValueError("action must be [B,blocks,16,62|64]") + return action + + def _canonicalize_state( + self, state: torch.Tensor, num_blocks: int + ) -> torch.Tensor: + if state.ndim == 3: + state = state.unsqueeze(2) + if ( + state.ndim != 4 + or state.shape[1] != num_blocks + or state.shape[2] != self.num_state_per_block + ): + raise ValueError("state must align to block/state-token axes") + return _pad_feature_width(state, self.max_state_dim, "state") + + @staticmethod + def _squeeze_rope(frequencies: torch.Tensor) -> torch.Tensor: + if frequencies.ndim == 3 and frequencies.shape[1] == 1: + return frequencies[:, 0] + return frequencies + + def _tokenize_latents( + self, + video_latents: torch.Tensor, + num_blocks: int, + conditioning_video_latents: torch.Tensor | None = None, + ) -> tuple[ + torch.Tensor | None, + torch.Tensor, + tuple[int, int, int], + torch.Tensor, + ]: + patched = self.patch_embedding(video_latents) + _, _, frames, height, width = patched.shape + spatial = height * width + if spatial != self.frame_seqlen: + raise ValueError( + f"patch grid has {spatial} tokens/frame, expected {self.frame_seqlen}" + ) + expected_future = num_blocks * self.num_frame_per_block + if conditioning_video_latents is not None: + conditioning_patched = self.patch_embedding(conditioning_video_latents) + if conditioning_patched.shape[-2:] != (height, width): + raise ValueError("conditioning/future latent patch grids differ") + if frames != expected_future: + raise ValueError( + "separate future video latents must match the action blocks" + ) + conditioning_frames = conditioning_patched.shape[2] + if conditioning_frames <= 0: + raise ValueError("conditioning video must contain at least one latent") + conditioning = conditioning_patched.flatten(2).transpose(1, 2) + future = patched + total_frames = conditioning_frames + frames + elif frames == expected_future + 1: + conditioning = patched[:, :, :1].flatten(2).transpose(1, 2) + future = patched[:, :, 1:] + total_frames = frames + elif frames == expected_future: + conditioning = None + future = patched + total_frames = frames + else: + raise ValueError("latent temporal length does not match action blocks") + full_grid = torch.tensor((total_frames, height, width), dtype=torch.long) + video_rope = self._squeeze_rope(self._create_freqs(full_grid, start_frame=0)) + future = ( + future.view( + future.shape[0], + self.dim, + num_blocks, + self.num_frame_per_block, + height, + width, + ) + .permute(0, 2, 3, 4, 5, 1) + .reshape( + future.shape[0], + num_blocks, + self.num_frame_per_block * spatial, + self.dim, + ) + ) + return conditioning, future, (expected_future, height, width), video_rope + + def _tokenize_observations( + self, + *, + num_blocks: int, + obs_tokens: torch.Tensor | None, + conditioning_obs_tokens: torch.Tensor | None, + video_latents: torch.Tensor | None, + conditioning_video_latents: torch.Tensor | None, + ) -> tuple[ + torch.Tensor | None, + torch.Tensor, + tuple[int, int, int] | None, + torch.Tensor, + ]: + if (obs_tokens is None) == (video_latents is None): + raise ValueError("provide exactly one observation representation") + if video_latents is not None: + return self._tokenize_latents( + video_latents, num_blocks, conditioning_video_latents + ) + if conditioning_video_latents is not None: + raise ValueError( + "conditioning_video_latents requires future video_latents" + ) + assert obs_tokens is not None + if obs_tokens.ndim != 4 or obs_tokens.shape[1] != num_blocks: + raise ValueError("obs_tokens must be [B,blocks,O,D]") + obs = self.obs_input_projection(obs_tokens) + conditioning = ( + None + if conditioning_obs_tokens is None + else self.obs_input_projection(conditioning_obs_tokens) + ) + total_video = obs.shape[1] * obs.shape[2] + ( + 0 if conditioning is None else conditioning.shape[1] + ) + if total_video > self.freqs_obs_1d.shape[0]: + raise ValueError("direct observation token sequence exceeds RoPE table") + return ( + conditioning, + obs, + None, + self._squeeze_rope(self.freqs_obs_1d[:total_video]).to(obs.device), + ) + + def _build_context( + self, + context: torch.Tensor | None, + text_context: torch.Tensor | None, + clip_context: torch.Tensor | None, + ) -> torch.Tensor | None: + groups: list[torch.Tensor] = [] + if context is not None: + groups.append(self.context_embedding(context)) + if clip_context is not None: + if not hasattr(self, "img_emb"): + raise ValueError("CLIP context requires i2v/ti2v WAN") + groups.append(self.img_emb(clip_context)) + if text_context is not None: + groups.append(self.text_embedding(text_context)) + return torch.cat(groups, dim=1) if groups else None + + def _encode_action_state( + self, + action: torch.Tensor, + state: torch.Tensor, + action_tau: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + batch, blocks = action.shape[:2] + category = torch.zeros(batch, dtype=torch.long, device=action.device) + flat_action = action.flatten(1, 2) + action_times = action_tau.repeat_interleave(self.action_horizon, dim=1) * 1000 + action_tokens = self.action_encoder( + flat_action, action_times, category + ).view(batch, blocks, self.action_horizon, self.dim) + flat_state = state.flatten(1, 2) + state_tokens = self.state_encoder(flat_state, category).view( + batch, blocks, self.num_state_per_block, self.dim + ) + return action_tokens, state_tokens + + def _pack_tokens( + self, + conditioning: torch.Tensor | None, + obs: torch.Tensor, + action: torch.Tensor, + state: torch.Tensor, + track_past: torch.Tensor, + track_future: torch.Tensor, + ) -> tuple[torch.Tensor, BlockTokenLayout]: + layout = BlockTokenLayout( + num_blocks=obs.shape[1], + obs_tokens_per_block=obs.shape[2], + action_tokens_per_block=self.action_horizon, + state_tokens_per_block=self.num_state_per_block, + track_points=self.track_points, + conditioning_obs_tokens=0 if conditioning is None else conditioning.shape[1], + ) + groups: list[torch.Tensor] = [] + if conditioning is not None: + groups.append( + conditioning + + self.modality_embedding.weight[ + int(TokenType.CONDITIONING_OBS) + ].view(1, 1, self.dim) + ) + for block_index in range(layout.num_blocks): + for values, token_type in ( + (obs[:, block_index], TokenType.OBS), + (action[:, block_index], TokenType.ACTION), + (state[:, block_index], TokenType.STATE), + (track_past[:, block_index], TokenType.TRACK_PAST), + (track_future[:, block_index], TokenType.TRACK_FUTURE), + ): + groups.append( + values + + self.modality_embedding.weight[int(token_type)].view( + 1, 1, self.dim + ) + ) + return torch.cat(groups, dim=1), layout + + def _pack_scalar_by_layout( + self, + layout: BlockTokenLayout, + video_tau: torch.Tensor, + action_tau: torch.Tensor, + track_tau: torch.Tensor, + ) -> torch.Tensor: + batch = video_tau.shape[0] + zeros = torch.zeros(batch, 1, device=video_tau.device) + groups: list[torch.Tensor] = [] + if layout.conditioning_obs_tokens: + groups.append(zeros.expand(-1, layout.conditioning_obs_tokens)) + for block in range(layout.num_blocks): + groups.extend( + ( + video_tau[:, block : block + 1].expand( + -1, layout.obs_tokens_per_block + ), + action_tau[:, block : block + 1].expand( + -1, layout.action_tokens_per_block + ), + zeros.expand(-1, layout.state_tokens_per_block), + zeros.expand(-1, layout.track_points), + track_tau[:, block : block + 1].expand( + -1, layout.track_points + ), + ) + ) + return torch.cat(groups, dim=1) + + def _pack_rope( + self, layout: BlockTokenLayout, video_rope: torch.Tensor, device: torch.device + ) -> torch.Tensor: + tables = (self.freqs_action, self.freqs_state, self.freqs_track) + self.freqs_action, self.freqs_state, self.freqs_track = ( + table.to(device) for table in tables + ) + video_rope = self._squeeze_rope(video_rope).to(device) + groups: list[torch.Tensor] = [] + video_cursor = 0 + if layout.conditioning_obs_tokens: + end = layout.conditioning_obs_tokens + groups.append(video_rope[:end]) + video_cursor = end + for block in range(layout.num_blocks): + obs_end = video_cursor + layout.obs_tokens_per_block + groups.append(video_rope[video_cursor:obs_end]) + video_cursor = obs_end + action_start = block * self.action_horizon + groups.append( + self._squeeze_rope( + self.freqs_action[ + action_start : action_start + self.action_horizon + ] + ) + ) + state_start = block * self.num_state_per_block + groups.append( + self._squeeze_rope( + self.freqs_state[ + state_start : state_start + self.num_state_per_block + ] + ) + ) + track_start = block * 2 * self.track_points + groups.append( + self._squeeze_rope( + self.freqs_track[track_start : track_start + self.track_points] + ) + ) + groups.append( + self._squeeze_rope( + self.freqs_track[ + track_start + + self.track_points : track_start + + 2 * self.track_points + ] + ) + ) + packed = torch.cat(groups, dim=0) + if packed.shape[0] != layout.sequence_length: + raise RuntimeError("packed RoPE does not match token layout") + return packed + + @staticmethod + def _unpack( + values: torch.Tensor, layout: BlockTokenLayout + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + obs, action, past, future = [], [], [], [] + for block in range(layout.num_blocks): + slices = layout.slices(block) + obs.append(values[:, slices[TokenType.OBS]]) + action.append(values[:, slices[TokenType.ACTION]]) + past.append(values[:, slices[TokenType.TRACK_PAST]]) + future.append(values[:, slices[TokenType.TRACK_FUTURE]]) + return tuple( + torch.stack(group, dim=1) for group in (obs, action, past, future) + ) # type: ignore[return-value] + + def _causal_coarse_memory( + self, + hidden: torch.Tensor, + layout: BlockTokenLayout, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Build padded per-block memory without exposing later AR blocks.""" + + per_block_indices: list[torch.Tensor] = [] + for query_block in range(layout.num_blocks): + groups: list[torch.Tensor] = [] + if layout.conditioning_obs_tokens: + groups.append( + torch.arange( + layout.conditioning_slice.start, + layout.conditioning_slice.stop, + device=hidden.device, + ) + ) + first_key_block = ( + 0 + if self.max_chunk_size < 0 + else max(0, query_block - self.max_chunk_size) + ) + for key_block in range(first_key_block, query_block + 1): + slices = layout.slices(key_block) + for token_type in ( + TokenType.OBS, + TokenType.ACTION, + TokenType.TRACK_PAST, + TokenType.TRACK_FUTURE, + ): + token_slice = slices[token_type] + groups.append( + torch.arange( + token_slice.start, + token_slice.stop, + device=hidden.device, + ) + ) + per_block_indices.append(torch.cat(groups).long()) + + max_length = max(indices.numel() for indices in per_block_indices) + memories: list[torch.Tensor] = [] + masks: list[torch.Tensor] = [] + for indices in per_block_indices: + memory = hidden.index_select(1, indices) + valid = indices.numel() + if valid < max_length: + memory = torch.nn.functional.pad( + memory, (0, 0, 0, max_length - valid) + ) + mask = torch.zeros( + max_length, + dtype=torch.bool, + device=hidden.device, + ) + mask[:valid] = True + memories.append(memory) + masks.append(mask) + return torch.stack(memories, dim=1), torch.stack(masks, dim=0) + + def forward( + self, + *, + action: torch.Tensor, + state: torch.Tensor, + past_track_xy: torch.Tensor, + past_track_visibility: torch.Tensor | None, + future_track_xy: torch.Tensor, + future_track_visibility: torch.Tensor | None, + tau: float | torch.Tensor, + obs_tokens: torch.Tensor | None = None, + conditioning_obs_tokens: torch.Tensor | None = None, + video_latents: torch.Tensor | None = None, + conditioning_video_latents: torch.Tensor | None = None, + video_tau: float | torch.Tensor | None = None, + action_tau: float | torch.Tensor | None = None, + track_tau: float | torch.Tensor | None = None, + context: torch.Tensor | None = None, + context_mask: torch.Tensor | None = None, + text_context: torch.Tensor | None = None, + clip_context: torch.Tensor | None = None, + return_hidden_states: bool = False, + ) -> CausalWanTrackForceOutput: + num_blocks = ( + action.shape[1] + if action.ndim == 4 + else action.shape[1] // self.action_horizon + ) + action = self._canonicalize_action(action, num_blocks) + state = self._canonicalize_state(state, num_blocks) + batch, device = action.shape[0], action.device + dtype = self.patch_embedding.weight.dtype + action = action.to(dtype=dtype) + state = state.to(device=device, dtype=dtype) + + conditioning, obs, video_grid, video_rope = self._tokenize_observations( + num_blocks=num_blocks, + obs_tokens=None if obs_tokens is None else obs_tokens.to(device, dtype), + conditioning_obs_tokens=( + None + if conditioning_obs_tokens is None + else conditioning_obs_tokens.to(device, dtype) + ), + video_latents=( + None if video_latents is None else video_latents.to(device, dtype) + ), + conditioning_video_latents=( + None + if conditioning_video_latents is None + else conditioning_video_latents.to(device, dtype) + ), + ) + base_tau = self._normalize_tau(tau, batch, num_blocks, device) + video_tau_tensor = self._normalize_tau( + tau if video_tau is None else video_tau, batch, num_blocks, device + ) + action_tau_tensor = self._normalize_tau( + tau if action_tau is None else action_tau, batch, num_blocks, device + ) + track_tau_tensor = self._normalize_tau( + base_tau if track_tau is None else track_tau, batch, num_blocks, device + ) + action_tokens, state_tokens = self._encode_action_state( + action, state, action_tau_tensor + ) + track_tokens = self.track_encoder( + past_track_xy.to(device, dtype), + None + if past_track_visibility is None + else past_track_visibility.to(device, dtype), + future_track_xy.to(device, dtype), + # Future visibility is a supervision mask, not an input. Feeding + # it here would leak target occlusion unavailable at inference. + None, + ) + hidden, layout = self._pack_tokens( + conditioning, + obs, + action_tokens, + state_tokens, + track_tokens.past_tokens, + track_tokens.future_tokens, + ) + + token_tau = self._pack_scalar_by_layout( + layout, video_tau_tensor, action_tau_tensor, track_tau_tensor + ) + time_embedding = self.time_embedding( + sinusoidal_embedding_1d( + self.freq_dim, (token_tau * 1000).flatten() + ).type_as(hidden) + ).view(batch, layout.sequence_length, self.dim) + modulation = self.time_projection(time_embedding).view( + batch, layout.sequence_length, 6, self.dim + ) + rope_frequencies = self._pack_rope(layout, video_rope, device) + context_tokens = self._build_context( + context, text_context, clip_context + ) + if context_tokens is not None: + context_tokens = context_tokens.to(device, dtype) + allow_matrix = build_dense_allow_matrix( + layout, + max_previous_blocks=( + None if self.max_chunk_size < 0 else self.max_chunk_size + ), + device=device, + ) + + for block in self.blocks: + def run_block( + values: torch.Tensor, + current_block: CausalWanTrackForceBlock = block, + ) -> torch.Tensor: + return current_block( + values, + e=modulation, + layout=layout, + rope_frequencies=rope_frequencies, + context=context_tokens, + context_mask=context_mask, + allow_matrix=allow_matrix, + ) + + if self.gradient_checkpointing and self.training and torch.is_grad_enabled(): + hidden = torch.utils.checkpoint.checkpoint( + run_block, hidden, use_reentrant=False + ) + else: + hidden = run_block(hidden) + + obs_hidden, action_hidden, _, future_hidden = self._unpack(hidden, layout) + category = torch.zeros(batch, dtype=torch.long, device=device) + action_flow = self.action_decoder( + action_hidden.flatten(1, 2), category + ).view(batch, num_blocks, self.action_horizon, self.action_dim) + track_flow = self.track_decoder(future_hidden) + obs_flow_tokens = self.obs_token_decoder(obs_hidden) + + memory_indices = layout.indices( + ( + TokenType.CONDITIONING_OBS, + TokenType.OBS, + TokenType.ACTION, + TokenType.TRACK_PAST, + TokenType.TRACK_FUTURE, + ), + device=device, + ) + coarse_memory = hidden.index_select(1, memory_indices) + coarse_memory_by_block, coarse_memory_mask_by_block = ( + self._causal_coarse_memory(hidden, layout) + ) + video_flow = None + if video_grid is not None: + obs_time, _, _, _ = self._unpack(time_embedding, layout) + video_tokens = self.head( + obs_hidden.flatten(1, 2), + obs_time.flatten(1, 2).unsqueeze(2), + ) + video_flow = self.unpatchify( + video_tokens, torch.tensor(video_grid, dtype=torch.long) + ) + return CausalWanTrackForceOutput( + action_flow=action_flow, + track_flow=track_flow, + obs_flow_tokens=obs_flow_tokens, + coarse_memory=coarse_memory, + coarse_memory_by_block=coarse_memory_by_block, + coarse_memory_mask_by_block=coarse_memory_mask_by_block, + layout=layout, + video_flow=video_flow, + hidden_states=hidden if return_hidden_states else None, + ) + + +CausalWanModel = CausalWanTrackForceModel +TrackForceCausalWanModel = CausalWanTrackForceModel + + +__all__ = [ + "CausalWanModel", + "CausalWanTrackForceModel", + "CausalWanTrackForceOutput", + "DreamZeroCausalWanModel", + "TrackForceCausalWanModel", + "WAN_SHARED_PREFIXES", +] diff --git a/groot/vla/model/trex_track_force/policy.py b/groot/vla/model/trex_track_force/policy.py new file mode 100644 index 0000000000000000000000000000000000000000..d713efe464a600ff2427ccdc7646f714ecdd43da --- /dev/null +++ b/groot/vla/model/trex_track_force/policy.py @@ -0,0 +1,1677 @@ +"""Independent WAN policy head for T-Rex track and force training/inference.""" + +from __future__ import annotations + +import json +import os +from collections.abc import Iterator +from typing import Any + +from hydra.utils import instantiate +import torch +from torch import nn +from transformers import PretrainedConfig +from transformers.feature_extraction_utils import BatchFeature + +from groot.vla.model.n1_5.action_head.base_action_head import ActionHead + +from .force import ( + ACTION_HORIZON, + FORCE_HISTORY_FRAMES, + FORCE_SENSOR_COUNT, + FORCE_SENSOR_DIM, + PADDED_ACTION_DIM, + ForceOnlyOutput, + ForceOnlyTransformer, + TwoStageFlowSchedule, + euler_flow_step, + pad_action_62_to_64, + sample_trex_timestep, +) +from .lora import inject_lora, set_lora_and_new_modules_trainable +from .model import CausalWanTrackForceModel, CausalWanTrackForceOutput +from .track import ( + NUM_TRACK_POINTS, + TRACK_HORIZON, + track_motion_weights, + visibility_masked_track_loss, + visibility_masked_track_smoothness_loss, + zero_invisible_track_xy, +) + + +class WANTrackForcePolicyHeadConfig(PretrainedConfig): + model_type = "trex_track_force_policy" + + def __init__( + self, + *, + diffusion_model_cfg: dict[str, Any] | None = None, + force_model_cfg: dict[str, Any] | None = None, + text_encoder_cfg: dict[str, Any] | None = None, + image_encoder_cfg: dict[str, Any] | None = None, + vae_cfg: dict[str, Any] | None = None, + load_pretrained_components: bool = False, + allow_checkpoint_download: bool = False, + action_dim: int = PADDED_ACTION_DIM, + action_horizon: int = ACTION_HORIZON, + physical_action_dim: int = 62, + track_points: int = NUM_TRACK_POINTS, + track_horizon: int = TRACK_HORIZON, + track_loss_weight: float = 1.0, + track_position_loss_weight: float = 0.0, + track_smoothness_loss_weight: float = 0.0, + track_motion_loss_boost: float = 0.0, + zero_invisible_track_xy: bool = True, + force_loss_weight: float = 1.0, + vq_loss_weight: float = 1.0, + vq_commitment_weight: float = 0.25, + tactile_dropout: float = 0.0, + dynamics_loss_weight: float = 1.0, + require_force_inputs: bool = False, + target_video_height: int = 160, + target_video_width: int = 320, + train_architecture: str = "lora", + lora_rank: int = 4, + lora_alpha: float = 4.0, + lora_target_modules: str = "q,k,v,o,k_img,v_img,ffn.0,ffn.2", + defer_lora_injection: bool = False, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self.diffusion_model_cfg = diffusion_model_cfg + self.force_model_cfg = force_model_cfg + self.text_encoder_cfg = text_encoder_cfg + self.image_encoder_cfg = image_encoder_cfg + self.vae_cfg = vae_cfg + self.load_pretrained_components = load_pretrained_components + self.allow_checkpoint_download = allow_checkpoint_download + self.action_dim = action_dim + self.action_horizon = action_horizon + self.physical_action_dim = physical_action_dim + self.track_points = track_points + self.track_horizon = track_horizon + self.track_loss_weight = track_loss_weight + if track_position_loss_weight < 0: + raise ValueError("track_position_loss_weight must be non-negative") + if track_smoothness_loss_weight < 0: + raise ValueError("track_smoothness_loss_weight must be non-negative") + if track_motion_loss_boost < 0: + raise ValueError("track_motion_loss_boost must be non-negative") + self.track_position_loss_weight = track_position_loss_weight + self.track_smoothness_loss_weight = track_smoothness_loss_weight + self.track_motion_loss_boost = track_motion_loss_boost + self.zero_invisible_track_xy = bool(zero_invisible_track_xy) + self.force_loss_weight = force_loss_weight + self.vq_loss_weight = vq_loss_weight + self.vq_commitment_weight = vq_commitment_weight + if not 0.0 <= tactile_dropout <= 1.0: + raise ValueError("tactile_dropout must be in [0,1]") + self.tactile_dropout = tactile_dropout + self.dynamics_loss_weight = dynamics_loss_weight + self.require_force_inputs = require_force_inputs + self.target_video_height = target_video_height + self.target_video_width = target_video_width + self.train_architecture = train_architecture + self.lora_rank = lora_rank + self.lora_alpha = lora_alpha + self.lora_target_modules = lora_target_modules + self.defer_lora_injection = defer_lora_injection + + +def _instantiate_or_default( + config: dict[str, Any] | None, + default_type: type[nn.Module], + **default_kwargs: Any, +) -> nn.Module: + if config is None: + return default_type(**default_kwargs) + if "_target_" in config: + return instantiate(config) + return default_type(**dict(config)) + + +class WANTrackForcePolicyHead(ActionHead): + """VLA-compatible action head with a runnable lightweight core interface.""" + + config_class = WANTrackForcePolicyHeadConfig + supports_gradient_checkpointing = True + + def __init__(self, config: WANTrackForcePolicyHeadConfig) -> None: + super().__init__() + if config.action_dim != PADDED_ACTION_DIM or config.action_horizon != ACTION_HORIZON: + raise ValueError("the model contract is action=[16,64] (62 physical + 2 pad)") + if config.track_points != NUM_TRACK_POINTS or config.track_horizon != TRACK_HORIZON: + raise ValueError("the model contract is 250 points x 16 future steps") + self.config = config + self.action_dim = config.action_dim + self.action_horizon = config.action_horizon + self.track_points = config.track_points + self.track_horizon = config.track_horizon + self.train_architecture = config.train_architecture.lower() + + self.model = _instantiate_or_default( + config.diffusion_model_cfg, + CausalWanTrackForceModel, + action_dim=config.action_dim, + action_horizon=config.action_horizon, + track_points=config.track_points, + track_horizon=config.track_horizon, + ) + if not isinstance(self.model, CausalWanTrackForceModel): + raise TypeError("diffusion_model_cfg must create CausalWanTrackForceModel") + + force_defaults = { + "action_dim": config.action_dim, + "action_horizon": config.action_horizon, + "coarse_memory_dim": self.model.dim, + } + self.force_transformer = _instantiate_or_default( + config.force_model_cfg, + ForceOnlyTransformer, + **force_defaults, + ) + if not isinstance(self.force_transformer, ForceOnlyTransformer): + raise TypeError("force_model_cfg must create ForceOnlyTransformer") + + self.text_encoder = ( + instantiate(config.text_encoder_cfg) + if config.text_encoder_cfg is not None + else None + ) + self.image_encoder = ( + instantiate(config.image_encoder_cfg) + if config.image_encoder_cfg is not None + else None + ) + self.vae = instantiate(config.vae_cfg) if config.vae_cfg is not None else None + self.schedule = TwoStageFlowSchedule() + self._lora_injected = False + + if config.load_pretrained_components: + self.load_pretrained_component_weights() + self.set_trainable_parameters() + + @property + def device(self) -> torch.device: + return next(self.model.parameters()).device + + @property + def dtype(self) -> torch.dtype: + return next(self.model.parameters()).dtype + + def prepare_input(self, batch: dict[str, Any]) -> BatchFeature: + return BatchFeature(data=batch) + + def _resolve_checkpoint_file( + self, + path: str | None, + filename: str, + *, + repo_id: str, + ) -> str: + if path and os.path.isfile(path): + return path + if path and os.path.isdir(path): + candidate = os.path.join(path, filename) + if os.path.isfile(candidate): + return candidate + if not self.config.allow_checkpoint_download: + raise FileNotFoundError(path or filename) + from huggingface_hub import hf_hub_download + + return hf_hub_download(repo_id=repo_id, filename=filename) + + @staticmethod + def _unwrap_state_dict(state: Any) -> dict[str, torch.Tensor]: + if not isinstance(state, dict): + raise TypeError("checkpoint must contain a state-dict mapping") + for key in ("state_dict", "model", "module"): + nested = state.get(key) + if isinstance(nested, dict) and nested: + state = nested + break + return {key: value for key, value in state.items() if torch.is_tensor(value)} + + @classmethod + def _load_component_state( + cls, + target: nn.Module, + path: str, + *, + minimum_numel_coverage: float, + ) -> dict[str, float | int]: + try: + checkpoint = torch.load( + path, map_location="cpu", weights_only=True, mmap=True + ) + except TypeError: + checkpoint = torch.load(path, map_location="cpu") + source = cls._unwrap_state_dict(checkpoint) + current = target.state_dict() + compatible: dict[str, torch.Tensor] = {} + target_numel = sum(value.numel() for value in current.values()) + matched_numel = 0 + for key, value in source.items(): + normalized = key.removeprefix("module.") + if normalized in current and current[normalized].shape == value.shape: + compatible[normalized] = value + matched_numel += value.numel() + coverage = matched_numel / max(target_numel, 1) + if coverage < minimum_numel_coverage: + raise ValueError( + f"component checkpoint coverage {coverage:.3f} " + f"< {minimum_numel_coverage:.3f}: {path}" + ) + target.load_state_dict(compatible, strict=False) + return { + "checkpoint_keys": len(source), + "matched_keys": len(compatible), + "numel_coverage": coverage, + } + + def _iter_diffusion_states( + self, checkpoint: str | None + ) -> Iterator[dict[str, torch.Tensor]]: + repo_id = "Wan-AI/Wan2.2-TI2V-5B" + filename = "diffusion_pytorch_model.safetensors" + if checkpoint and os.path.isfile(checkpoint): + if checkpoint.endswith(".safetensors"): + from safetensors.torch import load_file + + yield load_file(checkpoint) + else: + yield self._unwrap_state_dict( + torch.load(checkpoint, map_location="cpu") + ) + return + + directory = checkpoint if checkpoint and os.path.isdir(checkpoint) else None + index_name = "diffusion_pytorch_model.safetensors.index.json" + index_path = os.path.join(directory, index_name) if directory else None + single_path = os.path.join(directory, filename) if directory else None + if index_path and os.path.isfile(index_path): + pass + elif single_path and os.path.isfile(single_path): + from safetensors.torch import load_file + + yield load_file(single_path) + return + elif self.config.allow_checkpoint_download: + from huggingface_hub import hf_hub_download + + index_path = hf_hub_download(repo_id=repo_id, filename=index_name) + directory = os.path.dirname(index_path) + else: + raise FileNotFoundError(checkpoint or index_name) + + assert index_path is not None and directory is not None + with open(index_path, "r", encoding="utf-8") as handle: + index = json.load(handle) + shard_names = sorted(set(index["weight_map"].values())) + if self.config.allow_checkpoint_download: + from huggingface_hub import hf_hub_download + + shard_paths = [ + ( + os.path.join(directory, shard) + if os.path.isfile(os.path.join(directory, shard)) + else hf_hub_download(repo_id=repo_id, filename=shard) + ) + for shard in shard_names + ] + else: + shard_paths = [os.path.join(directory, shard) for shard in shard_names] + missing = [path for path in shard_paths if not os.path.isfile(path)] + if missing: + raise FileNotFoundError(missing[0]) + from safetensors.torch import load_file + + for path in shard_paths: + yield load_file(path) + + def load_pretrained_component_weights(self) -> None: + """Load standard Wan2.2 files; network access is separately opt-in.""" + + if self.text_encoder is not None: + path = self._resolve_checkpoint_file( + getattr(self.text_encoder, "text_encoder_pretrained_path", None), + "models_t5_umt5-xxl-enc-bf16.pth", + repo_id="Wan-AI/Wan2.2-TI2V-5B", + ) + self._load_component_state( + self.text_encoder, path, minimum_numel_coverage=0.99 + ) + if self.image_encoder is not None: + path = self._resolve_checkpoint_file( + getattr(self.image_encoder, "image_encoder_pretrained_path", None), + "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth", + repo_id="Wan-AI/Wan2.1-I2V-14B-480P", + ) + self._load_component_state( + self.image_encoder.model, path, minimum_numel_coverage=0.90 + ) + if self.vae is not None: + path = self._resolve_checkpoint_file( + getattr(self.vae, "vae_pretrained_path", None), + "Wan2.2_VAE.pth", + repo_id="Wan-AI/Wan2.2-TI2V-5B", + ) + self._load_component_state( + self.vae.model, path, minimum_numel_coverage=0.99 + ) + + checkpoint = self.model.diffusion_model_pretrained_path + if checkpoint is None and not self.config.allow_checkpoint_download: + return + for state in self._iter_diffusion_states(checkpoint): + self.model.load_pretrained_wan_state_dict( + state, minimum_numel_coverage=0.95 + ) + self._initialize_ti2v_image_attention() + + def _initialize_ti2v_image_attention(self) -> None: + """Seed Wan2.2's new CLIP path from pretrained text cross-attention.""" + + for block in self.model.blocks: + cross_attention = block.cross_attn + if not hasattr(cross_attention, "k_img"): + continue + cross_attention.k_img.load_state_dict(cross_attention.k.state_dict()) + cross_attention.v_img.load_state_dict(cross_attention.v.state_dict()) + if hasattr(cross_attention, "norm_k_img") and hasattr( + cross_attention, "norm_k" + ): + cross_attention.norm_k_img.load_state_dict( + cross_attention.norm_k.state_dict() + ) + + def _new_trainable_modules(self) -> list[nn.Module]: + modules = [ + self.model.action_encoder, + self.model.action_decoder, + self.model.state_encoder, + self.model.track_encoder, + self.model.track_decoder, + self.model.modality_embedding, + self.force_transformer, + ] + # The configured Wan path trains on video latents. The direct + # obs-token adapters are an alternate/testing path and receive no + # gradient from video_flow loss, so leave them frozen in LoRA mode. + if hasattr(self.model, "img_emb"): + # Wan2.2-TI2V-5B has no img_emb keys; this CLIP projection is new. + modules.append(self.model.img_emb) + return modules + + def inject_lora_after_loading(self) -> None: + if self._lora_injected or self.config.train_architecture != "lora": + return + inject_lora( + self.model.blocks, + target_modules=self.config.lora_target_modules, + rank=self.config.lora_rank, + alpha=self.config.lora_alpha, + ) + set_lora_and_new_modules_trainable( + self, new_modules=self._new_trainable_modules() + ) + self._freeze_shared_encoders() + self._lora_injected = True + + def _freeze_shared_encoders(self) -> None: + for module in (self.text_encoder, self.image_encoder, self.vae): + if module is not None: + module.requires_grad_(False) + module.eval() + + def set_trainable_parameters(self) -> None: + architecture = self.config.train_architecture.lower() + if architecture in {"full", "full_finetune"}: + self.model.requires_grad_(True) + self.force_transformer.requires_grad_(True) + self._freeze_shared_encoders() + return + if architecture != "lora": + raise ValueError(f"unsupported train_architecture={architecture!r}") + if self.config.defer_lora_injection: + self.requires_grad_(False) + for module in self._new_trainable_modules(): + module.requires_grad_(True) + self._freeze_shared_encoders() + else: + self.inject_lora_after_loading() + + def train(self, mode: bool = True): + super().train(mode) + self._freeze_shared_encoders() + return self + + @staticmethod + def _get( + data: dict[str, Any] | BatchFeature, + *names: str, + required: bool = True, + ) -> Any: + for name in names: + if name in data: + return data[name] + if required: + raise KeyError(f"missing required input; expected one of {names}") + return None + + def _reshape_action(self, action: torch.Tensor) -> tuple[torch.Tensor, int]: + action = pad_action_62_to_64(action) + if action.ndim == 4: + blocks = action.shape[1] + elif action.ndim == 3 and action.shape[1] % self.action_horizon == 0: + blocks = action.shape[1] // self.action_horizon + action = action.view( + action.shape[0], blocks, self.action_horizon, self.action_dim + ) + else: + raise ValueError("action must be [B,K,16,D] or [B,K*16,D]") + return action, blocks + + def _reshape_state(self, state: torch.Tensor, blocks: int) -> torch.Tensor: + if state.ndim == 2: + state = state[:, None].expand(-1, blocks, -1) + elif state.ndim == 3 and state.shape[1] == blocks: + pass + elif state.ndim == 3 and state.shape[1] % blocks == 0: + state = state.view(state.shape[0], blocks, -1, state.shape[-1]) + else: + raise ValueError("state cannot be aligned to action blocks") + return state + + @staticmethod + def _future_to_model(future: torch.Tensor) -> torch.Tensor: + if future.ndim != 5 or future.shape[-1] != 2: + raise ValueError("future track must be rank-5 coordinates") + if future.shape[2] == TRACK_HORIZON and future.shape[3] == NUM_TRACK_POINTS: + return future + if future.shape[2] == NUM_TRACK_POINTS and future.shape[3] == TRACK_HORIZON: + return future.transpose(2, 3) + raise ValueError("future track must contain 16 time and 250 point axes") + + @staticmethod + def _future_to_decoder(future: torch.Tensor) -> torch.Tensor: + return WANTrackForcePolicyHead._future_to_model(future).transpose(2, 3) + + @staticmethod + def _track_anchor(past: torch.Tensor) -> torch.Tensor: + """Return the current per-point position as ``[B,K,250,2]``.""" + + if past.ndim != 5 or past.shape[-1] != 2: + raise ValueError("past track must be rank-5 coordinates") + if past.shape[-2] == NUM_TRACK_POINTS: + return past[:, :, -1] + if past.shape[2] == NUM_TRACK_POINTS: + return past[:, :, :, -1] + raise ValueError("past track must contain a 250-point axis") + + def _prepare_absolute_track_target( + self, future: torch.Tensor, past: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return OpenPI-style absolute XY with target timestep zero anchored.""" + + absolute = self._future_to_model(future).to( + device=past.device, dtype=past.dtype + ) + anchor = self._track_anchor(past) + absolute = absolute.clone() + absolute[:, :, 0] = anchor + return absolute, anchor + + @staticmethod + def _anchor_track_prediction( + prediction: torch.Tensor, anchor: torch.Tensor + ) -> torch.Tensor: + """Make absolute prediction timestep zero exactly the current GT XY.""" + + if prediction.ndim != 5 or prediction.shape[2:4] != ( + TRACK_HORIZON, + NUM_TRACK_POINTS, + ): + raise ValueError("track prediction must be [B,K,16,250,2]") + prediction = prediction.clone() + prediction[:, :, 0] = anchor + return prediction + + @staticmethod + def _flow_noising( + clean: torch.Tensor, + tau: torch.Tensor, + *, + noise: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + noise = torch.randn_like(clean) if noise is None else noise + expanded_tau = tau + while expanded_tau.ndim < clean.ndim: + expanded_tau = expanded_tau.unsqueeze(-1) + noisy = (1 - expanded_tau) * clean + expanded_tau * noise + return noisy, noise - clean, noise + + def _extract_force_inputs( + self, data: dict[str, Any] | BatchFeature, blocks: int + ) -> tuple[ + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: + current_force = self._get( + data, + "current_force", + "force_current", + "observation.force_current", + required=False, + ) + history = self._get( + data, + "tactile_force_history", + "force_history_raw", + "observation.tactile_force", + "force_history_vq", + "force_vq_codes", + "observation.force_history_vq", + required=False, + ) + history_padding_mask = self._get( + data, + "tactile_force_history_padding_mask", + "force_history_padding_mask", + required=False, + ) + if ( + current_force is None + and history is not None + and history.is_floating_point() + ): + if history.ndim == 6: + current_force = history[:, :, :, -1] + elif history.ndim == 5: + current_force = history[:, :, -1] + elif history.ndim == 4: + current_force = history[:, -1] + if current_force is None or history is None: + if self.config.require_force_inputs: + raise KeyError( + "current force and raw [B,K,S,16,10,6] history " + "or precomputed [B,K,S,10] tactile codes are required" + ) + return None, None, None + if current_force.ndim == 3: + current_force = current_force[:, None, None].expand( + -1, blocks, 1, -1, -1 + ) + elif current_force.ndim == 4: + current_force = current_force[:, :, None] + if history.is_floating_point(): + if history.ndim == 4: + history = history[:, None, None].expand( + -1, blocks, 1, -1, -1, -1 + ) + elif history.ndim == 5: + history = history[:, :, None] + expected_history = ( + blocks, + history.shape[2], + FORCE_HISTORY_FRAMES, + FORCE_SENSOR_COUNT, + FORCE_SENSOR_DIM, + ) + if history.shape[1:] != expected_history: + raise ValueError( + "raw tactile history must be [B,K,S,16,10,6], " + f"got {tuple(history.shape)}" + ) + expected_padding = history.shape[:4] + if history_padding_mask is None: + history_valid_mask = torch.ones( + expected_padding, + device=history.device, + dtype=torch.bool, + ) + else: + padding = history_padding_mask.to( + device=history.device, dtype=torch.bool + ) + if padding.ndim == 2: + if padding.shape != ( + history.shape[0], + FORCE_HISTORY_FRAMES, + ): + raise ValueError( + "force history padding mask must be [B,16]" + ) + padding = padding[:, None, None].expand(expected_padding) + elif padding.ndim == 3: + if padding.shape[1:] == ( + blocks, + FORCE_HISTORY_FRAMES, + ): + padding = padding[:, :, None].expand(expected_padding) + elif blocks == 1 and padding.shape[1:] == ( + history.shape[2], + FORCE_HISTORY_FRAMES, + ): + padding = padding[:, None] + else: + raise ValueError( + "force history padding mask must align with " + "[B,K,S,16]" + ) + if padding.shape != expected_padding: + raise ValueError( + "force history padding mask must be [B,K,S,16], " + f"got {tuple(padding.shape)}" + ) + history_valid_mask = ~padding + else: + if history_padding_mask is not None: + raise ValueError( + "force history padding mask cannot accompany VQ codes" + ) + history_valid_mask = None + if history.ndim == 2: + history = history[:, None, None].expand(-1, blocks, 1, -1) + elif history.ndim == 3: + history = history[:, :, None] + expected_codes = (blocks, history.shape[2], FORCE_SENSOR_COUNT) + if history.shape[1:] != expected_codes: + raise ValueError( + "precomputed tactile codes must be [B,K,S,10], " + f"got {tuple(history.shape)}" + ) + expected_force = ( + blocks, + current_force.shape[2], + FORCE_SENSOR_COUNT, + FORCE_SENSOR_DIM, + ) + if current_force.shape[1:] != expected_force: + raise ValueError( + "current_force must be [B,K,S,10,6], got " + f"{tuple(current_force.shape)}" + ) + if current_force.shape[2] != history.shape[2]: + raise ValueError("current force and history update-slot counts differ") + if current_force.shape[2] not in (1, len(self.schedule.offsets)): + raise ValueError( + f"force update slots must be 1 or {len(self.schedule.offsets)}" + ) + return current_force, history, history_valid_mask + + def _prepare_observation_inputs( + self, + data: dict[str, Any] | BatchFeature, + tau: torch.Tensor, + blocks: int, + ) -> tuple[dict[str, torch.Tensor], torch.Tensor | None]: + obs_tokens = self._get(data, "obs_tokens", required=False) + video_latents = self._get(data, "video_latents", required=False) + if obs_tokens is not None: + if obs_tokens.ndim != 4 or obs_tokens.shape[1] != blocks: + raise ValueError("obs_tokens must be [B,blocks,O,D]") + noisy_obs, obs_target, _ = self._flow_noising(obs_tokens, tau) + result = {"obs_tokens": noisy_obs} + conditioning = self._get( + data, "conditioning_obs_tokens", required=False + ) + if conditioning is not None: + result["conditioning_obs_tokens"] = conditioning + return result, obs_target + conditioning_video_latents = self._get( + data, "conditioning_video_latents", "history_video_latents", required=False + ) + future_video_latents = self._get( + data, "future_video_latents", required=False + ) + history_images = self._get(data, "history_images", required=False) + future_images = self._get(data, "future_images", required=False) + if future_video_latents is not None or future_images is not None: + if future_video_latents is None: + if self.vae is None or history_images is None: + raise KeyError( + "future_images require history_images and a configured VAE" + ) + anchor_and_future = torch.cat( + (history_images[:, -1:], future_images), dim=1 + ) + encoded_future = self.encode_video(anchor_and_future) + future_video_latents = encoded_future[:, :, 1:] + if conditioning_video_latents is None: + if history_images is None or self.vae is None: + raise KeyError( + "separate future latents require history video latents/images" + ) + conditioning_video_latents = self.encode_video(history_images) + video_latents = future_video_latents + + if video_latents is not None: + future_frames = blocks * self.model.num_frame_per_block + if conditioning_video_latents is not None: + if video_latents.shape[2] != future_frames: + raise ValueError( + "future video latents temporal axis does not match blocks" + ) + conditioning = conditioning_video_latents + future = video_latents + elif video_latents.shape[2] == future_frames + 1: + conditioning = video_latents[:, :, :1] + future = video_latents[:, :, 1:] + elif video_latents.shape[2] == future_frames: + conditioning = None + future = video_latents + else: + raise ValueError("video_latents temporal axis does not match blocks") + per_frame_tau = tau.repeat_interleave( + self.model.num_frame_per_block, dim=1 + ) + noise = torch.randn_like(future) + expanded_tau = per_frame_tau[:, None, :, None, None] + noisy_future = (1 - expanded_tau) * future + expanded_tau * noise + target = noise - future + if conditioning_video_latents is not None: + return { + "conditioning_video_latents": conditioning, + "video_latents": noisy_future, + }, target + noisy_video = noisy_future if conditioning is None else torch.cat( + (conditioning, noisy_future), dim=2 + ) + return {"video_latents": noisy_video}, target + images = self._get(data, "images", required=False) + if images is None or self.vae is None: + raise KeyError( + "provide obs_tokens/video_latents, or configure VAE and provide images" + ) + latents = self.encode_video(images) + copied = dict(data) + copied["video_latents"] = latents + return self._prepare_observation_inputs(copied, tau, blocks) + + def _prepare_sampling_observation( + self, + data: dict[str, Any] | BatchFeature, + *, + blocks: int, + generator: torch.Generator | None, + ) -> dict[str, torch.Tensor]: + """Create the future-observation noise updated by the coarse stage.""" + + obs_tokens = self._get(data, "obs_tokens", required=False) + if obs_tokens is not None: + obs_tokens = obs_tokens.to(device=self.device, dtype=self.dtype) + noisy_obs = torch.randn( + obs_tokens.shape, + device=self.device, + dtype=self.dtype, + generator=generator, + ) + result = {"obs_tokens": noisy_obs} + conditioning = self._get( + data, "conditioning_obs_tokens", required=False + ) + if conditioning is not None: + result["conditioning_obs_tokens"] = conditioning.to( + device=self.device, dtype=self.dtype + ) + return result + + video_latents = self._get( + data, "future_video_latents", "video_latents", required=False + ) + conditioning_latent = self._get( + data, "conditioning_video_latents", "history_video_latents", required=False + ) + history_images = self._get(data, "history_images", required=False) + if conditioning_latent is None and history_images is not None: + if self.vae is None: + raise KeyError("history_images require a configured VAE") + conditioning_latent = self.encode_video(history_images) + if video_latents is None: + if conditioning_latent is not None: + video_latents = conditioning_latent + else: + images = ( + history_images + if history_images is not None + else self._get(data, "images", required=False) + ) + if images is None or self.vae is None: + raise KeyError( + "sampling requires obs token template, video latents, or VAE images" + ) + video_latents = self.encode_video(images) + video_latents = video_latents.to(device=self.device, dtype=self.dtype) + future_frames = blocks * self.model.num_frame_per_block + if conditioning_latent is not None: + if video_latents.shape[2] in (future_frames,): + future_shape = video_latents.shape + else: + # At inference, encoded history supplies only a spatial template. + future_shape = ( + video_latents.shape[0], + video_latents.shape[1], + future_frames, + video_latents.shape[3], + video_latents.shape[4], + ) + elif video_latents.shape[2] == future_frames + 1: + conditioning_latent = video_latents[:, :, :1] + future_shape = video_latents[:, :, 1:].shape + elif video_latents.shape[2] == 1: + conditioning_latent = video_latents + future_shape = ( + video_latents.shape[0], + video_latents.shape[1], + future_frames, + video_latents.shape[3], + video_latents.shape[4], + ) + elif video_latents.shape[2] == future_frames: + conditioning_latent = None + future_shape = video_latents.shape + else: + raise ValueError("video latent template cannot be aligned to coarse blocks") + noisy_future = torch.randn( + future_shape, + device=self.device, + dtype=self.dtype, + generator=generator, + ) + if conditioning_latent is not None and self._get( + data, + "conditioning_video_latents", + "history_video_latents", + "history_images", + required=False, + ) is None: + noisy_future = torch.cat((conditioning_latent, noisy_future), dim=2) + return {"video_latents": noisy_future} + result = {"video_latents": noisy_future} + if conditioning_latent is not None: + result["conditioning_video_latents"] = conditioning_latent.to( + device=self.device, dtype=self.dtype + ) + return result + + @staticmethod + def _advance_sampling_observation( + observation: dict[str, torch.Tensor], + output: CausalWanTrackForceOutput, + tau: float, + tau_next: float, + ) -> dict[str, torch.Tensor]: + """Euler-update only the generated observation suffix.""" + + updated = dict(observation) + if "obs_tokens" in updated: + updated["obs_tokens"] = euler_flow_step( + updated["obs_tokens"], + output.obs_flow_tokens, + tau, + tau_next, + ) + return updated + video = updated["video_latents"] + if output.video_flow is None: + raise RuntimeError("latent sampling requires video_flow from CausalWan") + if video.shape[2] == output.video_flow.shape[2] + 1: + updated_future = euler_flow_step( + video[:, :, 1:], output.video_flow, tau, tau_next + ) + updated["video_latents"] = torch.cat( + (video[:, :, :1], updated_future), dim=2 + ) + else: + updated["video_latents"] = euler_flow_step( + video, output.video_flow, tau, tau_next + ) + return updated + + def encode_video(self, images: torch.Tensor) -> torch.Tensor: + """Encode ``[B,T,H,W,C]`` images with the shared Wan VAE38.""" + + if self.vae is None: + raise RuntimeError("VAE is not configured") + if images.ndim == 5 and images.shape[-1] == 3: + images = images.permute(0, 4, 1, 2, 3) + if images.ndim != 5 or images.shape[1] != 3: + raise ValueError("images must be [B,T,H,W,3] or [B,3,T,H,W]") + images = images.to(device=self.device, dtype=self.dtype) + if images.max() > 1: + images = images / 255.0 + images = images * 2 - 1 + target_size = ( + int(self.config.target_video_height), + int(self.config.target_video_width), + ) + if images.shape[-2:] != target_size: + batch, channels, time, height, width = images.shape + images = torch.nn.functional.interpolate( + images.permute(0, 2, 1, 3, 4).reshape( + batch * time, channels, height, width + ), + size=target_size, + mode="bilinear", + align_corners=False, + ).view(batch, time, channels, *target_size).permute(0, 2, 1, 3, 4) + # WanVideoVAE stores normalization tensors as ordinary attributes, + # so Module.to() does not move them across runtime CUDA devices. + if hasattr(self.vae, "mean") and hasattr(self.vae, "std"): + self.vae.mean = self.vae.mean.to(images.device, images.dtype) + self.vae.std = self.vae.std.to(images.device, images.dtype) + self.vae.scale = [self.vae.mean, self.vae.std.reciprocal()] + with torch.no_grad(): + return self.vae.encode(images, tiled=False) + + def _context_kwargs( + self, data: dict[str, Any] | BatchFeature + ) -> dict[str, torch.Tensor]: + result: dict[str, torch.Tensor] = {} + context_mask: torch.Tensor | None = None + text_mask: torch.Tensor | None = None + clip_mask: torch.Tensor | None = None + context = self._get(data, "context_tokens", required=False) + if context is not None: + result["context"] = context + mask = self._get(data, "context_mask", required=False) + context_mask = ( + torch.ones(context.shape[:2], dtype=torch.bool, device=context.device) + if mask is None + else mask.to(dtype=torch.bool) + ) + if self.text_encoder is not None and "text" in data: + attention_mask = self._get(data, "text_attention_mask") + with torch.no_grad(): + result["text_context"] = self.text_encoder( + data["text"], attention_mask + ) + text_mask = attention_mask.to(dtype=torch.bool) + clip_context = self._get(data, "clip_context", required=False) + if clip_context is None and self.image_encoder is not None and ( + "history_images" in data or "images" in data + ): + images = data.get("history_images", data.get("images")) + if images.ndim == 5 and images.shape[-1] == 3: + first_image = images[:, -1:].permute(0, 1, 4, 2, 3) + elif images.ndim == 5 and images.shape[1] == 3: + first_image = images[:, :, -1:].transpose(1, 2) + else: + raise ValueError("images cannot be converted to CLIP input") + first_image = first_image.to(device=self.device, dtype=self.dtype) + if first_image.max() > 1: + first_image = first_image / 255.0 + first_image = first_image * 2 - 1 + with torch.no_grad(): + clip_context = self.image_encoder.encode_image(first_image) + if clip_context is not None: + result["clip_context"] = clip_context + clip_mask = torch.ones( + clip_context.shape[:2], + dtype=torch.bool, + device=clip_context.device, + ) + # Must match model._build_context order: generic, CLIP, then text. + masks = [ + mask + for mask in (context_mask, clip_mask, text_mask) + if mask is not None + ] + if masks: + result["context_mask"] = torch.cat(masks, dim=1) + return result + + def forward_core( + self, + data: dict[str, Any] | BatchFeature, + *, + tau: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: + action_clean, blocks = self._reshape_action(self._get(data, "action")) + state = self._reshape_state(self._get(data, "state"), blocks) + batch_size = action_clean.shape[0] + action_clean = action_clean.to(device=self.device, dtype=self.dtype) + state = state.to(device=self.device, dtype=self.dtype) + if tau is None: + # T-Rex trains the slow action expert over the full flow domain, + # even though cascaded inference only runs its upper segment. + tau = sample_trex_timestep( + (batch_size, blocks), + device=self.device, + dtype=self.dtype, + ) + else: + tau = tau.to(device=self.device, dtype=self.dtype) + if tau.ndim == 1: + tau = tau[:, None].expand(-1, blocks) + action_mask = self._get(data, "action_mask", required=False) + if action_mask is None: + action_mask = torch.zeros_like(action_clean, dtype=torch.bool) + action_mask[..., : self.config.physical_action_dim] = True + else: + action_mask, _ = self._reshape_action(action_mask) + action_mask = action_mask.to(device=self.device, dtype=self.dtype) + + past_xy = self._get(data, "track_past_xy", "past_track_xy").to( + device=self.device, dtype=self.dtype + ) + past_visibility = self._get( + data, + "track_past_visibility", + "past_track_visibility", + required=False, + ) + future_clean_model, track_anchor = self._prepare_absolute_track_target( + self._get(data, "track_future_xy", "future_track_xy"), past_xy + ) + future_clean_model = future_clean_model.to( + device=self.device, dtype=self.dtype + ) + future_visibility = self._get( + data, + "track_future_visibility", + "future_track_visibility", + ).to(device=self.device, dtype=self.dtype) + if past_visibility is not None: + past_visibility = past_visibility.to(device=self.device, dtype=self.dtype) + if self.config.zero_invisible_track_xy: + # Invisible coordinates are not supervised, but nonzero garbage XY + # still enters the track encoder via noisy future tokens. Zero them + # before CFM noising so the network never has to denoise junk. + future_clean_model = zero_invisible_track_xy( + future_clean_model, future_visibility + ) + past_xy = zero_invisible_track_xy(past_xy, past_visibility) + + action_noise = torch.randn_like(action_clean) * action_mask + action_noisy, action_target, _ = self._flow_noising( + action_clean, tau, noise=action_noise + ) + future_noisy, future_target_model, _ = self._flow_noising( + future_clean_model, tau + ) + observation_data: dict[str, Any] | BatchFeature = data + history_images = self._get(data, "history_images", required=False) + future_images = self._get(data, "future_images", required=False) + if ( + history_images is not None + and future_images is not None + and self._get(data, "future_video_latents", required=False) is None + ): + # Encode clean video once; the random-tau and split-tau training + # forwards share these immutable VAE targets. + observation_data = dict(data) + observation_data["conditioning_video_latents"] = self.encode_video( + history_images + ) + anchor_and_future = torch.cat( + (history_images[:, -1:], future_images), dim=1 + ) + observation_data["future_video_latents"] = self.encode_video( + anchor_and_future + )[:, :, 1:] + observation_kwargs, observation_target = self._prepare_observation_inputs( + observation_data, tau, blocks + ) + context_kwargs = self._context_kwargs(data) + output = self.model( + action=action_noisy, + state=state, + past_track_xy=past_xy, + past_track_visibility=past_visibility, + future_track_xy=future_noisy, + future_track_visibility=future_visibility, + tau=tau, + **observation_kwargs, + **context_kwargs, + ) + + action_error = (output.action_flow - action_target).square() * action_mask + action_denominator = action_mask.sum().clamp_min(1) + action_loss = action_error.sum() / action_denominator + + track_target = future_target_model.transpose(2, 3) + future_clean_decoder = future_clean_model.transpose(2, 3) + motion_weight = None + if self.config.track_motion_loss_boost > 0: + motion_weight = track_motion_weights( + future_clean_decoder, + future_visibility, + boost=self.config.track_motion_loss_boost, + ) + track_flow_loss = visibility_masked_track_loss( + output.track_flow, + track_target, + future_visibility, + weight=motion_weight, + ) + track_position_loss = action_loss.new_zeros(()) + track_smoothness_loss = action_loss.new_zeros(()) + if ( + self.config.track_position_loss_weight > 0 + or self.config.track_smoothness_loss_weight > 0 + ): + # One-step CFM reconstruction: x=(1-τ)c+τn and v*=n-c ⇒ c=x-τv. + expanded_tau = tau + while expanded_tau.ndim < future_noisy.ndim: + expanded_tau = expanded_tau.unsqueeze(-1) + pred_clean_model = future_noisy - expanded_tau * output.track_flow.transpose( + 2, 3 + ) + pred_clean_model = self._anchor_track_prediction( + pred_clean_model, track_anchor + ) + pred_clean_decoder = pred_clean_model.transpose(2, 3) + if self.config.track_position_loss_weight > 0: + track_position_loss = visibility_masked_track_loss( + pred_clean_decoder, + future_clean_decoder, + future_visibility, + weight=motion_weight, + ) + if self.config.track_smoothness_loss_weight > 0: + track_smoothness_loss = visibility_masked_track_smoothness_loss( + pred_clean_decoder, + future_clean_decoder, + future_visibility, + weight=motion_weight, + ) + track_loss = ( + track_flow_loss + + self.config.track_position_loss_weight * track_position_loss + + self.config.track_smoothness_loss_weight * track_smoothness_loss + ) + if observation_target is None: + dynamics_loss = action_loss.new_zeros(()) + elif output.video_flow is not None: + dynamics_loss = torch.nn.functional.mse_loss( + output.video_flow, observation_target + ) + else: + dynamics_loss = torch.nn.functional.mse_loss( + output.obs_flow_tokens, observation_target + ) + + current_force, force_history, force_history_valid = ( + self._extract_force_inputs(data, blocks) + ) + force_loss = action_loss.new_zeros(()) + vq_loss = action_loss.new_zeros(()) + commitment_loss = action_loss.new_zeros(()) + tactile_perplexity = action_loss.new_zeros(()) + tactile_active_codes = action_loss.new_zeros(()) + if current_force is not None and force_history is not None: + force_updates = current_force.shape[2] + force_tau = sample_trex_timestep( + (batch_size, blocks, force_updates), + device=self.device, + dtype=self.dtype, + upper=self.schedule.tau_split, + ) + force_clean = action_clean[:, :, None].expand( + -1, -1, force_updates, -1, -1 + ) + force_noise_base = torch.randn_like(action_clean) * action_mask + force_noise = force_noise_base[:, :, None].expand_as(force_clean) + force_noisy, force_target, _ = self._flow_noising( + force_clean, force_tau, noise=force_noise + ) + + # Force memory must represent the split state itself, not the + # unrelated random-tau forward used by the primary coarse losses. + split_tau = torch.full_like(tau, self.schedule.tau_split) + action_at_split, _, _ = self._flow_noising( + action_clean, split_tau, noise=force_noise_base + ) + future_at_split, _, _ = self._flow_noising( + future_clean_model, split_tau + ) + split_observation, _ = self._prepare_observation_inputs( + observation_data, split_tau, blocks + ) + with torch.no_grad(): + split_output = self.model( + action=action_at_split, + state=state, + past_track_xy=past_xy, + past_track_visibility=past_visibility, + future_track_xy=future_at_split, + future_track_visibility=future_visibility, + tau=split_tau, + **split_observation, + **context_kwargs, + ) + flat_action = force_noisy.flatten(0, 2) + flat_tau = force_tau.flatten() + flat_force = current_force.to( + device=self.device, dtype=self.dtype + ).flatten(0, 2) + flat_history = force_history.to(device=self.device).flatten(0, 2) + flat_history_valid = ( + None + if force_history_valid is None + else force_history_valid.to(device=self.device).flatten(0, 2) + ) + memory = ( + split_output.coarse_memory_by_block.detach()[:, :, None] + .expand(-1, -1, force_updates, -1, -1) + .flatten(0, 2) + ) + memory_mask = ( + split_output.coarse_memory_mask_by_block[None, :, None] + .expand(batch_size, -1, force_updates, -1) + .flatten(0, 2) + ) + update_offsets = torch.as_tensor( + self.schedule.offsets[:force_updates], + device=self.device, + dtype=torch.long, + ).view(1, 1, force_updates) + update_offsets = update_offsets.expand( + batch_size, blocks, -1 + ).flatten() + tactile_keep = ( + torch.rand( + batch_size * blocks * force_updates, device=self.device + ) + >= self.config.tactile_dropout + ) + force_output = self.force_transformer( + flat_action, + flat_tau, + flat_force, + flat_history, + coarse_memory=memory, + coarse_memory_mask=memory_mask, + update_offset=update_offsets, + tactile_keep_mask=tactile_keep, + tactile_history_valid_mask=flat_history_valid, + return_details=True, + ) + if not isinstance(force_output, ForceOnlyOutput): + raise RuntimeError("force transformer did not return training details") + force_prediction = force_output.flow.view_as(force_target) + action_positions = torch.arange( + self.action_horizon, device=self.device + ).view(1, 1, 1, self.action_horizon, 1) + suffix_mask = action_positions >= update_offsets.view( + batch_size, blocks, force_updates, 1, 1 + ) + force_mask = action_mask[:, :, None] * suffix_mask + force_error = (force_prediction - force_target).square() * force_mask + force_loss = force_error.sum() / force_mask.sum().clamp_min(1) + vq_loss = force_output.vq_loss + commitment_loss = force_output.commitment_loss + if force_output.tactile_perplexity is not None: + tactile_perplexity = force_output.tactile_perplexity + if force_output.tactile_active_codes is not None: + tactile_active_codes = force_output.tactile_active_codes.to( + dtype=action_loss.dtype + ) + + loss = ( + self.config.dynamics_loss_weight * dynamics_loss + + action_loss + + self.config.track_loss_weight * track_loss + + self.config.force_loss_weight * force_loss + + self.config.vq_loss_weight + * ( + vq_loss + + self.config.vq_commitment_weight * commitment_loss + ) + ) + return { + "loss": loss, + "dynamics_loss": dynamics_loss, + "action_loss": action_loss, + "track_loss": track_loss, + "track_flow_loss": track_flow_loss, + "track_position_loss": track_position_loss, + "track_smoothness_loss": track_smoothness_loss, + "force_loss": force_loss, + "vq_loss": vq_loss, + "commitment_loss": commitment_loss, + "tactile_perplexity": tactile_perplexity, + "tactile_active_codes": tactile_active_codes, + "action_flow": output.action_flow, + "track_flow": output.track_flow, + } + + def forward( + self, + backbone_output: BatchFeature, + action_input: BatchFeature, + ) -> BatchFeature: + del backbone_output + return BatchFeature(data=self.forward_core(action_input)) + + @torch.no_grad() + def sample( + self, + data: dict[str, Any] | BatchFeature, + *, + seed: int | None = None, + return_refinement_state: bool = False, + coarse_steps: int | None = None, + run_force_refinement: bool = True, + ) -> dict[str, torch.Tensor]: + if coarse_steps is not None and coarse_steps <= 0: + raise ValueError("coarse_steps must be positive") + if return_refinement_state and not run_force_refinement: + raise ValueError( + "return_refinement_state requires force refinement" + ) + coarse_schedule = ( + self.schedule + if coarse_steps is None + else TwoStageFlowSchedule( + coarse_steps=coarse_steps, + force_steps=self.schedule.force_steps, + tau_start=self.schedule.tau_start, + tau_split=self.schedule.tau_split, + tau_end=self.schedule.tau_end, + action_horizon=self.schedule.action_horizon, + force_stride=self.schedule.force_stride, + ) + ) + state = self._get(data, "state") + if state.ndim == 2: + blocks = 1 + else: + blocks = state.shape[1] + if blocks != 1: + raise ValueError("online sampling currently consumes exactly one block") + batch_size = state.shape[0] + generator = None + if seed is not None: + generator = torch.Generator(device=self.device).manual_seed(seed) + action = torch.randn( + batch_size, + 1, + self.action_horizon, + self.action_dim, + device=self.device, + dtype=self.dtype, + generator=generator, + ) + action[..., self.config.physical_action_dim :] = 0 + future_track = torch.randn( + batch_size, + 1, + self.track_horizon, + self.track_points, + 2, + device=self.device, + dtype=self.dtype, + generator=generator, + ) + past_xy = self._get(data, "track_past_xy", "past_track_xy").to( + device=self.device, dtype=self.dtype + ) + track_anchor = self._track_anchor(past_xy) + past_visibility = self._get( + data, + "track_past_visibility", + "past_track_visibility", + required=False, + ) + future_visibility = self._get( + data, + "track_future_visibility", + "future_track_visibility", + required=False, + ) + if future_visibility is None: + future_visibility = torch.ones( + batch_size, + 1, + self.track_horizon, + self.track_points, + device=self.device, + dtype=self.dtype, + ) + observation_kwargs = self._prepare_sampling_observation( + data, + blocks=1, + generator=generator, + ) + context_kwargs = self._context_kwargs(data) + coarse_output: CausalWanTrackForceOutput | None = None + for step in coarse_schedule.iter_steps("coarse"): + coarse_output = self.model( + action=action, + state=self._reshape_state(state, 1).to( + device=self.device, dtype=self.dtype + ), + past_track_xy=past_xy, + past_track_visibility=past_visibility, + future_track_xy=future_track, + future_track_visibility=future_visibility, + tau=step.tau, + **observation_kwargs, + **context_kwargs, + ) + action = euler_flow_step( + action, coarse_output.action_flow, step.tau, step.tau_next + ) + action[..., self.config.physical_action_dim :] = 0 + track_flow_model = coarse_output.track_flow.transpose(2, 3) + future_track = euler_flow_step( + future_track, track_flow_model, step.tau, step.tau_next + ) + observation_kwargs = self._advance_sampling_observation( + observation_kwargs, + coarse_output, + step.tau, + step.tau_next, + ) + assert coarse_output is not None + # The coarse state intentionally stops at tau=.4. Evaluate the model at + # that boundary once, then form clean video/track estimates with the + # flow-matching identity x_0 = x_tau - tau * v_theta(x_tau, tau). + coarse_output = self.model( + action=action, + state=self._reshape_state(state, 1).to( + device=self.device, dtype=self.dtype + ), + past_track_xy=past_xy, + past_track_visibility=past_visibility, + future_track_xy=future_track, + future_track_visibility=future_visibility, + tau=self.schedule.tau_split, + **observation_kwargs, + **context_kwargs, + ) + clean_future_track = euler_flow_step( + future_track, + coarse_output.track_flow.transpose(2, 3), + self.schedule.tau_split, + self.schedule.tau_end, + ) + clean_future_track = self._anchor_track_prediction( + clean_future_track, track_anchor + ) + clean_observation_kwargs = self._advance_sampling_observation( + observation_kwargs, + coarse_output, + self.schedule.tau_split, + self.schedule.tau_end, + ) + coarse_action = action[:, 0] + if run_force_refinement: + current_force, history, history_valid = self._extract_force_inputs( + data, 1 + ) + if current_force is None or history is None: + current_force = torch.zeros( + batch_size, + 1, + 1, + FORCE_SENSOR_COUNT, + FORCE_SENSOR_DIM, + device=self.device, + dtype=self.dtype, + ) + history = torch.zeros( + batch_size, + 1, + 1, + FORCE_SENSOR_COUNT, + device=self.device, + dtype=torch.long, + ) + history_valid = None + refined_action = self.refine_action_suffix( + coarse_action, + coarse_output.coarse_memory, + current_force[:, 0, 0], + history[:, 0, 0], + update_offset=0, + tactile_history_valid_mask=( + None if history_valid is None else history_valid[:, 0, 0] + ), + ) + else: + # Video/track monitoring does not consume the action prediction. + # Keep the clean-output boundary pass but skip all tactile passes. + refined_action = coarse_action + refined_action[..., self.config.physical_action_dim :] = 0 + result = { + "action_pred": refined_action, + "track_pred": clean_future_track[:, 0] + .transpose(1, 2) + .clamp(0.0, 1.0), + } + if "obs_tokens" in clean_observation_kwargs: + result["obs_pred"] = clean_observation_kwargs["obs_tokens"] + else: + result["video_latents_pred"] = clean_observation_kwargs[ + "video_latents" + ] + if "conditioning_video_latents" in clean_observation_kwargs: + result["conditioning_video_latents"] = clean_observation_kwargs[ + "conditioning_video_latents" + ] + if return_refinement_state: + result["coarse_action_at_split"] = coarse_action + result["coarse_memory"] = coarse_output.coarse_memory + return result + + @torch.no_grad() + def sample_visualization( + self, + data: dict[str, Any] | BatchFeature, + *, + coarse_steps: int = 1, + seed: int | None = None, + ) -> dict[str, torch.Tensor]: + """Generate video/tracks without action-only tactile refinement.""" + + return self.sample( + data, + seed=seed, + coarse_steps=coarse_steps, + run_force_refinement=False, + ) + + @torch.no_grad() + def refine_action_suffix( + self, + coarse_action_at_split: torch.Tensor, + coarse_memory: torch.Tensor, + current_force: torch.Tensor, + tactile_history_or_codes: torch.Tensor, + *, + update_offset: int, + tactile_history_valid_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Re-denoise only the unexecuted suffix at a 5 Hz tactile refresh.""" + + if update_offset not in self.schedule.offsets: + raise ValueError( + f"update_offset must be one of {self.schedule.offsets}" + ) + action = pad_action_62_to_64(coarse_action_at_split).to( + device=self.device, dtype=self.dtype + ).clone() + action[..., self.config.physical_action_dim :] = 0 + current_force = current_force.to(device=self.device, dtype=self.dtype) + tactile_history_or_codes = tactile_history_or_codes.to(device=self.device) + coarse_memory = coarse_memory.to(device=self.device, dtype=self.dtype) + for step in self.schedule.iter_steps("force"): + flow = self.force_transformer( + action, + step.tau, + current_force, + tactile_history_or_codes, + coarse_memory=coarse_memory, + update_offset=update_offset, + tactile_history_valid_mask=tactile_history_valid_mask, + ) + if not isinstance(flow, torch.Tensor): + raise RuntimeError("force inference expected a flow tensor") + proposed = euler_flow_step(action, flow, step.tau, step.tau_next) + action[:, update_offset:] = proposed[:, update_offset:] + action[..., self.config.physical_action_dim :] = 0 + return action + + def get_action( + self, + backbone_output: BatchFeature, + action_input: BatchFeature, + num_action_samples: int = 1, + inference_batch_size: int = 32, + ) -> BatchFeature: + del inference_batch_size + if num_action_samples != 1: + raise ValueError("track/force sampler currently returns one sample") + merged = dict(backbone_output) + merged.update(dict(action_input)) + detailed_sample = self.sample(merged) + if "action_pred" not in detailed_sample: + raise RuntimeError("sampler did not return action_pred") + return BatchFeature(data=detailed_sample) + + def get_visualization( + self, + backbone_output: BatchFeature, + action_input: BatchFeature, + *, + coarse_steps: int = 1, + seed: int | None = None, + ) -> BatchFeature: + merged = dict(backbone_output) + merged.update(dict(action_input)) + return BatchFeature( + data=self.sample_visualization( + merged, + coarse_steps=coarse_steps, + seed=seed, + ) + ) + + def post_initialize(self) -> None: + """Compatibility hook; placement remains controlled by VLA/Trainer.""" + + def parallelize(self, device_mesh=None) -> None: + """The independent force model currently uses ordinary data parallelism.""" + + del device_mesh + + +# Names requested by the config and by existing WAN integrations. +WANPolicyHeadConfig = WANTrackForcePolicyHeadConfig +WANPolicyHead = WANTrackForcePolicyHead +TrexTrackForceWANPolicyHead = WANTrackForcePolicyHead + + +__all__ = [ + "TrexTrackForceWANPolicyHead", + "WANPolicyHead", + "WANPolicyHeadConfig", + "WANTrackForcePolicyHead", + "WANTrackForcePolicyHeadConfig", +] diff --git a/groot/vla/model/trex_track_force/runtime.py b/groot/vla/model/trex_track_force/runtime.py new file mode 100644 index 0000000000000000000000000000000000000000..e02a97d0cefa8ffab719c3b4b8d333fe0b336f6a --- /dev/null +++ b/groot/vla/model/trex_track_force/runtime.py @@ -0,0 +1,184 @@ +"""Online 20 Hz / 5 Hz execution helpers for T-Rex Track-Force.""" + +from __future__ import annotations + +from dataclasses import dataclass +import json +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +from .force import FORCE_OFFSETS +from .policy import WANTrackForcePolicyHead +from .dataset import rotation_6d_to_matrix + + +def _matrix_to_rotation_6d(matrix: np.ndarray) -> np.ndarray: + return np.concatenate((matrix[..., :, 0], matrix[..., :, 1]), axis=-1) + + +def delta_base_to_absolute( + reference_state: np.ndarray, delta_base_action: np.ndarray +) -> np.ndarray: + """Invert T-Rex delta-base arm poses; hand targets are already absolute.""" + + reference = np.asarray(reference_state, dtype=np.float64) + delta = np.asarray(delta_base_action, dtype=np.float64) + if reference.shape != (62,) or delta.shape[-1] != 62: + raise ValueError("expected reference [62] and delta action [...,62]") + output = np.empty_like(delta) + for pose_slice, hand_slice in ( + (slice(0, 9), slice(9, 31)), + (slice(31, 40), slice(40, 62)), + ): + reference_pose = reference[pose_slice] + delta_pose = delta[..., pose_slice] + reference_rotation = rotation_6d_to_matrix(reference_pose[3:9]) + delta_rotation = rotation_6d_to_matrix(delta_pose[..., 3:9]) + target_translation = ( + reference_pose[:3] + + np.einsum("ij,...j->...i", reference_rotation, delta_pose[..., :3]) + ) + target_rotation = np.einsum( + "ij,...jk->...ik", reference_rotation, delta_rotation + ) + output[..., pose_slice] = np.concatenate( + (target_translation, _matrix_to_rotation_6d(target_rotation)), axis=-1 + ) + output[..., hand_slice] = delta[..., hand_slice] + return output.astype(np.float32) + + +@dataclass(frozen=True) +class TrexRuntimeStatistics: + action_q01: np.ndarray + action_q99: np.ndarray + state_q01: np.ndarray + state_q99: np.ndarray + force_q01: np.ndarray + force_q99: np.ndarray + + @classmethod + def from_dataset(cls, dataset_root: str | Path) -> "TrexRuntimeStatistics": + root = Path(dataset_root) + with (root / "meta" / "relative_stats_dreamzero.json").open() as handle: + action = json.load(handle)["eef62"] + with (root / "meta" / "stats.json").open() as handle: + stats = json.load(handle) + force = stats["observation.tactile_force"] + state = stats["observation.state_eef62"] + return cls( + action_q01=np.asarray(action["q01"], dtype=np.float32), + action_q99=np.asarray(action["q99"], dtype=np.float32), + state_q01=np.asarray(state["q01"], dtype=np.float32), + state_q99=np.asarray(state["q99"], dtype=np.float32), + force_q01=np.asarray(force["q01"], dtype=np.float32).reshape(10, 6), + force_q99=np.asarray(force["q99"], dtype=np.float32).reshape(10, 6), + ) + + def normalize_force(self, force: np.ndarray) -> np.ndarray: + values = np.asarray(force, dtype=np.float32) + scale = self.force_q99 - self.force_q01 + normalized = np.where( + scale != 0, + 2.0 * (values - self.force_q01) / np.where(scale == 0, 1.0, scale) + - 1.0, + values, + ) + return np.clip(normalized, -1.0, 1.0).astype(np.float32) + + def normalize_state(self, state: np.ndarray) -> np.ndarray: + values = np.asarray(state, dtype=np.float32) + scale = self.state_q99 - self.state_q01 + normalized = np.where( + scale != 0, + 2.0 * (values - self.state_q01) / np.where(scale == 0, 1.0, scale) + - 1.0, + values, + ) + return np.clip(normalized, -1.0, 1.0).astype(np.float32) + + def denormalize_action(self, normalized_action: np.ndarray) -> np.ndarray: + values = np.asarray(normalized_action, dtype=np.float32)[..., :62] + return ( + (values + 1.0) * 0.5 * (self.action_q99 - self.action_q01) + + self.action_q01 + ).astype(np.float32) + + +@dataclass +class TactileRefinementState: + coarse_action_at_split: torch.Tensor + coarse_memory: torch.Tensor + refined_action: torch.Tensor + latest_offset: int = 0 + + +class TrexTrackForceOnlineController: + """Cache coarse output and refresh only suffixes at steps 0/4/8/12.""" + + def __init__(self, policy: WANTrackForcePolicyHead) -> None: + self.policy = policy + self.state: TactileRefinementState | None = None + + @torch.no_grad() + def start_chunk( + self, + model_inputs: dict[str, Any], + *, + seed: int | None = None, + ) -> torch.Tensor: + result = self.policy.sample( + model_inputs, seed=seed, return_refinement_state=True + ) + self.state = TactileRefinementState( + coarse_action_at_split=result["coarse_action_at_split"].detach(), + coarse_memory=result["coarse_memory"].detach(), + refined_action=result["action_pred"].detach(), + latest_offset=0, + ) + return self.state.refined_action + + @torch.no_grad() + def refresh_tactile( + self, + current_force: torch.Tensor, + tactile_history_or_codes: torch.Tensor, + *, + executed_steps: int, + ) -> torch.Tensor: + if self.state is None: + raise RuntimeError("start_chunk must be called before tactile refresh") + if executed_steps not in FORCE_OFFSETS: + raise ValueError(f"executed_steps must be one of {FORCE_OFFSETS}") + if executed_steps < self.state.latest_offset: + raise ValueError("tactile refresh offsets must be monotonic") + refreshed = self.policy.refine_action_suffix( + self.state.coarse_action_at_split, + self.state.coarse_memory, + current_force, + tactile_history_or_codes, + update_offset=executed_steps, + ) + # Preserve commands already emitted from earlier refinements. + refreshed[:, :executed_steps] = self.state.refined_action[:, :executed_steps] + self.state.refined_action = refreshed.detach() + self.state.latest_offset = executed_steps + return self.state.refined_action + + def action_at(self, step: int) -> torch.Tensor: + if self.state is None: + raise RuntimeError("no active action chunk") + if not 0 <= step < self.state.refined_action.shape[1]: + raise IndexError(step) + return self.state.refined_action[:, step] + + +__all__ = [ + "TactileRefinementState", + "TrexRuntimeStatistics", + "TrexTrackForceOnlineController", + "delta_base_to_absolute", +] diff --git a/groot/vla/model/trex_track_force/tactile_vq.py b/groot/vla/model/trex_track_force/tactile_vq.py new file mode 100644 index 0000000000000000000000000000000000000000..08ec09b37cdbd6479bae94cf8080b6f63aff4afa --- /dev/null +++ b/groot/vla/model/trex_track_force/tactile_vq.py @@ -0,0 +1,452 @@ +"""Force-only adaptation of T-Rex's per-finger tactile VQ-VAE. + +The original implementation encodes each five-finger hand separately with +shared temporal convolutions. Here the same weights process both hands, +yielding ten discrete codes for a 16-sample F6 history. The codebook uses +distributed EMA updates and periodic dead-code revival. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import torch.distributed as dist +from torch import nn +from torch.nn import functional as F + + +def _is_distributed() -> bool: + return dist.is_available() and dist.is_initialized() + + +def _all_reduce_sum(values: torch.Tensor) -> torch.Tensor: + if _is_distributed(): + dist.all_reduce(values, op=dist.ReduceOp.SUM) + return values + + +def _all_gather_2d(values: torch.Tensor) -> torch.Tensor: + if not _is_distributed(): + return values + outputs = [torch.zeros_like(values) for _ in range(dist.get_world_size())] + dist.all_gather(outputs, values.contiguous()) + return torch.cat(outputs, dim=0) + + +def _conv_block( + input_channels: int, output_channels: int, *, stride: int +) -> nn.Sequential: + return nn.Sequential( + nn.Conv1d( + input_channels, + output_channels, + kernel_size=5, + stride=stride, + padding=2, + ), + nn.GroupNorm(min(8, output_channels), output_channels), + nn.GELU(), + ) + + +def _upconv_block( + input_channels: int, output_channels: int, *, stride: int +) -> nn.Sequential: + if stride > 1: + layer: nn.Module = nn.ConvTranspose1d( + input_channels, + output_channels, + kernel_size=5, + stride=stride, + padding=2, + output_padding=stride - 1, + ) + else: + layer = nn.Conv1d( + input_channels, output_channels, kernel_size=5, padding=2 + ) + return nn.Sequential( + layer, + nn.GroupNorm(min(8, output_channels), output_channels), + nn.GELU(), + ) + + +class PerFingerF6Encoder(nn.Module): + def __init__( + self, + *, + window: int = 16, + force_dim: int = 6, + fingers_per_hand: int = 5, + hidden_channels: int = 128, + bottleneck_channels: int = 256, + embed_dim: int = 256, + strided_blocks: int = 2, + ) -> None: + super().__init__() + self.window = window + self.fingers_per_hand = fingers_per_hand + self.embed_dim = embed_dim + self.stem = _conv_block(force_dim, hidden_channels, stride=1) + self.finger_embedding = nn.Embedding(fingers_per_hand, hidden_channels) + blocks: list[nn.Module] = [] + time = window + channels = hidden_channels + for index in range(strided_blocks): + stride = 2 if time >= 4 else 1 + output_channels = ( + bottleneck_channels + if index == strided_blocks - 1 + else hidden_channels + ) + blocks.append( + _conv_block(channels, output_channels, stride=stride) + ) + channels = output_channels + time //= stride + self.strided = nn.Sequential(*blocks) + self.projection = nn.Conv1d(channels, embed_dim, kernel_size=3, padding=1) + + def forward(self, values: torch.Tensor) -> torch.Tensor: + batch, time, fingers, force_dim = values.shape + if time != self.window or fingers != self.fingers_per_hand: + raise ValueError( + f"F6 encoder expects [B,{self.window},{self.fingers_per_hand},D]" + ) + encoded = values.permute(0, 2, 3, 1).reshape( + batch * fingers, force_dim, time + ) + encoded = self.stem(encoded) + finger_ids = torch.arange(fingers, device=values.device).repeat(batch) + encoded = encoded + self.finger_embedding(finger_ids).unsqueeze(-1) + encoded = self.strided(encoded) + encoded = self.projection(encoded).mean(dim=-1) + return encoded.view(batch, fingers, self.embed_dim) + + +class PerFingerF6Decoder(nn.Module): + def __init__( + self, + *, + window: int = 16, + force_dim: int = 6, + fingers_per_hand: int = 5, + hidden_channels: int = 128, + bottleneck_channels: int = 256, + embed_dim: int = 256, + strided_blocks: int = 2, + ) -> None: + super().__init__() + self.window = window + self.fingers_per_hand = fingers_per_hand + time = window + strides: list[int] = [] + channel_chain = [hidden_channels] + for index in range(strided_blocks): + stride = 2 if time >= 4 else 1 + strides.append(stride) + channel_chain.append( + bottleneck_channels + if index == strided_blocks - 1 + else hidden_channels + ) + time //= stride + self.bottleneck_time = time + self.from_embedding = nn.Conv1d( + embed_dim, bottleneck_channels, kernel_size=3, padding=1 + ) + self.finger_embedding = nn.Embedding( + fingers_per_hand, bottleneck_channels + ) + blocks: list[nn.Module] = [] + channels = bottleneck_channels + reversed_chain = list(reversed(channel_chain)) + for index, stride in enumerate(reversed(strides)): + output_channels = reversed_chain[index + 1] + blocks.append( + _upconv_block(channels, output_channels, stride=stride) + ) + channels = output_channels + self.up_strided = nn.Sequential(*blocks) + self.head = nn.Conv1d( + hidden_channels, force_dim, kernel_size=5, padding=2 + ) + + def forward(self, quantized: torch.Tensor) -> torch.Tensor: + batch, fingers, embed_dim = quantized.shape + values = quantized.reshape(batch * fingers, embed_dim) + values = values.unsqueeze(-1).expand(-1, -1, self.bottleneck_time) + values = self.from_embedding(values) + finger_ids = torch.arange(fingers, device=quantized.device).repeat(batch) + values = values + self.finger_embedding(finger_ids).unsqueeze(-1) + values = self.up_strided(values) + values = self.head(values) + if values.shape[-1] > self.window: + values = values[..., : self.window] + elif values.shape[-1] < self.window: + values = F.pad(values, (0, self.window - values.shape[-1])) + return ( + values.transpose(1, 2) + .reshape(batch, fingers, self.window, -1) + .permute(0, 2, 1, 3) + .contiguous() + ) + + +@dataclass +class EMAQuantizerOutput: + quantized: torch.Tensor + codes: torch.Tensor + commitment_loss: torch.Tensor + perplexity: torch.Tensor + active_codes: torch.Tensor + + +class EMAQuantizer(nn.Module): + def __init__( + self, + *, + codebook_size: int = 64, + embed_dim: int = 256, + decay: float = 0.99, + epsilon: float = 1e-5, + revive_frequency: int = 200, + revive_threshold: float = 1.0, + ) -> None: + super().__init__() + self.codebook_size = codebook_size + self.embed_dim = embed_dim + self.decay = decay + self.epsilon = epsilon + self.revive_frequency = revive_frequency + self.revive_threshold = revive_threshold + embedding = torch.randn(codebook_size, embed_dim) * 0.02 + self.register_buffer("embedding", embedding) + self.register_buffer("cluster_size", torch.zeros(codebook_size)) + self.register_buffer("embedding_average", embedding.clone()) + self.register_buffer("step", torch.zeros((), dtype=torch.long)) + + @torch.no_grad() + def _update(self, encoded: torch.Tensor, codes: torch.Tensor) -> None: + one_hot = F.one_hot(codes, self.codebook_size).type_as(encoded) + counts = _all_reduce_sum(one_hot.sum(dim=0).contiguous()) + sums = _all_reduce_sum((one_hot.t() @ encoded).contiguous()) + self.cluster_size.mul_(self.decay).add_(counts, alpha=1 - self.decay) + self.embedding_average.mul_(self.decay).add_( + sums, alpha=1 - self.decay + ) + total = self.cluster_size.sum() + smoothed = ( + (self.cluster_size + self.epsilon) + / (total + self.codebook_size * self.epsilon) + * total + ) + self.embedding.copy_( + self.embedding_average / smoothed.unsqueeze(-1).clamp_min(self.epsilon) + ) + + @torch.no_grad() + def _revive(self, encoded: torch.Tensor) -> None: + dead = self.cluster_size < self.revive_threshold + count = int(dead.sum()) + if count == 0: + return + pool = _all_gather_2d(encoded) + if pool.shape[0] == 0: + return + if not _is_distributed() or dist.get_rank() == 0: + selected = torch.randint(0, pool.shape[0], (count,), device=pool.device) + else: + selected = torch.zeros(count, dtype=torch.long, device=pool.device) + if _is_distributed(): + dist.broadcast(selected, src=0) + replacements = pool[selected].to(self.embedding.dtype) + dead_indices = dead.nonzero(as_tuple=False).flatten() + self.embedding[dead_indices] = replacements + self.embedding_average[dead_indices] = replacements + self.cluster_size[dead_indices] = self.revive_threshold * 2 + + def lookup(self, codes: torch.Tensor) -> torch.Tensor: + if codes.min() < 0 or codes.max() >= self.codebook_size: + raise ValueError("precomputed tactile code is outside the codebook") + return self.embedding[codes.long()] + + def forward(self, encoded: torch.Tensor) -> EMAQuantizerOutput: + if encoded.ndim != 2 or encoded.shape[-1] != self.embed_dim: + raise ValueError(f"EMA quantizer expects [N,{self.embed_dim}]") + # Keep nearest-neighbour search and EMA accumulation in the codebook + # buffer dtype (normally FP32 under BF16 mixed-precision training). + # The straight-through value is converted back to the encoder dtype so + # downstream BF16 projections receive matching inputs. + lookup_encoded = encoded.to(dtype=self.embedding.dtype) + distances = ( + lookup_encoded.square().sum(-1, keepdim=True) + - 2 * lookup_encoded @ self.embedding.t() + + self.embedding.square().sum(-1) + ) + codes = distances.argmin(dim=-1) + quantized = self.embedding[codes] + commitment = F.mse_loss(lookup_encoded, quantized.detach()) + quantized_for_encoder = quantized.to(dtype=encoded.dtype) + straight_through = encoded + (quantized_for_encoder - encoded).detach() + if self.training: + self._update(lookup_encoded.detach(), codes) + self.step += 1 + if int(self.step) % self.revive_frequency == 0: + self._revive(lookup_encoded.detach()) + with torch.no_grad(): + counts = _all_reduce_sum( + F.one_hot(codes, self.codebook_size) + .type_as(lookup_encoded) + .sum(dim=0) + ) + probabilities = counts / counts.sum().clamp_min(1) + perplexity = torch.exp( + -(probabilities * probabilities.clamp_min(1e-12).log()).sum() + ) + active = (counts > 0).sum() + return EMAQuantizerOutput( + straight_through, codes, commitment, perplexity, active + ) + + +@dataclass +class PerFingerVQVAEOutput: + tokens: torch.Tensor + codes: torch.Tensor + reconstruction: torch.Tensor | None + reconstruction_loss: torch.Tensor + commitment_loss: torch.Tensor + perplexity: torch.Tensor + active_codes: torch.Tensor + + +class PerFingerForceVQVAE(nn.Module): + """Two hands × five fingers, sharing T-Rex encoder/decoder/codebook.""" + + def __init__( + self, + *, + history_frames: int = 16, + num_fingers: int = 10, + force_dim: int = 6, + codebook_size: int = 64, + code_dim: int = 256, + ) -> None: + super().__init__() + if num_fingers != 10: + raise ValueError("force VQ-VAE requires ten fingers") + self.history_frames = history_frames + self.num_fingers = num_fingers + self.force_dim = force_dim + self.codebook_size = codebook_size + self.code_dim = code_dim + kwargs = { + "window": history_frames, + "force_dim": force_dim, + "fingers_per_hand": 5, + "embed_dim": code_dim, + } + self.encoder = PerFingerF6Encoder(**kwargs) + self.decoder = PerFingerF6Decoder(**kwargs) + self.quantizer = EMAQuantizer( + codebook_size=codebook_size, embed_dim=code_dim + ) + self.side_embedding = nn.Embedding(2, code_dim) + self.output_finger_embedding = nn.Embedding(5, code_dim) + + def _identity(self, device: torch.device) -> torch.Tensor: + ids = torch.arange(self.num_fingers, device=device) + return self.side_embedding(ids // 5) + self.output_finger_embedding(ids % 5) + + def forward( + self, + history_or_codes: torch.Tensor, + valid_history_mask: torch.Tensor | None = None, + ) -> PerFingerVQVAEOutput: + if not history_or_codes.is_floating_point(): + if history_or_codes.ndim != 2 or history_or_codes.shape[1] != 10: + raise ValueError("precomputed tactile codes must be [B,10]") + if valid_history_mask is not None: + raise ValueError("history validity mask is only valid for raw force") + codes = history_or_codes.long() + quantized = self.quantizer.lookup(codes) + tokens = quantized + self._identity(codes.device).unsqueeze(0) + zero = tokens.new_zeros(()) + return PerFingerVQVAEOutput( + tokens, codes, None, zero, zero, zero, zero + ) + + expected = ( + history_or_codes.shape[0], + self.history_frames, + self.num_fingers, + self.force_dim, + ) + if history_or_codes.shape != expected: + raise ValueError( + "raw tactile history must be " + f"[B,{self.history_frames},10,{self.force_dim}]" + ) + batch = history_or_codes.shape[0] + hands = ( + history_or_codes.view( + batch, self.history_frames, 2, 5, self.force_dim + ) + .permute(0, 2, 1, 3, 4) + .reshape(batch * 2, self.history_frames, 5, self.force_dim) + ) + encoded = self.encoder(hands) + quantizer_output = self.quantizer( + encoded.reshape(batch * self.num_fingers, self.code_dim) + ) + quantized_hands = quantizer_output.quantized.view(batch * 2, 5, self.code_dim) + reconstruction_hands = self.decoder(quantized_hands) + reconstruction = ( + reconstruction_hands.view( + batch, 2, self.history_frames, 5, self.force_dim + ) + .permute(0, 2, 1, 3, 4) + .reshape(batch, self.history_frames, self.num_fingers, self.force_dim) + ) + if valid_history_mask is None: + reconstruction_loss = F.mse_loss(reconstruction, history_or_codes) + else: + if valid_history_mask.shape != (batch, self.history_frames): + raise ValueError( + "raw tactile validity mask must be [B,history_frames]" + ) + valid = valid_history_mask.to( + device=history_or_codes.device, + dtype=history_or_codes.dtype, + )[:, :, None, None] + squared_error = (reconstruction - history_or_codes).square() * valid + denominator = ( + valid.sum() * self.num_fingers * self.force_dim + ).clamp_min(1) + reconstruction_loss = squared_error.sum() / denominator + codes = quantizer_output.codes.view(batch, self.num_fingers) + tokens = quantizer_output.quantized.view( + batch, self.num_fingers, self.code_dim + ) + self._identity(history_or_codes.device).unsqueeze(0) + return PerFingerVQVAEOutput( + tokens=tokens, + codes=codes, + reconstruction=reconstruction, + reconstruction_loss=reconstruction_loss, + commitment_loss=quantizer_output.commitment_loss, + perplexity=quantizer_output.perplexity, + active_codes=quantizer_output.active_codes, + ) + + +__all__ = [ + "EMAQuantizer", + "EMAQuantizerOutput", + "PerFingerF6Decoder", + "PerFingerF6Encoder", + "PerFingerForceVQVAE", + "PerFingerVQVAEOutput", +] diff --git a/groot/vla/model/trex_track_force/track.py b/groot/vla/model/trex_track_force/track.py new file mode 100644 index 0000000000000000000000000000000000000000..68e7d35af085f6ba5f9d7326dbfd8730add3ad32 --- /dev/null +++ b/groot/vla/model/trex_track_force/track.py @@ -0,0 +1,538 @@ +"""Track layout, embedding, decoding, and visibility-masked flow loss.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import IntEnum + +import torch +from torch import nn + + +NUM_TRACK_POINTS = 250 +TRACK_HORIZON = 16 + + +class TrackView(IntEnum): + HEAD = 0 + LEFT_WRIST = 1 + RIGHT_WRIST = 2 + + +class TrackHand(IntEnum): + NONE = 0 + LEFT = 1 + RIGHT = 2 + + +class TrackRole(IntEnum): + HEAD_HAND = 0 + WRIST_BACKGROUND = 1 + WRIST_HAND = 2 + + +@dataclass(frozen=True) +class TrackSegment: + name: str + start: int + stop: int + view: TrackView + hand: TrackHand + role: TrackRole + + @property + def size(self) -> int: + return self.stop - self.start + + +# Stable identity contract shared with the 250-point data builder. +TRACK_SEGMENTS = ( + TrackSegment( + "head_left_hand", 0, 50, TrackView.HEAD, TrackHand.LEFT, TrackRole.HEAD_HAND + ), + TrackSegment( + "head_right_hand", + 50, + 100, + TrackView.HEAD, + TrackHand.RIGHT, + TrackRole.HEAD_HAND, + ), + TrackSegment( + "left_wrist_background", + 100, + 125, + TrackView.LEFT_WRIST, + TrackHand.NONE, + TrackRole.WRIST_BACKGROUND, + ), + TrackSegment( + "left_wrist_hand", + 125, + 175, + TrackView.LEFT_WRIST, + TrackHand.LEFT, + TrackRole.WRIST_HAND, + ), + TrackSegment( + "right_wrist_background", + 175, + 200, + TrackView.RIGHT_WRIST, + TrackHand.NONE, + TrackRole.WRIST_BACKGROUND, + ), + TrackSegment( + "right_wrist_hand", + 200, + 250, + TrackView.RIGHT_WRIST, + TrackHand.RIGHT, + TrackRole.WRIST_HAND, + ), +) +TRACK_POINT_SLICES = tuple((segment.start, segment.stop) for segment in TRACK_SEGMENTS) + + +def validate_track_layout() -> None: + """Validate that the canonical segments cover every point exactly once.""" + + cursor = 0 + for segment in TRACK_SEGMENTS: + if segment.start != cursor or segment.stop <= segment.start: + raise RuntimeError(f"invalid track segment ordering at {segment}") + cursor = segment.stop + if cursor != NUM_TRACK_POINTS: + raise RuntimeError( + f"track segments cover {cursor} points, expected {NUM_TRACK_POINTS}" + ) + + +def canonical_point_metadata( + *, device: torch.device | str | None = None +) -> dict[str, torch.Tensor]: + """Return canonical per-point view, hand, role, and point identities.""" + + validate_track_layout() + view = torch.empty(NUM_TRACK_POINTS, dtype=torch.long, device=device) + hand = torch.empty_like(view) + role = torch.empty_like(view) + for segment in TRACK_SEGMENTS: + view[segment.start : segment.stop] = int(segment.view) + hand[segment.start : segment.stop] = int(segment.hand) + role[segment.start : segment.stop] = int(segment.role) + return { + "view": view, + "hand": hand, + "role": role, + "point": torch.arange(NUM_TRACK_POINTS, device=device), + } + + +def _canonicalize_xy( + xy: torch.Tensor, + *, + points: int = NUM_TRACK_POINTS, + name: str, +) -> torch.Tensor: + """Canonicalize track coordinates to ``[B, blocks, time, points, 2]``.""" + + if xy.ndim != 5 or xy.shape[-1] != 2: + raise ValueError( + f"{name} must be [B,blocks,time,points,2] or " + f"[B,blocks,points,time,2], got {tuple(xy.shape)}" + ) + if xy.shape[-2] == points: + return xy + if xy.shape[2] == points: + return xy.transpose(2, 3) + raise ValueError(f"{name} does not contain a {points}-point axis: {tuple(xy.shape)}") + + +def _canonicalize_visibility( + visibility: torch.Tensor | None, + reference_xy: torch.Tensor, + *, + points: int = NUM_TRACK_POINTS, + name: str, +) -> torch.Tensor: + """Canonicalize visibility to ``[B, blocks, time, points]``.""" + + expected = reference_xy.shape[:-1] + if visibility is None: + return torch.ones(expected, dtype=reference_xy.dtype, device=reference_xy.device) + if visibility.ndim != 4: + raise ValueError(f"{name} must be rank 4, got {tuple(visibility.shape)}") + if visibility.shape == expected: + result = visibility + elif ( + visibility.shape[:2] == expected[:2] + and visibility.shape[2] == points + and visibility.shape[3] == expected[2] + ): + result = visibility.transpose(2, 3) + else: + raise ValueError( + f"{name} shape {tuple(visibility.shape)} does not match {tuple(expected)}" + ) + return result.to(device=reference_xy.device, dtype=reference_xy.dtype) + + +@dataclass +class TrackEmbeddingOutput: + """One clean-past and one noisy-future token per canonical point.""" + + past_tokens: torch.Tensor + future_tokens: torch.Tensor + + @property + def packed(self) -> torch.Tensor: + return torch.cat((self.past_tokens, self.future_tokens), dim=-2) + + def __iter__(self): + yield self.past_tokens + yield self.future_tokens + + +class TrackEncoder(nn.Module): + """Embed 250 point trajectories into one past/future token per point. + + Every token contains additive view, hand, semantic-role, point, block, and + temporal-role embeddings. The trajectory projection consumes ``xy``, + visibility, and first-order motion at each time index. + """ + + def __init__( + self, + dim: int, + *, + num_points: int = NUM_TRACK_POINTS, + future_horizon: int = TRACK_HORIZON, + max_history_frames: int = 32, + max_blocks: int = 32, + ) -> None: + super().__init__() + if num_points != NUM_TRACK_POINTS: + raise ValueError( + "T-Rex identity metadata is fixed at exactly 250 points; " + f"got {num_points}" + ) + self.dim = dim + self.num_points = num_points + self.future_horizon = future_horizon + self.max_history_frames = max_history_frames + + self.view_embedding = nn.Embedding(len(TrackView), dim) + self.hand_embedding = nn.Embedding(len(TrackHand), dim) + self.role_embedding = nn.Embedding(len(TrackRole), dim) + self.point_embedding = nn.Embedding(num_points, dim) + self.block_embedding = nn.Embedding(max_blocks, dim) + self.token_role_embedding = nn.Embedding(2, dim) # past / future + self.sequence_frames = max(max_history_frames, future_horizon) + + # Keep the complete temporal sequence inside each point token. The + # previous visible-frame average discarded which independent diffusion + # noise belonged to which future timestep, making 16-step flow + # reconstruction mathematically underdetermined. + self.trajectory_projection = nn.Sequential( + nn.Linear(self.sequence_frames * 5, dim), + nn.SiLU(), + nn.Linear(dim, dim), + ) + self.output_norm = nn.LayerNorm(dim) + + metadata = canonical_point_metadata() + self.register_buffer("point_view_ids", metadata["view"], persistent=False) + self.register_buffer("point_hand_ids", metadata["hand"], persistent=False) + self.register_buffer("point_role_ids", metadata["role"], persistent=False) + self.register_buffer("point_ids", metadata["point"], persistent=False) + + def _embed_trajectory( + self, + xy: torch.Tensor, + visibility: torch.Tensor, + ) -> torch.Tensor: + batch_size, num_blocks, time, points, _ = xy.shape + if time > self.sequence_frames: + raise ValueError( + f"trajectory has {time} frames, maximum is " + f"{self.sequence_frames}" + ) + visible = visibility.clamp(0, 1) + valid_xy = torch.where( + visible.unsqueeze(-1).bool(), xy, torch.zeros_like(xy) + ) + motion = torch.zeros_like(xy) + if time > 1: + pair_visible = (visible[:, :, 1:] * visible[:, :, :-1]).unsqueeze(-1) + delta = valid_xy[:, :, 1:] - valid_xy[:, :, :-1] + motion[:, :, 1:] = torch.where( + pair_visible.bool(), delta, torch.zeros_like(delta) + ) + sample = torch.cat((valid_xy, visible.unsqueeze(-1), motion), dim=-1) + # [B,K,T,N,5] -> [B,K,N,max(T)*5]. Padding is deterministic and only + # supports shorter history windows used by online inference/tests. + ordered = sample.permute(0, 1, 3, 2, 4) + padded = sample.new_zeros( + batch_size, num_blocks, points, self.sequence_frames, 5 + ) + padded[:, :, :, :time] = ordered + return self.trajectory_projection(padded.flatten(-2)) + + def _identity_embedding( + self, num_blocks: int, *, device: torch.device + ) -> torch.Tensor: + if num_blocks > self.block_embedding.num_embeddings: + raise ValueError( + f"num_blocks={num_blocks} exceeds {self.block_embedding.num_embeddings}" + ) + point_identity = ( + self.view_embedding(self.point_view_ids) + + self.hand_embedding(self.point_hand_ids) + + self.role_embedding(self.point_role_ids) + + self.point_embedding(self.point_ids) + ) + block_ids = torch.arange(num_blocks, device=device) + return point_identity.view(1, 1, self.num_points, self.dim) + self.block_embedding( + block_ids + ).view(1, num_blocks, 1, self.dim) + + def forward( + self, + past_xy: torch.Tensor, + past_visibility: torch.Tensor | None, + future_xy: torch.Tensor, + future_visibility: torch.Tensor | None, + ) -> TrackEmbeddingOutput: + past_xy = _canonicalize_xy(past_xy, name="past_xy") + future_xy = _canonicalize_xy(future_xy, name="future_xy") + if past_xy.shape[:2] != future_xy.shape[:2]: + raise ValueError("past and future batch/block axes must match") + if future_xy.shape[2] != self.future_horizon: + raise ValueError( + f"future trajectory must have {self.future_horizon} steps, " + f"got {future_xy.shape[2]}" + ) + + past_visibility = _canonicalize_visibility( + past_visibility, past_xy, name="past_visibility" + ) + future_visibility = _canonicalize_visibility( + future_visibility, future_xy, name="future_visibility" + ) + identity = self._identity_embedding( + past_xy.shape[1], device=past_xy.device + ).to(dtype=past_xy.dtype) + past = self._embed_trajectory(past_xy, past_visibility) + future = self._embed_trajectory(future_xy, future_visibility) + past_role = self.token_role_embedding.weight[0].view(1, 1, 1, self.dim) + future_role = self.token_role_embedding.weight[1].view(1, 1, 1, self.dim) + return TrackEmbeddingOutput( + past_tokens=self.output_norm(past + identity + past_role), + future_tokens=self.output_norm(future + identity + future_role), + ) + + +class TrackDecoder(nn.Module): + """Decode absolute-XY CFM velocity to ``[B,blocks,250,16,2]``.""" + + def __init__( + self, + dim: int, + *, + num_points: int = NUM_TRACK_POINTS, + future_horizon: int = TRACK_HORIZON, + ) -> None: + super().__init__() + self.num_points = num_points + self.future_horizon = future_horizon + self.norm = nn.LayerNorm(dim) + self.projection = nn.Linear(dim, future_horizon * 2) + + def forward(self, future_tokens: torch.Tensor) -> torch.Tensor: + if future_tokens.ndim != 4 or future_tokens.shape[-2] != self.num_points: + raise ValueError( + "future_tokens must have shape [B,blocks,250,D], got " + f"{tuple(future_tokens.shape)}" + ) + batch_size, blocks, points, _ = future_tokens.shape + decoded = self.projection(self.norm(future_tokens)) + return decoded.view( + batch_size, blocks, points, self.future_horizon, 2 + ) + + +def _canonicalize_visibility_layout( + visibility: torch.Tensor, + expected: torch.Size | tuple[int, ...], +) -> torch.Tensor: + """Map visibility to ``[B,blocks,250,T]`` matching decoder-order targets.""" + + expected_shape = tuple(expected) + if visibility.shape == expected_shape: + return visibility + if ( + visibility.ndim == 4 + and visibility.shape[:2] == expected_shape[:2] + and visibility.shape[2] == expected_shape[3] + and visibility.shape[3] == expected_shape[2] + ): + return visibility.transpose(2, 3) + raise ValueError( + f"visibility {tuple(visibility.shape)} does not match {expected_shape}" + ) + + +def zero_invisible_track_xy( + xy: torch.Tensor, + visibility: torch.Tensor | None, +) -> torch.Tensor: + """Zero XY where visibility is 0 so invisible coords never enter noising.""" + + if visibility is None: + return xy + if xy.ndim != 5 or xy.shape[-1] != 2: + raise ValueError(f"xy must be [...,2] rank-5, got {tuple(xy.shape)}") + if visibility.shape == xy.shape[:-1]: + mask = visibility + elif ( + visibility.ndim == 4 + and visibility.shape[:2] == xy.shape[:2] + and visibility.shape[2] == xy.shape[3] + and visibility.shape[3] == xy.shape[2] + ): + mask = visibility.transpose(2, 3) + else: + raise ValueError( + f"visibility {tuple(visibility.shape)} does not match xy {tuple(xy.shape)}" + ) + return xy * mask.to(device=xy.device, dtype=xy.dtype).clamp(0, 1).unsqueeze(-1) + + +def track_motion_weights( + clean_xy: torch.Tensor, + visibility: torch.Tensor, + *, + boost: float = 2.0, +) -> torch.Tensor: + """Per-point weights that upweight trajectories with larger GT motion. + + ``clean_xy`` / returned weights use ``[B,blocks,250,T]`` (decoder layout). + """ + + if clean_xy.ndim != 5 or clean_xy.shape[-1] != 2: + raise ValueError("clean_xy must be [B,K,N,T,2]") + visible = _canonicalize_visibility_layout(visibility, clean_xy.shape[:-1]) + visible = visible.to(device=clean_xy.device, dtype=clean_xy.dtype).clamp(0, 1) + deltas = clean_xy[..., 1:, :] - clean_xy[..., :-1, :] + delta_vis = visible[..., 1:] * visible[..., :-1] + motion = (deltas.norm(dim=-1) * delta_vis).sum(dim=-1) # [B,K,N] + point_weight = 1.0 + float(boost) * motion + return point_weight.unsqueeze(-1).expand_as(visible) + + +def visibility_masked_track_loss( + prediction: torch.Tensor, + target: torch.Tensor, + visibility: torch.Tensor, + *, + weight: torch.Tensor | None = None, + reduction: str = "mean", +) -> torch.Tensor: + """MSE track-flow loss evaluated only at visible target points. + + ``prediction`` and ``target`` use ``[B,blocks,250,16,2]``. Visibility may + be ``[B,blocks,250,16]`` or ``[B,blocks,16,250]``. + """ + + if prediction.shape != target.shape or prediction.ndim != 5: + raise ValueError( + f"prediction/target must share [B,K,N,T,2], got " + f"{tuple(prediction.shape)} and {tuple(target.shape)}" + ) + expected = prediction.shape[:-1] + visible = _canonicalize_visibility_layout(visibility, expected) + + finite = torch.isfinite(target).all(dim=-1) & torch.isfinite(prediction).all(dim=-1) + mask = visible.to(device=prediction.device, dtype=prediction.dtype).clamp(0, 1) + mask = mask * finite.to(mask.dtype) + squared_error = (prediction - target).square() + if weight is not None: + while weight.ndim < squared_error.ndim: + weight = weight.unsqueeze(-1) + squared_error = squared_error * weight.to( + device=prediction.device, dtype=prediction.dtype + ) + masked = squared_error * mask.unsqueeze(-1) + if reduction == "none": + return masked + if reduction == "sum": + return masked.sum() + if reduction != "mean": + raise ValueError(f"unsupported reduction: {reduction}") + denominator = (mask.sum() * prediction.shape[-1]).clamp_min(1.0) + return masked.sum() / denominator + + +def visibility_masked_track_smoothness_loss( + prediction: torch.Tensor, + target: torch.Tensor, + visibility: torch.Tensor, + *, + weight: torch.Tensor | None = None, +) -> torch.Tensor: + """Temporal first-difference MSE on visible consecutive track positions.""" + + if prediction.shape != target.shape or prediction.ndim != 5: + raise ValueError("prediction/target must share [B,K,N,T,2]") + if prediction.shape[3] < 2: + return prediction.new_zeros(()) + pred_delta = prediction[..., 1:, :] - prediction[..., :-1, :] + target_delta = target[..., 1:, :] - target[..., :-1, :] + visible = _canonicalize_visibility_layout(visibility, prediction.shape[:-1]) + delta_visibility = visible[..., 1:] * visible[..., :-1] + delta_weight = None + if weight is not None: + aligned = weight + while aligned.ndim < visible.ndim: + aligned = aligned.unsqueeze(-1) + if aligned.shape != visible.shape: + raise ValueError( + f"weight {tuple(weight.shape)} incompatible with " + f"{tuple(visible.shape)}" + ) + delta_weight = aligned[..., 1:] + return visibility_masked_track_loss( + pred_delta, + target_delta, + delta_visibility, + weight=delta_weight, + ) + + +# Concise aliases used by integrations. +TrackEmbedding = TrackEncoder +masked_track_flow_loss = visibility_masked_track_loss + + +__all__ = [ + "NUM_TRACK_POINTS", + "TRACK_HORIZON", + "TRACK_POINT_SLICES", + "TRACK_SEGMENTS", + "TrackDecoder", + "TrackEmbedding", + "TrackEmbeddingOutput", + "TrackEncoder", + "TrackHand", + "TrackRole", + "TrackSegment", + "TrackView", + "canonical_point_metadata", + "masked_track_flow_loss", + "track_motion_weights", + "validate_track_layout", + "visibility_masked_track_loss", + "visibility_masked_track_smoothness_loss", + "zero_invisible_track_xy", +] diff --git a/groot/vla/model/trex_track_force/transform.py b/groot/vla/model/trex_track_force/transform.py new file mode 100644 index 0000000000000000000000000000000000000000..bf795bc63cab23c18e2e367c02b50146bd477d72 --- /dev/null +++ b/groot/vla/model/trex_track_force/transform.py @@ -0,0 +1,402 @@ +"""Model-side transform helpers for fixed-size track/force columns.""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +from pydantic import Field + +from groot.vla.model.dreamzero.transform.dreamzero_cotrain import ( + DefaultDataCollator, + DreamTransform, +) + +from .force import ( + ACTION_HORIZON, + FORCE_HISTORY_FRAMES, + FORCE_OFFSETS, + FORCE_SENSOR_COUNT, + FORCE_SENSOR_DIM, +) +from .track import NUM_TRACK_POINTS, TRACK_HORIZON + + +TRACK_XY_ALIASES = ( + "track_xy", + "tracks", + "track.xy", + "observation.track_xy", + "observation.tracks_xy", +) +TRACK_VISIBILITY_ALIASES = ( + "track_visibility", + "visibility", + "vis", + "track.visibility", + "observation.track_visibility", + "observation.tracks_visibility", +) +FORCE_CURRENT_ALIASES = ( + "current_force", + "force_current", + "observation.force_current", +) +FORCE_HISTORY_ALIASES = ( + "force_history", + "tactile_force_history", + "force_history_raw", + "observation.tactile_force", + "force_history_vq", + "force_vq_codes", + "observation.force_history_vq", +) +FORCE_HISTORY_PADDING_MASK_ALIASES = ( + "force_history_padding_mask", + "tactile_force_history_padding_mask", +) + + +def _first_present(data: dict[str, Any], names: tuple[str, ...]) -> Any | None: + for name in names: + if name in data: + return data[name] + return None + + +def _reshape_track_xy(values: Any) -> np.ndarray: + xy = np.asarray(values, dtype=np.float32) + if xy.ndim == 2 and xy.shape[-1] == NUM_TRACK_POINTS * 2: + xy = xy.reshape(xy.shape[0], NUM_TRACK_POINTS, 2) + if xy.ndim != 3 or xy.shape[1:] != (NUM_TRACK_POINTS, 2): + raise ValueError( + "track xy must be [T,250,2] or flattened [T,500], got " + f"{xy.shape}" + ) + if not np.isfinite(xy).all(): + raise ValueError("track xy contains NaN or infinity") + if xy.min(initial=0.0) < -1e-4 or xy.max(initial=1.0) > 1.0001: + raise ValueError("track xy must be normalized to [0,1]") + return xy + + +def _reshape_track_visibility(values: Any, time: int) -> np.ndarray: + visibility = np.asarray(values) + if visibility.ndim == 1 and visibility.size == time * NUM_TRACK_POINTS: + visibility = visibility.reshape(time, NUM_TRACK_POINTS) + if visibility.shape != (time, NUM_TRACK_POINTS): + raise ValueError( + f"track visibility must be [T,250], got {visibility.shape}" + ) + return visibility.astype(np.float32) + + +def prepare_track_windows( + track_xy: Any, + track_visibility: Any, + *, + num_blocks: int, + history_frames: int = FORCE_HISTORY_FRAMES, + future_horizon: int = TRACK_HORIZON, +) -> dict[str, np.ndarray]: + """Split a contiguous track timeline into clean-past and future windows. + + The first future step immediately follows the final clean-past step. + Boundary indices are clamped, which is deterministic for online prefixes. + """ + + if num_blocks <= 0 or history_frames <= 0: + raise ValueError("num_blocks and history_frames must be positive") + xy = _reshape_track_xy(track_xy) + visibility = _reshape_track_visibility(track_visibility, xy.shape[0]) + total_time = xy.shape[0] + minimum = history_frames + num_blocks * future_horizon + if total_time < future_horizon: + raise ValueError( + f"track timeline has {total_time} frames; need at least {future_horizon}" + ) + + # A full training window is [H past][K*16 future]. For a shorter online + # prefix, anchor at its first frame and use deterministic edge padding. + first_anchor = history_frames - 1 if total_time >= minimum else 0 + past_xy, past_vis, future_xy, future_vis = [], [], [], [] + for block_index in range(num_blocks): + anchor = first_anchor + block_index * future_horizon + past_indices = np.arange( + anchor - history_frames + 1, anchor + 1, dtype=np.int64 + ) + future_indices = np.arange( + anchor + 1, anchor + 1 + future_horizon, dtype=np.int64 + ) + past_indices = np.clip(past_indices, 0, total_time - 1) + future_indices = np.clip(future_indices, 0, total_time - 1) + past_xy.append(xy[past_indices]) + past_vis.append(visibility[past_indices]) + future_xy.append(xy[future_indices]) + future_vis.append(visibility[future_indices]) + return { + "track_past_xy": np.stack(past_xy), + "track_past_visibility": np.stack(past_vis), + "track_future_xy": np.stack(future_xy), + "track_future_visibility": np.stack(future_vis), + } + + +class TrexTrackForceTransform(DreamTransform): + """DreamTransform extension that preserves track and force model inputs.""" + + track_history_frames: int = FORCE_HISTORY_FRAMES + track_horizon: int = TRACK_HORIZON + track_points: int = NUM_TRACK_POINTS + require_track_inputs: bool = True + require_force_inputs: bool = True + track_xy_aliases: tuple[str, ...] = Field(default=TRACK_XY_ALIASES) + track_visibility_aliases: tuple[str, ...] = Field( + default=TRACK_VISIBILITY_ALIASES + ) + force_current_aliases: tuple[str, ...] = Field(default=FORCE_CURRENT_ALIASES) + force_history_aliases: tuple[str, ...] = Field(default=FORCE_HISTORY_ALIASES) + force_history_padding_mask_aliases: tuple[str, ...] = Field( + default=FORCE_HISTORY_PADDING_MASK_ALIASES + ) + + def _num_blocks(self, data: dict[str, Any]) -> int: + if "action" not in data: + return 1 + action = np.asarray(data["action"]) + if action.shape[0] % ACTION_HORIZON: + raise ValueError("action timeline must contain complete 16-step blocks") + return max(action.shape[0] // ACTION_HORIZON, 1) + + def _prepare_tracks(self, data: dict[str, Any], blocks: int) -> dict[str, np.ndarray]: + explicit_keys = ( + "track_past_xy", + "track_past_visibility", + "track_future_xy", + "track_future_visibility", + ) + if all(key in data for key in explicit_keys): + result = { + key: np.asarray(data[key], dtype=np.float32) for key in explicit_keys + } + expected = { + "track_past_xy": ( + blocks, + self.track_history_frames, + self.track_points, + 2, + ), + "track_past_visibility": ( + blocks, + self.track_history_frames, + self.track_points, + ), + "track_future_xy": ( + blocks, + self.track_horizon, + self.track_points, + 2, + ), + "track_future_visibility": ( + blocks, + self.track_horizon, + self.track_points, + ), + } + for key, shape in expected.items(): + if result[key].shape != shape: + raise ValueError(f"{key} is {result[key].shape}, expected {shape}") + return result + + xy = _first_present(data, self.track_xy_aliases) + visibility = _first_present(data, self.track_visibility_aliases) + if xy is None or visibility is None: + if self.require_track_inputs: + raise KeyError( + "250-point track xy/visibility columns were not provided" + ) + return {} + return prepare_track_windows( + xy, + visibility, + num_blocks=blocks, + history_frames=self.track_history_frames, + future_horizon=self.track_horizon, + ) + + def _prepare_force(self, data: dict[str, Any], blocks: int) -> dict[str, np.ndarray]: + current = _first_present(data, self.force_current_aliases) + history = _first_present(data, self.force_history_aliases) + history_padding_mask = _first_present( + data, self.force_history_padding_mask_aliases + ) + if history is None: + if self.require_force_inputs: + raise KeyError( + "raw 16x10x6 tactile history or 10 precomputed codes are required" + ) + return {} + history_array = np.asarray(history) + is_raw = np.issubdtype(history_array.dtype, np.floating) + force_updates = len(FORCE_OFFSETS) + if is_raw: + if history_array.shape == ( + FORCE_HISTORY_FRAMES, + FORCE_SENSOR_COUNT * FORCE_SENSOR_DIM, + ): + history_array = history_array.reshape( + FORCE_HISTORY_FRAMES, FORCE_SENSOR_COUNT, FORCE_SENSOR_DIM + ) + if history_array.shape == ( + blocks, + FORCE_HISTORY_FRAMES, + FORCE_SENSOR_COUNT * FORCE_SENSOR_DIM, + ): + history_array = history_array.reshape( + blocks, + FORCE_HISTORY_FRAMES, + FORCE_SENSOR_COUNT, + FORCE_SENSOR_DIM, + ) + if history_array.shape == ( + FORCE_HISTORY_FRAMES, + FORCE_SENSOR_COUNT, + FORCE_SENSOR_DIM, + ): + history_array = np.broadcast_to( + history_array, + ( + blocks, + FORCE_HISTORY_FRAMES, + FORCE_SENSOR_COUNT, + FORCE_SENSOR_DIM, + ), + ).copy() + expected_history = ( + blocks, + FORCE_HISTORY_FRAMES, + FORCE_SENSOR_COUNT, + FORCE_SENSOR_DIM, + ) + expected_update_history = ( + blocks, + force_updates, + FORCE_HISTORY_FRAMES, + FORCE_SENSOR_COUNT, + FORCE_SENSOR_DIM, + ) + if history_array.shape not in (expected_history, expected_update_history): + raise ValueError( + "raw tactile history must be [blocks,16,10,6] or " + f"[blocks,{force_updates},16,10,6], " + f"got {history_array.shape}" + ) + history_array = history_array.astype(np.float32) + expected_padding_shape = history_array.shape[:-2] + if history_padding_mask is None: + padding_array = np.zeros(expected_padding_shape, dtype=bool) + else: + padding_array = np.asarray(history_padding_mask, dtype=bool) + if padding_array.shape == (FORCE_HISTORY_FRAMES,): + padding_array = np.broadcast_to( + padding_array, expected_padding_shape + ).copy() + if padding_array.shape != expected_padding_shape: + raise ValueError( + "force history padding mask must match the raw history " + f"time axes {expected_padding_shape}, got " + f"{padding_array.shape}" + ) + else: + if history_array.shape == (FORCE_SENSOR_COUNT,): + history_array = np.broadcast_to( + history_array, (blocks, FORCE_SENSOR_COUNT) + ).copy() + expected_codes = (blocks, FORCE_SENSOR_COUNT) + expected_update_codes = (blocks, force_updates, FORCE_SENSOR_COUNT) + if history_array.shape not in (expected_codes, expected_update_codes): + raise ValueError( + "precomputed tactile codes must be [blocks,10] or " + f"[blocks,{force_updates},10], " + f"got {history_array.shape}" + ) + history_array = history_array.astype(np.int64) + + if current is None and is_raw: + current = ( + history_array[:, :, -1] + if history_array.ndim == 6 + else history_array[:, -1] + ) + if current is None: + if self.require_force_inputs: + raise KeyError("current 10x6 force is required with precomputed codes") + return {} + current_array = np.asarray(current, dtype=np.float32) + if current_array.shape == (FORCE_SENSOR_COUNT * FORCE_SENSOR_DIM,): + current_array = current_array.reshape(FORCE_SENSOR_COUNT, FORCE_SENSOR_DIM) + if current_array.shape == (FORCE_SENSOR_COUNT, FORCE_SENSOR_DIM): + current_array = np.broadcast_to( + current_array, (blocks, FORCE_SENSOR_COUNT, FORCE_SENSOR_DIM) + ).copy() + expected_current = ( + blocks, + FORCE_SENSOR_COUNT, + FORCE_SENSOR_DIM, + ) + expected_update_current = ( + blocks, + force_updates, + FORCE_SENSOR_COUNT, + FORCE_SENSOR_DIM, + ) + if current_array.shape not in (expected_current, expected_update_current): + raise ValueError( + "current force must be [blocks,10,6] or " + f"[blocks,{force_updates},10,6], got {current_array.shape}" + ) + + result = {"current_force": current_array} + result[ + "tactile_force_history" if is_raw else "force_history_vq" + ] = history_array + if is_raw: + result["tactile_force_history_padding_mask"] = padding_array + return result + + def apply_single(self, data: dict) -> dict: + blocks = self._num_blocks(data) + track_data = self._prepare_tracks(data, blocks) + force_data = self._prepare_force(data, blocks) + transformed = super().apply_single(data) + images = np.asarray(transformed["images"]) + expected_training_frames = 1 + blocks * 8 + if self.training and images.shape[0] == expected_training_frames: + # Preserve DreamZero's 4-block/33-frame autoregressive geometry: + # one clean conditioning frame plus eight 10 Hz targets per block. + transformed["history_images"] = images[:1] + transformed["future_images"] = images[1:] + elif images.shape[0] >= 1 and not self.training: + transformed["history_images"] = images + else: + raise ValueError( + "T-Rex Track-Force expects one conditioning frame plus " + f"eight frames per autoregressive block ({expected_training_frames} " + "training frames), or one or more history frames for inference; " + f"got {images.shape[0]}" + ) + transformed.update(track_data) + transformed.update(force_data) + return transformed + + +TrackForceTransform = TrexTrackForceTransform +TrackForceDataCollator = DefaultDataCollator + + +__all__ = [ + "TrackForceDataCollator", + "TrackForceTransform", + "TrexTrackForceTransform", + "prepare_track_windows", +] diff --git a/groot/vla/model/trex_track_force/vla.py b/groot/vla/model/trex_track_force/vla.py new file mode 100644 index 0000000000000000000000000000000000000000..f86cd03aeb2c5766deb8e28011ef7e0ca8cf812e --- /dev/null +++ b/groot/vla/model/trex_track_force/vla.py @@ -0,0 +1,137 @@ +"""Distinct Hugging Face/Hydra identity while reusing DreamZero's VLA shell.""" + +from __future__ import annotations + +import json +import os + +from safetensors.torch import load_file +from transformers import AutoConfig, AutoModel + +from groot.vla.model.dreamzero.base_vla import VLA, VLAConfig + + +class TrexTrackForceVLAConfig(VLAConfig): + model_type = "trex_track_force_vla" + + +class TrexTrackForceVLA(VLA): + """Unmodified VLA orchestration with the independent track/force head.""" + + config_class = TrexTrackForceVLAConfig + + def get_visualization( + self, + inputs: dict, + *, + coarse_steps: int = 1, + seed: int | None = None, + ): + """Fast video/track sampling without the action-only tactile stage.""" + + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + outputs = self.action_head.get_visualization( + backbone_outputs, + action_inputs, + coarse_steps=coarse_steps, + seed=seed, + ) + self.validate_data(outputs, backbone_outputs, is_training=False) + return outputs + + def get_trainable_checkpoint_state_dict( + self, state_dict: dict | None = None + ) -> dict: + """Save trainable weights plus EMA VQ buffers in LoRA-only checkpoints.""" + + source = self.state_dict() if state_dict is None else state_dict + trainable = { + name for name, parameter in self.named_parameters() if parameter.requires_grad + } + ema_prefix = "action_head.force_transformer.tactile_vq.quantizer." + return { + name: value + for name, value in source.items() + if name in trainable or name.startswith(ema_prefix) + } + + @classmethod + def load_lora(cls, pretrained_model_name_or_path: str) -> "TrexTrackForceVLA": + """Load base Wan components, custom LoRA adapters, and all new modules.""" + + config_path = os.path.join(pretrained_model_name_or_path, "config.json") + with open(config_path, "r", encoding="utf-8") as handle: + config = TrexTrackForceVLAConfig(**json.load(handle)) + action_head_cfg = config.action_head_cfg + inner = ( + action_head_cfg.get("config", action_head_cfg) + if isinstance(action_head_cfg, dict) + else None + ) + if not isinstance(inner, dict): + raise TypeError("checkpoint action_head_cfg is not a mapping") + inner["defer_lora_injection"] = False + inner["load_pretrained_components"] = True + model = cls(config) + current = model.state_dict() + + index_path = os.path.join( + pretrained_model_name_or_path, "model.safetensors.index.json" + ) + single_path = os.path.join(pretrained_model_name_or_path, "model.safetensors") + if os.path.isfile(index_path): + with open(index_path, "r", encoding="utf-8") as handle: + index = json.load(handle) + paths = [ + os.path.join(pretrained_model_name_or_path, filename) + for filename in sorted(set(index["weight_map"].values())) + ] + elif os.path.isfile(single_path): + paths = [single_path] + else: + raise FileNotFoundError( + f"no model.safetensors checkpoint in {pretrained_model_name_or_path}" + ) + + loaded: set[str] = set() + for path in paths: + state = load_file(path) + unexpected = [key for key in state if key not in current] + mismatched = [ + key + for key, value in state.items() + if key in current and current[key].shape != value.shape + ] + if unexpected or mismatched: + raise ValueError( + "incompatible T-Rex Track-Force checkpoint: " + f"unexpected={unexpected[:10]}, mismatched={mismatched[:10]}" + ) + model.load_state_dict(state, strict=False) + loaded.update(state) + if not loaded: + raise ValueError("T-Rex Track-Force checkpoint contains no parameters") + expected = set(model.get_trainable_checkpoint_state_dict(current)) + missing = sorted(expected - loaded) + if missing: + raise ValueError( + "T-Rex Track-Force checkpoint is incomplete; " + f"missing trainable/EMA keys={missing[:20]}" + ) + return model + + +AutoConfig.register( + TrexTrackForceVLAConfig.model_type, + TrexTrackForceVLAConfig, + exist_ok=True, +) +AutoModel.register( + TrexTrackForceVLAConfig, + TrexTrackForceVLA, + exist_ok=True, +) + + +__all__ = ["TrexTrackForceVLA", "TrexTrackForceVLAConfig"] diff --git a/groot/vla/utils/__init__.py b/groot/vla/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/groot/vla/utils/action_args_override_utils.py b/groot/vla/utils/action_args_override_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..67ecdbb307bf02ebef0282a1d5065f4d3bd8b24f --- /dev/null +++ b/groot/vla/utils/action_args_override_utils.py @@ -0,0 +1,68 @@ +""" +Utility functions for automatically updating action horizon and related configurations. +""" + +from omegaconf import DictConfig, open_dict + + +def update_action_horizon_configs(cfg: DictConfig, action_horizon: int) -> DictConfig: + """ + Automatically update action_horizon and corresponding delta_indices in all modality configs. + + Args: + cfg: The hydra configuration + action_horizon: The desired action horizon (e.g., 30) + + Returns: + Updated configuration with action_horizon and delta_indices set appropriately + """ + # Generate delta_indices for the given action_horizon [0, 1, 2, ..., action_horizon-1] + delta_indices = list(range(action_horizon)) + + # Update the global action_horizon + with open_dict(cfg): + cfg.action_horizon = action_horizon + if hasattr(cfg.model, "vla_override_kwargs"): + cfg.model.vla_override_kwargs.action_horizon = action_horizon + if hasattr(cfg.model, "action_head_override_kwargs"): + cfg.model.action_head_override_kwargs.action_horizon = action_horizon + + # Update delta_indices for all action modalities in modality_configs + if hasattr(cfg, "modality_configs"): + for embodiment_name, modality_config in cfg.modality_configs.items(): + if hasattr(modality_config, "action"): + # Update the action delta_indices + modality_config.action.delta_indices = delta_indices + print(f"Updated {embodiment_name}.action.delta_indices to {delta_indices}") + + return cfg + + +def update_action_dim_configs(cfg: DictConfig, new_action_dim: int) -> DictConfig: + """ + Update the action dimension in all modality configs. + """ + with open_dict(cfg): + cfg.max_action_dim = new_action_dim + return cfg + + +def apply_action_overrides(cfg: DictConfig) -> DictConfig: + """ + Apply action horizon overrides if action_horizon is specified in config. + This function should be called after the config is loaded but before model instantiation. + """ + if hasattr(cfg.model, "action_head_override_kwargs"): + action_horizon = cfg.action_horizon + print(f"Applying action_horizon={action_horizon} overrides...") + cfg = update_action_horizon_configs(cfg, action_horizon) + + if hasattr(cfg.model, "expand_action_head_kwargs"): + expand_action_head_kwargs = cfg.model.expand_action_head_kwargs + if "expand_action_dim" in expand_action_head_kwargs: + old_action_dim = expand_action_head_kwargs.expand_action_dim.old_action_dim + new_action_dim = expand_action_head_kwargs.expand_action_dim.new_action_dim + print(f"Applying expand_action_dim={old_action_dim}->{new_action_dim} overrides...") + cfg = update_action_dim_configs(cfg, new_action_dim) + + return cfg diff --git a/groot/vla/utils/timer.py b/groot/vla/utils/timer.py new file mode 100644 index 0000000000000000000000000000000000000000..d4e17daa2caee54c85fd7869195136cd2b098a87 --- /dev/null +++ b/groot/vla/utils/timer.py @@ -0,0 +1,42 @@ +import time + + +class ContextTimer: + + def __init__(self, trainer): + self.last_key = None + self.trainer = trainer + self.start_times = {} + self.key_stack = [] + + def with_label(self, key): + self.last_key = key + return self + + def __enter__(self): + self.key_stack.append(self.last_key) # Push key to stack + self.start_times[self.last_key] = time.time() # Start timing for this key + return self + + def __exit__(self, exc_type, exc_value, traceback): + key = self.key_stack.pop() # Pop key from stack + diff = time.time() - self.start_times[key] + self.trainer.log({f"{key}_time": diff}) + # print(f"{key}: {diff:.2f} seconds") + + +if __name__ == "__main__": + + class MockTrainer: + def log(self, data): + print("Logging:", data) + + trainer = MockTrainer() + my_timer = ContextTimer(trainer) + + with my_timer.with_label("outer"): + time.sleep(1) + with my_timer.with_label("inner"): + time.sleep(2) + with my_timer.with_label("inner"): # Another inner block + time.sleep(1) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..cec38dd81094702b690711ddb4ba69e8c3034d96 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,103 @@ +[build-system] +requires = ["setuptools>=67", "wheel", "pip"] +build-backend = "setuptools.build_meta" + +[project] +name = "dreamzero" +description = "DreamZero: World Action Models Are Zero-Shot Policies" +version="1.0.0" +readme = "README.md" +classifiers = [ + "Intended Audience :: Science/Research", + "Development Status :: 3 - Alpha", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3", + "Topic :: Scientific/Engineering :: Artificial Intelligence", +] +authors = [ + {name = "NVIDIA Gear Lab"} +] +requires-python = "~=3.11,<3.13" +dependencies = [ + "av==15.0.0", + "torch==2.8.0", + "torchvision==0.23.0", + "torchaudio==2.8.0", + "pyttsx3==2.90", + "scipy==1.15.3", + "numpy==1.26.4", + "matplotlib", + "hydra-core", + "ray[default]==2.47.1", + "click", + "gymnasium", + "mujoco", + "termcolor", + "flask", + "python-socketio>=5.13.0", + "flask_socketio", + "loguru", + "lmdb", + "meshcat", + "meshcat-shapes", + "rerun-sdk==0.21.0", + "pygame", + "sshkeyboard", + "msgpack", + "msgpack-numpy", + "peft==0.5.0", + "pyzmq", + "PyQt6; platform_machine != 'aarch64'", + "pin", + "pin-pink", + "timm", + "tyro", + "redis", + "lark", + "datasets==3.6.0", + "pandas", + "evdev; sys_platform == 'linux'", + "pybullet; sys_platform == 'linux'", + "gear", + # Pin the MSC dependency version to prevent potential issues caused by publishing new releases. + "multi-storage-client[boto3,msal,observability-otel]==0.33.0", + "dm_tree", + "openai", + "transformers==4.51.3", + "albumentations==1.4.18", + "einops==0.8.1", + "tyro", + "tianshou==0.5.1", + "imageio==2.34.2", + "imageio-ffmpeg", + "wandb", + "opencv-python==4.8.0.74", + "diffusers==0.30.2", + "ftfy", + "nvidia-modelopt", + "nvidia-modelopt-core", + "tensorrt", + "openpi-client==0.1.1", + "huggingface_hub", + "decord2", + "deepspeed", + "tiktoken", + "sentencepiece" +] + +[project.optional-dependencies] +dev = [ + "pytest>=7.0.0", + "black>=23.0.0", + "isort>=5.12.0", +] + +[project.scripts] +dreamzero-server = "socket_test_optimized_AR:main" + +[tool.setuptools.packages.find] +where = ["."] +include = ["groot*"] + +[tool.setuptools] +include-package-data = true diff --git a/scripts/compare_loss.py b/scripts/compare_loss.py new file mode 100644 index 0000000000000000000000000000000000000000..7b79ee1f8a9b900d34bd80ec3e8da428cc8b91d4 --- /dev/null +++ b/scripts/compare_loss.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python3 +"""Compare loss curves between LoRA and full fine-tuning runs. + +Usage: + python scripts/compare_loss.py \ + --lora-log ./checkpoints/dreamzero_droid_lora/loss_log.jsonl \ + --full-log ./checkpoints/dreamzero_droid_full_finetune/loss_log.jsonl \ + [--plot loss_comparison.png] +""" + +import argparse +import json + + +def load_loss_log(path): + entries = [] + with open(path) as f: + for line in f: + line = line.strip() + if line: + entries.append(json.loads(line)) + return entries + + +def print_comparison_table(lora_entries, full_entries): + # Index by step + lora_by_step = {e["step"]: e for e in lora_entries} + full_by_step = {e["step"]: e for e in full_entries} + all_steps = sorted(set(lora_by_step.keys()) | set(full_by_step.keys())) + + header = f"{'Step':>6} {'LoRA Loss':>10} {'Full Loss':>10} {'LoRA Dyn':>10} {'Full Dyn':>10} {'LoRA Act':>10} {'Full Act':>10}" + print(header) + print("-" * len(header)) + + for step in all_steps: + lora = lora_by_step.get(step, {}) + full = full_by_step.get(step, {}) + + def fmt(d, key): + v = d.get(key) + return f"{v:10.4f}" if v is not None else f"{'—':>10}" + + print( + f"{step:>6} " + f"{fmt(lora, 'loss')} {fmt(full, 'loss')} " + f"{fmt(lora, 'dynamics_loss_avg')} {fmt(full, 'dynamics_loss_avg')} " + f"{fmt(lora, 'action_loss_avg')} {fmt(full, 'action_loss_avg')}" + ) + + +def plot_comparison(lora_entries, full_entries, output_path): + try: + import matplotlib.pyplot as plt + except ImportError: + print("matplotlib not installed, skipping plot generation.") + print("Install with: pip install matplotlib") + return + + metrics = [ + ("loss", "Total Loss"), + ("dynamics_loss_avg", "Dynamics Loss"), + ("action_loss_avg", "Action Loss"), + ] + + fig, axes = plt.subplots(1, len(metrics), figsize=(5 * len(metrics), 4)) + if len(metrics) == 1: + axes = [axes] + + for ax, (key, title) in zip(axes, metrics): + lora_steps = [e["step"] for e in lora_entries if key in e] + lora_vals = [e[key] for e in lora_entries if key in e] + full_steps = [e["step"] for e in full_entries if key in e] + full_vals = [e[key] for e in full_entries if key in e] + + if lora_steps: + ax.plot(lora_steps, lora_vals, label="LoRA", marker="o", markersize=3) + if full_steps: + ax.plot(full_steps, full_vals, label="Full FT", marker="s", markersize=3) + + ax.set_title(title) + ax.set_xlabel("Step") + ax.set_ylabel("Loss") + ax.legend() + ax.grid(True, alpha=0.3) + + fig.tight_layout() + fig.savefig(output_path, dpi=150) + print(f"Plot saved to {output_path}") + + +def main(): + parser = argparse.ArgumentParser(description="Compare LoRA vs full fine-tuning loss curves") + parser.add_argument("--lora-log", required=True, help="Path to LoRA run loss_log.jsonl") + parser.add_argument("--full-log", required=True, help="Path to full FT run loss_log.jsonl") + parser.add_argument("--plot", default=None, help="Output path for comparison plot (e.g., loss_comparison.png)") + args = parser.parse_args() + + lora_entries = load_loss_log(args.lora_log) + full_entries = load_loss_log(args.full_log) + + print(f"LoRA: {len(lora_entries)} log entries") + print(f"Full: {len(full_entries)} log entries") + print() + + print_comparison_table(lora_entries, full_entries) + + if args.plot: + plot_comparison(lora_entries, full_entries, args.plot) + + +if __name__ == "__main__": + main() diff --git a/scripts/data/audit_trex_tracks.py b/scripts/data/audit_trex_tracks.py new file mode 100644 index 0000000000000000000000000000000000000000..b03f95072580de4617f1d294ea0ba8b9d587f12f --- /dev/null +++ b/scripts/data/audit_trex_tracks.py @@ -0,0 +1,647 @@ +#!/usr/bin/env python3 +"""Audit all T-Rex track caches and investigate near-static wrist tracks.""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import json +import math +from collections import Counter +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path + +import cv2 +import numpy as np +import pyarrow.parquet as pq + +VIEW_SLICES = { + "head_left": (0, 100), + "left_wrist": (100, 175), + "right_wrist": (175, 250), +} +WRIST_GROUPS = { + "left_wrist": {"background": (100, 125), "hand": (125, 175)}, + "right_wrist": {"background": (175, 200), "hand": (200, 250)}, +} +IMAGE_SCALE = np.array([320.0, 180.0], dtype=np.float32) +WINDOW_FRAMES = 768 +WINDOW_OVERLAP = 64 +WINDOW_STEP = WINDOW_FRAMES - WINDOW_OVERLAP + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--dataset-root", + type=Path, + default=Path("/scratch1/home/zhicao/dreamzero/data/trex_full_force"), + ) + parser.add_argument( + "--output-dir", + type=Path, + default=Path( + "/scratch1/home/zhicao/dreamzero/data/" + "trex_full_force/audit/track_quality" + ), + ) + parser.add_argument("--static-span-px", type=float, default=2.0) + parser.add_argument("--video-samples", type=int, default=5) + parser.add_argument("--video-workers", type=int, default=8) + parser.add_argument("--skip-video-check", action="store_true") + return parser.parse_args() + + +def _quantiles(values: list[float]) -> dict[str, float]: + array = np.asarray(values, dtype=np.float64) + return { + name: float(np.percentile(array, percentile)) + for name, percentile in ( + ("min", 0), + ("p01", 1), + ("p05", 5), + ("median", 50), + ("p95", 95), + ("p99", 99), + ("max", 100), + ) + } + + +def _point_span_px(tracks: np.ndarray) -> np.ndarray: + pixel_tracks = np.asarray(tracks, dtype=np.float32) * IMAGE_SCALE + return np.sqrt( + np.ptp(pixel_tracks[..., 0], axis=0) ** 2 + + np.ptp(pixel_tracks[..., 1], axis=0) ** 2 + ) + + +def _seam_metrics( + tracks: np.ndarray, + visibility: np.ndarray, +) -> tuple[float, float, float]: + frames = int(tracks.shape[0]) + boundaries = list(range(WINDOW_FRAMES, frames, WINDOW_STEP)) + if not boundaries: + return math.nan, math.nan, math.nan + + pixel_tracks = np.asarray(tracks, dtype=np.float32) * IMAGE_SCALE + delta = np.linalg.norm(np.diff(pixel_tracks, axis=0), axis=-1) + visible_pair = (visibility[1:] > 0.5) & (visibility[:-1] > 0.5) + seam_indices = np.asarray([boundary - 1 for boundary in boundaries], dtype=np.int64) + seam_values = delta[seam_indices][visible_pair[seam_indices]] + + regular_indices = np.unique( + np.rint(np.linspace(0, max(0, frames - 2), min(128, frames - 1))).astype( + np.int64 + ) + ) + regular_indices = regular_indices[ + ~np.isin(regular_indices, seam_indices) + ] + regular_values = delta[regular_indices][visible_pair[regular_indices]] + seam_p95 = ( + float(np.percentile(seam_values, 95)) if seam_values.size else math.nan + ) + regular_p95 = ( + float(np.percentile(regular_values, 95)) + if regular_values.size + else math.nan + ) + ratio = ( + seam_p95 / max(regular_p95, 0.1) + if np.isfinite(seam_p95) and np.isfinite(regular_p95) + else math.nan + ) + return seam_p95, regular_p95, ratio + + +def _audit_track( + path: Path, + episode_index: int, + expected_frames: int, +) -> tuple[dict[str, object], list[str]]: + errors: list[str] = [] + with np.load(path, allow_pickle=False) as payload: + tracks = np.asarray(payload["tracks"], dtype=np.float32) + visibility = np.asarray(payload["vis"], dtype=np.float32) + if tracks.shape != (expected_frames, 250, 2): + errors.append(f"tracks shape {tracks.shape} != {(expected_frames, 250, 2)}") + if visibility.shape != (expected_frames, 250): + errors.append( + f"visibility shape {visibility.shape} != {(expected_frames, 250)}" + ) + if errors: + return {"episode_index": episode_index, "frames": expected_frames}, errors + + finite = np.isfinite(tracks).all(axis=-1) + in_frame = ((tracks >= 0.0) & (tracks <= 1.0)).all(axis=-1) + binary_visibility = (visibility == 0.0) | (visibility == 1.0) + if not finite.all(): + errors.append("non-finite track coordinates") + if not in_frame.all(): + errors.append("out-of-range normalized track coordinates") + if not binary_visibility.all(): + errors.append("non-binary visibility") + + row: dict[str, object] = { + "episode_index": episode_index, + "frames": expected_frames, + "finite_fraction": float(finite.mean()), + "in_frame_fraction": float(in_frame.mean()), + "binary_visibility_fraction": float(binary_visibility.mean()), + } + for view, (start, end) in VIEW_SLICES.items(): + view_tracks = tracks[:, start:end] + view_visibility = visibility[:, start:end] + span = _point_span_px(view_tracks) + row[f"{view}_visibility"] = float(view_visibility.mean()) + row[f"{view}_span_median_px"] = float(np.median(span)) + row[f"{view}_span_p95_px"] = float(np.percentile(span, 95)) + row[f"{view}_static_point_fraction"] = float(np.mean(span < 1.0)) + if episode_index >= 1737: + seam_p95, regular_p95, seam_ratio = _seam_metrics( + view_tracks, + view_visibility, + ) + else: + seam_p95, regular_p95, seam_ratio = math.nan, math.nan, math.nan + row[f"{view}_seam_jump_p95_px"] = seam_p95 + row[f"{view}_regular_jump_p95_px"] = regular_p95 + row[f"{view}_seam_jump_ratio"] = seam_ratio + + for view, groups in WRIST_GROUPS.items(): + for group, (start, end) in groups.items(): + span = _point_span_px(tracks[:, start:end]) + row[f"{view}_{group}_span_median_px"] = float(np.median(span)) + row[f"{view}_{group}_span_p95_px"] = float( + np.percentile(span, 95) + ) + return row, errors + + +def _video_path(root: Path, episode_index: int, view: str) -> Path: + return ( + root + / "videos" + / f"chunk-{episode_index // 1000:03d}" + / f"observation.images.{view}" + / f"episode_{episode_index:06d}.mp4" + ) + + +def _sample_video_motion( + root: Path, + episode_index: int, + view: str, + sample_count: int, +) -> dict[str, object]: + path = _video_path(root, episode_index, view) + cap = cv2.VideoCapture(str(path)) + frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + frames: list[np.ndarray] = [] + if frame_count > 0: + sample_indices = np.rint( + np.linspace(0, frame_count - 1, sample_count) + ).astype(int) + for frame_index in sample_indices: + cap.set(cv2.CAP_PROP_POS_FRAMES, int(frame_index)) + ok, frame = cap.read() + if not ok: + continue + gray = cv2.resize( + cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY), + (80, 45), + interpolation=cv2.INTER_AREA, + ).astype(np.float32) + frames.append(gray) + cap.release() + if len(frames) < 2: + return { + "episode_index": episode_index, + "view": view, + "video_error": f"decoded only {len(frames)} sampled frame(s)", + } + adjacent = [ + float(np.mean(np.abs(second - first))) + for first, second in zip(frames, frames[1:]) + ] + return { + "episode_index": episode_index, + "view": view, + "video_frames": frame_count, + "video_adjacent_mad_mean": float(np.mean(adjacent)), + "video_adjacent_mad_max": float(np.max(adjacent)), + "video_first_last_mad": float(np.mean(np.abs(frames[-1] - frames[0]))), + "video_first_frame_hash": hashlib.sha256( + frames[0].astype(np.uint8).tobytes() + ).hexdigest()[:16], + "video_sample_hash": hashlib.sha256( + np.stack(frames).astype(np.uint8).tobytes() + ).hexdigest()[:16], + } + + +def _arm_motion(root: Path, episode_index: int) -> dict[str, float]: + path = ( + root + / "data" + / f"chunk-{episode_index // 1000:03d}" + / f"episode_{episode_index:06d}.parquet" + ) + column = pq.read_table(path, columns=["observation.state"]).column(0) + array = column.combine_chunks() + values = np.asarray(array.values.to_numpy(zero_copy_only=False)).reshape( + len(array), + -1, + ) + result: dict[str, float] = {} + for name, selection in ( + ("left_arm", slice(0, 7)), + ("right_arm", slice(29, 36)), + ): + arm = values[:, selection].astype(np.float32, copy=False) + result[f"{name}_joint_path_l2"] = float( + np.linalg.norm(np.diff(arm, axis=0), axis=1).sum() + ) + result[f"{name}_end_delta_l2"] = float( + np.linalg.norm(arm[-1] - arm[0]) + ) + result[f"{name}_max_joint_range"] = float( + np.ptp(arm, axis=0).max() + ) + return result + + +def main() -> int: + args = _parse_args() + root = args.dataset_root.expanduser().resolve() + output_dir = args.output_dir.expanduser().resolve() + output_dir.mkdir(parents=True, exist_ok=True) + + episodes = [ + json.loads(line) + for line in (root / "meta" / "episodes.jsonl").read_text().splitlines() + if line.strip() + ] + rows: list[dict[str, object]] = [] + errors: list[dict[str, object]] = [] + for position, episode in enumerate(episodes, start=1): + episode_index = int(episode["episode_index"]) + path = ( + root + / "tracks_trex_track_force_v2" + / f"episode_{episode_index:06d}.npz" + ) + try: + row, track_errors = _audit_track( + path, + episode_index, + int(episode["length"]), + ) + except Exception as exc: # noqa: BLE001 + row = { + "episode_index": episode_index, + "frames": int(episode["length"]), + } + track_errors = [f"{type(exc).__name__}: {exc}"] + row["task"] = " | ".join(episode.get("tasks", [])) + rows.append(row) + if track_errors: + errors.append( + { + "episode_index": episode_index, + "errors": track_errors, + } + ) + if position % 250 == 0 or position == len(episodes): + print(f"Audited tracks: {position}/{len(episodes)}", flush=True) + + static_threshold = float(args.static_span_px) + candidates: list[tuple[int, str]] = [] + for row in rows: + for view in ("left_wrist", "right_wrist"): + value = row.get(f"{view}_background_span_median_px") + if isinstance(value, float) and value < static_threshold: + candidates.append((int(row["episode_index"]), view)) + candidate_keys = set(candidates) + candidate_episode_indices = sorted({episode for episode, _ in candidates}) + + video_results: list[dict[str, object]] = [] + video_targets = [ + (int(row["episode_index"]), view) + for row in rows + for view in ("left_wrist", "right_wrist") + ] + if not args.skip_video_check: + with ThreadPoolExecutor(max_workers=max(1, int(args.video_workers))) as pool: + futures = { + pool.submit( + _sample_video_motion, + root, + episode_index, + view, + max(2, int(args.video_samples)), + ): (episode_index, view) + for episode_index, view in video_targets + } + for position, future in enumerate(as_completed(futures), start=1): + video_results.append(future.result()) + if position % 500 == 0 or position == len(futures): + print( + f"Checked wrist videos: " + f"{position}/{len(futures)}", + flush=True, + ) + + video_by_key = { + (int(result["episode_index"]), str(result["view"])): result + for result in video_results + } + rows_by_episode = { + int(row["episode_index"]): row + for row in rows + } + video_frame_count_mismatches = [ + { + "episode_index": int(result["episode_index"]), + "view": str(result["view"]), + "expected_frames": int( + rows_by_episode[int(result["episode_index"])]["frames"] + ), + "actual_frames": int(result["video_frames"]), + } + for result in video_results + if "video_frames" in result + and int(result["video_frames"]) + != int(rows_by_episode[int(result["episode_index"])]["frames"]) + ] + video_near_static_keys = { + key + for key, video in video_by_key.items() + if not video.get("video_error") + and float(video.get("video_adjacent_mad_mean", math.inf)) < 1.0 + and float(video.get("video_first_last_mad", math.inf)) < 2.0 + } + video_error_keys = { + key + for key, video in video_by_key.items() + if video.get("video_error") + } + relevant_keys = candidate_keys | video_near_static_keys | video_error_keys + relevant_episode_indices = sorted( + {episode for episode, _ in relevant_keys} + ) + arm_by_episode: dict[int, dict[str, float]] = {} + for position, episode_index in enumerate( + relevant_episode_indices, + start=1, + ): + arm_by_episode[episode_index] = _arm_motion(root, episode_index) + if position % 250 == 0 or position == len(relevant_episode_indices): + print( + f"Checked relevant arm motion: " + f"{position}/{len(relevant_episode_indices)}", + flush=True, + ) + + classifications: dict[str, list[dict[str, object]]] = { + "source_video_near_static_with_moving_arm": [], + "track_near_static_with_moving_video_and_arm": [], + "stationary_arm_or_low_motion": [], + "video_check_error": [], + } + for episode_index, view in sorted(relevant_keys): + row = rows[episode_index] + video = video_by_key.get((episode_index, view), {}) + arm_name = "left_arm" if view == "left_wrist" else "right_arm" + arm = arm_by_episode[episode_index] + max_joint_range = float(arm[f"{arm_name}_max_joint_range"]) + video_error = video.get("video_error") + key = (episode_index, view) + video_near_static = key in video_near_static_keys + track_near_static = key in candidate_keys + arm_moving = max_joint_range >= 0.1 + record = { + "episode_index": episode_index, + "view": view, + "frames": int(row["frames"]), + "task": row["task"], + "background_span_median_px": row[ + f"{view}_background_span_median_px" + ], + "hand_span_median_px": row[f"{view}_hand_span_median_px"], + "video_adjacent_mad_mean": video.get( + "video_adjacent_mad_mean" + ), + "video_first_last_mad": video.get("video_first_last_mad"), + "video_first_frame_hash": video.get("video_first_frame_hash"), + "video_sample_hash": video.get("video_sample_hash"), + "arm_max_joint_range": max_joint_range, + "track_near_static": track_near_static, + "video_near_static": video_near_static, + } + if video_error: + record["video_error"] = video_error + classifications["video_check_error"].append(record) + elif video_near_static and arm_moving: + classifications[ + "source_video_near_static_with_moving_arm" + ].append(record) + elif track_near_static and not video_near_static and arm_moving: + classifications[ + "track_near_static_with_moving_video_and_arm" + ].append(record) + else: + classifications["stationary_arm_or_low_motion"].append(record) + + seam_anomalies: list[dict[str, object]] = [] + for row in rows: + for view in VIEW_SLICES: + ratio = row.get(f"{view}_seam_jump_ratio") + seam_p95 = row.get(f"{view}_seam_jump_p95_px") + if ( + isinstance(ratio, float) + and isinstance(seam_p95, float) + and np.isfinite(ratio) + and np.isfinite(seam_p95) + and ratio > 5.0 + and seam_p95 > 5.0 + ): + seam_anomalies.append( + { + "episode_index": int(row["episode_index"]), + "view": view, + "seam_jump_p95_px": seam_p95, + "regular_jump_p95_px": row[ + f"{view}_regular_jump_p95_px" + ], + "ratio": ratio, + } + ) + + metric_distributions: dict[str, dict[str, float]] = {} + for view in VIEW_SLICES: + for suffix in ("visibility", "span_median_px", "seam_jump_ratio"): + key = f"{view}_{suffix}" + values = [ + float(row[key]) + for row in rows + if isinstance(row.get(key), float) + and np.isfinite(float(row[key])) + ] + if values: + metric_distributions[key] = _quantiles(values) + for view in WRIST_GROUPS: + for group in ("background", "hand"): + key = f"{view}_{group}_span_median_px" + metric_distributions[key] = _quantiles( + [float(row[key]) for row in rows] + ) + + csv_path = output_dir / "episode_metrics.csv" + fieldnames = sorted({key for row in rows for key in row}) + with csv_path.open("w", newline="") as file: + writer = csv.DictWriter(file, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(rows) + + video_csv_path = output_dir / "wrist_video_metrics.csv" + if video_results: + video_fieldnames = sorted( + {key for result in video_results for key in result} + ) + with video_csv_path.open("w", newline="") as file: + writer = csv.DictWriter(file, fieldnames=video_fieldnames) + writer.writeheader() + writer.writerows( + sorted( + video_results, + key=lambda result: ( + int(result["episode_index"]), + str(result["view"]), + ), + ) + ) + + frozen_records = classifications[ + "source_video_near_static_with_moving_arm" + ] + frozen_episode_indices = sorted( + {int(record["episode_index"]) for record in frozen_records} + ) + frozen_views_by_episode: dict[int, set[str]] = {} + for record in frozen_records: + frozen_views_by_episode.setdefault( + int(record["episode_index"]), + set(), + ).add(str(record["view"])) + frozen_episode_breakdown = Counter( + "both" if len(views) == 2 else next(iter(views)) + for views in frozen_views_by_episode.values() + ) + frozen_frame_count = sum( + int(rows_by_episode[episode_index]["frames"]) + for episode_index in frozen_episode_indices + ) + first_frame_hash_counts = Counter( + str(record["video_first_frame_hash"]) + for record in frozen_records + if record.get("video_first_frame_hash") + ) + repeated_frozen_frames = [ + {"first_frame_hash": frame_hash, "view_count": count} + for frame_hash, count in first_frame_hash_counts.most_common() + if count > 1 + ] + all_classification_records = [ + {"classification": name, **record} + for name, records in classifications.items() + for record in records + ] + classification_path = output_dir / "wrist_static_classifications.json" + classification_path.write_text( + json.dumps(all_classification_records, indent=2) + "\n" + ) + blacklist_path = output_dir / "frozen_wrist_episode_indices.json" + blacklist_path.write_text( + json.dumps(frozen_episode_indices, indent=2) + "\n" + ) + + summary = { + "dataset_root": str(root), + "total_episodes": len(rows), + "total_frames": int(sum(int(row["frames"]) for row in rows)), + "track_integrity_error_count": len(errors), + "track_integrity_errors": errors[:100], + "static_background_threshold_px": static_threshold, + "near_static_wrist_view_count": len(candidates), + "near_static_episode_count": len(candidate_episode_indices), + "wrist_video_check_count": len(video_results), + "wrist_video_check_error_count": len(video_error_keys), + "wrist_video_frame_count_mismatch_count": len( + video_frame_count_mismatches + ), + "wrist_video_frame_count_mismatches": ( + video_frame_count_mismatches[:100] + ), + "source_video_near_static_view_count": len( + video_near_static_keys + ), + "source_video_near_static_episode_count": len( + {episode for episode, _ in video_near_static_keys} + ), + "classification_counts": { + name: len(records) + for name, records in classifications.items() + }, + "classification_episode_counts": { + name: len( + { + int(record["episode_index"]) + for record in records + } + ) + for name, records in classifications.items() + }, + "classification_examples": { + name: records[:30] + for name, records in classifications.items() + if records + }, + "focus_episode_5463": [ + record + for record in all_classification_records + if int(record["episode_index"]) == 5463 + ], + "repeated_frozen_first_frame_groups": repeated_frozen_frames[:30], + "unique_frozen_first_frames": len(first_frame_hash_counts), + "frozen_view_breakdown": dict( + Counter(str(record["view"]) for record in frozen_records) + ), + "frozen_episode_breakdown": dict(frozen_episode_breakdown), + "frozen_episode_frame_count": frozen_frame_count, + "frozen_episode_frame_fraction": ( + frozen_frame_count + / sum(int(row["frames"]) for row in rows) + ), + "seam_anomaly_count": len(seam_anomalies), + "seam_anomalies": seam_anomalies[:100], + "metric_distributions": metric_distributions, + "episode_metrics_csv": str(csv_path), + "wrist_video_metrics_csv": ( + str(video_csv_path) if video_results else None + ), + "wrist_static_classifications_json": str(classification_path), + "frozen_wrist_episode_indices_json": str(blacklist_path), + } + summary_path = output_dir / "summary.json" + summary_path.write_text(json.dumps(summary, indent=2) + "\n") + print(f"Wrote {summary_path}") + print(f"Wrote {csv_path}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/data/build_trex_mini_force.py b/scripts/data/build_trex_mini_force.py new file mode 100644 index 0000000000000000000000000000000000000000..6776f94209574dea24fcf6138284ea889e9bb3d9 --- /dev/null +++ b/scripts/data/build_trex_mini_force.py @@ -0,0 +1,214 @@ +#!/usr/bin/env python3 +"""Build one train-ready T-Rex Track-Force demo with fresh statistics.""" + +from __future__ import annotations + +import argparse +import copy +import hashlib +import json +import os +import shutil +import sys +from pathlib import Path +from typing import Sequence + +_ROOT = Path(__file__).resolve().parents[2] +_DATA_SCRIPTS = _ROOT / "scripts" / "data" +if str(_DATA_SCRIPTS) not in sys.path: + sys.path.insert(0, str(_DATA_SCRIPTS)) + +import build_trex_track_force_v2 as force_builder # noqa: E402 +import check_trex_dataset_ready as ready_check # noqa: E402 +import rebuild_trex_dataset_variants as variants # noqa: E402 + + +def _read_json(path: Path) -> dict: + return json.loads(path.read_text()) + + +def _read_jsonl(path: Path) -> list[dict]: + return [ + json.loads(line) + for line in path.read_text().splitlines() + if line.strip() + ] + + +def _sha256_or_empty(path: Path) -> str: + if path.is_file(): + return variants._sha256(path) + return hashlib.sha256(b"[]\n").hexdigest() + + +def build_mini_dataset( + *, + source_root: Path, + output_root: Path, + source_episode: int, + workers: int, +) -> dict: + source_root = source_root.expanduser().resolve() + output_root = output_root.expanduser().resolve() + if source_root == output_root: + raise ValueError("source and output dataset roots must differ") + + ready_check.check_dataset(source_root, require_force=True) + if output_root.exists(): + result = ready_check.check_dataset(output_root, require_force=True) + variant = _read_json(output_root / "meta" / "dataset_variant.json") + if int(variant.get("source_episode", -1)) != source_episode: + raise ValueError( + f"{output_root} already uses source episode " + f"{variant.get('source_episode')}, not {source_episode}" + ) + return result + + source_info = _read_json(source_root / "meta" / "info.json") + source_modality = _read_json(source_root / "meta" / "modality.json") + source_embodiment = _read_json(source_root / "meta" / "embodiment.json") + source_episodes = _read_jsonl(source_root / "meta" / "episodes.jsonl") + source_tasks = _read_jsonl(source_root / "meta" / "tasks.jsonl") + selected = [ + episode + for episode in source_episodes + if int(episode["episode_index"]) == source_episode + ] + if len(selected) != 1: + raise ValueError( + f"source episode {source_episode} was not found exactly once" + ) + registry, tasks = variants._build_registry( + selected, + source_tasks, + excluded=set(), + source_limit=None, + ) + if len(registry) != 1 or registry[0].episode_index != 0: + raise AssertionError("mini dataset must contain dense episode 0") + + stage_root = output_root.parent / f".{output_root.name}.staging" + if stage_root.exists(): + shutil.rmtree(stage_root) + video_keys = variants._video_keys(source_info) + blacklist = source_root / "meta" / "excluded_source_episode_indices.json" + variants._write_variant_metadata_skeleton( + variant_root=stage_root, + final_root=output_root, + source_info=source_info, + source_modality=source_modality, + source_embodiment=source_embodiment, + registry=registry, + tasks=tasks, + video_keys=video_keys, + excluded_source_indices=[], + force=True, + source_dataset=source_root, + blacklist_sha256=_sha256_or_empty(blacklist), + ) + variant_path = stage_root / "meta" / "dataset_variant.json" + variant = _read_json(variant_path) + variant["source_episode"] = source_episode + variants._atomic_write_json(variant_path, variant) + + variants._rewrite_parquets( + source_root=source_root, + destination_root=stage_root, + registry=registry, + workers=workers, + force=True, + ) + variants._link_videos( + source_root=source_root, + destination_root=stage_root, + registry=registry, + video_keys=video_keys, + source_uses_new_indices=False, + workers=workers, + ) + track_hashes = variants._rewrite_tracks( + source_root=source_root, + destination_root=stage_root, + registry=registry, + workers=workers, + ) + + source_manifest = _read_json( + source_root / "meta" / "trex_track_force_manifest.json" + ) + source_entries = { + int(index): copy.deepcopy(entry) + for index, entry in source_manifest["episodes"].items() + } + variants._build_force_manifest( + variant_root=stage_root, + final_root=output_root, + registry=registry, + source_entries=source_entries, + track_hashes=track_hashes, + ) + + variants._atomic_write_json( + stage_root / "meta" / "stats.json", + variants._compute_base_stats(stage_root, registry), + ) + metadata_result = force_builder.update_metadata( + stage_root, + assume_all_converted=True, + ) + manifest_path = stage_root / "meta" / "trex_track_force_manifest.json" + manifest = _read_json(manifest_path) + manifest["metadata"] = metadata_result + manifest["updated_at"] = variants._utc_now() + variants._atomic_write_json(manifest_path, manifest) + force_builder.validate_metadata(stage_root) + for backup in (stage_root / "meta").glob("*.trex_track_force.bak"): + backup.unlink() + + variants._validate_variant( + root=stage_root, + expected_registry=registry, + video_keys=video_keys, + force=True, + ) + os.replace(stage_root, output_root) + variants._validate_variant( + root=output_root, + expected_registry=registry, + video_keys=video_keys, + force=True, + ) + return ready_check.check_dataset(output_root, require_force=True) + + +def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--source-root", + type=Path, + default=_ROOT / "data" / "trex_full_force", + ) + parser.add_argument( + "--output-root", + type=Path, + default=_ROOT / "data" / "trex_mini_force", + ) + parser.add_argument("--source-episode", type=int, default=49) + parser.add_argument("--workers", type=int, default=4) + return parser.parse_args(argv) + + +def main(argv: Sequence[str] | None = None) -> int: + args = _parse_args(argv) + result = build_mini_dataset( + source_root=args.source_root, + output_root=args.output_root, + source_episode=args.source_episode, + workers=max(1, int(args.workers)), + ) + print(json.dumps(result)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/data/build_trex_track_force_v2.py b/scripts/data/build_trex_track_force_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..e0eb1e5db6c3af2d1dbf5fb3f78ffe62d4ab259b --- /dev/null +++ b/scripts/data/build_trex_track_force_v2.py @@ -0,0 +1,1936 @@ +"""Build the T-Rex track/EEF extension in a LeRobot v2 dataset. + +This builder preserves every existing parquet column and adds: + +* model-facing ``observation.track_xy`` / ``observation.track_visibility``; +* view-preserving ``observation.tracks.{head_left,left_wrist,right_wrist}``, + each frame stored as fixed-size ``[x, y, visibility]`` values; +* ``observation.state_eef62`` and ``action.eef62_absolute`` using T-Rex FK and + the canonical ``translation + rotation-6D + hand`` representation. + +Writes are resumable and atomic. Existing valid episode outputs are skipped, +and the original parquet/metadata files receive one-time ``.trex_track_force.bak`` +backups before their first replacement. +""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.util +import json +import os +import shutil +import sys +import tempfile +from datetime import datetime, timezone +from functools import lru_cache +from pathlib import Path +from typing import Callable, Iterable, Sequence + +import numpy as np + +_DATA_SCRIPT_DIR = Path(__file__).resolve().parent +_SCRIPT_DIR = _DATA_SCRIPT_DIR.parent +_DREAMZERO_ROOT = _SCRIPT_DIR.parent +if str(_SCRIPT_DIR) not in sys.path: + sys.path.insert(0, str(_SCRIPT_DIR)) + +from trex_track.layout import ( # noqa: E402 + NUM_COMBINED_POINTS, + POINT_SLICES, + TRACK_LAYOUT_VERSION, + VIEW_ORDER, + VIEW_POINT_COUNTS, + VIEW_SLICES, + identity_metadata, + layout_metadata, +) + +SCHEMA_VERSION = "trex_track_force_v2.3" +BACKUP_SUFFIX = ".trex_track_force.bak" +DEFAULT_DATASET_ROOT = _DREAMZERO_ROOT / "data" / "trex_small_force" +DEFAULT_TREX_ROOT = Path("/scratch1/home/zhicao/T-Rex") +PARQUET_SCHEMA_METADATA_KEY = b"trex_track_force_schema_version" + +TARGET_RATE_HZ = 20.0 +ACTION_CHUNK_STEPS = 16 +ACTION_CHUNK_DURATION_SECONDS = ACTION_CHUNK_STEPS / TARGET_RATE_HZ +ACTION_CHUNK_TIMESTAMP_SPAN_SECONDS = (ACTION_CHUNK_STEPS - 1) / TARGET_RATE_HZ +AUTOREGRESSIVE_BLOCKS = 4 +VIDEO_FRAMES_PER_BLOCK = 8 +TRAINING_VIDEO_FRAMES = 1 + AUTOREGRESSIVE_BLOCKS * VIDEO_FRAMES_PER_BLOCK +FORCE_COLUMN = "observation.tactile_force" +FORCE_FLAT_DIM = 60 +FORCE_SENSOR_COUNT = 10 +FORCE_SENSOR_DIM = 6 +FORCE_HISTORY_FRAMES = 16 +RELATIVE_ACTION_STATS_FILENAME = "relative_stats_dreamzero.json" + +STATE_EEF_COLUMN = "observation.state_eef62" +ACTION_EEF_COLUMN = "action.eef62_absolute" +TRACK_XY_COLUMN = "observation.track_xy" +TRACK_VISIBILITY_COLUMN = "observation.track_visibility" +TRACK_COLUMNS = { + "head_left": "observation.tracks.head_left", + "left_wrist": "observation.tracks.left_wrist", + "right_wrist": "observation.tracks.right_wrist", +} +NEW_COLUMNS = ( + TRACK_XY_COLUMN, + TRACK_VISIBILITY_COLUMN, + *TRACK_COLUMNS.values(), + STATE_EEF_COLUMN, + ACTION_EEF_COLUMN, +) + +LEFT_EEF = slice(0, 9) +LEFT_HAND_EEF = slice(9, 31) +RIGHT_EEF = slice(31, 40) +RIGHT_HAND_EEF = slice(40, 62) + +EefConverter = Callable[[np.ndarray], np.ndarray] + + +class DatasetSchemaError(RuntimeError): + """Raised when an episode cannot satisfy the track/EEF schema.""" + + +def sample_timestamps_nearest( + source_timestamps: np.ndarray | Sequence[float], + target_rate_hz: float = TARGET_RATE_HZ, + *, + anchor_index: int | None = None, + anchor_timestamp: float | None = None, + offsets: Sequence[int] | np.ndarray | None = None, + alignment_tolerance: float = 1e-6, +) -> dict[str, object]: + """Deterministically align a target-rate grid to nearest source frames. + + Ties choose the earlier source frame. Non-padding source indices must be + unique, and every non-padding alignment error is bounded by half the + median source period plus ``alignment_tolerance``. Queries outside the + source interval clamp to an endpoint and are explicitly marked in + ``padding_mask``. + """ + + source = np.asarray(source_timestamps, dtype=np.float64) + if source.ndim != 1 or source.size < 2: + raise DatasetSchemaError( + f"source_timestamps must be a 1D array with >=2 values, got {source.shape}" + ) + if not np.isfinite(source).all(): + raise DatasetSchemaError("source_timestamps contain NaN/Inf") + source_deltas = np.diff(source) + if not np.all(source_deltas > 0.0): + raise DatasetSchemaError("source_timestamps must be strictly increasing") + source_period = float(np.median(source_deltas)) + if not np.isfinite(source_period) or source_period <= 0.0: + raise DatasetSchemaError("could not infer a positive source period") + max_source_period = float(source_deltas.max()) + if max_source_period > 1.5 * source_period: + raise DatasetSchemaError( + "source timestamps contain a dropped-frame gap: " + f"max={max_source_period:.9f}s median={source_period:.9f}s" + ) + target_rate = float(target_rate_hz) + if not np.isfinite(target_rate) or target_rate <= 0.0: + raise DatasetSchemaError("target_rate_hz must be finite and positive") + tolerance = float(alignment_tolerance) + if not np.isfinite(tolerance) or tolerance < 0.0: + raise DatasetSchemaError("alignment_tolerance must be finite and non-negative") + + if anchor_index is not None and anchor_timestamp is not None: + raise DatasetSchemaError("set only anchor_index or anchor_timestamp") + if anchor_timestamp is None: + index = 0 if anchor_index is None else int(anchor_index) + if index < 0 or index >= source.size: + raise DatasetSchemaError( + f"anchor_index {index} is outside [0,{source.size})" + ) + anchor = float(source[index]) + else: + anchor = float(anchor_timestamp) + if not np.isfinite(anchor): + raise DatasetSchemaError("anchor_timestamp must be finite") + + if offsets is None: + last_offset = int( + np.floor((float(source[-1]) - anchor) * target_rate + tolerance * target_rate) + ) + if last_offset < 0: + raise DatasetSchemaError("anchor is after the source timestamp interval") + offset_array = np.arange(last_offset + 1, dtype=np.int64) + else: + raw_offsets = np.asarray(offsets) + if raw_offsets.ndim != 1 or raw_offsets.size == 0: + raise DatasetSchemaError("offsets must be a non-empty 1D sequence") + offset_array = raw_offsets.astype(np.int64) + if not np.array_equal(raw_offsets, offset_array): + raise DatasetSchemaError("offsets must contain integer target steps") + if not np.all(np.diff(offset_array) > 0): + raise DatasetSchemaError("offsets must be strictly increasing and unique") + + target = anchor + offset_array.astype(np.float64) / target_rate + if not np.all(np.diff(target) > 0.0): + raise DatasetSchemaError("target timestamps must be strictly increasing") + padding = (target < source[0] - tolerance) | (target > source[-1] + tolerance) + + insertion = np.searchsorted(source, target, side="left") + lower = np.clip(insertion - 1, 0, source.size - 1) + upper = np.clip(insertion, 0, source.size - 1) + lower_error = np.abs(target - source[lower]) + upper_error = np.abs(source[upper] - target) + # Differences within tolerance count as midpoint ties and choose earlier. + choose_upper = upper_error < (lower_error - tolerance) + indices = np.where(choose_upper, upper, lower).astype(np.int64) + indices[target < source[0]] = 0 + indices[target > source[-1]] = source.size - 1 + alignment_errors = np.abs(source[indices] - target) + + non_padding = ~padding + # Real MP4 timestamps have small per-frame jitter. Nearest-neighbour error + # is bounded by half the local gap, not half the median source period. + max_allowed_error = max_source_period / 2.0 + tolerance + if non_padding.any() and np.any( + alignment_errors[non_padding] > max_allowed_error + ): + worst = float(alignment_errors[non_padding].max()) + raise DatasetSchemaError( + f"timestamp alignment error {worst:.9f}s exceeds " + f"source_period/2+tolerance={max_allowed_error:.9f}s" + ) + selected = indices[non_padding] + if np.unique(selected).size != selected.size: + raise DatasetSchemaError( + "nearest timestamp alignment selected duplicate non-padding source frames" + ) + + return { + "indices": indices, + "target_timestamps": target, + "offsets": offset_array, + "padding_mask": padding.astype(bool), + "alignment_errors": alignment_errors, + "source_period_seconds": source_period, + "max_source_period_seconds": max_source_period, + "source_rate_hz": 1.0 / source_period, + "target_rate_hz": target_rate, + "max_allowed_alignment_error_seconds": max_allowed_error, + } + + +def summarize_timestamp_sampling( + source_timestamps: np.ndarray | Sequence[float], + *, + target_rate_hz: float = TARGET_RATE_HZ, + action_chunk_steps: int = ACTION_CHUNK_STEPS, +) -> dict[str, object]: + """Return a JSON-safe per-episode 20 Hz coverage/chunk validation summary.""" + + source = np.asarray(source_timestamps, dtype=np.float64) + coverage = sample_timestamps_nearest( + source, + target_rate_hz=target_rate_hz, + anchor_index=0, + ) + chunk_steps = int(action_chunk_steps) + if chunk_steps <= 0: + raise DatasetSchemaError("action_chunk_steps must be positive") + chunk = sample_timestamps_nearest( + source, + target_rate_hz=target_rate_hz, + anchor_index=0, + offsets=np.arange(chunk_steps, dtype=np.int64), + ) + coverage_indices = np.asarray(coverage["indices"], dtype=np.int64) + coverage_targets = np.asarray(coverage["target_timestamps"], dtype=np.float64) + coverage_errors = np.asarray(coverage["alignment_errors"], dtype=np.float64) + coverage_padding = np.asarray(coverage["padding_mask"], dtype=bool) + chunk_padding = np.asarray(chunk["padding_mask"], dtype=bool) + target_rate = float(target_rate_hz) + chunk_duration = chunk_steps / target_rate + chunk_timestamp_span = (chunk_steps - 1) / target_rate + if chunk_steps == ACTION_CHUNK_STEPS and np.isclose( + target_rate, TARGET_RATE_HZ + ) and ( + not np.isclose(chunk_duration, ACTION_CHUNK_DURATION_SECONDS, atol=1e-12) + or not np.isclose( + chunk_timestamp_span, + ACTION_CHUNK_TIMESTAMP_SPAN_SECONDS, + atol=1e-12, + ) + ): + raise DatasetSchemaError("invalid 16-step/20 Hz action chunk definition") + + return { + "source_frame_count": int(source.size), + "source_start_timestamp": float(source[0]), + "source_end_timestamp": float(source[-1]), + "source_timestamp_span_seconds": float(source[-1] - source[0]), + "source_period_seconds": float(coverage["source_period_seconds"]), + "source_rate_hz": float(coverage["source_rate_hz"]), + "target_rate_hz": target_rate, + "target_sample_count": int(coverage_indices.size), + "target_start_timestamp": float(coverage_targets[0]), + "target_end_timestamp": float(coverage_targets[-1]), + "first_source_index": int(coverage_indices[0]), + "last_source_index": int(coverage_indices[-1]), + "max_alignment_error_seconds": float(coverage_errors.max(initial=0.0)), + "max_allowed_alignment_error_seconds": float( + coverage["max_allowed_alignment_error_seconds"] + ), + "coverage_padding_count": int(coverage_padding.sum()), + "action_chunk_steps": chunk_steps, + # A 16-step control horizon is the half-open interval [t, t+0.8s). + "action_chunk_duration_seconds": chunk_duration, + # The first/last sampled timestamps in that horizon are 15/20=0.75s apart. + "action_chunk_timestamp_span_seconds": chunk_timestamp_span, + "action_chunk_fully_covered": not bool(chunk_padding.any()), + "action_chunk_padding_count": int(chunk_padding.sum()), + "action_chunk_padding_mask": chunk_padding.tolist(), + "complete_action_chunks": int(coverage_indices.size // chunk_steps), + } + + +def _utc_now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _load_json(path: Path) -> dict: + if not path.is_file(): + raise FileNotFoundError(path) + with path.open("r", encoding="utf-8") as file: + value = json.load(file) + if not isinstance(value, dict): + raise DatasetSchemaError(f"expected JSON object in {path}") + return value + + +def _fsync_directory(path: Path) -> None: + try: + fd = os.open(path, os.O_RDONLY) + except OSError: + return + try: + os.fsync(fd) + finally: + os.close(fd) + + +def _atomic_write_json(path: Path, value: dict) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fd, tmp_name = tempfile.mkstemp( + prefix=f".{path.name}.", + suffix=".tmp", + dir=path.parent, + ) + try: + with os.fdopen(fd, "w", encoding="utf-8") as file: + json.dump(value, file, indent=2, sort_keys=False) + file.write("\n") + file.flush() + os.fsync(file.fileno()) + os.replace(tmp_name, path) + _fsync_directory(path.parent) + except BaseException: + try: + os.unlink(tmp_name) + except FileNotFoundError: + pass + raise + + +def backup_path(path: Path) -> Path: + return path.with_name(path.name + BACKUP_SUFFIX) + + +def _atomic_backup(path: Path) -> Path | None: + """Create a one-time atomic backup, never replacing an existing backup.""" + + if not path.exists(): + return None + destination = backup_path(path) + if destination.exists(): + return destination + fd, tmp_name = tempfile.mkstemp( + prefix=f".{destination.name}.", + suffix=".tmp", + dir=path.parent, + ) + os.close(fd) + try: + shutil.copy2(path, tmp_name) + with open(tmp_name, "rb") as file: + os.fsync(file.fileno()) + # A concurrent builder may have completed the backup while we copied. + if destination.exists(): + os.unlink(tmp_name) + return destination + os.replace(tmp_name, destination) + _fsync_directory(path.parent) + return destination + except BaseException: + try: + os.unlink(tmp_name) + except FileNotFoundError: + pass + raise + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as file: + for block in iter(lambda: file.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _import_pyarrow(): + try: + import pyarrow as pa + import pyarrow.parquet as pq + except ImportError as exc: + raise RuntimeError("pyarrow is required to build LeRobot parquet files") from exc + return pa, pq + + +def _fixed_size_array(values: np.ndarray): + """Convert ``[rows, *shape]`` to nested Arrow FixedSizeListArray.""" + + pa, _ = _import_pyarrow() + array = np.asarray(values, dtype=np.float32) + if array.ndim < 2: + raise ValueError(f"fixed-size feature must have at least 2 dims, got {array.shape}") + result = pa.array(array.reshape(-1), type=pa.float32()) + for size in reversed(array.shape[1:]): + result = pa.FixedSizeListArray.from_arrays(result, int(size)) + if len(result) != array.shape[0]: + raise AssertionError(f"Arrow rows {len(result)} != numpy rows {array.shape[0]}") + return result + + +def _set_or_append_column(table, name: str, values: np.ndarray): + column = _fixed_size_array(values) + index = table.schema.get_field_index(name) + if index >= 0: + return table.set_column(index, name, column) + return table.append_column(name, column) + + +def _column_to_numpy(table, name: str, dtype=np.float32) -> np.ndarray: + if name not in table.column_names: + raise DatasetSchemaError(f"missing parquet column {name!r}") + values = table[name].combine_chunks().to_pylist() + try: + return np.asarray(values, dtype=dtype) + except (TypeError, ValueError) as exc: + raise DatasetSchemaError(f"column {name!r} is not a dense numeric array") from exc + + +def validate_tactile_force(table) -> dict[str, object]: + """Validate the existing force-only source without creating VQ-code columns.""" + + force = _column_to_numpy(table, FORCE_COLUMN, dtype=np.float32) + expected = (int(table.num_rows), FORCE_FLAT_DIM) + if force.shape != expected: + raise DatasetSchemaError( + f"{FORCE_COLUMN} must have shape {expected}, got {force.shape}" + ) + if not np.isfinite(force).all(): + raise DatasetSchemaError(f"{FORCE_COLUMN} contains NaN/Inf") + return { + "source_column": FORCE_COLUMN, + "stored_shape": [FORCE_FLAT_DIM], + "reshape": [FORCE_SENSOR_COUNT, FORCE_SENSOR_DIM], + "history_frames": FORCE_HISTORY_FRAMES, + "history_encoding": "online_model_encoder", + "vq_codes_on_disk": False, + "finite": True, + } + + +def _is_fixed_shape(field_type, shape: Sequence[int]) -> bool: + pa, _ = _import_pyarrow() + current = field_type + for size in shape: + if not pa.types.is_fixed_size_list(current) or current.list_size != int(size): + return False + current = current.value_type + return pa.types.is_float32(current) + + +def _atomic_write_parquet(table, path: Path, *, expected_rows: int) -> None: + _, pq = _import_pyarrow() + path.parent.mkdir(parents=True, exist_ok=True) + fd, tmp_name = tempfile.mkstemp( + prefix=f".{path.name}.", + suffix=".tmp.parquet", + dir=path.parent, + ) + os.close(fd) + tmp_path = Path(tmp_name) + try: + pq.write_table(table, tmp_path, compression="zstd") + with tmp_path.open("rb") as file: + os.fsync(file.fileno()) + validate_episode_parquet( + tmp_path, + expected_frames=expected_rows, + verify_source_fk=False, + ) + _atomic_backup(path) + os.replace(tmp_path, path) + _fsync_directory(path.parent) + except BaseException: + tmp_path.unlink(missing_ok=True) + raise + + +def _scalar_text(value: np.ndarray) -> str: + scalar = np.asarray(value) + if scalar.shape != (): + raise DatasetSchemaError(f"expected scalar string, got shape {scalar.shape}") + return str(scalar.item()) + + +def _validate_unit_interval(name: str, values: np.ndarray) -> None: + array = np.asarray(values) + if not np.isfinite(array).all(): + raise DatasetSchemaError(f"{name} contains NaN/Inf") + if array.size and (float(array.min()) < 0.0 or float(array.max()) > 1.0): + raise DatasetSchemaError( + f"{name} must be in [0,1], got [{array.min()}, {array.max()}]" + ) + + +def load_track_payload( + path: Path, + *, + expected_frames: int | None = None, + episode_index: int | None = None, +) -> dict[str, np.ndarray]: + """Load and strictly validate a canonical extraction NPZ.""" + + if not path.is_file(): + raise FileNotFoundError(path) + with np.load(path, allow_pickle=False) as archive: + forbidden = [name for name in archive.files if name.startswith("images_")] + if forbidden: + raise DatasetSchemaError( + f"{path} embeds full RGB arrays ({forbidden}); regenerate with the new extractor" + ) + payload = {name: np.asarray(archive[name]).copy() for name in archive.files} + + required = { + "tracks", + "vis", + "tracks_head_left", + "tracks_left_wrist", + "tracks_right_wrist", + "vis_head_left", + "vis_left_wrist", + "vis_right_wrist", + "episode_index", + "num_steps", + "point_slices", + "track_layout_version", + "point_view_ids", + "point_hand_ids", + "point_role_ids", + "point_local_ids", + "point_global_ids", + "point_names", + } + missing = sorted(required.difference(payload)) + if missing: + raise DatasetSchemaError(f"{path} is missing keys: {missing}") + + tracks = np.asarray(payload["tracks"], dtype=np.float32) + visibility = np.asarray(payload["vis"], dtype=np.float32) + if tracks.ndim != 3 or tracks.shape[1:] != (NUM_COMBINED_POINTS, 2): + raise DatasetSchemaError( + f"{path}: tracks must be (T,{NUM_COMBINED_POINTS},2), got {tracks.shape}" + ) + if visibility.shape != tracks.shape[:2]: + raise DatasetSchemaError( + f"{path}: visibility {visibility.shape} != {tracks.shape[:2]}" + ) + num_frames = int(tracks.shape[0]) + if int(np.asarray(payload["num_steps"]).item()) != num_frames: + raise DatasetSchemaError(f"{path}: num_steps does not match tracks") + if expected_frames is not None and num_frames != int(expected_frames): + raise DatasetSchemaError( + f"{path}: {num_frames} track frames != {expected_frames} parquet frames" + ) + stored_episode = int(np.asarray(payload["episode_index"]).item()) + if episode_index is not None and stored_episode != int(episode_index): + raise DatasetSchemaError( + f"{path}: episode_index={stored_episode}, expected {episode_index}" + ) + if _scalar_text(payload["track_layout_version"]) != TRACK_LAYOUT_VERSION: + raise DatasetSchemaError(f"{path}: unsupported track layout version") + if not np.array_equal( + np.asarray(payload["point_slices"], dtype=np.int32), + np.asarray(POINT_SLICES, dtype=np.int32), + ): + raise DatasetSchemaError(f"{path}: point_slices do not match canonical layout") + + expected_ids = identity_metadata() + identity_keys = { + "point_view_ids": "view_ids", + "point_hand_ids": "hand_ids", + "point_role_ids": "role_ids", + "point_local_ids": "local_ids", + "point_global_ids": "global_ids", + } + for stored_key, expected_key in identity_keys.items(): + if not np.array_equal( + np.asarray(payload[stored_key], dtype=np.int64), + np.asarray(expected_ids[expected_key], dtype=np.int64), + ): + raise DatasetSchemaError(f"{path}: unstable identity metadata in {stored_key}") + if not np.array_equal( + np.asarray(payload["point_names"]).astype(str), + np.asarray(expected_ids["point_names"]).astype(str), + ): + raise DatasetSchemaError(f"{path}: unstable identity metadata in point_names") + + _validate_unit_interval("tracks", tracks) + _validate_unit_interval("visibility", visibility) + if not np.all((visibility == 0.0) | (visibility == 1.0)): + raise DatasetSchemaError(f"{path}: visibility must be binary") + + view_tracks: list[np.ndarray] = [] + view_visibility: list[np.ndarray] = [] + for view in VIEW_ORDER: + count = VIEW_POINT_COUNTS[view] + track_key = f"tracks_{view}" + vis_key = f"vis_{view}" + track = np.asarray(payload[track_key], dtype=np.float32) + vis = np.asarray(payload[vis_key], dtype=np.float32) + if track.shape != (num_frames, count, 2): + raise DatasetSchemaError(f"{path}: {track_key} has shape {track.shape}") + if vis.shape != (num_frames, count): + raise DatasetSchemaError(f"{path}: {vis_key} has shape {vis.shape}") + _validate_unit_interval(track_key, track) + _validate_unit_interval(vis_key, vis) + view_tracks.append(track) + view_visibility.append(vis) + if not np.array_equal(np.concatenate(view_tracks, axis=1), tracks): + raise DatasetSchemaError(f"{path}: combined tracks differ from per-view tracks") + if not np.array_equal(np.concatenate(view_visibility, axis=1), visibility): + raise DatasetSchemaError(f"{path}: combined visibility differs from per-view visibility") + return payload + + +def track_features_from_payload(payload: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + features: dict[str, np.ndarray] = { + TRACK_XY_COLUMN: np.asarray(payload["tracks"], dtype=np.float32), + TRACK_VISIBILITY_COLUMN: np.asarray(payload["vis"], dtype=np.float32), + } + for view, column in TRACK_COLUMNS.items(): + xy = np.asarray(payload[f"tracks_{view}"], dtype=np.float32) + vis = np.asarray(payload[f"vis_{view}"], dtype=np.float32)[..., None] + features[column] = np.concatenate([xy, vis], axis=-1).astype(np.float32) + return features + + +@lru_cache(maxsize=1) +def _load_lerobot_common(): + path = DEFAULT_TREX_ROOT / "utils" / "lerobot_common.py" + if not path.is_file(): + raise FileNotFoundError(f"T-Rex pose semantics module not found: {path}") + spec = importlib.util.spec_from_file_location("_trex_lerobot_common_schema", path) + if spec is None or spec.loader is None: + raise ImportError(f"cannot load {path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + for name in ("pose_matrix_to_9d", "get_rot_mat"): + if not hasattr(module, name): + raise ImportError(f"{path} does not define {name}") + return module + + +def _validate_transform(matrix: np.ndarray, *, label: str) -> None: + transform = np.asarray(matrix, dtype=np.float64) + if transform.shape != (4, 4) or not np.isfinite(transform).all(): + raise DatasetSchemaError(f"{label}: FK returned an invalid transform") + if not np.allclose(transform[3], [0.0, 0.0, 0.0, 1.0], atol=1e-8): + raise DatasetSchemaError(f"{label}: FK transform has an invalid homogeneous row") + rotation = transform[:3, :3] + if not np.allclose(rotation.T @ rotation, np.eye(3), atol=1e-5): + raise DatasetSchemaError(f"{label}: FK rotation is not orthonormal") + if not np.isclose(np.linalg.det(rotation), 1.0, atol=1e-5): + raise DatasetSchemaError(f"{label}: FK rotation determinant is not +1") + + +def joint58_to_eef62_batch(joints: np.ndarray) -> np.ndarray: + """Convert joint-space state/action rows to absolute 62-D EEF semantics. + + This function deliberately has no fallback. Missing robot assets, invalid + joints, or unreliable FK raise an exception rather than fabricating poses. + """ + + source = np.asarray(joints, dtype=np.float64) + if source.ndim != 2 or source.shape[1] != 58: + raise DatasetSchemaError(f"FK expects (T,58), got {source.shape}") + if not np.isfinite(source).all(): + raise DatasetSchemaError("FK input contains NaN/Inf") + + try: + from trex_track.trex_fk import ( + frame_pose_matrix, + get_bimanual_robot, + state_to_components, + ) + + robot, assemble_qpos, _ = get_bimanual_robot() + common = _load_lerobot_common() + output = np.empty((source.shape[0], 62), dtype=np.float32) + for index, row in enumerate(source): + components = state_to_components(row) + qpos = assemble_qpos( + { + "left_arm": components["left_arm"], + "right_arm": components["right_arm"], + } + ) + left_pose = frame_pose_matrix(robot, qpos, "L_ee") + right_pose = frame_pose_matrix(robot, qpos, "R_ee") + _validate_transform(left_pose, label=f"row {index} left") + _validate_transform(right_pose, label=f"row {index} right") + left_9d = common.pose_matrix_to_9d(left_pose[None])[0] + right_9d = common.pose_matrix_to_9d(right_pose[None])[0] + output[index] = np.concatenate( + [ + left_9d, + components["left_hand"], + right_9d, + components["right_hand"], + ] + ) + except DatasetSchemaError: + raise + except Exception as exc: + raise DatasetSchemaError( + "reliable T-Rex FK failed; refusing to synthesize EEF values" + ) from exc + validate_eef62(output, source_joint58=source, label="FK output") + return output + + +def validate_eef62( + values: np.ndarray, + *, + source_joint58: np.ndarray | None = None, + label: str, +) -> None: + """Validate shape, hand preservation, rotations, and pose/rot6d roundtrip.""" + + array = np.asarray(values, dtype=np.float64) + if array.ndim != 2 or array.shape[1] != 62: + raise DatasetSchemaError(f"{label}: expected (T,62), got {array.shape}") + if not np.isfinite(array).all(): + raise DatasetSchemaError(f"{label}: contains NaN/Inf") + common = _load_lerobot_common() + + for side, arm_slice in (("left", LEFT_EEF), ("right", RIGHT_EEF)): + arm = array[:, arm_slice] + for row_index, pose9 in enumerate(arm): + rotation = np.asarray(common.get_rot_mat(pose9[3:9]), dtype=np.float64) + if not np.allclose(rotation.T @ rotation, np.eye(3), atol=2e-5): + raise DatasetSchemaError( + f"{label}: {side} row {row_index} rot6d is not orthonormal" + ) + if not np.isclose(np.linalg.det(rotation), 1.0, atol=2e-5): + raise DatasetSchemaError( + f"{label}: {side} row {row_index} rotation determinant is not +1" + ) + transform = np.eye(4, dtype=np.float64) + transform[:3, :3] = rotation + transform[:3, 3] = pose9[:3] + roundtrip = common.pose_matrix_to_9d(transform[None])[0] + if not np.allclose(roundtrip, pose9, atol=2e-5, rtol=1e-5): + raise DatasetSchemaError( + f"{label}: {side} row {row_index} pose/rot6d roundtrip failed" + ) + + if source_joint58 is not None: + source = np.asarray(source_joint58, dtype=np.float64) + if source.shape != (array.shape[0], 58): + raise DatasetSchemaError( + f"{label}: source shape {source.shape} does not match EEF rows" + ) + if not np.allclose(array[:, LEFT_HAND_EEF], source[:, 7:29], atol=1e-6): + raise DatasetSchemaError(f"{label}: left hand values were not preserved") + if not np.allclose(array[:, RIGHT_HAND_EEF], source[:, 36:58], atol=1e-6): + raise DatasetSchemaError(f"{label}: right hand values were not preserved") + + +def convert_eef_columns( + state58: np.ndarray, + action58: np.ndarray, + *, + converter: EefConverter | None = None, +) -> tuple[np.ndarray, np.ndarray]: + convert = converter or joint58_to_eef62_batch + state = np.asarray(state58, dtype=np.float64) + action = np.asarray(action58, dtype=np.float64) + if state.ndim != 2 or state.shape[1] != 58: + raise DatasetSchemaError(f"observation.state must be (T,58), got {state.shape}") + if action.shape != state.shape: + raise DatasetSchemaError(f"action shape {action.shape} != state shape {state.shape}") + state_eef = np.asarray(convert(state), dtype=np.float32) + action_eef = np.asarray(convert(action), dtype=np.float32) + validate_eef62(state_eef, source_joint58=state, label=STATE_EEF_COLUMN) + validate_eef62(action_eef, source_joint58=action, label=ACTION_EEF_COLUMN) + return state_eef, action_eef + + +def episode_parquet_path(dataset_root: Path, episode_index: int, info: dict | None = None) -> Path: + metadata = info or _load_json(dataset_root / "meta" / "info.json") + pattern = metadata.get( + "data_path", + "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet", + ) + chunk_size = int(metadata.get("chunks_size", 1000)) + return dataset_root / pattern.format( + episode_chunk=int(episode_index) // chunk_size, + episode_index=int(episode_index), + ) + + +def default_track_cache(dataset_root: Path) -> Path: + return dataset_root.with_name(dataset_root.name + "_tracks") + + +def track_npz_path(track_cache: Path, episode_index: int) -> Path: + return track_cache / f"episode_{int(episode_index):06d}.npz" + + +def validate_episode_parquet( + path: Path, + *, + expected_frames: int | None = None, + verify_source_fk: bool = False, + converter: EefConverter | None = None, +) -> dict[str, object]: + """Validate fixed-size Arrow types, values, frame count, and EEF semantics.""" + + _, pq = _import_pyarrow() + if not path.is_file(): + raise FileNotFoundError(path) + table = pq.read_table(path) + schema_version = (table.schema.metadata or {}).get(PARQUET_SCHEMA_METADATA_KEY) + if schema_version != SCHEMA_VERSION.encode("utf-8"): + found = schema_version.decode("utf-8") if schema_version is not None else None + raise DatasetSchemaError( + f"{path}: parquet schema version {found!r} != {SCHEMA_VERSION!r}; rebuild required" + ) + if expected_frames is not None and table.num_rows != int(expected_frames): + raise DatasetSchemaError( + f"{path}: {table.num_rows} rows != expected {expected_frames}" + ) + for old_column in ("observation.state", "action"): + if old_column not in table.column_names: + raise DatasetSchemaError(f"{path}: original column {old_column!r} is missing") + timestamps = _column_to_numpy(table, "timestamp", dtype=np.float64) + if timestamps.shape == (table.num_rows, 1): + timestamps = timestamps[:, 0] + if timestamps.shape != (table.num_rows,): + raise DatasetSchemaError( + f"{path}: timestamp must have shape ({table.num_rows},), got {timestamps.shape}" + ) + sampling_summary = summarize_timestamp_sampling(timestamps) + force_summary = validate_tactile_force(table) + + combined_xy_field = ( + table.schema.field(TRACK_XY_COLUMN) + if TRACK_XY_COLUMN in table.column_names + else None + ) + if combined_xy_field is None or not _is_fixed_shape( + combined_xy_field.type, (NUM_COMBINED_POINTS, 2) + ): + raise DatasetSchemaError( + f"{path}: {TRACK_XY_COLUMN} must be Arrow fixed-size float32 " + f"({NUM_COMBINED_POINTS}, 2)" + ) + combined_visibility_field = ( + table.schema.field(TRACK_VISIBILITY_COLUMN) + if TRACK_VISIBILITY_COLUMN in table.column_names + else None + ) + if combined_visibility_field is None or not _is_fixed_shape( + combined_visibility_field.type, (NUM_COMBINED_POINTS,) + ): + raise DatasetSchemaError( + f"{path}: {TRACK_VISIBILITY_COLUMN} must be Arrow fixed-size float32 " + f"({NUM_COMBINED_POINTS},)" + ) + combined_xy = _column_to_numpy(table, TRACK_XY_COLUMN) + combined_visibility = _column_to_numpy(table, TRACK_VISIBILITY_COLUMN) + _validate_unit_interval(TRACK_XY_COLUMN, combined_xy) + _validate_unit_interval(TRACK_VISIBILITY_COLUMN, combined_visibility) + if not np.all( + (combined_visibility == 0.0) | (combined_visibility == 1.0) + ): + raise DatasetSchemaError( + f"{path}: {TRACK_VISIBILITY_COLUMN} visibility is not binary" + ) + + view_xy: list[np.ndarray] = [] + view_visibility: list[np.ndarray] = [] + for view, column in TRACK_COLUMNS.items(): + field = table.schema.field(column) if column in table.column_names else None + shape = (VIEW_POINT_COUNTS[view], 3) + if field is None or not _is_fixed_shape(field.type, shape): + raise DatasetSchemaError( + f"{path}: {column} must be Arrow fixed-size float32 {shape}" + ) + values = _column_to_numpy(table, column) + if values.shape != (table.num_rows, *shape): + raise DatasetSchemaError(f"{path}: {column} has shape {values.shape}") + _validate_unit_interval(column, values) + visibility = values[..., 2] + if not np.all((visibility == 0.0) | (visibility == 1.0)): + raise DatasetSchemaError(f"{path}: {column} visibility is not binary") + view_xy.append(values[..., :2]) + view_visibility.append(visibility) + if not np.array_equal(np.concatenate(view_xy, axis=1), combined_xy): + raise DatasetSchemaError(f"{path}: combined and per-view track XY differ") + if not np.array_equal( + np.concatenate(view_visibility, axis=1), combined_visibility + ): + raise DatasetSchemaError(f"{path}: combined and per-view visibility differ") + + for column in (STATE_EEF_COLUMN, ACTION_EEF_COLUMN): + field = table.schema.field(column) if column in table.column_names else None + if field is None or not _is_fixed_shape(field.type, (62,)): + raise DatasetSchemaError( + f"{path}: {column} must be Arrow fixed-size float32 (62,)" + ) + + source_state = _column_to_numpy(table, "observation.state") + source_action = _column_to_numpy(table, "action") + state_eef = _column_to_numpy(table, STATE_EEF_COLUMN) + action_eef = _column_to_numpy(table, ACTION_EEF_COLUMN) + validate_eef62(state_eef, source_joint58=source_state, label=f"{path}:{STATE_EEF_COLUMN}") + validate_eef62(action_eef, source_joint58=source_action, label=f"{path}:{ACTION_EEF_COLUMN}") + + if verify_source_fk: + convert = converter or joint58_to_eef62_batch + expected_state = np.asarray(convert(source_state), dtype=np.float32) + expected_action = np.asarray(convert(source_action), dtype=np.float32) + if not np.allclose(state_eef, expected_state, atol=2e-5, rtol=1e-5): + raise DatasetSchemaError(f"{path}: state EEF does not roundtrip through FK") + if not np.allclose(action_eef, expected_action, atol=2e-5, rtol=1e-5): + raise DatasetSchemaError(f"{path}: action EEF does not roundtrip through FK") + return { + "path": str(path), + "num_frames": int(table.num_rows), + "schema_version": SCHEMA_VERSION, + "sampling_20hz": sampling_summary, + "force_only": force_summary, + } + + +def output_is_valid(path: Path) -> tuple[bool, str]: + try: + validate_episode_parquet(path, verify_source_fk=False) + except Exception as exc: # Validation intentionally collapses to a skip decision. + return False, str(exc) + return True, "valid" + + +def build_episode( + *, + dataset_root: Path, + episode_index: int, + track_path: Path, + converter: EefConverter | None = None, + verify_source_fk: bool = True, +) -> dict[str, object]: + """Merge one validated track cache and reliable FK columns into a parquet.""" + + _, pq = _import_pyarrow() + info = _load_json(dataset_root / "meta" / "info.json") + parquet_path = episode_parquet_path(dataset_root, episode_index, info) + if not parquet_path.is_file(): + raise FileNotFoundError(parquet_path) + table = pq.read_table(parquet_path) + old_names = tuple(table.column_names) + expected_frames = int(table.num_rows) + payload = load_track_payload( + track_path, + expected_frames=expected_frames, + episode_index=episode_index, + ) + state58 = _column_to_numpy(table, "observation.state") + action58 = _column_to_numpy(table, "action") + state_eef, action_eef = convert_eef_columns( + state58, + action58, + converter=converter, + ) + + output = table + for column, values in track_features_from_payload(payload).items(): + output = _set_or_append_column(output, column, values) + output = _set_or_append_column(output, STATE_EEF_COLUMN, state_eef) + output = _set_or_append_column(output, ACTION_EEF_COLUMN, action_eef) + if any(name not in output.column_names for name in old_names): + raise AssertionError("an original parquet column was dropped") + schema_metadata = dict(output.schema.metadata or {}) + schema_metadata[PARQUET_SCHEMA_METADATA_KEY] = SCHEMA_VERSION.encode("utf-8") + output = output.replace_schema_metadata(schema_metadata) + _atomic_write_parquet(output, parquet_path, expected_rows=expected_frames) + validation = validate_episode_parquet( + parquet_path, + expected_frames=expected_frames, + verify_source_fk=verify_source_fk, + converter=converter, + ) + return { + "episode_index": int(episode_index), + "num_frames": expected_frames, + "parquet": str(parquet_path.relative_to(dataset_root)), + "track_npz": str(track_path), + "track_sha256": _sha256(track_path), + "sampling_20hz": validation["sampling_20hz"], + "force_only": validation["force_only"], + "validated_at": _utc_now(), + } + + +def _eef_feature_names(prefix: str) -> list[str]: + rotation_names = [ + "rot6d_col1_x", + "rot6d_col1_y", + "rot6d_col1_z", + "rot6d_col2_x", + "rot6d_col2_y", + "rot6d_col2_z", + ] + names = [ + f"left_{prefix}_x", + f"left_{prefix}_y", + f"left_{prefix}_z", + *(f"left_{prefix}_{name}" for name in rotation_names), + *(f"left_hand_q_{index}" for index in range(22)), + f"right_{prefix}_x", + f"right_{prefix}_y", + f"right_{prefix}_z", + *(f"right_{prefix}_{name}" for name in rotation_names), + *(f"right_hand_q_{index}" for index in range(22)), + ] + if len(names) != 62: + raise AssertionError("EEF feature names must have length 62") + return names + + +def _new_feature_metadata() -> dict[str, dict]: + features: dict[str, dict] = { + TRACK_XY_COLUMN: { + "dtype": "float32", + "shape": [NUM_COMBINED_POINTS, 2], + "names": None, + }, + TRACK_VISIBILITY_COLUMN: { + "dtype": "float32", + "shape": [NUM_COMBINED_POINTS], + "names": None, + }, + } + features.update( + { + column: { + "dtype": "float32", + "shape": [VIEW_POINT_COUNTS[view], 3], + "names": None, + } + for view, column in TRACK_COLUMNS.items() + } + ) + features[STATE_EEF_COLUMN] = { + "dtype": "float32", + "shape": [62], + "names": _eef_feature_names("eef"), + } + features[ACTION_EEF_COLUMN] = { + "dtype": "float32", + "shape": [62], + "names": _eef_feature_names("eef_target"), + } + return features + + +def _force_only_metadata() -> dict[str, object]: + return { + "source_column": FORCE_COLUMN, + "stored_shape": [FORCE_FLAT_DIM], + "reshape": [FORCE_SENSOR_COUNT, FORCE_SENSOR_DIM], + "target_rate_hz": 5.0, + "action_rate_hz": TARGET_RATE_HZ, + "action_update_stride": 4, + "action_chunk_offsets": [0, 4, 8, 12], + "history_frames": FORCE_HISTORY_FRAMES, + "history_duration_seconds": FORCE_HISTORY_FRAMES / 5.0, + "history_encoding": "online_model_encoder", + "vq_codes_on_disk": False, + "deformation_maps_used": False, + } + + +def _autoregressive_metadata() -> dict[str, object]: + return { + "blocks": AUTOREGRESSIVE_BLOCKS, + "action_steps_per_block": ACTION_CHUNK_STEPS, + "action_steps_per_sample": AUTOREGRESSIVE_BLOCKS * ACTION_CHUNK_STEPS, + "video_conditioning_frames": 1, + "video_frames_per_block": VIDEO_FRAMES_PER_BLOCK, + "video_frames_per_sample": TRAINING_VIDEO_FRAMES, + } + + +def _eef_modality_entries(original_key: str, *, action: bool) -> dict[str, dict]: + prefix = "eef62_absolute" if action else "eef62" + + def entry(start: int, end: int, rotation_type: str | None = None) -> dict: + return { + "original_key": original_key, + "start": start, + "end": end, + "rotation_type": rotation_type, + "absolute": True, + "dtype": "float32", + "range": None, + } + + return { + prefix: entry(0, 62), + f"left_{prefix}_position": entry(0, 3), + f"left_{prefix}_rotation_6d": entry(3, 9, "rotation_6d"), + f"left_{prefix}_hand": entry(9, 31), + f"right_{prefix}_position": entry(31, 34), + f"right_{prefix}_rotation_6d": entry(34, 40, "rotation_6d"), + f"right_{prefix}_hand": entry(40, 62), + } + + +def _statistics(values: np.ndarray) -> dict[str, list]: + array = np.asarray(values, dtype=np.float64) + if array.ndim < 2 or not np.isfinite(array).all(): + raise DatasetSchemaError(f"cannot compute stats for shape {array.shape}") + return { + "mean": np.mean(array, axis=0).tolist(), + "std": np.std(array, axis=0).tolist(), + "min": np.min(array, axis=0).tolist(), + "max": np.max(array, axis=0).tolist(), + "q01": np.quantile(array, 0.01, axis=0).tolist(), + "q99": np.quantile(array, 0.99, axis=0).tolist(), + } + + +def _rotation_6d_to_matrix(rotation_6d: np.ndarray) -> np.ndarray: + values = np.asarray(rotation_6d, dtype=np.float64) + if values.shape[-1] != 6: + raise DatasetSchemaError("rotation_6d must end in six values") + first = values[..., :3] + first /= np.linalg.norm(first, axis=-1, keepdims=True).clip(min=1e-8) + second = values[..., 3:6] + second = second - np.sum(first * second, axis=-1, keepdims=True) * first + second /= np.linalg.norm(second, axis=-1, keepdims=True).clip(min=1e-8) + third = np.cross(first, second) + return np.stack((first, second, third), axis=-1) + + +def eef62_delta_base( + reference_state: np.ndarray, absolute_targets: np.ndarray +) -> np.ndarray: + """T-Rex chunk-start-frame action: relative EEF pose + absolute hand joints.""" + + reference = np.asarray(reference_state, dtype=np.float64) + targets = np.asarray(absolute_targets, dtype=np.float64) + if reference.shape != (62,) or targets.shape[-1] != 62: + raise DatasetSchemaError("delta-base conversion expects [62] and [...,62]") + output = np.empty_like(targets, dtype=np.float64) + for pose_slice, hand_slice in ( + (LEFT_EEF, LEFT_HAND_EEF), + (RIGHT_EEF, RIGHT_HAND_EEF), + ): + reference_pose = reference[pose_slice] + target_pose = targets[..., pose_slice] + reference_rotation = _rotation_6d_to_matrix(reference_pose[3:9]) + target_rotation = _rotation_6d_to_matrix(target_pose[..., 3:9]) + delta_xyz = np.einsum( + "ji,...j->...i", + reference_rotation, + target_pose[..., :3] - reference_pose[:3], + ) + delta_rotation = np.einsum( + "ji,...jk->...ik", reference_rotation, target_rotation + ) + output[..., pose_slice] = np.concatenate( + ( + delta_xyz, + delta_rotation[..., :, 0], + delta_rotation[..., :, 1], + ), + axis=-1, + ) + output[..., hand_slice] = targets[..., hand_slice] + return output.astype(np.float32) + + +def compute_delta_base_stats(parquet_paths: Iterable[Path]) -> dict[str, list]: + """Pool all complete 16-step 20 Hz chunks for action normalization.""" + + _, pq = _import_pyarrow() + chunks: list[np.ndarray] = [] + boundary_fallback_chunks: list[np.ndarray] = [] + offsets = np.arange(ACTION_CHUNK_STEPS, dtype=np.int64) + for path in parquet_paths: + table = pq.read_table( + path, + columns=["timestamp", STATE_EEF_COLUMN, ACTION_EEF_COLUMN], + ) + timestamps = np.asarray(table["timestamp"].to_numpy(), dtype=np.float64) + state = _column_to_numpy(table, STATE_EEF_COLUMN) + action = _column_to_numpy(table, ACTION_EEF_COLUMN) + anchors = sample_timestamps_nearest( + timestamps, + target_rate_hz=TARGET_RATE_HZ, + anchor_index=0, + ) + anchor_indices = np.asarray(anchors["indices"], dtype=np.int64) + anchor_times = np.asarray(anchors["target_timestamps"], dtype=np.float64) + for anchor_index, anchor_time in zip(anchor_indices, anchor_times): + selection = sample_timestamps_nearest( + timestamps, + target_rate_hz=TARGET_RATE_HZ, + anchor_timestamp=float(anchor_time), + offsets=offsets, + ) + target_indices = np.asarray(selection["indices"], dtype=np.int64) + delta_chunk = eef62_delta_base( + state[int(anchor_index)], action[target_indices] + ) + if np.asarray(selection["padding_mask"], dtype=bool).any(): + boundary_fallback_chunks.append(delta_chunk) + else: + chunks.append(delta_chunk) + if not chunks: + # Tiny schema fixtures and unusually short episodes cannot contain a + # complete horizon. Edge-clamped values keep metadata writable, while + # the runtime loader still excludes these anchors from training. + chunks = boundary_fallback_chunks + if not chunks: + raise DatasetSchemaError("no rows are available for delta-base stats") + return _statistics(np.concatenate(chunks, axis=0)) + + +def compute_new_stats(parquet_paths: Iterable[Path]) -> dict[str, dict]: + _, pq = _import_pyarrow() + stat_columns = (*NEW_COLUMNS, FORCE_COLUMN) + buffers: dict[str, list[np.ndarray]] = {column: [] for column in stat_columns} + for path in parquet_paths: + table = pq.read_table(path, columns=list(stat_columns)) + for column in stat_columns: + buffers[column].append(_column_to_numpy(table, column)) + if not all(buffers.values()): + raise DatasetSchemaError("no valid converted episodes are available for stats") + return { + column: _statistics(np.concatenate(parts, axis=0)) + for column, parts in buffers.items() + } + + +def _valid_converted_episodes( + dataset_root: Path, + *, + info: dict, +) -> tuple[list[int], list[Path]]: + indices: list[int] = [] + paths: list[Path] = [] + for episode_index in range(int(info["total_episodes"])): + path = episode_parquet_path(dataset_root, episode_index, info) + valid, _ = output_is_valid(path) + if valid: + indices.append(episode_index) + paths.append(path) + return indices, paths + + +def update_metadata( + dataset_root: Path, + *, + assume_all_converted: bool = False, +) -> dict[str, object]: + """Atomically update info/modality/stats while preserving all old entries.""" + + meta_dir = dataset_root / "meta" + info_path = meta_dir / "info.json" + modality_path = meta_dir / "modality.json" + stats_path = meta_dir / "stats.json" + info = _load_json(info_path) + modality = _load_json(modality_path) + stats = _load_json(stats_path) if stats_path.exists() else {} + + if assume_all_converted: + converted_indices = list(range(int(info["total_episodes"]))) + converted_paths = [ + episode_parquet_path(dataset_root, episode_index, info) + for episode_index in converted_indices + ] + missing = [path for path in converted_paths if not path.is_file()] + if missing: + raise FileNotFoundError(missing[0]) + else: + converted_indices, converted_paths = _valid_converted_episodes( + dataset_root, + info=info, + ) + if not converted_paths: + raise DatasetSchemaError("metadata cannot be updated before one valid episode exists") + new_stats = compute_new_stats(converted_paths) + + features = info.setdefault("features", {}) + force_feature = features.get(FORCE_COLUMN) + if not isinstance(force_feature, dict) or force_feature.get("shape") != [ + FORCE_FLAT_DIM + ]: + raise DatasetSchemaError( + f"info.json must declare existing {FORCE_COLUMN} with shape [{FORCE_FLAT_DIM}]" + ) + if "float" not in str(force_feature.get("dtype", "")): + raise DatasetSchemaError(f"info.json {FORCE_COLUMN} must be floating-point") + features.update(_new_feature_metadata()) + info["trex_track_force"] = { + "schema_version": SCHEMA_VERSION, + "track_layout": layout_metadata(), + "sampling_20hz": { + "source_column": "timestamp", + "target_rate_hz": TARGET_RATE_HZ, + "method": "deterministic_nearest_earlier_on_tie", + "source_data_overwritten": False, + "action_chunk_steps": ACTION_CHUNK_STEPS, + "action_chunk_duration_seconds": ACTION_CHUNK_DURATION_SECONDS, + "action_chunk_timestamp_span_seconds": ( + ACTION_CHUNK_TIMESTAMP_SPAN_SECONDS + ), + }, + "autoregressive_training": _autoregressive_metadata(), + "force_only": _force_only_metadata(), + "eef62_layout": { + "order": [ + "left_eef_pose9", + "left_hand22", + "right_eef_pose9", + "right_hand22", + ], + "slices": { + "left_eef_pose9": [0, 9], + "left_hand22": [9, 31], + "right_eef_pose9": [31, 40], + "right_hand22": [40, 62], + }, + "pose9": "translation_xyz + rotation_matrix_column_1 + rotation_matrix_column_2", + "source": "T-Rex trex_fk + utils/lerobot_common.py", + }, + "converted_episode_indices": converted_indices, + "complete": len(converted_indices) == int(info["total_episodes"]), + "updated_at": _utc_now(), + } + + modality.setdefault("state", {}).update( + _eef_modality_entries(STATE_EEF_COLUMN, action=False) + ) + modality.setdefault("action", {}).update( + _eef_modality_entries(ACTION_EEF_COLUMN, action=True) + ) + # The dedicated loader exposes this alias after applying T-Rex delta-base + # conversion. The source parquet remains absolute and is never overwritten. + modality["action"]["eef62"] = { + "original_key": ACTION_EEF_COLUMN, + "start": 0, + "end": 62, + "rotation_type": None, + "absolute": False, + "dtype": "float32", + "range": None, + } + modality["track"] = { + "xy": { + "original_key": TRACK_XY_COLUMN, + "shape": [NUM_COMBINED_POINTS, 2], + "coordinate_space": "normalized_xy_div_wh", + }, + "visibility": { + "original_key": TRACK_VISIBILITY_COLUMN, + "shape": [NUM_COMBINED_POINTS], + "range": [0.0, 1.0], + }, + "views": { + view: { + "original_key": column, + "shape": [VIEW_POINT_COUNTS[view], 3], + "value_order": ["x", "y", "visibility"], + "coordinate_space": "normalized_xy_div_wh", + "slice": list(VIEW_SLICES[view]), + } + for view, column in TRACK_COLUMNS.items() + }, + } + modality["force"] = { + "current": { + "original_key": FORCE_COLUMN, + "stored_shape": [FORCE_FLAT_DIM], + "reshape": [FORCE_SENSOR_COUNT, FORCE_SENSOR_DIM], + }, + "history": { + "original_key": FORCE_COLUMN, + "frames": FORCE_HISTORY_FRAMES, + "target_rate_hz": 5.0, + "action_update_stride": 4, + "encoding": "online_model_encoder", + "vq_codes_on_disk": False, + }, + } + stats.update(new_stats) + + relative_stats_path = meta_dir / RELATIVE_ACTION_STATS_FILENAME + relative_stats = {"eef62": compute_delta_base_stats(converted_paths)} + for path in (info_path, modality_path, stats_path, relative_stats_path): + _atomic_backup(path) + _atomic_write_json(info_path, info) + _atomic_write_json(modality_path, modality) + _atomic_write_json(stats_path, stats) + _atomic_write_json(relative_stats_path, relative_stats) + return { + "converted_episode_indices": converted_indices, + "complete": info["trex_track_force"]["complete"], + "stats_episode_count": len(converted_indices), + } + + +def validate_metadata(dataset_root: Path) -> None: + info = _load_json(dataset_root / "meta" / "info.json") + modality = _load_json(dataset_root / "meta" / "modality.json") + stats = _load_json(dataset_root / "meta" / "stats.json") + relative_stats = _load_json( + dataset_root / "meta" / RELATIVE_ACTION_STATS_FILENAME + ) + feature_specs = _new_feature_metadata() + for column, expected in feature_specs.items(): + if info.get("features", {}).get(column) != expected: + raise DatasetSchemaError(f"info.json has invalid feature metadata for {column}") + if column not in stats: + raise DatasetSchemaError(f"stats.json is missing {column}") + schema_block = info.get("trex_track_force", {}) + if schema_block.get("schema_version") != SCHEMA_VERSION: + raise DatasetSchemaError("info.json is missing the track-force schema version") + if schema_block.get("track_layout") != layout_metadata(): + raise DatasetSchemaError("info.json has unstable track identity metadata") + if schema_block.get("sampling_20hz", {}).get("target_rate_hz") != TARGET_RATE_HZ: + raise DatasetSchemaError("info.json is missing the 20 Hz sampling contract") + if schema_block.get("force_only") != _force_only_metadata(): + raise DatasetSchemaError("info.json has invalid force-only metadata") + if schema_block.get("autoregressive_training") != _autoregressive_metadata(): + raise DatasetSchemaError("info.json has invalid autoregressive training metadata") + force_feature = info.get("features", {}).get(FORCE_COLUMN, {}) + if force_feature.get("shape") != [FORCE_FLAT_DIM]: + raise DatasetSchemaError(f"info.json has invalid {FORCE_COLUMN} shape") + if FORCE_COLUMN not in stats or any( + len(stats[FORCE_COLUMN].get(name, [])) != FORCE_FLAT_DIM + for name in ("mean", "std", "min", "max", "q01", "q99") + ): + raise DatasetSchemaError(f"stats.json is missing 60-D {FORCE_COLUMN} stats") + if "observation.force_history_vq" in info.get("features", {}): + raise DatasetSchemaError("metadata must not declare fabricated force VQ codes") + for name in _eef_modality_entries(STATE_EEF_COLUMN, action=False): + if name not in modality.get("state", {}): + raise DatasetSchemaError(f"modality.json is missing state.{name}") + for name in _eef_modality_entries(ACTION_EEF_COLUMN, action=True): + if name not in modality.get("action", {}): + raise DatasetSchemaError(f"modality.json is missing action.{name}") + if modality.get("action", {}).get("eef62", {}).get("absolute") is not False: + raise DatasetSchemaError("modality.json is missing delta-base action.eef62") + delta_stats = relative_stats.get("eef62", {}) + if set(delta_stats) != {"mean", "std", "min", "max", "q01", "q99"}: + raise DatasetSchemaError("relative action stats are missing action.eef62") + if any(len(delta_stats[name]) != 62 for name in delta_stats): + raise DatasetSchemaError("relative action.eef62 stats must have 62 values") + track_meta = modality.get("track", {}) + if track_meta.get("xy", {}).get("original_key") != TRACK_XY_COLUMN: + raise DatasetSchemaError("modality.json has invalid track XY mapping") + if ( + track_meta.get("visibility", {}).get("original_key") + != TRACK_VISIBILITY_COLUMN + ): + raise DatasetSchemaError("modality.json has invalid track visibility mapping") + if set(track_meta.get("views", {})) != set(VIEW_ORDER): + raise DatasetSchemaError("modality.json has invalid track views") + force_meta = modality.get("force", {}) + if force_meta.get("current", {}).get("original_key") != FORCE_COLUMN: + raise DatasetSchemaError("modality.json has invalid current force source") + if force_meta.get("history", {}).get("encoding") != "online_model_encoder": + raise DatasetSchemaError("modality.json must encode force history online") + if force_meta.get("history", {}).get("vq_codes_on_disk") is not False: + raise DatasetSchemaError("modality.json must not claim on-disk VQ codes") + + +def _new_manifest(dataset_root: Path, track_cache: Path) -> dict: + return { + "schema_version": SCHEMA_VERSION, + "track_layout_version": TRACK_LAYOUT_VERSION, + "track_layout": layout_metadata(), + "sampling_contract": { + "source_column": "timestamp", + "target_rate_hz": TARGET_RATE_HZ, + "action_chunk_steps": ACTION_CHUNK_STEPS, + "action_chunk_duration_seconds": ACTION_CHUNK_DURATION_SECONDS, + "action_chunk_timestamp_span_seconds": ( + ACTION_CHUNK_TIMESTAMP_SPAN_SECONDS + ), + }, + "autoregressive_training": _autoregressive_metadata(), + "force_only": _force_only_metadata(), + "dataset_root": str(dataset_root), + "track_cache": str(track_cache), + "created_at": _utc_now(), + "updated_at": _utc_now(), + "episodes": {}, + } + + +def load_manifest(path: Path, *, dataset_root: Path, track_cache: Path) -> dict: + if not path.exists(): + return _new_manifest(dataset_root, track_cache) + manifest = _load_json(path) + if manifest.get("schema_version") != SCHEMA_VERSION: + fresh = _new_manifest(dataset_root, track_cache) + fresh["supersedes_schema_version"] = manifest.get("schema_version") + fresh["stale_episode_entries_discarded"] = len(manifest.get("episodes", {})) + return fresh + if manifest.get("track_layout") != layout_metadata(): + raise DatasetSchemaError(f"{path}: manifest point layout is not canonical") + manifest["autoregressive_training"] = _autoregressive_metadata() + manifest.setdefault("episodes", {}) + return manifest + + +def write_manifest(path: Path, manifest: dict) -> None: + manifest["updated_at"] = _utc_now() + _atomic_backup(path) + _atomic_write_json(path, manifest) + + +def select_episode_indices( + total_episodes: int, + *, + episode_index: int | None = None, + episode_range: Sequence[int] | None = None, + all_episodes: bool = False, +) -> list[int]: + modes = int(episode_index is not None) + int(episode_range is not None) + int(all_episodes) + if modes != 1: + raise ValueError("select exactly one of episode_index, episode_range, or all_episodes") + if episode_index is not None: + result = [int(episode_index)] + elif episode_range is not None: + if len(episode_range) != 2: + raise ValueError("episode_range must contain START END") + start, end = map(int, episode_range) + if end <= start: + raise ValueError("episode range is half-open and requires END > START") + result = list(range(start, end)) + else: + result = list(range(int(total_episodes))) + invalid = [index for index in result if index < 0 or index >= int(total_episodes)] + if invalid: + raise ValueError( + f"episode indices out of range [0,{total_episodes}): {invalid[:5]}" + ) + return result + + +def validate_dataset( + *, + dataset_root: Path, + episode_indices: Sequence[int], + manifest_path: Path, + verify_fk: bool, +) -> list[dict[str, object]]: + info = _load_json(dataset_root / "meta" / "info.json") + manifest = load_manifest( + manifest_path, + dataset_root=dataset_root, + track_cache=default_track_cache(dataset_root), + ) + results: list[dict[str, object]] = [] + for episode_index in episode_indices: + path = episode_parquet_path(dataset_root, episode_index, info) + result = validate_episode_parquet( + path, + verify_source_fk=verify_fk, + ) + manifest_entry = manifest.get("episodes", {}).get(f"{episode_index:06d}") + if not manifest_entry or manifest_entry.get("status") != "complete": + raise DatasetSchemaError( + f"manifest has no complete entry for episode {episode_index}" + ) + recorded_sampling = manifest_entry.get("sampling_20hz", {}) + current_sampling = result["sampling_20hz"] + for name in ( + "source_frame_count", + "target_rate_hz", + "target_sample_count", + "action_chunk_steps", + "action_chunk_duration_seconds", + "action_chunk_timestamp_span_seconds", + ): + if recorded_sampling.get(name) != current_sampling[name]: + raise DatasetSchemaError( + f"manifest episode {episode_index} has stale sampling field {name}" + ) + if not np.isclose( + float(recorded_sampling.get("source_rate_hz", np.nan)), + float(current_sampling["source_rate_hz"]), + rtol=1e-9, + atol=1e-9, + ): + raise DatasetSchemaError( + f"manifest episode {episode_index} has stale source_rate_hz" + ) + if manifest_entry.get("force_only") != result["force_only"]: + raise DatasetSchemaError( + f"manifest episode {episode_index} has stale force-only metadata" + ) + results.append(result) + validate_metadata(dataset_root) + return results + + +def _ensure_track_npz( + *, + dataset_root: Path, + track_cache: Path, + episode_index: int, + expected_frames: int, + args: argparse.Namespace, + runtime_holder: dict[str, object], +) -> Path: + path = track_npz_path(track_cache, episode_index) + try: + load_track_payload( + path, + expected_frames=expected_frames, + episode_index=episode_index, + ) + return path + except (FileNotFoundError, DatasetSchemaError) as exc: + if not args.extract_missing: + raise DatasetSchemaError( + f"episode {episode_index}: no valid track cache and extraction is disabled" + ) from exc + print(f"episode {episode_index}: extracting tracks ({exc})") + + if "runtime" not in runtime_holder: + import extract_track + + runtime_holder["module"] = extract_track + runtime_holder["runtime"] = extract_track.create_tracking_runtime( + calib_path=args.calib_path, + openpi_root=args.openpi_root, + cotracker_checkpoint=args.cotracker_checkpoint, + cotracker_device=args.cotracker_device, + sam2_model=args.sam2_model, + sam2_device=args.sam2_device, + sam2_libs=args.sam2_libs, + image_height=args.image_height, + image_width=args.image_width, + ) + module = runtime_holder["module"] + runtime = runtime_holder["runtime"] + viz_dir = track_cache / "viz_tracks" + masks_dir = track_cache / "sam2_masks" + module.process_episode( + dataset_root=dataset_root, + episode_index=episode_index, + output_path=track_cache, + calib=runtime.calib, + out_hw=runtime.out_hw, + cotracker_model=runtime.cotracker_model, + cotracker_device=runtime.cotracker_device, + save_viz=bool(args.save_viz), + viz_out_dir=viz_dir, + viz_fps=int(args.viz_fps), + viz_trail=int(args.viz_trail), + sam2_predictor=runtime.sam2_predictor, + sam2_seed=int(args.sam2_seed), + save_sam2_masks_flag=bool(args.save_sam2_masks), + sam2_masks_dir=masks_dir, + ) + load_track_payload( + path, + expected_frames=expected_frames, + episode_index=episode_index, + ) + return path + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--dataset-root", type=Path, default=DEFAULT_DATASET_ROOT) + selection = parser.add_mutually_exclusive_group(required=True) + selection.add_argument("--episode-index", type=int) + selection.add_argument( + "--episode-range", + type=int, + nargs=2, + metavar=("START", "END"), + help="Half-open episode range [START, END)", + ) + selection.add_argument("--all", dest="all_episodes", action="store_true") + parser.add_argument("--track-cache", type=Path, default=None) + parser.add_argument("--manifest-path", type=Path, default=None) + parser.add_argument("--dry-run", action="store_true") + parser.add_argument("--validate-only", action="store_true") + parser.add_argument("--force", action="store_true", help="Rebuild even valid parquets") + parser.add_argument( + "--extract-missing", + action=argparse.BooleanOptionalAction, + default=True, + ) + parser.add_argument( + "--update-metadata", + action=argparse.BooleanOptionalAction, + default=True, + ) + parser.add_argument( + "--verify-fk", + action=argparse.BooleanOptionalAction, + default=True, + ) + + # Heavy extraction dependencies are only imported if a cache is missing. + parser.add_argument( + "--calib-path", + type=Path, + default=_DREAMZERO_ROOT / "assets" / "trex_camera_calib.json", + ) + parser.add_argument( + "--openpi-root", + type=Path, + default=Path("/scratch2/home/zhicao/openpi"), + ) + parser.add_argument("--cotracker-checkpoint", type=str, default="") + parser.add_argument("--cotracker-device", type=str, default="") + parser.add_argument( + "--sam2-model", + type=str, + default=os.environ.get("SAM2_MODEL", "facebook/sam2-hiera-large"), + ) + parser.add_argument("--sam2-device", type=str, default="") + parser.add_argument("--sam2-seed", type=int, default=0) + parser.add_argument( + "--sam2-libs", + type=Path, + default=Path(os.environ.get("SAM2_LIBS", "/scratch1/home/zhicao/physctrl/libs")), + ) + parser.add_argument("--image-height", type=int, default=0) + parser.add_argument("--image-width", type=int, default=0) + parser.add_argument( + "--save-viz", + action=argparse.BooleanOptionalAction, + default=False, + ) + parser.add_argument("--viz-fps", type=int, default=10) + parser.add_argument("--viz-trail", type=int, default=15) + parser.add_argument( + "--save-sam2-masks", + action=argparse.BooleanOptionalAction, + default=False, + ) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + args = _build_parser().parse_args(argv) + if args.dry_run and args.validate_only: + raise ValueError("--dry-run and --validate-only are mutually exclusive") + dataset_root = args.dataset_root.expanduser().resolve() + info = _load_json(dataset_root / "meta" / "info.json") + episode_indices = select_episode_indices( + int(info["total_episodes"]), + episode_index=args.episode_index, + episode_range=args.episode_range, + all_episodes=bool(args.all_episodes), + ) + track_cache = ( + args.track_cache.expanduser().resolve() + if args.track_cache is not None + else default_track_cache(dataset_root) + ) + manifest_path = ( + args.manifest_path.expanduser().resolve() + if args.manifest_path is not None + else dataset_root / "meta" / "trex_track_force_manifest.json" + ) + + if args.validate_only: + results = validate_dataset( + dataset_root=dataset_root, + episode_indices=episode_indices, + manifest_path=manifest_path, + verify_fk=bool(args.verify_fk), + ) + for result in results: + sampling = result["sampling_20hz"] + print( + f"{Path(str(result['path'])).name}: " + f"source={sampling['source_rate_hz']:.6f}Hz " + f"target={sampling['target_rate_hz']:.1f}Hz " + f"samples={sampling['target_sample_count']} " + f"chunk={sampling['action_chunk_steps']} steps/" + f"{sampling['action_chunk_duration_seconds']:.1f}s " + f"(timestamp span " + f"{sampling['action_chunk_timestamp_span_seconds']:.2f}s)" + ) + print(f"validated {len(results)} episode(s)") + return 0 + + if args.dry_run: + _, pq = _import_pyarrow() + for episode_index in episode_indices: + parquet_path = episode_parquet_path(dataset_root, episode_index, info) + valid, reason = output_is_valid(parquet_path) + cache_path = track_npz_path(track_cache, episode_index) + cache_valid = False + cache_reason = "missing" + if cache_path.exists() and parquet_path.exists(): + try: + expected_frames = int(pq.read_metadata(parquet_path).num_rows) + load_track_payload( + cache_path, + expected_frames=expected_frames, + episode_index=episode_index, + ) + cache_valid = True + cache_reason = "valid" + except Exception as exc: + cache_reason = str(exc) + if valid and not args.force: + action = "skip valid output" + elif cache_valid: + action = "merge cache + FK" + elif args.extract_missing: + action = f"extract SAM2/CoTracker, merge + FK (cache: {cache_reason})" + else: + action = f"FAIL: no valid track cache ({cache_reason})" + print(f"[dry-run] episode {episode_index:06d}: {action} ({reason})") + print("[dry-run] no files were changed") + return 0 + + manifest = load_manifest( + manifest_path, + dataset_root=dataset_root, + track_cache=track_cache, + ) + runtime_holder: dict[str, object] = {} + _, pq = _import_pyarrow() + for episode_index in episode_indices: + key = f"{episode_index:06d}" + parquet_path = episode_parquet_path(dataset_root, episode_index, info) + valid, reason = output_is_valid(parquet_path) + try: + validated_summary: dict[str, object] | None = None + if valid and not args.force: + try: + validated_summary = validate_episode_parquet( + parquet_path, + verify_source_fk=bool(args.verify_fk), + ) + except DatasetSchemaError as exc: + valid = False + reason = f"deep validation failed: {exc}" + if valid and not args.force: + if validated_summary is None: + raise AssertionError("valid output was not validated") + summary = validated_summary + summary.update( + { + "episode_index": episode_index, + "status": "complete", + "skipped": True, + "validated_at": _utc_now(), + } + ) + print(f"episode {episode_index:06d}: skip valid output") + else: + if not parquet_path.is_file(): + raise FileNotFoundError(parquet_path) + expected_frames = int(pq.read_metadata(parquet_path).num_rows) + cache_path = _ensure_track_npz( + dataset_root=dataset_root, + track_cache=track_cache, + episode_index=episode_index, + expected_frames=expected_frames, + args=args, + runtime_holder=runtime_holder, + ) + summary = build_episode( + dataset_root=dataset_root, + episode_index=episode_index, + track_path=cache_path, + verify_source_fk=bool(args.verify_fk), + ) + summary["status"] = "complete" + summary["skipped"] = False + print(f"episode {episode_index:06d}: built and validated ({reason})") + manifest["episodes"][key] = summary + write_manifest(manifest_path, manifest) + except Exception as exc: + manifest["episodes"][key] = { + "episode_index": episode_index, + "status": "failed", + "error": f"{type(exc).__name__}: {exc}", + "failed_at": _utc_now(), + } + write_manifest(manifest_path, manifest) + raise + + if args.update_metadata: + manifest["metadata"] = update_metadata(dataset_root) + write_manifest(manifest_path, manifest) + + if args.update_metadata: + results = validate_dataset( + dataset_root=dataset_root, + episode_indices=episode_indices, + manifest_path=manifest_path, + # Every built/skipped episode was already deep-FK checked above. + verify_fk=False, + ) + else: + results = [ + validate_episode_parquet( + episode_parquet_path(dataset_root, episode_index, info), + verify_source_fk=False, + ) + for episode_index in episode_indices + ] + print( + f"completed {len(results)} episode(s); manifest={manifest_path}; " + f"metadata={'updated' if args.update_metadata else 'unchanged'}" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/data/check_trex_dataset_ready.py b/scripts/data/check_trex_dataset_ready.py new file mode 100644 index 0000000000000000000000000000000000000000..1a2085b10ee2440f182da4d26e8a6bf6cf2ecfb7 --- /dev/null +++ b/scripts/data/check_trex_dataset_ready.py @@ -0,0 +1,140 @@ +#!/usr/bin/env python3 +"""Fast startup check for a previously validated T-Rex dataset variant.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Sequence + +import numpy as np +import pyarrow.parquet as pq + +TRACK_CACHE_NAME = "tracks_trex_track_force_v2" +FORCE_SCHEMA_METADATA_KEY = b"trex_track_force_schema_version" + + +def _read_json(path: Path) -> dict: + if not path.is_file(): + raise FileNotFoundError(path) + return json.loads(path.read_text()) + + +def _jsonl_count(path: Path) -> int: + if not path.is_file(): + raise FileNotFoundError(path) + return sum(bool(line.strip()) for line in path.read_text().splitlines()) + + +def _episode_path(root: Path, episode_index: int) -> Path: + return ( + root + / "data" + / f"chunk-{episode_index // 1000:03d}" + / f"episode_{episode_index:06d}.parquet" + ) + + +def _track_path(root: Path, episode_index: int) -> Path: + return root / TRACK_CACHE_NAME / f"episode_{episode_index:06d}.npz" + + +def check_dataset(root: Path, *, require_force: bool) -> dict: + root = root.expanduser().resolve() + ready = _read_json(root / "meta" / "dataset_ready.json") + info = _read_json(root / "meta" / "info.json") + episodes = int(info["total_episodes"]) + frames = int(info["total_frames"]) + tasks = int(info["total_tasks"]) + videos = int(info["total_videos"]) + if episodes <= 0 or frames <= 0: + raise ValueError(f"{root}: empty dataset") + if int(ready.get("episodes", -1)) != episodes: + raise ValueError(f"{root}: stale dataset_ready episode count") + if int(ready.get("frames", -1)) != frames: + raise ValueError(f"{root}: stale dataset_ready frame count") + if bool(ready.get("force")) != require_force: + raise ValueError( + f"{root}: force={ready.get('force')} but require_force={require_force}" + ) + if _jsonl_count(root / "meta" / "episodes.jsonl") != episodes: + raise ValueError(f"{root}: episodes.jsonl count mismatch") + if _jsonl_count(root / "meta" / "tasks.jsonl") != tasks: + raise ValueError(f"{root}: tasks.jsonl count mismatch") + for required in ( + root / "meta" / "stats.json", + root / "meta" / "relative_stats_dreamzero.json", + root / "meta" / "source_episode_index_map.json", + ): + if not required.is_file(): + raise FileNotFoundError(required) + + manifest = None + if require_force: + manifest = _read_json( + root / "meta" / "trex_track_force_manifest.json" + ) + if len(manifest.get("episodes", {})) != episodes: + raise ValueError(f"{root}: force manifest count mismatch") + if not (root / TRACK_CACHE_NAME).is_dir(): + raise FileNotFoundError(root / TRACK_CACHE_NAME) + + sample_indices = sorted({0, episodes // 2, episodes - 1}) + for episode_index in sample_indices: + parquet_path = _episode_path(root, episode_index) + parquet_file = pq.ParquetFile(parquet_path) + if int(parquet_file.metadata.num_rows) <= 0: + raise ValueError(f"{parquet_path}: empty parquet") + if require_force: + metadata = parquet_file.schema_arrow.metadata or {} + if FORCE_SCHEMA_METADATA_KEY not in metadata: + raise ValueError(f"{parquet_path}: force schema metadata missing") + entry = manifest["episodes"].get(f"{episode_index:06d}", {}) + if entry.get("status") != "complete": + raise ValueError( + f"{root}: force manifest sample {episode_index} incomplete" + ) + track_path = _track_path(root, episode_index) + with np.load(track_path, allow_pickle=False) as payload: + if int(np.asarray(payload["episode_index"]).item()) != episode_index: + raise ValueError(f"{track_path}: episode_index mismatch") + if int(np.asarray(payload["num_steps"]).item()) != int( + parquet_file.metadata.num_rows + ): + raise ValueError(f"{track_path}: frame count mismatch") + + return { + "dataset": str(root), + "episodes": episodes, + "frames": frames, + "tasks": tasks, + "videos": videos, + "force": require_force, + "validated_at": ready.get("validated_at"), + } + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dataset-root", type=Path, required=True) + parser.add_argument( + "--require-force", + action=argparse.BooleanOptionalAction, + default=False, + ) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + args = _build_parser().parse_args(argv) + result = check_dataset( + args.dataset_root, + require_force=args.require_force, + ) + print(json.dumps(result)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/data/convert_agibot.py b/scripts/data/convert_agibot.py new file mode 100644 index 0000000000000000000000000000000000000000..8d888cc17e402ff7045060f200438bc3536fd7a0 --- /dev/null +++ b/scripts/data/convert_agibot.py @@ -0,0 +1,1271 @@ +import argparse +from copy import deepcopy +from functools import partial +import gc +import json +import logging +from math import ceil +import os +from pathlib import Path +import shutil +from typing import Callable + +import einops +import h5py +from lerobot.common.datasets.lerobot_dataset import LeRobotDataset +from lerobot.common.datasets.utils import ( + STATS_PATH, + check_timestamps_sync, + get_episode_data_index, + serialize_dict, + write_json, +) +import numpy as np +import torch +from tqdm import tqdm +from tqdm.contrib.concurrent import process_map + + +def generate_modality_json(output_dir: str) -> None: + """Generate modality.json file defining field mappings for the dataset.""" + modality_config = { + "state": { + "left_arm_joint_position": { + "original_key": "observation.state", + "start": 0, + "end": 7, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "right_arm_joint_position": { + "original_key": "observation.state", + "start": 7, + "end": 14, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "left_effector_position": { + "original_key": "observation.state", + "start": 14, + "end": 15, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "right_effector_position": { + "original_key": "observation.state", + "start": 15, + "end": 16, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "head_position": { + "original_key": "observation.state", + "start": 16, + "end": 18, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "waist_pitch": { + "original_key": "observation.state", + "start": 18, + "end": 19, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "waist_lift": { + "original_key": "observation.state", + "start": 19, + "end": 20, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + }, + "action": { + "left_arm_joint_position": { + "original_key": "action", + "start": 0, + "end": 7, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "right_arm_joint_position": { + "original_key": "action", + "start": 7, + "end": 14, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "left_effector_position": { + "original_key": "action", + "start": 14, + "end": 15, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "right_effector_position": { + "original_key": "action", + "start": 15, + "end": 16, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "head_position": { + "original_key": "action", + "start": 16, + "end": 18, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "waist_pitch": { + "original_key": "action", + "start": 18, + "end": 19, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "waist_lift": { + "original_key": "action", + "start": 19, + "end": 20, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + "robot_velocity": { + "original_key": "action", + "start": 20, + "end": 22, + "rotation_type": None, + "absolute": True, + "dtype": "float64", + "range": None + }, + }, + "video": { + "top_head": { + "original_key": "observation.images.top_head" + }, + "hand_left": { + "original_key": "observation.images.hand_left" + }, + "hand_right": { + "original_key": "observation.images.hand_right" + }, + }, + "annotation": { + "language.action_text": { + "original_key": "task_index" + }, + "agibot.sub_task": { + "original_key": "annotation.agibot.sub_task" + }, + "frame_type": { + "original_key": "annotation.frame_type" + }, + }, + } + + modality_path = os.path.join(output_dir, "modality.json") + with open(modality_path, "w") as f: + json.dump(modality_config, f, indent=4) + print(f"Generated modality.json at {modality_path}") + + +HEAD_COLOR = "head_color.mp4" +HAND_LEFT_COLOR = "hand_left_color.mp4" +HAND_RIGHT_COLOR = "hand_right_color.mp4" + +FEATURES = { + "observation.images.top_head": { + "dtype": "video", + "shape": [480, 640, 3], + "names": ["height", "width", "channel"], + "video_info": { + "video.fps": 30.0, + "video.codec": "av1", + "video.pix_fmt": "yuv420p", + "video.is_depth_map": False, + "has_audio": False, + }, + }, + "observation.images.hand_left": { + "dtype": "video", + "shape": [480, 640, 3], + "names": ["height", "width", "channel"], + "video_info": { + "video.fps": 30.0, + "video.codec": "av1", + "video.pix_fmt": "yuv420p", + "video.is_depth_map": False, + "has_audio": False, + }, + }, + "observation.images.hand_right": { + "dtype": "video", + "shape": [480, 640, 3], + "names": ["height", "width", "channel"], + "video_info": { + "video.fps": 30.0, + "video.codec": "av1", + "video.pix_fmt": "yuv420p", + "video.is_depth_map": False, + "has_audio": False, + }, + }, + "observation.state": { + "dtype": "float32", + "shape": [20], + }, + "action": { + "dtype": "float32", + "shape": [22], + }, + "annotation.language.action_text": { + "dtype": "int64", + "shape": [1], + "names": None, + }, + "annotation.agibot.tasks": { + "dtype": "int64", + "shape": [1], + "names": None, + }, + "episode_index": { + "dtype": "int64", + "shape": [1], + "names": None, + }, + "frame_index": { + "dtype": "int64", + "shape": [1], + "names": None, + }, + "index": { + "dtype": "int64", + "shape": [1], + "names": None, + }, + "task_index": { + "dtype": "int64", + "shape": [1], + "names": None, + }, +} + + +def get_stats_einops_patterns(dataset, num_workers=0): + """These einops patterns will be used to aggregate batches and compute statistics. + + Note: We assume the images are in channel first format + """ + + dataloader = torch.utils.data.DataLoader( + dataset, + num_workers=num_workers, + batch_size=2, + shuffle=False, + ) + batch = next(iter(dataloader)) + + stats_patterns = {} + + for key in dataset.features: + # sanity check that tensors are not float64 + assert batch[key].dtype != torch.float64 + + # if isinstance(feats_type, (VideoFrame, Image)): + if key in dataset.meta.camera_keys: + # sanity check that images are channel first + _, c, h, w = batch[key].shape + assert c < h and c < w, f"expect channel first images, but instead {batch[key].shape}" + assert ( + batch[key].dtype == torch.float32 + ), f"expect torch.float32, but instead {batch[key].dtype=}" + # assert batch[key].max() <= 1, f"expect pixels lower than 1, but instead {batch[key].max()=}" + # assert batch[key].min() >= 0, f"expect pixels greater than 1, but instead {batch[key].min()=}" + stats_patterns[key] = "b c h w -> c 1 1" + elif batch[key].ndim == 2: + stats_patterns[key] = "b c -> c " + elif batch[key].ndim == 1: + stats_patterns[key] = "b -> 1" + else: + raise ValueError(f"{key}, {batch[key].shape}") + + return stats_patterns + + +def compute_stats(dataset, batch_size=8, num_workers=4, max_num_samples=None): + """Compute mean/std and min/max statistics of all data keys in a LeRobotDataset.""" + if max_num_samples is None: + max_num_samples = len(dataset) + + # for more info on why we need to set the same number of workers, see `load_from_videos` + stats_patterns = get_stats_einops_patterns(dataset, num_workers) + + # mean and std will be computed incrementally while max and min will track the running value. + mean, std, max, min = {}, {}, {}, {} + for key in stats_patterns: + mean[key] = torch.tensor(0.0).float() + std[key] = torch.tensor(0.0).float() + max[key] = torch.tensor(-float("inf")).float() + min[key] = torch.tensor(float("inf")).float() + + def create_seeded_dataloader(dataset, batch_size, seed): + generator = torch.Generator() + generator.manual_seed(seed) + dataloader = torch.utils.data.DataLoader( + dataset, + num_workers=num_workers, + batch_size=batch_size, + shuffle=True, + drop_last=False, + generator=generator, + ) + return dataloader + + # Note: Due to be refactored soon. The point of storing `first_batch` is to make sure we don't get + # surprises when rerunning the sampler. + first_batch = None + running_item_count = 0 # for online mean computation + dataloader = create_seeded_dataloader(dataset, batch_size, seed=1337) + for i, batch in enumerate( + tqdm( + dataloader, + total=ceil(max_num_samples / batch_size), + desc="Compute mean, min, max", + ) + ): + this_batch_size = len(batch["index"]) + running_item_count += this_batch_size + if first_batch is None: + first_batch = deepcopy(batch) + for key, pattern in stats_patterns.items(): + batch[key] = batch[key].float() + # Numerically stable update step for mean computation. + batch_mean = einops.reduce(batch[key], pattern, "mean") + # Hint: to update the mean we need x̄ₙ = (Nₙ₋₁x̄ₙ₋₁ + Bₙxₙ) / Nₙ, where the subscript represents + # the update step, N is the running item count, B is this batch size, x̄ is the running mean, + # and x is the current batch mean. Some rearrangement is then required to avoid risking + # numerical overflow. Another hint: Nₙ₋₁ = Nₙ - Bₙ. Rearrangement yields + # x̄ₙ = x̄ₙ₋₁ + Bₙ * (xₙ - x̄ₙ₋₁) / Nₙ + mean[key] = mean[key] + this_batch_size * (batch_mean - mean[key]) / running_item_count + max[key] = torch.maximum(max[key], einops.reduce(batch[key], pattern, "max")) + min[key] = torch.minimum(min[key], einops.reduce(batch[key], pattern, "min")) + + if i == ceil(max_num_samples / batch_size) - 1: + break + + first_batch_ = None + running_item_count = 0 # for online std computation + dataloader = create_seeded_dataloader(dataset, batch_size, seed=1337) + for i, batch in enumerate( + tqdm(dataloader, total=ceil(max_num_samples / batch_size), desc="Compute std") + ): + this_batch_size = len(batch["index"]) + running_item_count += this_batch_size + # Sanity check to make sure the batches are still in the same order as before. + if first_batch_ is None: + first_batch_ = deepcopy(batch) + for key in stats_patterns: + assert torch.equal(first_batch_[key], first_batch[key]) + for key, pattern in stats_patterns.items(): + batch[key] = batch[key].float() + # Numerically stable update step for mean computation (where the mean is over squared + # residuals).See notes in the mean computation loop above. + batch_std = einops.reduce((batch[key] - mean[key]) ** 2, pattern, "mean") + std[key] = std[key] + this_batch_size * (batch_std - std[key]) / running_item_count + + if i == ceil(max_num_samples / batch_size) - 1: + break + + for key in stats_patterns: + std[key] = torch.sqrt(std[key]) + + stats = {} + for key in stats_patterns: + stats[key] = { + "mean": mean[key], + "std": std[key], + "max": max[key], + "min": min[key], + } + return stats + + +class AgiBotDataset(LeRobotDataset): + def __init__( + self, + repo_id: str, + root: str | Path | None = None, + episodes: list[int] | None = None, + image_transforms: Callable | None = None, + delta_timestamps: dict[list[float]] | None = None, + tolerance_s: float = 1e-4, + download_videos: bool = True, + local_files_only: bool = False, + video_backend: str | None = None, + ): + super().__init__( + repo_id=repo_id, + root=root, + episodes=episodes, + image_transforms=image_transforms, + delta_timestamps=delta_timestamps, + tolerance_s=tolerance_s, + download_videos=download_videos, + local_files_only=local_files_only, + video_backend=video_backend, + ) + + def save_episode( + self, task: str, episode_data: dict | None = None, videos: dict | None = None + ) -> None: + """ + We rewrite this method to copy mp4 videos to the target position + """ + if not episode_data: + episode_buffer = self.episode_buffer + + episode_length = episode_buffer.pop("size") + episode_index = episode_buffer["episode_index"] + if episode_index != self.meta.total_episodes: + # TODO(aliberts): Add option to use existing episode_index + raise NotImplementedError( + "You might have manually provided the episode_buffer with an episode_index that doesn't " + "match the total number of episodes in the dataset. This is not supported for now." + ) + + if episode_length == 0: + raise ValueError( + "You must add one or several frames with `add_frame` before calling `add_episode`." + ) + + # Use our custom task indexing instead of LeRobot's built-in mechanism + task_index = getattr(self, "_custom_task_to_index", {}).get(task, 0) + + # Remove the 'task' key if it exists (it's passed as a parameter, not needed in buffer) + episode_buffer.pop("task", None) + + if not set(episode_buffer.keys()) == set(self.features): + raise ValueError() + + for key, ft in self.features.items(): + if key == "index": + episode_buffer[key] = np.arange( + self.meta.total_frames, self.meta.total_frames + episode_length + ) + elif key == "episode_index": + episode_buffer[key] = np.full((episode_length,), episode_index) + elif key == "task_index": + episode_buffer[key] = np.full((episode_length,), task_index) + elif ft["dtype"] in ["image", "video"]: + continue + elif ft["dtype"] == "string": + pass + elif len(ft["shape"]) == 1 and ft["shape"][0] == 1: + episode_buffer[key] = np.array(episode_buffer[key], dtype=ft["dtype"]) + elif len(ft["shape"]) == 1 and ft["shape"][0] > 1: + episode_buffer[key] = np.stack(episode_buffer[key]) + else: + raise ValueError(key) + + self._wait_image_writer() + self._save_episode_table(episode_buffer, episode_index) + + # Copy videos first before calling meta.save_episode which might try to read them + for key in self.meta.video_keys: + video_path = self.root / self.meta.get_video_file_path(episode_index, key) + episode_buffer[key] = video_path + video_path.parent.mkdir(parents=True, exist_ok=True) + # Copy video files to target location + shutil.copyfile(str(videos[key]), str(video_path)) + + try: + # FIX: Call meta.save_episode with correct parameters + # Note: We pass an empty task list to prevent duplicate entries in tasks.jsonl + # since we create our own tasks.jsonl file with proper indexing + self.meta.save_episode(episode_index, episode_length, [], {}) + except AttributeError as e: + if "'NoneType' object has no attribute 'items'" in str(e): + # Handle the episode stats issue - skip episode stats for now + print( + f"Warning: Episode stats computation failed, proceeding without stats " + f"for episode {episode_index}" + ) + # Just skip the problematic save_episode call - the core data is already saved + pass + else: + raise + if not episode_data: # Reset the buffer + self.episode_buffer = self.create_episode_buffer() + self.consolidated = False + + def consolidate(self, run_compute_stats: bool = True, keep_image_files: bool = False) -> None: + self.hf_dataset = self.load_hf_dataset() + self.episode_data_index = get_episode_data_index(self.meta.episodes, self.episodes) + check_timestamps_sync(self.hf_dataset, self.episode_data_index, self.fps, self.tolerance_s) + if len(self.meta.video_keys) > 0: + self.meta.write_video_info() + + if not keep_image_files: + img_dir = self.root / "images" + if img_dir.is_dir(): + shutil.rmtree(self.root / "images") + video_files = list(self.root.rglob("*.mp4")) + assert len(video_files) == self.num_episodes * len(self.meta.video_keys) + + parquet_files = list(self.root.rglob("*.parquet")) + assert len(parquet_files) == self.num_episodes + + if run_compute_stats: + self.stop_image_writer() + self.meta.stats = compute_stats(self) + serialized_stats = serialize_dict(self.meta.stats) + write_json(serialized_stats, self.root / STATS_PATH) + self.consolidated = True + else: + logging.warning( + "Skipping computation of the dataset statistics, dataset is not fully consolidated." + ) + + def add_frame(self, frame: dict) -> None: + """ + This function only adds the frame to the episode_buffer. Apart from images — which are written in a + temporary directory — nothing is written to disk. To save those frames, the 'save_episode()' method + then needs to be called. + """ + # TODO(aliberts, rcadene): Add sanity check for the input, check it's numpy or torch, + # check the dtype and shape matches, etc. + + if self.episode_buffer is None: + self.episode_buffer = self.create_episode_buffer() + + frame_index = self.episode_buffer["size"] + timestamp = frame.pop("timestamp") if "timestamp" in frame else frame_index / self.fps + self.episode_buffer["frame_index"].append(frame_index) + self.episode_buffer["timestamp"].append(timestamp) + + for key in frame: + if key not in self.features: + raise ValueError(key) + item = frame[key].numpy() if isinstance(frame[key], torch.Tensor) else frame[key] + self.episode_buffer[key].append(item) + + self.episode_buffer["size"] += 1 + + +def detect_dataset_format(src_path: str) -> str: + """Detect whether the dataset follows old or new format structure""" + src_path = Path(src_path) + + # Check for old format indicators + if (src_path / "task_info").exists() and (src_path / "proprio_stats").exists(): + return "old" + + # Check for new format indicators + # Look for pattern: job_id/robot_id/episode_id with aligned_joints.h5 + # The test_data/2810125 directory contains job_id directories + subdirs = [d for d in src_path.iterdir() if d.is_dir()] + if subdirs: + # Check if we have nested structure with aligned_joints.h5 + for job_dir in subdirs: # These are job_id directories like 3335477 + if not job_dir.is_dir(): + continue + for ( + robot_dir + ) in job_dir.iterdir(): # These are robot_id directories like A2D0015AB00061 + if not robot_dir.is_dir(): + continue + for ( + episode_dir + ) in robot_dir.iterdir(): # These are episode_id directories like 12052353 + if episode_dir.is_dir() and (episode_dir / "aligned_joints.h5").exists(): + return "new" + + return "unknown" + + +def load_local_dataset_old_format(episode_id: int, src_path: str, task_id: int) -> list | None: + """Load local dataset from old format and return a dict with observations and actions""" + + # --- Load task info for this specific call --- + task_json_path = Path(src_path) / f"task_info/task_{task_id}.json" + task_info_list = None + if task_json_path.exists(): + try: + with open(task_json_path, "r") as f: + task_info_list = json.load(f) + except json.JSONDecodeError: + print(f"Warning: Failed to decode JSON {task_json_path} for episode {episode_id}") + task_info_list = [] # Treat as empty if decode fails + else: + print(f"Warning: Task info JSON not found at {task_json_path} for episode {episode_id}") + task_info_list = [] # Treat as empty if not found + + # --- Find action_config for this episode_id --- + episode_action_config = None + if isinstance(task_info_list, list): # Check if loading succeeded and it's a list + for item in task_info_list: + # Ensure episode_id exists and compare as int + if "episode_id" in item and int(item["episode_id"]) == episode_id: + episode_action_config = item.get("label_info", {}).get("action_config") + break + + default_action_text = "N/A" + + ob_dir = Path(src_path) / f"observations/{task_id}/{episode_id}" + proprio_dir = Path(src_path) / f"proprio_stats/{task_id}/{episode_id}" + + with h5py.File(proprio_dir / "proprio_stats.h5") as f: + state_joint = np.array(f["state/joint/position"]) + state_effector = np.clip((np.array(f["state/effector/position"]) - 35.0) / (120.0 - 35.0), 0.0, 1.0) + state_head = np.array(f["state/head/position"]) + state_waist = np.array(f["state/waist/position"]) + action_joint = np.array(f["action/joint/position"]) + action_effector = np.clip((np.array(f["action/effector/position"]) - 35.0) / (120.0 - 35.0), 0.0, 1.0) + action_head = np.array(f["action/head/position"]) + action_waist = np.array(f["action/waist/position"]) + action_velocity = np.array(f["action/robot/velocity"]) + + # State (20 DOF): joint(14) + effector(2) + head(2) + waist(2) + states_value = np.hstack( + [ + state_joint, + state_effector, + state_head, + state_waist, + ] + ).astype(np.float32) + assert ( + action_joint.shape[0] == action_effector.shape[0] + ), f"shape of action_joint:{action_joint.shape};shape of action_effector:{action_effector.shape}" + # Action (22 DOF): joint(14) + effector(2) + head(2) + waist(2) + velocity(2) + action_value = np.hstack( + [ + action_joint, + action_effector, + action_head, + action_waist, + action_velocity, + ] + ).astype(np.float32) + + num_frames = len(states_value) + + # --- Create frame -> action_text mapping --- + frame_action_texts = [default_action_text] * num_frames + if episode_action_config: # Only proceed if config was found + for action in episode_action_config: + start = action.get("start_frame") + end = action.get("end_frame") + text = action.get("action_text", default_action_text) + + if start is None or end is None: + continue + + clamped_start = max(0, start) + clamped_end = min(num_frames, end) + for i in range(clamped_start, clamped_end): + frame_action_texts[i] = text + + frames = [ + { + "observation.state": states_value[i], + "action": action_value[i], + "annotation.language.action_text": [frame_action_texts[i]], # Add action_text here + } + for i in range(num_frames) + ] + + v_path = ob_dir / "videos" + videos = { + "observation.images.top_head": v_path / HEAD_COLOR, + "observation.images.hand_left": v_path / HAND_LEFT_COLOR, + "observation.images.hand_right": v_path / HAND_RIGHT_COLOR, + } + return frames, videos + + +def load_local_dataset_new_format(episode_path: str) -> list | None: + """Load local dataset from new format and return a dict with observations and actions""" + + episode_dir = Path(episode_path) + + # Load data info JSON + data_info_path = episode_dir / "data_info.json" + episode_action_config = None + default_action_text = "N/A" + + if data_info_path.exists(): + try: + with open(data_info_path, "r") as f: + data_info = json.load(f) + episode_action_config = data_info.get("label_info", {}).get("action_config") + except json.JSONDecodeError: + print(f"Warning: Failed to decode JSON {data_info_path}") + + # Load aligned joints data + joints_path = episode_dir / "aligned_joints.h5" + if not joints_path.exists(): + print(f"Warning: aligned_joints.h5 not found at {joints_path}") + return None + + with h5py.File(joints_path) as f: + # Extract state data - using same mapping as old format for compatibility + state_joint = np.array(f["state/joint/position"]) + state_head = np.array(f["state/head/position"]) + state_waist = np.array(f["state/waist/position"]) + + # For new format, use separate left/right effector position data + # Normalize from raw [35, 120] to [0, 1] + state_left_effector = np.array(f["state/left_effector/position"]) # Shape: (N, 1) + state_right_effector = np.array(f["state/right_effector/position"]) # Shape: (N, 1) + state_effector = np.clip( + np.column_stack([state_left_effector.flatten(), state_right_effector.flatten()]) - 35.0, + 0.0, 85.0 + ) / 85.0 # Shape: (N, 2), range [0, 1] + + # Extract action data + action_joint = np.array(f["action/joint/position"]) + action_head = np.array(f["action/head/position"]) + action_waist = np.array(f["action/waist/position"]) + + # For new format, use separate left/right effector position data for actions + # Normalize from raw [35, 120] to [0, 1] + action_left_effector = np.array(f["action/left_effector/position"]) # Shape: (N, 1) + action_right_effector = np.array(f["action/right_effector/position"]) # Shape: (N, 1) + action_effector = np.clip( + np.column_stack([action_left_effector.flatten(), action_right_effector.flatten()]) - 35.0, + 0.0, 85.0 + ) / 85.0 # Shape: (N, 2), range [0, 1] + + # Get robot velocity (N,) or (N, 2) + action_velocity_raw = np.array(f["action/robot/velocity"]) + if action_velocity_raw.ndim == 1: + # Scalar velocity: pad second component with zeros (linear_x only) + action_velocity = np.column_stack( + [action_velocity_raw, np.zeros_like(action_velocity_raw)] + ) + else: + action_velocity = action_velocity_raw[:, :2] + + # State (20 DOF): joint(14) + effector(2) + head(2) + waist(2) + states_value = np.hstack( + [ + state_joint, + state_effector, + state_head, + state_waist, + ] + ).astype(np.float32) + # Action (22 DOF): joint(14) + effector(2) + head(2) + waist(2) + velocity(2) + action_value = np.hstack( + [ + action_joint, + action_effector, + action_head, + action_waist, + action_velocity, + ] + ).astype(np.float32) + + num_frames = len(states_value) + + # --- Create frame -> action_text mapping --- + frame_action_texts = [default_action_text] * num_frames + if episode_action_config: # Only proceed if config was found + for action in episode_action_config: + start = action.get("start_frame") + end = action.get("end_frame") + # Use English action text if available, otherwise use Chinese + text = action.get("english_action_text") or action.get( + "action_text", default_action_text + ) + + if start is None or end is None: + continue + + clamped_start = max(0, start) + clamped_end = min(num_frames, end) + for i in range(clamped_start, clamped_end): + frame_action_texts[i] = text + + frames = [ + { + "observation.state": states_value[i], + "action": action_value[i], + "annotation.language.action_text": [frame_action_texts[i]], + } + for i in range(num_frames) + ] + + # Videos are at episode level in new format + videos = { + "observation.images.top_head": episode_dir / HEAD_COLOR, + "observation.images.hand_left": episode_dir / HAND_LEFT_COLOR, + "observation.images.hand_right": episode_dir / HAND_RIGHT_COLOR, + } + return frames, videos + + +def load_local_dataset( + episode_id: int, + src_path: str, + task_id: int = None, + episode_path: str = None, + format_type: str = "old", +) -> list | None: + """Load local dataset and return a dict with observations and actions + + Args: + episode_id: Episode ID (used for old format) + src_path: Source path (used for old format) + task_id: Task ID (used for old format) + episode_path: Full path to episode directory (used for new format) + format_type: "old" or "new" format + """ + if format_type == "old": + return load_local_dataset_old_format(episode_id, src_path, task_id) + elif format_type == "new": + return load_local_dataset_new_format(episode_path) + else: + raise ValueError(f"Unknown format type: {format_type}") + + +def get_task_instruction_old_format(task_json_path: str) -> str: + """Get task language instruction from old format""" + with open(task_json_path, "r") as f: + task_info = json.load(f) + task_name = task_info[0]["task_name"] + task_init_scene = task_info[0]["init_scene_text"] + task_instruction = f"{task_name}.{task_init_scene}" + print(f"Get Task Instruction <{task_instruction}>") + return task_instruction + + +def get_task_instruction_new_format(episode_paths: list) -> str: + """Get task language instruction from new format - use first episode's data_info.json""" + if not episode_paths: + return "Unknown Task" + + first_episode_path = Path(episode_paths[0]) + data_info_path = first_episode_path / "data_info.json" + + if data_info_path.exists(): + try: + with open(data_info_path, "r") as f: + data_info = json.load(f) + # Use English task name if available, otherwise use Chinese + task_name = data_info.get("english_task_name") or data_info.get( + "task_name", "Unknown Task" + ) + task_instruction = task_name + print( + f"Get Task Instruction <{task_instruction}> " + f"(english_task_name: {data_info.get('english_task_name')}, " + f"task_name: {data_info.get('task_name')})" + ) + return task_instruction + except json.JSONDecodeError: + print(f"Warning: Failed to decode JSON {data_info_path}") + + return "Unknown Task" + + +def get_task_instruction( + task_json_path: str = None, episode_paths: list = None, format_type: str = "old" +) -> str: + """Get task language instruction""" + if format_type == "old": + return get_task_instruction_old_format(task_json_path) + elif format_type == "new": + return get_task_instruction_new_format(episode_paths) + else: + raise ValueError(f"Unknown format type: {format_type}") + + +def load_new_format_episode(episode_path): + """Helper function for multiprocessing - load new format episode""" + return load_local_dataset( + episode_id=0, src_path="", episode_path=episode_path, format_type="new" + ) + + +def create_tasks_jsonl(tgt_path: str, repo_id: str, task_name: str, all_action_texts: set) -> dict: + """Create tasks.jsonl file with unique task names and action texts.""" + meta_path = os.path.join(tgt_path, repo_id, "meta") + os.makedirs(meta_path, exist_ok=True) + + tasks_jsonl_path = os.path.join(meta_path, "tasks.jsonl") + + # Create a list of unique tasks combining task name and action texts + # Remove task_name from action texts if it exists to avoid duplicates + unique_action_texts = all_action_texts - {task_name} + + tasks = [task_name] # First entry is the main task + tasks.extend(sorted(unique_action_texts)) # Then all unique action texts (excluding task name) + + # Check if file already exists and has correct content + should_write = True + if os.path.exists(tasks_jsonl_path): + try: + with open(tasks_jsonl_path, "r") as f: + existing_content = f.read().strip() + + # Generate expected content + expected_lines = [] + for i, task in enumerate(tasks): + task_entry = {"task_index": i, "task": task} + expected_lines.append(json.dumps(task_entry)) + expected_content = "\n".join(expected_lines) + + # If content matches, don't rewrite + if existing_content == expected_content: + should_write = False + except Exception as e: + print(f"Warning: Failed to read tasks.jsonl: {e}") + # If there's any issue reading, we'll rewrite + pass + + if should_write: + # Write tasks.jsonl (overwrite to ensure clean content) + with open(tasks_jsonl_path, "w") as f: + for i, task in enumerate(tasks): + task_entry = {"task_index": i, "task": task} + f.write(json.dumps(task_entry) + "\n") + + print(f"Created tasks.jsonl with {len(tasks)} entries at {tasks_jsonl_path}") + else: + print(f"tasks.jsonl already exists with correct content at {tasks_jsonl_path}") + + # Create mapping for lookups + task_to_index = {task: i for i, task in enumerate(tasks)} + + return task_to_index + + +def main( + src_path: str, + tgt_path: str, + task_id: int = None, + repo_id: str = None, + task_info_json: str = None, + debug: bool = False, + chunk_size: int = 10, +): + # Detect dataset format + format_type = detect_dataset_format(src_path) + print(f"Detected dataset format: {format_type}") + + if format_type == "unknown": + raise ValueError(f"Unable to detect dataset format for path: {src_path}") + + # Collect all unique action texts first + all_action_texts = set() + + # Initialize dataset + if not repo_id: + if format_type == "old": + repo_id = f"agibotworld/task_{task_id}" + else: # new format + # Use the top-level directory name as task_id + task_id = Path(src_path).name + repo_id = f"agibotworld/task_{task_id}" + + dataset = AgiBotDataset.create( + repo_id=repo_id, + root=f"{tgt_path}/{repo_id}", + fps=30, + robot_type="a2d", + features=FEATURES, + ) + + if format_type == "old": + # Old format processing + task_name = get_task_instruction(task_json_path=task_info_json, format_type="old") + + all_subdir = sorted( + [f.as_posix() for f in Path(src_path).glob(f"observations/{task_id}/*") if f.is_dir()] + ) + + if debug: + all_subdir = all_subdir[:2] + + # Get all episode id + all_subdir_eids = [int(Path(path).name) for path in all_subdir] + all_subdir_episode_desc = [task_name] * len(all_subdir_eids) + + # First pass: collect all unique action texts + print("Collecting unique action texts...") + for episode_id in tqdm(all_subdir_eids, desc="Scanning for action texts"): + frames_data, _ = load_local_dataset( + episode_id, src_path=src_path, task_id=task_id, format_type="old" + ) + if frames_data: + for frame in frames_data: + action_text = frame["annotation.language.action_text"][0] + all_action_texts.add(action_text) + + # Create tasks.jsonl with all unique texts + task_to_index = create_tasks_jsonl(tgt_path, repo_id, task_name, all_action_texts) + + # Pass the task mapping to the dataset for custom indexing + dataset._custom_task_to_index = task_to_index + + # Process in chunks to reduce memory usage + for chunk_start in tqdm( + range(0, len(all_subdir_eids), chunk_size), desc="Processing chunks" + ): + chunk_end = min(chunk_start + chunk_size, len(all_subdir_eids)) + chunk_eids = all_subdir_eids[chunk_start:chunk_end] + chunk_descs = all_subdir_episode_desc[chunk_start:chunk_end] + + # Process only this chunk + if debug: + raw_datasets_chunk = [ + load_local_dataset( + subdir, src_path=src_path, task_id=task_id, format_type="old" + ) + for subdir in tqdm(chunk_eids, desc="Loading chunk data") + ] + else: + raw_datasets_chunk = process_map( + partial( + load_local_dataset, src_path=src_path, task_id=task_id, format_type="old" + ), + chunk_eids, + max_workers=os.cpu_count() // 2, + desc=f"Loading chunk {chunk_start//chunk_size + 1}/" + f"{(len(all_subdir_eids) + chunk_size - 1)//chunk_size}", + ) + + # Filter out None results + valid_datasets = [ + (ds, desc) for ds, desc in zip(raw_datasets_chunk, chunk_descs) if ds is not None + ] + + # Process each dataset in the chunk + for raw_dataset, episode_desc in tqdm( + valid_datasets, desc="Processing episodes in chunk" + ): + for raw_dataset_sub in tqdm(raw_dataset[0], desc="Processing frames", leave=False): + # Convert string annotation to int index + action_text = raw_dataset_sub["annotation.language.action_text"][0] + raw_dataset_sub["annotation.language.action_text"] = [ + task_to_index[action_text] + ] + raw_dataset_sub["annotation.agibot.tasks"] = [task_to_index[episode_desc]] + dataset.add_frame(raw_dataset_sub) + dataset.save_episode(task=episode_desc, videos=raw_dataset[1]) + + # Clear memory after each chunk + raw_datasets_chunk = None + valid_datasets = None + gc.collect() + + else: # new format + # Find all episode directories + all_episode_paths = [] + src_path = Path(src_path) + + # Walk through job_id/robot_id/episode_id structure (test_data/2810125 contains job_id directories) + for job_dir in src_path.iterdir(): # These are job_id directories like 3335477 + if not job_dir.is_dir(): + continue + for ( + robot_dir + ) in job_dir.iterdir(): # These are robot_id directories like A2D0015AB00061 + if not robot_dir.is_dir(): + continue + for ( + episode_dir + ) in robot_dir.iterdir(): # These are episode_id directories like 12052353 + if episode_dir.is_dir() and (episode_dir / "aligned_joints.h5").exists(): + all_episode_paths.append(str(episode_dir)) + + all_episode_paths = sorted(all_episode_paths) + + if debug: + all_episode_paths = all_episode_paths[:2] + + # Get task name from first episode + task_name = get_task_instruction(episode_paths=all_episode_paths, format_type="new") + all_episode_descs = [task_name] * len(all_episode_paths) + + # First pass: collect all unique action texts + print("Collecting unique action texts...") + for episode_path in tqdm(all_episode_paths, desc="Scanning for action texts"): + frames_data, _ = load_local_dataset( + episode_id=0, src_path="", episode_path=episode_path, format_type="new" + ) + if frames_data: + for frame in frames_data: + action_text = frame["annotation.language.action_text"][0] + all_action_texts.add(action_text) + + # Create tasks.jsonl with all unique texts + task_to_index = create_tasks_jsonl(tgt_path, repo_id, task_name, all_action_texts) + + # Pass the task mapping to the dataset for custom indexing + dataset._custom_task_to_index = task_to_index + + # Process in chunks to reduce memory usage + for chunk_start in tqdm( + range(0, len(all_episode_paths), chunk_size), desc="Processing chunks" + ): + chunk_end = min(chunk_start + chunk_size, len(all_episode_paths)) + chunk_paths = all_episode_paths[chunk_start:chunk_end] + chunk_descs = all_episode_descs[chunk_start:chunk_end] + + # Process only this chunk + if debug: + raw_datasets_chunk = [ + load_local_dataset( + episode_id=0, src_path="", episode_path=episode_path, format_type="new" + ) + for episode_path in tqdm(chunk_paths, desc="Loading chunk data") + ] + else: + raw_datasets_chunk = process_map( + load_new_format_episode, + chunk_paths, + max_workers=os.cpu_count() // 2, + desc=f"Loading chunk {chunk_start//chunk_size + 1}/" + f"{(len(all_episode_paths) + chunk_size - 1)//chunk_size}", + ) + + # Filter out None results + valid_datasets = [ + (ds, desc) for ds, desc in zip(raw_datasets_chunk, chunk_descs) if ds is not None + ] + + # Process each dataset in the chunk + for raw_dataset, episode_desc in tqdm( + valid_datasets, desc="Processing episodes in chunk" + ): + for raw_dataset_sub in tqdm(raw_dataset[0], desc="Processing frames", leave=False): + # Convert string annotation to int index + action_text = raw_dataset_sub["annotation.language.action_text"][0] + raw_dataset_sub["annotation.language.action_text"] = [ + task_to_index[action_text] + ] + raw_dataset_sub["annotation.agibot.tasks"] = [task_to_index[episode_desc]] + dataset.add_frame(raw_dataset_sub) + dataset.save_episode(task=episode_desc, videos=raw_dataset[1]) + + # Clear memory after each chunk + raw_datasets_chunk = None + valid_datasets = None + gc.collect() + + # Only consolidate at the end + try: + dataset.consolidate(run_compute_stats=False) + except Exception as e: + print(f"Warning: Consolidation failed with error: {e}") + print("Dataset conversion completed but may not be fully consolidated.") + + # Generate modality.json file in the dataset meta directory + meta_path = os.path.join(tgt_path, repo_id, "meta") + os.makedirs(meta_path, exist_ok=True) + generate_modality_json(meta_path) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Convert AgiBot dataset to LeRobot format. Supports both old and new format datasets." + ) + parser.add_argument( + "--src_path", type=str, required=True, help="Path to source dataset directory" + ) + parser.add_argument( + "--task_id", + type=str, + required=False, + help="Task ID (required for old format, optional for new format)", + ) + parser.add_argument( + "--tgt_path", type=str, required=True, help="Path to target output directory" + ) + parser.add_argument( + "--repo_id", + type=str, + required=False, + help="Repository ID for the dataset (auto-generated if not provided)", + ) + parser.add_argument( + "--debug", action="store_true", help="Process only first 2 episodes for debugging" + ) + parser.add_argument( + "--chunk_size", type=int, default=10, help="Number of episodes to process at once" + ) + args = parser.parse_args() + + # Detect format first + format_type = detect_dataset_format(args.src_path) + print(f"Auto-detected format: {format_type}") + + # Validate arguments based on format + if format_type == "old": + if not args.task_id: + parser.error("--task_id is required for old format datasets") + + task_id = int(args.task_id) + json_file = f"{args.src_path}/task_info/task_{args.task_id}.json" + + if not Path(json_file).exists(): + parser.error(f"Cannot find task info file: {json_file}") + + main( + src_path=args.src_path, + tgt_path=args.tgt_path, + task_id=task_id, + repo_id=args.repo_id, + task_info_json=json_file, + debug=args.debug, + chunk_size=args.chunk_size, + ) + + elif format_type == "new": + main( + src_path=args.src_path, + tgt_path=args.tgt_path, + task_id=args.task_id, + repo_id=args.repo_id, + task_info_json=None, + debug=args.debug, + chunk_size=args.chunk_size, + ) + + else: + parser.error( + f"Unknown dataset format. Please check the directory structure at: {args.src_path}" + ) + diff --git a/scripts/data/convert_droid.py b/scripts/data/convert_droid.py new file mode 100644 index 0000000000000000000000000000000000000000..2a548e196f6912b341c3906c02b9dcd89584f5af --- /dev/null +++ b/scripts/data/convert_droid.py @@ -0,0 +1,679 @@ +""" +Convert DROID 1.0.1 (RLDS/TFDS format) to LeRobot format with idle frame filtering. + +This script takes the raw DROID dataset in RLDS format, applies idle frame filtering +using a pre-computed JSON file of non-idle frame ranges, filters out failed episodes +and episodes without language annotations, and outputs the dataset in LeRobot v2.0 format. + +The idle filtering is based on Physical Intelligence's approach (see openpi): + https://github.com/Physical-Intelligence/openpi/blob/main/examples/droid/README_train.md + +The pre-computed idle filter ranges can be downloaded from: + gsutil cp gs://openpi-assets/droid/droid_sample_ranges_v1_0_1.json + +Usage: + python scripts/data/convert_droid.py \\ + --keep-ranges-path \\ + [--fps 15] [--first-n N] [-n 16] [--filter-failed] + +Example: + # Download DROID 1.0.1 raw dataset + gsutil -m cp -r gs://gresearch/robotics/droid/1.0.1 ./data/droid/1.0.1 + + # Download idle filter ranges from openpi + gsutil cp gs://openpi-assets/droid/droid_sample_ranges_v1_0_1.json ./data/keep_ranges.json + + # Run conversion + python scripts/data/convert_droid.py ./data/droid/1.0.1 ./data/droid_lerobot \\ + --keep-ranges-path ./data/keep_ranges.json --filter-failed + +Original dataset structure (RLDS): + - 3 camera views: exterior_image_1_left, exterior_image_2_left, wrist_image_left + - State: cartesian_position (6), gripper_position (1), joint_position (7) + - Action: cartesian_position (6), cartesian_velocity (6), gripper_position (1), + gripper_velocity (1), joint_position (7), joint_velocity (7) + - Language instructions (up to 3 per episode) + +Credits: + - Original conversion script by Loic Magne (NVIDIA) + - Idle filtering by Scott Reed (NVIDIA), based on Physical Intelligence's approach +""" + +from concurrent.futures import ProcessPoolExecutor, as_completed +import json +import multiprocessing as mp +import os +from pathlib import Path + +import av +import numpy as np +import polars as pl +import tensorflow as tf +import tensorflow_datasets as tfds +import torch +import tqdm + +# Limit thread counts to avoid oversubscription in multiprocessing +os.environ["MKL_NUM_THREADS"] = "1" +os.environ["NUMEXPR_NUM_THREADS"] = "1" +os.environ["OMP_NUM_THREADS"] = "1" +os.environ["MPI_NUM_THREADS"] = "1" +os.environ["TF_NUM_INTRAOP_THREADS"] = "1" +os.environ["TF_NUM_INTEROP_THREADS"] = "1" +os.environ["OPENBLAS_NUM_THREADS"] = "1" +os.environ["VECLIB_MAXIMUM_THREADS"] = "1" + +tf.get_logger().setLevel("WARN") +tf.config.threading.set_inter_op_parallelism_threads(1) +tf.config.threading.set_intra_op_parallelism_threads(1) +tf.config.set_soft_device_placement(True) + + +def tf_to_torch(data): + return torch.from_numpy(data.numpy()) + + +def tf_img_convert(img): + if img.dtype == tf.string: + img = tf.io.decode_image(img, expand_animations=False, dtype=tf.uint8) + elif img.dtype != tf.uint8: + raise ValueError(f"Unsupported image dtype: found with dtype {img.dtype}") + return img.numpy() + + +def _broadcast_metadata_rlds(i: tf.Tensor, traj: dict) -> dict: + steps = traj.pop("steps") + traj_len = tf.shape(tf.nest.flatten(steps)[0])[0] + metadata = tf.nest.map_structure(lambda x: tf.repeat(x, traj_len), traj) + traj = {**steps, "traj_metadata": metadata} + traj["_len"] = tf.repeat(traj_len, traj_len) + traj["_traj_index"] = tf.repeat(i, traj_len) + traj["_frame_index"] = tf.range(traj_len) + return traj + + +def concat_state_or_action(modality_dict, keys, compute_concat_info=False): + arrays = [] + if compute_concat_info: + concat_info = {} + start_index = 0 + for key in keys: + array = tf_to_torch(modality_dict[key]) + arrays.append(array) + if compute_concat_info: + D = array.shape[1] + data_dtype = array.numpy().dtype + if np.issubdtype(data_dtype, bool): + data_dtype = "int64" + data_range = [0, 1] + else: + data_dtype = data_dtype.name + data_range = None + concat_info[key] = { + "start": start_index, + "end": start_index + D, + } + if data_dtype != "float64": + concat_info[key]["dtype"] = data_dtype + if data_range is not None: + concat_info[key]["range"] = data_range + start_index += D + concatenated = torch.cat(arrays, dim=1) + ret_dict = {} + if compute_concat_info: + ret_dict["concat_info"] = concat_info + ret_dict["concatenated"] = concatenated + return ret_dict + + +def encode_video(frames: np.ndarray, output_path: Path, fps: int) -> None: + """Encode a sequence of frames to a video file using PyAV.""" + options = { + "threads": "1", + "thread_type": "slice", + "preset": "ultrafast", + "tune": "zerolatency", + "crf": "23", + } + + container = av.open(str(output_path), mode="w") + stream = container.add_stream("h264", rate=fps, options=options) + stream.width = frames.shape[2] + stream.height = frames.shape[1] + stream.pix_fmt = "yuv420p" + + video_frame = av.VideoFrame(width=stream.width, height=stream.height, format="rgb24") + frame_array = video_frame.to_ndarray(format="rgb24") + + for frame in frames: + frame_array[:] = frame + packet = stream.encode(video_frame) + container.mux(packet) + + packet = stream.encode(None) + container.mux(packet) + container.close() + + +def process_tfrecord( + ith_shard, + raw_dir, + output_path, + fps, + all_tasks, + state_keys, + action_keys, + lang_keys, + image_keys, + start_episode_idx, + kept_registry, + keep_ranges_path, +): + config = tfds.ReadConfig( + try_autocache=False, + num_parallel_calls_for_decode=1, + num_parallel_calls_for_interleave_files=1, + interleave_cycle_length=1, + shuffle_reshuffle_each_iteration=False, + ) + + ds_builder = tfds.builder_from_directory(str(raw_dir)) + dataset = ds_builder.as_dataset( + split=f"train[{ith_shard}shard]", + decoders={"steps": tfds.decode.SkipDecoding()}, + read_config=config, + ) + + dataset = dataset.enumerate().map(_broadcast_metadata_rlds) + all_keep_ranges = json.load(open(keep_ranges_path, "r")) + + episodes_data = [] + for local_idx, episode in enumerate(dataset): + + # Add keep frame info to episode. + file_path = ( + episode["traj_metadata"]["episode_metadata"]["file_path"][0].numpy().decode("utf-8") + ) + recording_folderpath = ( + episode["traj_metadata"]["episode_metadata"]["recording_folderpath"][0] + .numpy() + .decode("utf-8") + ) + idle_key = f"{recording_folderpath}--{file_path}" + keep_ranges = all_keep_ranges[idle_key] + + global_episode_idx = start_episode_idx + local_idx + + # check if the episode has been filtered + if global_episode_idx not in kept_registry: + continue + + episode_idx = kept_registry[global_episode_idx] + episode_data = process_sample( + episode_idx, + episode, + output_path, + fps, + all_tasks, + state_keys, + action_keys, + lang_keys, + image_keys, + keep_ranges, + ) + episodes_data.append(episode_data) + return episodes_data + + +def process_sample( + ep_idx, + episode, + output_path, + fps, + all_tasks, + state_keys, + action_keys, + lang_keys, + image_keys, + keep_ranges, +): + chunk_idx = ep_idx // 1000 + + # Create chunk directory + (output_path / f"data/chunk-{chunk_idx:03d}").mkdir(parents=True, exist_ok=True) + for img_key in image_keys: + (output_path / f"videos/chunk-{chunk_idx:03d}/observation.images.{img_key}").mkdir( + parents=True, exist_ok=True + ) + + # Use concat_state_or_action for state and action + state_dict = concat_state_or_action(episode["observation"], state_keys) + action_dict = concat_state_or_action(episode["action_dict"], action_keys) + + # Count number of non-idle frames. + num_frames = len(episode["observation"][state_keys[0]]) + actual_num_frames = 0 + for start_ix, end_ix in keep_ranges: + actual_num_frames += end_ix - start_ix + + # Build episode data dictionary + episode_dict = { + "observation.state": state_dict["concatenated"].numpy(), + "action": action_dict["concatenated"].numpy(), + "next.reward": tf_to_torch(episode["reward"]).numpy(), + "next.done": tf_to_torch(episode["is_last"]).numpy(), + "is_terminal": tf_to_torch(episode["is_terminal"]).numpy(), + "is_first": tf_to_torch(episode["is_first"]).numpy(), + "discount": tf_to_torch(episode["discount"]).numpy(), + "timestamp": np.arange(actual_num_frames) / fps, + "episode_index": np.full(actual_num_frames, ep_idx), + "frame_index": np.arange(actual_num_frames), + } + + # Initialize all annotation columns with default value + for lang_key in lang_keys: + episode_dict[f"annotation.language.{lang_key}"] = np.full( + num_frames, all_tasks["not provided"], dtype=np.int64 + ) + + # Add language instruction indices to parquet + episode_tasks = [] + for lang_key in lang_keys: + if lang_key in episode: + task = episode[lang_key][0].numpy().decode("utf-8") + if task and len(task) > 1: + episode_tasks.append(task) + task_idx = all_tasks[task] + episode_dict[f"annotation.language.{lang_key}"] = np.full( + num_frames, task_idx, dtype=np.int64 + ) + + # Set task_index to match the first language instruction annotation + episode_dict["task_index"] = episode_dict[f"annotation.language.{lang_keys[0]}"].copy() + + # Filter idle frames from episode_dict. + for key in episode_dict: + if key in ["timestamp", "episode_index", "frame_index"]: + continue + tensor_parts = [] + for start_ix, end_ix in keep_ranges: + tensor_parts.append(episode_dict[key][start_ix:end_ix]) + episode_dict[key] = np.concatenate(tensor_parts, axis=0) + + # Filter idle frames from observation images. + for img_key in image_keys: + video_parts = [] + all_frames = np.stack( + [tf_img_convert(episode["observation"][img_key][i]) for i in range(num_frames)] + ) + for start_ix, end_ix in keep_ranges: + video_parts.append(all_frames[start_ix:end_ix]) + new_video = np.concatenate(video_parts, axis=0) + assert new_video.shape[0] == actual_num_frames + episode["observation"][img_key] = new_video + + # Save to parquet using polars + df = pl.DataFrame(episode_dict) + parquet_path = output_path / f"data/chunk-{chunk_idx:03d}/episode_{ep_idx:06d}.parquet" + df.write_parquet(parquet_path) + + # Process videos for each image key + for img_key in image_keys: + frames = episode["observation"][img_key] + video_path = ( + output_path + / f"videos/chunk-{chunk_idx:03d}/observation.images.{img_key}/episode_{ep_idx:06d}.mp4" + ) + encode_video(frames, video_path, fps) + + episode_data = { + "episode_index": ep_idx, + "tasks": episode_tasks, + "length": actual_num_frames, + "success": bool(np.any(tf_to_torch(episode["reward"]).numpy() != 0)), + } + return episode_data + + +def convert_droid_dataset( + raw_dir: str, + output_dir: str, + keep_ranges_path: str, + fps: int = 15, + first_n: int | None = None, + max_workers: int = 16, + filter_failed: bool = False, +): + """ + Convert DROID 1.0.1 RLDS dataset to LeRobot format with idle filtering. + + Args: + raw_dir: Path to raw DROID RLDS dataset (e.g., ./data/droid/1.0.1) + output_dir: Path to output directory for LeRobot dataset + keep_ranges_path: Path to JSON file containing idle filter ranges. + Download from: gsutil cp gs://openpi-assets/droid/droid_sample_ranges_v1_0_1.json + fps: Frames per second for output videos + first_n: Only process the first N tfrecord shards (for debugging) + max_workers: Max workers for multiprocessing + filter_failed: Whether to filter out failed episodes (all zero rewards) + """ + output_path = Path(output_dir) + + # Validate keep_ranges_path exists + if not os.path.exists(keep_ranges_path): + raise FileNotFoundError( + f"Keep ranges file not found: {keep_ranges_path}\n" + "Download it with: gsutil cp gs://openpi-assets/droid/droid_sample_ranges_v1_0_1.json " + ) + + # Load dataset + config = tfds.ReadConfig( + try_autocache=False, + num_parallel_calls_for_decode=1, + num_parallel_calls_for_interleave_files=1, + interleave_cycle_length=1, + shuffle_reshuffle_each_iteration=False, + ) + ds_builder = tfds.builder_from_directory(str(raw_dir)) + split_str = f"train[:{first_n}shard]" if first_n is not None else "train" + dataset = ds_builder.as_dataset( + split=split_str, + decoders={"steps": tfds.decode.SkipDecoding()}, + read_config=config, + ) + dataset_info = ds_builder.info + dataset = dataset.enumerate().map(_broadcast_metadata_rlds) + + # Extract keys + image_keys = [] + state_keys = [ + "cartesian_position", + "gripper_position", + "joint_position", + ] + action_keys = [ + "cartesian_position", + "cartesian_velocity", + "gripper_position", + "gripper_velocity", + "joint_position", + "joint_velocity", + ] + lang_keys = [ + "language_instruction", + "language_instruction_2", + "language_instruction_3", + ] + + observation_info = dataset_info.features["steps"]["observation"] + for key in observation_info: + if len(observation_info[key].shape) == 3: + if observation_info[key].dtype == tf.uint8: + image_keys.append(key) + else: + assert key in state_keys, f"{key=}, {state_keys=}" + + print(f"Found image keys: {image_keys}") + print(f"Using state keys: {state_keys}") + print(f"Using action keys: {action_keys}") + + (output_path / "meta").mkdir(parents=True, exist_ok=True) + + # Get concat info for modality.json from first episode + first_episode = next(iter(dataset)) + state_info = concat_state_or_action( + first_episode["observation"], state_keys, compute_concat_info=True + ) + action_info = concat_state_or_action( + first_episode["action_dict"], action_keys, compute_concat_info=True + ) + + # Generate modality.json + modality_config = { + "state": state_info["concat_info"], + "action": action_info["concat_info"], + "video": {k: {"original_key": f"observation.images.{k}"} for k in image_keys}, + "annotation": {f"language.{lang_key}": {} for lang_key in lang_keys}, + } + + with open(output_path / "meta/modality.json", "w") as f: + json.dump(modality_config, f, indent=4) + + # Get file instructions from TFDS + ds_builder = tfds.builder_from_directory(str(raw_dir)) + file_instructions = ds_builder.info.splits["train"].file_instructions + if first_n is not None: + file_instructions = file_instructions[:first_n] + + # First pass: collect unique tasks and determine which episodes to keep + all_tasks = {} # task string -> task index + task_counter = 0 + print(f"First pass: collecting unique tasks from {len(dataset)} episodes") + # Add a default "not provided" task + all_tasks["not provided"] = task_counter + task_counter += 1 + + # kept_registry maps global episode index -> filtered episode index + kept_registry = {} + kept_count = 0 + all_keep_ranges = json.load(open(keep_ranges_path, "r")) + + for i, episode in enumerate(tqdm.tqdm(dataset)): + # filter out failed episodes + filtered = False + if filter_failed: + if not np.any(tf_to_torch(episode["reward"]).numpy() != 0): + filtered = True + + # Check language annotations + has_lang = False + for lang_key in lang_keys: + if lang_key in episode: + task = episode[lang_key][0].numpy().decode("utf-8") + if task and (len(task) > 1) and task not in all_tasks: + has_lang = True + all_tasks[task] = task_counter + task_counter += 1 + if not has_lang: + # Do not include episodes missing language annotations + filtered = True + + # Filter out episodes that are only idle + file_path = ( + episode["traj_metadata"]["episode_metadata"]["file_path"][0].numpy().decode("utf-8") + ) + recording_folderpath = ( + episode["traj_metadata"]["episode_metadata"]["recording_folderpath"][0] + .numpy() + .decode("utf-8") + ) + idle_key = f"{recording_folderpath}--{file_path}" + keep_ranges = all_keep_ranges[idle_key] + if len(keep_ranges) == 0: + filtered = True + + if not filtered: + kept_registry[i] = kept_count + kept_count += 1 + + print(f"Kept {len(kept_registry)}/{len(dataset)} episodes") + + # Write tasks.jsonl + with open(output_path / "meta/tasks.jsonl", "w") as f: + for task, task_idx in all_tasks.items(): + f.write(json.dumps({"task_index": task_idx, "task": task}) + "\n") + + if max_workers > 1: + # Calculate process args with cumulative indices + cumsum = 0 + process_args = [] + for i, instruction in enumerate(file_instructions): + args = ( + i, + raw_dir, + output_path, + fps, + all_tasks, + state_keys, + action_keys, + lang_keys, + image_keys, + cumsum, + kept_registry, + keep_ranges_path, + ) + process_args.append(args) + cumsum += instruction.examples_in_shard + + ctx = mp.get_context("spawn") + with ProcessPoolExecutor(mp_context=ctx, max_workers=max_workers) as executor: + futures = [executor.submit(process_tfrecord, *args) for args in process_args] + episodes_data = [] + for future in tqdm.tqdm(as_completed(futures), total=len(futures)): + episodes_data.extend(future.result()) + else: + episodes_data = [] + cumsum = 0 + for i, instruction in enumerate(file_instructions): + episodes_data.extend( + process_tfrecord( + i, + raw_dir, + output_path, + fps, + all_tasks, + state_keys, + action_keys, + lang_keys, + image_keys, + cumsum, + kept_registry, + keep_ranges_path, + ) + ) + cumsum += instruction.examples_in_shard + + # Order episodes by episode index + episodes_data = sorted(episodes_data, key=lambda x: x["episode_index"]) + + # Generate episodes.jsonl + with open(output_path / "meta/episodes.jsonl", "w") as f: + for episode in episodes_data: + f.write(json.dumps(episode) + "\n") + + # Generate info.json + ds_length = len(episodes_data) + num_chunks = (ds_length // 1000) + (1 if ds_length % 1000 else 0) + info = { + "codebase_version": "v2.0", + "robot_type": "droid", + "total_episodes": ds_length, + "total_frames": sum(ep["length"] for ep in episodes_data), + "total_tasks": len(all_tasks), + "total_videos": len(image_keys), + "total_chunks": num_chunks, + "chunks_size": 1000, + "fps": fps, + "splits": {"train": "0:100"}, + "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": { + # Video features + **{ + f"observation.images.{k}": { + "dtype": "video", + "shape": list(tf_img_convert(first_episode["observation"][k][0]).shape), + "names": ["height", "width", "channel"], + "video_info": { + "video.fps": fps, + "video.codec": "h264", + "video.pix_fmt": "yuv420p", + "video.is_depth_map": False, + "has_audio": False, + }, + } + for k in image_keys + }, + # State feature + "observation.state": { + "dtype": "float64", + "shape": [state_info["concatenated"].shape[1]], + "names": state_keys, + }, + # Action feature + "action": { + "dtype": "float64", + "shape": [action_info["concatenated"].shape[1]], + "names": action_keys, + }, + # Single value features + "timestamp": {"dtype": "float64", "shape": [1]}, + "task_index": {"dtype": "int64", "shape": [1]}, + "episode_index": {"dtype": "int64", "shape": [1]}, + "index": {"dtype": "int64", "shape": [1]}, + "next.reward": {"dtype": "float64", "shape": [1]}, + "next.done": {"dtype": "bool", "shape": [1]}, + "is_terminal": {"dtype": "bool", "shape": [1]}, + "is_first": {"dtype": "bool", "shape": [1]}, + "discount": {"dtype": "float64", "shape": [1]}, + # Language annotation features + **{f"annotation.language.{k}": {"dtype": "int64", "shape": [1]} for k in lang_keys}, + }, + } + + with open(output_path / "meta/info.json", "w") as f: + json.dump(info, f, indent=4) + + # Sanity check: chunk directories should contain exactly 1000 episodes (except last) + for i in range(num_chunks): + chunk_path = output_path / f"data/chunk-{i:03d}" + episodes = list(chunk_path.glob("episode_*.parquet")) + assert ( + len(episodes) == 1000 if i != num_chunks - 1 else len(episodes) <= 1000 + ), f"chunk-{i:03d} contains {len(episodes)} episodes" + + for img_key in image_keys: + img_path = output_path / f"videos/chunk-{i:03d}/observation.images.{img_key}" + episodes = list(img_path.glob("episode_*.mp4")) + assert ( + len(episodes) == 1000 if i != num_chunks - 1 else len(episodes) <= 1000 + ), f"{img_path} contains {len(episodes)} episodes" + + print("Sanity check passed.") + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser( + description="Convert DROID 1.0.1 (RLDS) to LeRobot format with idle filtering." + ) + parser.add_argument("raw_dir", help="Path to raw DROID RLDS dataset (e.g., ./data/droid/1.0.1)") + parser.add_argument("output_dir", help="Path to output directory for LeRobot dataset") + parser.add_argument( + "--keep-ranges-path", + required=True, + help="Path to idle filter JSON file. Download with: " + "gsutil cp gs://openpi-assets/droid/droid_sample_ranges_v1_0_1.json ", + ) + parser.add_argument("--fps", type=int, default=15, help="Frames per second for videos") + parser.add_argument( + "--first-n", type=int, help="Only convert first N tfrecord shards (for debugging)" + ) + parser.add_argument("-n", type=int, default=16, help="Max workers for multiprocessing") + parser.add_argument( + "--filter-failed", + action="store_true", + help="Whether to filter out failed episodes (i.e., episodes with all zero rewards)", + ) + args = parser.parse_args() + + convert_droid_dataset( + args.raw_dir, + args.output_dir, + args.keep_ranges_path, + args.fps, + args.first_n, + args.n, + args.filter_failed, + ) \ No newline at end of file diff --git a/scripts/data/convert_lerobot_to_gear.py b/scripts/data/convert_lerobot_to_gear.py new file mode 100644 index 0000000000000000000000000000000000000000..2e1542e1c7a50bdcb471626d183d992dd3f5b16c --- /dev/null +++ b/scripts/data/convert_lerobot_to_gear.py @@ -0,0 +1,634 @@ +""" +Convert a standard LeRobot v2 dataset to the GEAR/DreamZero training format. + +This script takes a dataset collected with LeRobot v2 and generates/augments the +metadata files required by DreamZero's training pipeline: + + - meta/modality.json (state/action/video/annotation key mapping) + - meta/embodiment.json (embodiment tag for the training pipeline) + - meta/stats.json (dataset-level statistics: mean, std, min, max, q01, q99) + - meta/relative_stats_dreamzero.json (relative action statistics) + - meta/tasks.jsonl (task descriptions) + - meta/episodes.jsonl (episode-level metadata) + +The script does NOT modify parquet files or videos -- it only creates metadata. + +Usage: + # Auto-detect state/action structure, default embodiment tag 'xdof': + python scripts/data/convert_lerobot_to_gear.py --dataset-path ./Dataset/my_robot_data + + # Explicit modality mapping via JSON: + python scripts/data/convert_lerobot_to_gear.py \\ + --dataset-path ./Dataset/my_robot_data \\ + --embodiment-tag xdof \\ + --state-keys '{"joint_pos": [0, 6], "gripper_pos": [6, 7]}' \\ + --action-keys '{"joint_pos": [0, 6], "gripper_pos": [6, 7]}' \\ + --relative-action-keys joint_pos \\ + --task-key annotation.task + + # Copy to a new output directory instead of modifying in-place: + python scripts/data/convert_lerobot_to_gear.py \\ + --dataset-path ./Dataset/my_robot_data \\ + --output-path ./Dataset/my_robot_data_gear +""" + +from __future__ import annotations + +import argparse +import json +import logging +import shutil +import sys +from pathlib import Path + +import numpy as np +import pandas as pd +from tqdm import tqdm + +logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") +log = logging.getLogger(__name__) + +VALID_EMBODIMENT_TAGS = [ + "real_gr1_arms_only", "real_gr1_arms_only_annotated", + "real_gr1_arms_waist", "real_gr1_arms_waist_annotated", + "dexmg_gr1_arms_only_inspire", "dexmg_gr1_arms_only_fourier", + "dexmg_gr1_arms_waist_fourier", + "robocasa_single_arm", "onex_eve_gripper", + "robocasa_gr1_arms_only_inspire_hands", "robocasa_gr1_arms_only_fourier_hands", + "robocasa_gr1_fixed_lower_body_inspire_hands", "robocasa_gr1_fixed_lower_body_fourier_hands", + "robocasa_panda_omron", + "robocasa_bimanual_panda_parallel_gripper", "robocasa_bimanual_panda_inspire_hand", + "oxe_droid", "oxe_fractal", "oxe_language_table", "oxe_bridge", + "real_panda_single_arm", "hot3d_hands_only", + "gr1_unified", "robocasa_gr1_arms_waist_fourier_hands", + "agibot", "lapa", "oxe_mutex", "oxe_roboset", "oxe_plex", + "dream", "yam", "xdof", "trex", + "gr1_unified_segmentation", "language_table_sim", "gr1_isaac", + "sim_behavior_r1_pro", "mecka_hands", "real_r1_pro_sharpa", +] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def load_info(dataset_path: Path) -> dict: + info_path = dataset_path / "meta" / "info.json" + if not info_path.exists(): + log.error("meta/info.json not found at %s", info_path) + sys.exit(1) + with open(info_path) as f: + return json.load(f) + + +def get_parquet_paths(dataset_path: Path, info: dict) -> list[Path]: + pattern = info.get("data_path", "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet") + total_episodes = info["total_episodes"] + chunks_size = info.get("chunks_size", 1000) + paths = [] + for ep_idx in range(total_episodes): + chunk_idx = ep_idx // chunks_size + p = dataset_path / pattern.format(episode_chunk=chunk_idx, episode_index=ep_idx) + if p.exists(): + paths.append(p) + return sorted(paths) + + +def detect_features(info: dict) -> dict: + """Return categorised feature names from info.json.""" + features = info.get("features", {}) + state_keys = [k for k in features if k.startswith("observation.state")] + action_keys = [k for k in features if k == "action" or k.startswith("action.")] + video_keys = [k for k in features if features[k].get("dtype") == "video"] + annotation_keys = [k for k in features if k.startswith("annotation")] + return { + "state": state_keys, + "action": action_keys, + "video": video_keys, + "annotation": annotation_keys, + "features": features, + } + + +def parse_key_mapping(raw: str | None) -> dict[str, list[int]] | None: + """Parse a JSON string like '{"joint_pos": [0, 6], "gripper": [6, 7]}'.""" + if raw is None: + return None + try: + mapping = json.loads(raw) + except json.JSONDecodeError as e: + log.error("Invalid JSON for key mapping: %s", e) + sys.exit(1) + for name, bounds in mapping.items(): + if not isinstance(bounds, list) or len(bounds) != 2: + log.error("Each entry must be [start, end]. Got %s for '%s'", bounds, name) + sys.exit(1) + return mapping + + +# --------------------------------------------------------------------------- +# Modality JSON +# --------------------------------------------------------------------------- + +def build_modality_json( + info: dict, + detected: dict, + state_mapping: dict[str, list[int]] | None, + action_mapping: dict[str, list[int]] | None, + task_key: str | None, +) -> dict: + """Build the modality.json structure expected by GEAR/DreamZero.""" + features = detected["features"] + modality: dict = {"state": {}, "action": {}, "video": {}, "annotation": {}} + + # --- State --- + state_col = detected["state"][0] if detected["state"] else None + if state_col and state_mapping: + for name, (start, end) in state_mapping.items(): + dtype = features[state_col].get("dtype", "float64") + modality["state"][name] = { + "original_key": state_col, + "start": start, + "end": end, + "rotation_type": None, + "absolute": True, + "dtype": dtype, + "range": None, + } + elif state_col: + shape = features[state_col].get("shape", [1]) + dim = shape[0] if isinstance(shape, list) else shape + dtype = features[state_col].get("dtype", "float64") + modality["state"]["state"] = { + "original_key": state_col, + "start": 0, + "end": dim, + "rotation_type": None, + "absolute": True, + "dtype": dtype, + "range": None, + } + + # --- Action --- + action_col = detected["action"][0] if detected["action"] else None + if action_col and action_mapping: + for name, (start, end) in action_mapping.items(): + dtype = features[action_col].get("dtype", "float64") + modality["action"][name] = { + "original_key": action_col, + "start": start, + "end": end, + "rotation_type": None, + "absolute": True, + "dtype": dtype, + "range": None, + } + elif action_col: + shape = features[action_col].get("shape", [1]) + dim = shape[0] if isinstance(shape, list) else shape + dtype = features[action_col].get("dtype", "float64") + modality["action"]["action"] = { + "original_key": action_col, + "start": 0, + "end": dim, + "rotation_type": None, + "absolute": True, + "dtype": dtype, + "range": None, + } + + # --- Video --- + for vk in detected["video"]: + short_name = vk.replace("observation.images.", "") + modality["video"][short_name] = {"original_key": vk} + + # --- Annotation --- + if task_key: + short = task_key.replace("annotation.", "") + modality["annotation"][short] = {"original_key": task_key} + else: + for ak in detected["annotation"]: + short = ak.replace("annotation.", "") + modality["annotation"][short] = {"original_key": ak} + + return modality + + +# --------------------------------------------------------------------------- +# Stats computation +# --------------------------------------------------------------------------- + +def compute_stats(parquet_paths: list[Path], columns: list[str]) -> dict: + """Compute mean/std/min/max/q01/q99 for numeric columns across all episodes.""" + all_data: dict[str, list] = {col: [] for col in columns} + for pp in tqdm(parquet_paths, desc="Computing stats"): + df = pd.read_parquet(pp) + for col in columns: + if col not in df.columns: + continue + arr = np.stack(df[col].values) + if arr.ndim == 1: + arr = arr.reshape(-1, 1) + all_data[col].append(arr) + + stats = {} + for col in columns: + if not all_data[col]: + continue + data = np.concatenate(all_data[col], axis=0).astype(np.float64) + stats[col] = { + "mean": np.mean(data, axis=0).tolist(), + "std": np.std(data, axis=0).tolist(), + "min": np.min(data, axis=0).tolist(), + "max": np.max(data, axis=0).tolist(), + "q01": np.quantile(data, 0.01, axis=0).tolist(), + "q99": np.quantile(data, 0.99, axis=0).tolist(), + } + return stats + + +def compute_relative_stats( + parquet_paths: list[Path], + modality: dict, + relative_action_keys: list[str], + action_horizon: int = 24, +) -> dict: + """Compute relative-action statistics: (action - reference_state) for each key. + + This replicates the logic in groot/vla/data/dataset/lerobot.py + _calculate_relative_stats_for_key. + """ + stats: dict = {} + for rel_key in relative_action_keys: + if rel_key not in modality["action"]: + log.warning("Relative action key '%s' not found in action modality, skipping", rel_key) + continue + if rel_key not in modality["state"]: + log.warning( + "Relative action key '%s' has no matching state key -- " + "relative stats require a corresponding state key with the same name. Skipping.", + rel_key, + ) + continue + + action_meta = modality["action"][rel_key] + state_meta = modality["state"][rel_key] + + all_relative = [] + for pp in tqdm(parquet_paths, desc=f"Relative stats [{rel_key}]"): + df = pd.read_parquet(pp) + action_col = action_meta["original_key"] + state_col = state_meta["original_key"] + if action_col not in df.columns or state_col not in df.columns: + continue + + action_data = np.stack(df[action_col].values).astype(np.float64) + state_data = np.stack(df[state_col].values).astype(np.float64) + if action_data.ndim == 1: + action_data = action_data.reshape(-1, 1) + if state_data.ndim == 1: + state_data = state_data.reshape(-1, 1) + + a_start, a_end = action_meta["start"], action_meta["end"] + s_start, s_end = state_meta["start"], state_meta["end"] + + action_slice = action_data[:, a_start:a_end] + state_slice = state_data[:, s_start:s_end] + + traj_len = len(df) + usable = traj_len - action_horizon + for i in range(max(usable, 0)): + ref_state = state_slice[i] + chunk_end = min(i + action_horizon, traj_len) + actions = action_slice[i:chunk_end] + relative = actions - ref_state + all_relative.extend(relative) + + if not all_relative: + log.warning("No relative actions computed for '%s'", rel_key) + continue + + data = np.array(all_relative) + stats[rel_key] = { + "max": np.max(data, axis=0).tolist(), + "min": np.min(data, axis=0).tolist(), + "mean": np.mean(data, axis=0).tolist(), + "std": np.std(data, axis=0).tolist(), + "q01": np.quantile(data, 0.01, axis=0).tolist(), + "q99": np.quantile(data, 0.99, axis=0).tolist(), + } + + return stats + + +# --------------------------------------------------------------------------- +# Tasks & episodes +# --------------------------------------------------------------------------- + +def build_tasks(parquet_paths: list[Path], task_key: str | None) -> list[dict]: + """Build tasks.jsonl entries from the dataset.""" + if task_key is None: + return [{"task_index": 0, "task": ""}] + + task_set: dict[str, int] = {} + for pp in tqdm(parquet_paths, desc="Extracting tasks"): + df = pd.read_parquet(pp) + if task_key not in df.columns: + continue + for val in df[task_key].unique(): + text = str(val) if not isinstance(val, str) else val + if text not in task_set: + task_set[text] = len(task_set) + + if not task_set: + return [{"task_index": 0, "task": ""}] + + return [{"task_index": idx, "task": text} for text, idx in sorted(task_set.items(), key=lambda x: x[1])] + + +def build_episodes(parquet_paths: list[Path], info: dict, task_key: str | None, tasks: list[dict]) -> list[dict]: + """Build episodes.jsonl entries.""" + task_text_to_idx = {t["task"]: t["task_index"] for t in tasks} + episodes = [] + for ep_idx, pp in enumerate(tqdm(parquet_paths, desc="Building episodes")): + df = pd.read_parquet(pp) + length = len(df) + + ep_tasks: list[str] = [] + if task_key and task_key in df.columns: + unique_tasks = df[task_key].unique() + for t in unique_tasks: + text = str(t) if not isinstance(t, str) else t + if text and text in task_text_to_idx: + ep_tasks.append(text) + if not ep_tasks: + ep_tasks = [""] + + episodes.append({ + "episode_index": ep_idx, + "tasks": ep_tasks, + "length": length, + }) + + return episodes + + +# --------------------------------------------------------------------------- +# Validation +# --------------------------------------------------------------------------- + +def validate_dataset(dataset_path: Path, info: dict, modality: dict) -> list[str]: + """Run basic validation and return a list of warnings.""" + warnings = [] + + # Check required directories + for subdir in ["data", "videos", "meta"]: + if not (dataset_path / subdir).exists(): + warnings.append(f"Missing directory: {subdir}/") + + # Check at least one video key exists + if not modality["video"]: + warnings.append("No video features detected -- DreamZero requires at least one camera view") + + # Check state/action exist + if not modality["state"]: + warnings.append("No state modality keys defined") + if not modality["action"]: + warnings.append("No action modality keys defined") + + # Check total_episodes > 0 + if info.get("total_episodes", 0) == 0: + warnings.append("total_episodes is 0 in info.json") + + # Check FPS + if info.get("fps") is None: + warnings.append("fps not set in info.json") + + return warnings + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main(): + parser = argparse.ArgumentParser( + description="Convert a LeRobot v2 dataset to GEAR/DreamZero training format.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=__doc__, + ) + parser.add_argument("--dataset-path", type=str, required=True, help="Path to the LeRobot v2 dataset") + parser.add_argument("--output-path", type=str, default=None, help="Output path (default: modify in-place)") + parser.add_argument( + "--embodiment-tag", type=str, default="xdof", + help=f"Embodiment tag (default: xdof). Valid: {', '.join(sorted(set(VALID_EMBODIMENT_TAGS)))}" + ) + parser.add_argument( + "--state-keys", type=str, default=None, + help='JSON mapping of state sub-keys to [start, end] index ranges, ' + 'e.g. \'{"joint_pos": [0, 6], "gripper_pos": [6, 7]}\'' + ) + parser.add_argument( + "--action-keys", type=str, default=None, + help='JSON mapping of action sub-keys to [start, end] index ranges' + ) + parser.add_argument( + "--relative-action-keys", type=str, nargs="*", default=None, + help="Action sub-key names to compute relative stats for (e.g. joint_pos gripper_pos). " + "Each key must also exist in --state-keys. If omitted, skips relative stats." + ) + parser.add_argument("--task-key", type=str, default=None, help="Column name for language annotations (auto-detected if not set)") + parser.add_argument("--fps", type=float, default=None, help="Override FPS (default: use dataset FPS from info.json)") + parser.add_argument("--action-horizon", type=int, default=24, help="Action horizon for relative stats (default: 24)") + parser.add_argument("--force", action="store_true", help="Overwrite existing GEAR metadata files") + + args = parser.parse_args() + + dataset_path = Path(args.dataset_path).resolve() + if not dataset_path.exists(): + log.error("Dataset path does not exist: %s", dataset_path) + sys.exit(1) + + # Validate embodiment tag + if args.embodiment_tag not in VALID_EMBODIMENT_TAGS: + log.error( + "Invalid embodiment tag '%s'. Valid tags:\n %s", + args.embodiment_tag, + "\n ".join(sorted(set(VALID_EMBODIMENT_TAGS))), + ) + sys.exit(1) + + # Output path handling + if args.output_path: + output_path = Path(args.output_path).resolve() + if output_path != dataset_path: + log.info("Copying dataset to %s", output_path) + if output_path.exists(): + if not args.force: + log.error("Output path already exists. Use --force to overwrite.") + sys.exit(1) + shutil.rmtree(output_path) + shutil.copytree(dataset_path, output_path) + dataset_path = output_path + else: + output_path = dataset_path + + meta_dir = output_path / "meta" + meta_dir.mkdir(parents=True, exist_ok=True) + + # 1. Load info.json + info = load_info(dataset_path) + detected = detect_features(info) + + log.info("Dataset: %s", dataset_path.name) + log.info(" Episodes: %d", info.get("total_episodes", 0)) + log.info(" FPS: %s", info.get("fps", "not set")) + log.info(" State columns: %s", detected["state"]) + log.info(" Action columns: %s", detected["action"]) + log.info(" Video features: %d camera(s)", len(detected["video"])) + log.info(" Annotation columns: %s", detected["annotation"]) + + if args.fps is not None: + info["fps"] = args.fps + with open(output_path / "meta" / "info.json", "w") as f: + json.dump(info, f, indent=4) + log.info(" Overriding FPS to %s", args.fps) + + # Parse user-provided key mappings + state_mapping = parse_key_mapping(args.state_keys) + action_mapping = parse_key_mapping(args.action_keys) + + # Auto-detect task key if not provided + task_key = args.task_key + if task_key is None and detected["annotation"]: + for candidate in ["annotation.task", "annotation.language.language_instruction"]: + if candidate in detected["annotation"]: + task_key = candidate + break + if task_key is None: + task_key = detected["annotation"][0] + log.info(" Auto-detected task key: %s", task_key) + + # 2. Build modality.json + modality = build_modality_json(info, detected, state_mapping, action_mapping, task_key) + + modality_path = meta_dir / "modality.json" + if modality_path.exists() and not args.force: + log.info(" modality.json already exists, skipping (use --force to overwrite)") + else: + with open(modality_path, "w") as f: + json.dump(modality, f, indent=4) + log.info(" Wrote modality.json (%d state keys, %d action keys, %d video keys)", + len(modality["state"]), len(modality["action"]), len(modality["video"])) + + # 3. Write embodiment.json + embodiment = {"robot_type": args.embodiment_tag, "embodiment_tag": args.embodiment_tag} + embodiment_path = meta_dir / "embodiment.json" + if embodiment_path.exists() and not args.force: + log.info(" embodiment.json already exists, skipping") + else: + with open(embodiment_path, "w") as f: + json.dump(embodiment, f, indent=4) + log.info(" Wrote embodiment.json (tag=%s)", args.embodiment_tag) + + # 4. Get parquet file paths + parquet_paths = get_parquet_paths(output_path, info) + if not parquet_paths: + log.error("No parquet files found. Check dataset structure.") + sys.exit(1) + log.info(" Found %d parquet files", len(parquet_paths)) + + # 5. Compute stats.json + stats_path = meta_dir / "stats.json" + numeric_cols = detected["state"] + detected["action"] + if "timestamp" in info.get("features", {}): + numeric_cols.append("timestamp") + + if stats_path.exists() and not args.force: + log.info(" stats.json already exists, skipping") + else: + log.info(" Computing dataset statistics...") + stats = compute_stats(parquet_paths, numeric_cols) + with open(stats_path, "w") as f: + json.dump(stats, f, indent=4) + log.info(" Wrote stats.json (%d features)", len(stats)) + + # 6. Compute relative_stats_dreamzero.json + rel_stats_path = meta_dir / "relative_stats_dreamzero.json" + if args.relative_action_keys: + if rel_stats_path.exists() and not args.force: + log.info(" relative_stats_dreamzero.json already exists, skipping") + else: + log.info(" Computing relative action statistics for keys: %s", args.relative_action_keys) + rel_stats = compute_relative_stats( + parquet_paths, modality, args.relative_action_keys, + action_horizon=args.action_horizon, + ) + if rel_stats: + with open(rel_stats_path, "w") as f: + json.dump(rel_stats, f, indent=4) + log.info(" Wrote relative_stats_dreamzero.json (%d keys)", len(rel_stats)) + else: + log.warning(" No relative stats computed (check key names match between state and action)") + else: + log.info(" Skipping relative stats (no --relative-action-keys provided)") + + # 7. Build tasks.jsonl + tasks_path = meta_dir / "tasks.jsonl" + if tasks_path.exists() and not args.force: + log.info(" tasks.jsonl already exists, skipping") + else: + tasks = build_tasks(parquet_paths, task_key) + with open(tasks_path, "w") as f: + for t in tasks: + f.write(json.dumps(t) + "\n") + log.info(" Wrote tasks.jsonl (%d tasks)", len(tasks)) + + # 8. Build episodes.jsonl + episodes_path = meta_dir / "episodes.jsonl" + if episodes_path.exists() and not args.force: + log.info(" episodes.jsonl already exists, skipping") + else: + tasks = [] + if tasks_path.exists(): + with open(tasks_path) as f: + for line in f: + tasks.append(json.loads(line.strip())) + if not tasks: + tasks = [{"task_index": 0, "task": ""}] + episodes = build_episodes(parquet_paths, info, task_key, tasks) + with open(episodes_path, "w") as f: + for ep in episodes: + f.write(json.dumps(ep) + "\n") + log.info(" Wrote episodes.jsonl (%d episodes)", len(episodes)) + + # 9. Validation + warnings = validate_dataset(output_path, info, modality) + if warnings: + log.warning("Validation warnings:") + for w in warnings: + log.warning(" - %s", w) + else: + log.info("Validation passed -- no warnings") + + # Summary + print("\n" + "=" * 60) + print("Conversion complete!") + print(f" Output: {output_path}") + print(f" Embodiment tag: {args.embodiment_tag}") + print(f" State keys: {list(modality['state'].keys())}") + print(f" Action keys: {list(modality['action'].keys())}") + print(f" Video keys: {list(modality['video'].keys())}") + print(f" Task key: {task_key or '(none)'}") + if args.relative_action_keys: + print(f" Relative action keys: {args.relative_action_keys}") + print("=" * 60) + print("\nNext steps:") + print(" 1. Create a YAML data config in groot/vla/configs/data/dreamzero/") + print(" 2. Add modality configs to base_48_wan_fine_aug_relative.yaml") + print(" 3. Create a training script in scripts/train/") + print(" See docs/CUSTOM_EMBODIMENT_TRAINING.md for the full guide.") + + +if __name__ == "__main__": + main() diff --git a/scripts/data/convert_trex_v3_to_v2.py b/scripts/data/convert_trex_v3_to_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..134a69e0bdddf15d82a50c3da56ab014fec0e728 --- /dev/null +++ b/scripts/data/convert_trex_v3_to_v2.py @@ -0,0 +1,483 @@ +""" +Convert the T-Rex LeRobot v3.0 dataset to LeRobot v2.1 layout for DreamZero. + +LeRobot v3 packs many episodes into shared parquet/video files: + data/chunk-XXX/file-XXX.parquet (rows of many episodes) + videos/{video_key}/chunk-XXX/file-XXX.mp4 (concatenated episodes) + meta/episodes/chunk-XXX/file-XXX.parquet (episode metadata) + meta/tasks.parquet + +DreamZero's loader (groot/vla/data/dataset/lerobot.py) expects v2 layout: + data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet + videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4 + meta/episodes.jsonl, meta/tasks.jsonl, meta/info.json + +RGB cameras are re-encoded to 320x180. Tactile videos keep native resolution +(raw 320x240, deform 240x240) and are re-encoded with libx264 for smaller size. + +Usage: + python scripts/data/convert_trex_v3_to_v2.py --phase data + python scripts/data/convert_trex_v3_to_v2.py --phase videos + python scripts/data/convert_trex_v3_to_v2.py --phase videos --include-tactile + python scripts/data/convert_trex_v3_to_v2.py --phase meta --include-tactile + python scripts/data/convert_trex_v3_to_v2.py --phase verify --include-tactile + +All phases are resumable: existing valid outputs are skipped. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import subprocess +from concurrent.futures import ProcessPoolExecutor, as_completed +from pathlib import Path + +import numpy as np +import pandas as pd +from tqdm import tqdm + +logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") +log = logging.getLogger(__name__) + +DEFAULT_SRC = Path("/scratch1/home/zhicao/dreamzero/data/trex_dataset") +DEFAULT_DST = Path("/scratch1/home/zhicao/dreamzero/data/trex_datasetv2") +SRC = DEFAULT_SRC +DST = DEFAULT_DST + +RGB_VIDEO_KEYS = [ + "observation.images.head_left", + "observation.images.left_wrist", + "observation.images.right_wrist", +] +RGB_OUT_W, RGB_OUT_H = 320, 180 +FPS = 30 +CHUNKS_SIZE = 1000 +RGB_CRF = 23 +TACTILE_CRF = 28 + + +def load_src_info() -> dict: + return json.loads((SRC / "meta" / "info.json").read_text()) + + +def get_tactile_video_keys(src_info: dict | None = None) -> list[str]: + src_info = src_info or load_src_info() + return sorted( + k + for k, v in src_info["features"].items() + if v.get("dtype") == "video" and "tactile" in k + ) + + +def get_output_size(video_key: str, feature: dict) -> tuple[int, int]: + """Return (width, height) for ffmpeg scale filter.""" + if video_key in RGB_VIDEO_KEYS: + return RGB_OUT_W, RGB_OUT_H + shape = feature.get("shape", []) + if len(shape) >= 2: + height, width = int(shape[0]), int(shape[1]) + return width, height + info = feature.get("info", {}) + return int(info["video.width"]), int(info["video.height"]) + + +def get_crf(video_key: str) -> int: + return RGB_CRF if video_key in RGB_VIDEO_KEYS else TACTILE_CRF + + +def resolve_video_keys( + include_tactile: bool, + tactile_only: bool, + explicit_keys: list[str] | None, +) -> list[str]: + if explicit_keys: + return explicit_keys + if tactile_only: + return get_tactile_video_keys() + keys = list(RGB_VIDEO_KEYS) + if include_tactile: + keys.extend(get_tactile_video_keys()) + return keys + + +def load_episode_meta() -> pd.DataFrame: + files = sorted(SRC.glob("meta/episodes/chunk-*/file-*.parquet")) + if not files: + raise FileNotFoundError(f"No episode metadata under {SRC / 'meta/episodes'}") + df = pd.concat([pd.read_parquet(f) for f in files], ignore_index=True) + return df.sort_values("episode_index").reset_index(drop=True) + + +def load_task_map() -> dict[int, str]: + t = pd.read_parquet(SRC / "meta" / "tasks.parquet") + return {int(row.task_index): str(idx) for idx, row in t.iterrows()} + + +def ep_parquet_path(ep_idx: int) -> Path: + return DST / f"data/chunk-{ep_idx // CHUNKS_SIZE:03d}/episode_{ep_idx:06d}.parquet" + + +def ep_video_path(ep_idx: int, video_key: str) -> Path: + return DST / f"videos/chunk-{ep_idx // CHUNKS_SIZE:03d}/{video_key}/episode_{ep_idx:06d}.mp4" + + +def build_video_features(src_info: dict, video_keys: list[str]) -> dict: + features: dict = {} + for k, v in src_info["features"].items(): + if v.get("dtype") != "video" or k not in video_keys: + continue + v = dict(v) + out_w, out_h = get_output_size(k, v) + v["shape"] = [out_h, out_w, 3] + info_blk = dict(v.get("info", {})) + info_blk.update( + { + "video.height": out_h, + "video.width": out_w, + "video.codec": "h264", + "video.pix_fmt": "yuv420p", + "video.fps": FPS, + "video.channels": 3, + "has_audio": False, + } + ) + v["info"] = info_blk + features[k] = v + return features + + +def write_info_json( + src_info: dict, + video_keys: list[str], + *, + preserve_existing_features: bool = False, +) -> None: + meta_dir = DST / "meta" + meta_dir.mkdir(parents=True, exist_ok=True) + + existing = {} + info_path = meta_dir / "info.json" + if preserve_existing_features and info_path.exists(): + existing = json.loads(info_path.read_text()) + + features = dict(existing.get("features", {})) + for k, v in src_info["features"].items(): + if v.get("dtype") != "video": + features[k] = v + features.update(build_video_features(src_info, video_keys)) + features["annotation.task"] = {"dtype": "string", "shape": [1], "names": None} + + total_episodes = int(src_info["total_episodes"]) + all_video_keys = [k for k, v in features.items() if v.get("dtype") == "video"] + info = { + "codebase_version": "v2.1", + "robot_type": src_info.get("robot_type", "dexmate_vega1_and_sharpa_wave"), + "total_episodes": total_episodes, + "total_frames": int(src_info["total_frames"]), + "total_tasks": int(src_info["total_tasks"]), + "total_videos": total_episodes * len(all_video_keys), + "total_chunks": (total_episodes + CHUNKS_SIZE - 1) // CHUNKS_SIZE, + "chunks_size": CHUNKS_SIZE, + "fps": FPS, + "splits": {"train": f"0:{total_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, + } + with open(info_path, "w") as f: + json.dump(info, f, indent=4) + log.info( + "Wrote meta/info.json with %d video keys (%d total videos)", + len(all_video_keys), + info["total_videos"], + ) + + +# --------------------------------------------------------------------------- +# Phase: data (parquets + meta) +# --------------------------------------------------------------------------- + +def convert_data( + ep_meta: pd.DataFrame, + task_map: dict[int, str], + video_keys: list[str], +) -> None: + src_info = load_src_info() + + n_done = 0 + groups = ep_meta.groupby(["data/chunk_index", "data/file_index"]) + for (chunk_idx, file_idx), eps in tqdm(groups, desc="Converting data files"): + src_pq = SRC / f"data/chunk-{int(chunk_idx):03d}/file-{int(file_idx):03d}.parquet" + if not src_pq.exists(): + log.warning("Missing source parquet: %s", src_pq) + continue + if all(ep_parquet_path(int(r.episode_index)).exists() for r in eps.itertuples()): + n_done += len(eps) + continue + df = pd.read_parquet(src_pq) + for r in eps.itertuples(): + ep_idx = int(r.episode_index) + out = ep_parquet_path(ep_idx) + if out.exists(): + n_done += 1 + continue + ep_df = df[df["episode_index"] == ep_idx].copy() + assert len(ep_df) == int(r.length), ( + f"episode {ep_idx}: rows {len(ep_df)} != meta length {r.length}" + ) + task_texts = [str(t) for t in r.tasks] + ep_df["annotation.task"] = task_texts[0] if task_texts else "" + out.parent.mkdir(parents=True, exist_ok=True) + ep_df.to_parquet(out, index=False) + n_done += 1 + log.info("Data phase done: %d episode parquets", n_done) + + meta_dir = DST / "meta" + meta_dir.mkdir(parents=True, exist_ok=True) + with open(meta_dir / "tasks.jsonl", "w") as f: + for idx in sorted(task_map): + f.write(json.dumps({"task_index": idx, "task": task_map[idx]}) + "\n") + + with open(meta_dir / "episodes.jsonl", "w") as f: + for r in ep_meta.itertuples(): + f.write( + json.dumps( + { + "episode_index": int(r.episode_index), + "tasks": [str(t) for t in r.tasks], + "length": int(r.length), + } + ) + + "\n" + ) + + write_info_json(src_info, video_keys, preserve_existing_features=False) + log.info("Wrote meta/episodes.jsonl, meta/tasks.jsonl") + + +def update_meta(video_keys: list[str]) -> None: + src_info = load_src_info() + write_info_json(src_info, video_keys, preserve_existing_features=True) + + +# --------------------------------------------------------------------------- +# Phase: videos +# --------------------------------------------------------------------------- + +def _cut_one(job: tuple) -> tuple[int, str, bool, str]: + ep_idx, video_key, src_mp4, from_ts, n_frames, out_path, out_w, out_h, crf = job + out = Path(out_path) + out.parent.mkdir(parents=True, exist_ok=True) + tmp = out.with_suffix(".tmp.mp4") + ss = max(0.0, from_ts - 0.5 / FPS) + cmd = [ + "ffmpeg", + "-y", + "-loglevel", + "error", + "-ss", + f"{ss:.6f}", + "-i", + src_mp4, + "-frames:v", + str(n_frames), + "-vf", + f"scale={out_w}:{out_h}", + "-c:v", + "libx264", + "-preset", + "veryfast", + "-crf", + str(crf), + "-pix_fmt", + "yuv420p", + "-movflags", + "+faststart", + "-an", + "-threads", + "2", + str(tmp), + ] + try: + res = subprocess.run(cmd, capture_output=True, text=True, timeout=600) + if res.returncode != 0: + tmp.unlink(missing_ok=True) + return ep_idx, video_key, False, res.stderr[-500:] + tmp.rename(out) + return ep_idx, video_key, True, "" + except Exception as e: # noqa: BLE001 + tmp.unlink(missing_ok=True) + return ep_idx, video_key, False, str(e) + + +def convert_videos(ep_meta: pd.DataFrame, video_keys: list[str], workers: int) -> None: + src_info = load_src_info() + jobs = [] + missing_src = set() + for vk in video_keys: + feature = src_info["features"][vk] + out_w, out_h = get_output_size(vk, feature) + crf = get_crf(vk) + for r in ep_meta.itertuples(): + ep_idx = int(r.episode_index) + out = ep_video_path(ep_idx, vk) + if out.exists(): + continue + chunk_i = int(ep_meta.loc[r.Index, f"videos/{vk}/chunk_index"]) + file_i = int(ep_meta.loc[r.Index, f"videos/{vk}/file_index"]) + from_ts = float(ep_meta.loc[r.Index, f"videos/{vk}/from_timestamp"]) + src_mp4 = SRC / f"videos/{vk}/chunk-{chunk_i:03d}/file-{file_i:03d}.mp4" + if not src_mp4.exists(): + missing_src.add(str(src_mp4)) + continue + jobs.append( + ( + ep_idx, + vk, + str(src_mp4), + from_ts, + int(r.length), + str(out), + out_w, + out_h, + crf, + ) + ) + + if missing_src: + log.warning( + "%d source videos missing (not yet downloaded?), e.g. %s", + len(missing_src), + sorted(missing_src)[0], + ) + log.info("Cutting %d episode videos with %d workers", len(jobs), workers) + + failures = [] + with ProcessPoolExecutor(max_workers=workers) as pool: + futs = [pool.submit(_cut_one, j) for j in jobs] + for fut in tqdm(as_completed(futs), total=len(futs), desc="Cutting videos"): + ep_idx, vk, ok, err = fut.result() + if not ok: + failures.append((ep_idx, vk, err)) + if failures: + log.error("%d failures, first: %s", len(failures), failures[0]) + else: + log.info("Video phase done, no failures") + + +# --------------------------------------------------------------------------- +# Phase: verify +# --------------------------------------------------------------------------- + +def _probe_frames(path: Path) -> int: + res = subprocess.run( + [ + "ffprobe", + "-v", + "error", + "-count_frames", + "-select_streams", + "v:0", + "-show_entries", + "stream=nb_read_frames", + "-of", + "csv=p=0", + str(path), + ], + capture_output=True, + text=True, + timeout=120, + ) + return int(res.stdout.strip()) + + +def verify(ep_meta: pd.DataFrame, video_keys: list[str], n_samples: int) -> None: + rng = np.random.default_rng(0) + total = len(ep_meta) + + missing_pq = [ + int(r.episode_index) + for r in ep_meta.itertuples() + if not ep_parquet_path(int(r.episode_index)).exists() + ] + log.info("Parquets: %d/%d present", total - len(missing_pq), total) + + for vk in video_keys: + missing = [ + int(r.episode_index) + for r in ep_meta.itertuples() + if not ep_video_path(int(r.episode_index), vk).exists() + ] + log.info("Videos [%s]: %d/%d present", vk, total - len(missing), total) + + sample = rng.choice(total, size=min(n_samples, total), replace=False) + for ep_idx in sample: + ep_idx = int(ep_idx) + row = ep_meta[ep_meta["episode_index"] == ep_idx].iloc[0] + length = int(row["length"]) + pq = ep_parquet_path(ep_idx) + if pq.exists(): + df = pd.read_parquet(pq) + assert len(df) == length, f"ep {ep_idx}: parquet {len(df)} != {length}" + assert np.asarray(df["action"].iloc[0]).shape == (58,) + assert df["annotation.task"].iloc[0] == str(row["tasks"][0]) + for vk in video_keys: + vp = ep_video_path(ep_idx, vk) + if vp.exists(): + n = _probe_frames(vp) + assert n == length, f"ep {ep_idx} {vk}: video {n} frames != {length}" + log.info("ep %06d OK (length=%d)", ep_idx, length) + log.info("Verification passed on %d sampled episodes", len(sample)) + + +def main() -> None: + global SRC, DST + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--phase", + choices=["data", "videos", "meta", "verify"], + required=True, + ) + parser.add_argument("--src", type=Path, default=DEFAULT_SRC, help="v3 dataset root") + parser.add_argument("--dst", type=Path, default=DEFAULT_DST, help="v2 output root") + parser.add_argument( + "--include-tactile", + action="store_true", + help="Include all 20 tactile video streams", + ) + parser.add_argument( + "--tactile-only", + action="store_true", + help="Convert/update only tactile video streams (skip RGB)", + ) + parser.add_argument( + "--video-keys", + nargs="+", + default=None, + help="Explicit video keys to convert (overrides --include-tactile default set)", + ) + parser.add_argument("--workers", type=int, default=16) + parser.add_argument("--verify-samples", type=int, default=20) + args = parser.parse_args() + SRC = args.src + DST = args.dst + + ep_meta = load_episode_meta() + log.info("Loaded %d episodes from v3 metadata", len(ep_meta)) + video_keys = resolve_video_keys(args.include_tactile, args.tactile_only, args.video_keys) + log.info("Video keys (%d): %s", len(video_keys), ", ".join(video_keys)) + + if args.phase == "data": + convert_data(ep_meta, load_task_map(), video_keys) + elif args.phase == "meta": + update_meta(video_keys) + elif args.phase == "videos": + convert_videos(ep_meta, video_keys, args.workers) + elif args.phase == "verify": + verify(ep_meta, video_keys, args.verify_samples) + + +if __name__ == "__main__": + main() diff --git a/scripts/data/download_droid_hf.py b/scripts/data/download_droid_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..0862e39d15acd3414afe183ef566951b9bb568f9 --- /dev/null +++ b/scripts/data/download_droid_hf.py @@ -0,0 +1,81 @@ +#!/usr/bin/env python3 +""" +Download DreamZero DROID dataset from Hugging Face with rate-limit handling. + +Use this script when `huggingface-cli download` hits 429 (Too Many Requests). +It uses a single worker and retries with backoff to stay within the 3000 req/5min limit. + +Usage: + python scripts/data/download_droid_hf.py [--local-dir ./data/droid_lerobot] [--max-workers 1] + + # Or with env (same as CLI default): + DROID_DATA_ROOT=./data/droid_lerobot python scripts/data/download_droid_hf.py +""" + +from __future__ import annotations + +import argparse +import os +import sys +import time + +try: + from huggingface_hub import snapshot_download +except ImportError: + print("Install huggingface_hub: pip install huggingface_hub", file=sys.stderr) + sys.exit(1) + +REPO_ID = "GEAR-Dreams/DreamZero-DROID-Data" +REPO_TYPE = "dataset" + + +def main() -> None: + p = argparse.ArgumentParser(description="Download DreamZero DROID dataset with rate-limit handling.") + p.add_argument( + "--local-dir", + default=os.environ.get("DROID_DATA_ROOT", "./data/droid_lerobot"), + help="Local directory to download into (default: DROID_DATA_ROOT or ./data/droid_lerobot)", + ) + p.add_argument( + "--max-workers", + type=int, + default=1, + help="Concurrent download threads (default 1 to reduce API requests and avoid 429)", + ) + p.add_argument( + "--retry-wait", + type=int, + default=320, + help="Seconds to wait on 429 before retry (default 320 ≈ 5min)", + ) + args = p.parse_args() + + local_dir = os.path.abspath(args.local_dir) + os.makedirs(local_dir, exist_ok=True) + + attempt = 0 + while True: + attempt += 1 + try: + print(f"Download attempt {attempt} (max_workers={args.max_workers})...") + snapshot_download( + repo_id=REPO_ID, + repo_type=REPO_TYPE, + local_dir=local_dir, + local_dir_use_symlinks=False, + max_workers=args.max_workers, + resume_download=True, + ) + print(f"Done. Dataset at: {local_dir}") + return + except Exception as e: + err_str = str(e).lower() + if "429" in err_str or "too many requests" in err_str or "rate limit" in err_str: + print(f"Rate limited (429). Waiting {args.retry_wait}s before retry...", file=sys.stderr) + time.sleep(args.retry_wait) + continue + raise + + +if __name__ == "__main__": + main() diff --git a/scripts/data/prepare_trex_track_force_training.sh b/scripts/data/prepare_trex_track_force_training.sh new file mode 100644 index 0000000000000000000000000000000000000000..1af39c409f9c7c706ad0eef89e08e00370006291 --- /dev/null +++ b/scripts/data/prepare_trex_track_force_training.sh @@ -0,0 +1,835 @@ +#!/usr/bin/env bash +set -Eeuo pipefail + +# Prepare the complete LeRobot-v2 T-Rex dataset for VLTA/T-Rex Track-Force +# training. The source dataset is preserved: by default a hard-link clone is +# created, so RGB/tactile videos do not consume a second copy of their blocks. +# +# The Python builder then: +# 1. extracts canonical 250-point SAM2 + CoTracker tracks; +# 2. adds track, visibility, and absolute EEF62 parquet columns; +# 3. validates the existing 60-D (10 x 6) tactile-force column; +# 4. writes force/track/action statistics and the training metadata contract; +# 5. validates every episode before optionally starting training. + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" + +SOURCE_ROOT="${SOURCE_TREX_DATA_ROOT:-$REPO_ROOT/data/trex_full}" +DATASET_ROOT="${TREX_DATA_ROOT:-$REPO_ROOT/data/trex_full_force}" +TRACK_CACHE="${TRACK_CACHE:-}" +DEFAULT_PYTHON="$REPO_ROOT/../miniconda3/envs/dreamzero/bin/python" +if [[ ! -x "$DEFAULT_PYTHON" ]]; then + DEFAULT_PYTHON="python" +fi +PYTHON_BIN="${PYTHON_BIN:-$DEFAULT_PYTHON}" +GPU_LIST="${GPUS:-6,7}" +EPISODES_PER_JOB="${EPISODES_PER_JOB:-32}" +GPU_MIN_FREE_MIB="${GPU_MIN_FREE_MIB:-78000}" +GPU_RETRY_FREE_MIB="${GPU_RETRY_FREE_MIB:-78000}" +GPU_POLL_SECONDS="${GPU_POLL_SECONDS:-60}" +GPU_REQUIRE_IDLE="${GPU_REQUIRE_IDLE:-6,7}" +START_EPISODE="${START_EPISODE:-auto}" + +OPENPI_ROOT="${OPENPI_ROOT:-/scratch2/home/zhicao/openpi}" +SAM2_LIBS="${SAM2_LIBS:-/scratch1/home/zhicao/physctrl/libs}" +SAM2_MODEL="${SAM2_MODEL:-facebook/sam2-hiera-large}" +CALIB_PATH="${CALIB_PATH:-$REPO_ROOT/assets/trex_camera_calib.json}" +COTRACKER_CHECKPOINT="${COTRACKER_CHECKPOINT:-}" +T_REX_ROOT="${T_REX_ROOT:-/scratch1/home/zhicao/T-Rex}" + +RUN_TRAINING=0 +VALIDATE_ONLY=0 +FORCE_REBUILD=0 +ALLOW_IN_PLACE=0 +VERIFY_FK="${VERIFY_FK:-0}" +TRAIN_ARGS=() + +usage() { + cat <&2 + exit 2 + fi +} + +while [[ $# -gt 0 ]]; do + case "$1" in + --source-root) + need_value "$@" + SOURCE_ROOT="$2" + shift 2 + ;; + --dataset-root) + need_value "$@" + DATASET_ROOT="$2" + shift 2 + ;; + --track-cache) + need_value "$@" + TRACK_CACHE="$2" + shift 2 + ;; + --gpus) + need_value "$@" + GPU_LIST="$2" + shift 2 + ;; + --start-episode) + need_value "$@" + START_EPISODE="$2" + shift 2 + ;; + --episodes-per-job) + need_value "$@" + EPISODES_PER_JOB="$2" + shift 2 + ;; + --min-free-gpu-mib) + need_value "$@" + GPU_MIN_FREE_MIB="$2" + shift 2 + ;; + --require-idle-gpus) + need_value "$@" + GPU_REQUIRE_IDLE="$2" + shift 2 + ;; + --python) + need_value "$@" + PYTHON_BIN="$2" + shift 2 + ;; + --force) + FORCE_REBUILD=1 + shift + ;; + --verify-fk) + VERIFY_FK=1 + shift + ;; + --validate-only) + VALIDATE_ONLY=1 + shift + ;; + --in-place) + ALLOW_IN_PLACE=1 + shift + ;; + --train) + RUN_TRAINING=1 + shift + ;; + --) + shift + TRAIN_ARGS=("$@") + break + ;; + -h|--help) + usage + exit 0 + ;; + *) + echo "ERROR: unknown option: $1" >&2 + usage >&2 + exit 2 + ;; + esac +done + +if [[ "$START_EPISODE" != "auto" && ! "$START_EPISODE" =~ ^[0-9]+$ ]]; then + echo "ERROR: START_EPISODE must be 'auto' or a non-negative integer" >&2 + exit 2 +fi + +for numeric_setting in \ + "EPISODES_PER_JOB=$EPISODES_PER_JOB" \ + "GPU_MIN_FREE_MIB=$GPU_MIN_FREE_MIB" \ + "GPU_RETRY_FREE_MIB=$GPU_RETRY_FREE_MIB" \ + "GPU_POLL_SECONDS=$GPU_POLL_SECONDS"; do + setting_name="${numeric_setting%%=*}" + setting_value="${numeric_setting#*=}" + if [[ ! "$setting_value" =~ ^[1-9][0-9]*$ ]]; then + echo "ERROR: $setting_name must be a positive integer, got: $setting_value" >&2 + exit 2 + fi +done + +if ! command -v "$PYTHON_BIN" >/dev/null 2>&1; then + echo "ERROR: Python executable not found: $PYTHON_BIN" >&2 + exit 1 +fi + +canonical_path() { + "$PYTHON_BIN" - "$1" <<'PY' +import sys +from pathlib import Path + +print(Path(sys.argv[1]).expanduser().resolve()) +PY +} + +SOURCE_ROOT="$(canonical_path "$SOURCE_ROOT")" +DATASET_ROOT="$(canonical_path "$DATASET_ROOT")" +if [[ -z "$TRACK_CACHE" ]]; then + TRACK_CACHE="$DATASET_ROOT/tracks_trex_track_force_v2" +fi +TRACK_CACHE="$(canonical_path "$TRACK_CACHE")" +OPENPI_ROOT="$(canonical_path "$OPENPI_ROOT")" +SAM2_LIBS="$(canonical_path "$SAM2_LIBS")" +CALIB_PATH="$(canonical_path "$CALIB_PATH")" +T_REX_ROOT="$(canonical_path "$T_REX_ROOT")" +if [[ -z "$COTRACKER_CHECKPOINT" ]]; then + COTRACKER_CHECKPOINT="$OPENPI_ROOT/co-tracker/checkpoints/scaled_offline.pth" +else + COTRACKER_CHECKPOINT="$(canonical_path "$COTRACKER_CHECKPOINT")" +fi + +BUILDER="$REPO_ROOT/scripts/data/build_trex_track_force_v2.py" +TRAIN_SCRIPT="$REPO_ROOT/scripts/train/trex_track_force_training_wan22.sh" +CANONICAL_MANIFEST="$DATASET_ROOT/meta/trex_track_force_manifest.json" + +for required_file in "$BUILDER" "$TRAIN_SCRIPT"; do + if [[ ! -f "$required_file" ]]; then + echo "ERROR: required script is missing: $required_file" >&2 + exit 1 + fi +done + +if [[ "$SOURCE_ROOT" == "$DATASET_ROOT" && "$ALLOW_IN_PLACE" != "1" ]]; then + echo "ERROR: refusing to modify the source dataset in place." >&2 + echo " Use a different --dataset-root or explicitly pass --in-place." >&2 + exit 1 +fi + +if [[ "$VALIDATE_ONLY" == "1" ]]; then + if [[ ! -f "$DATASET_ROOT/meta/info.json" ]]; then + echo "ERROR: prepared dataset does not exist: $DATASET_ROOT" >&2 + exit 1 + fi +else + if [[ ! -f "$SOURCE_ROOT/meta/info.json" ]]; then + echo "ERROR: source is not a LeRobot-v2 dataset: $SOURCE_ROOT" >&2 + exit 1 + fi + + if [[ ! -e "$DATASET_ROOT" ]]; then + mkdir -p "$(dirname "$DATASET_ROOT")" + echo "Creating hard-link clone (source remains unchanged):" + echo " source: $SOURCE_ROOT" + echo " output: $DATASET_ROOT" + if ! cp -al -- "$SOURCE_ROOT" "$DATASET_ROOT"; then + echo "ERROR: hard-link clone failed." >&2 + echo " Source and output must be on the same filesystem." >&2 + echo " Remove the incomplete output or provide an existing copied dataset." >&2 + exit 1 + fi + elif [[ ! -f "$DATASET_ROOT/meta/info.json" ]]; then + echo "ERROR: output exists but is not a resumable dataset: $DATASET_ROOT" >&2 + exit 1 + fi +fi + +# Check the complete disk contract before loading expensive tracking models. +TOTAL_EPISODES="$( + "$PYTHON_BIN" - "$SOURCE_ROOT" "$DATASET_ROOT" <<'PY' +import json +import sys +from pathlib import Path + +source = Path(sys.argv[1]) +output = Path(sys.argv[2]) +output_info = json.loads((output / "meta" / "info.json").read_text()) +source_info_path = source / "meta" / "info.json" +source_info = ( + json.loads(source_info_path.read_text()) + if source_info_path.is_file() + else output_info +) + +if source_info.get("codebase_version") != "v2.1": + raise SystemExit( + f"source codebase_version must be v2.1, got {source_info.get('codebase_version')!r}" + ) +if output_info.get("total_episodes") != source_info.get("total_episodes"): + raise SystemExit("source/output total_episodes mismatch") + +required_metadata = ( + "modality.json", + "stats.json", + "episodes.jsonl", + "tasks.jsonl", + "embodiment.json", +) +missing_metadata = [ + name for name in required_metadata if not (output / "meta" / name).is_file() +] +if missing_metadata: + raise SystemExit(f"dataset is missing metadata files: {missing_metadata}") + +features = output_info.get("features", {}) +required_shapes = { + "observation.state": [58], + "action": [58], + "observation.tactile_force": [60], +} +for name, shape in required_shapes.items(): + actual = features.get(name, {}).get("shape") + if actual != shape: + raise SystemExit(f"{name} shape must be {shape}, got {actual}") + +video_keys = ( + "observation.images.head_left", + "observation.images.left_wrist", + "observation.images.right_wrist", +) +for key in video_keys: + if features.get(key, {}).get("dtype") != "video": + raise SystemExit(f"missing RGB video feature {key}") + +total = int(output_info["total_episodes"]) +chunk_size = int(output_info.get("chunks_size", 1000)) +data_pattern = output_info.get( + "data_path", + "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet", +) +video_pattern = output_info.get( + "video_path", + "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4", +) +missing = [] +for episode in range(total): + values = {"episode_chunk": episode // chunk_size, "episode_index": episode} + parquet = output / data_pattern.format(**values) + if not parquet.is_file(): + missing.append(str(parquet)) + for key in video_keys: + video = output / video_pattern.format(video_key=key, **values) + if not video.is_file(): + missing.append(str(video)) + if len(missing) >= 20: + break +if missing: + preview = "\n".join(f" - {path}" for path in missing) + raise SystemExit(f"dataset is incomplete; missing required files:\n{preview}") + +print(total) +PY +)" + +if [[ ! "$TOTAL_EPISODES" =~ ^[1-9][0-9]*$ ]]; then + echo "ERROR: invalid episode count: $TOTAL_EPISODES" >&2 + exit 1 +fi + +FK_FLAG="--no-verify-fk" +if [[ "$VERIFY_FK" == "1" ]]; then + FK_FLAG="--verify-fk" +fi + +if [[ "$VALIDATE_ONLY" == "1" ]]; then + echo "Validating $TOTAL_EPISODES prepared episodes..." + "$PYTHON_BIN" "$BUILDER" \ + --dataset-root "$DATASET_ROOT" \ + --track-cache "$TRACK_CACHE" \ + --manifest-path "$CANONICAL_MANIFEST" \ + --all \ + --validate-only \ + "$FK_FLAG" + echo "Dataset is ready for training: $DATASET_ROOT" +else + for required_path in \ + "$OPENPI_ROOT/droid" \ + "$COTRACKER_CHECKPOINT" \ + "$SAM2_LIBS/sam2" \ + "$CALIB_PATH" \ + "$T_REX_ROOT/utils/lerobot_common.py"; do + if [[ ! -e "$required_path" ]]; then + echo "ERROR: preprocessing dependency is missing: $required_path" >&2 + exit 1 + fi + done + + "$PYTHON_BIN" - "$SAM2_LIBS" "$OPENPI_ROOT" <<'PY' +import sys +from pathlib import Path + +missing = [] +for module in ("cv2", "hydra", "numpy", "pandas", "pinocchio", "pyarrow", "torch"): + try: + __import__(module) + except Exception as exc: # noqa: BLE001 + missing.append(f"{module}: {exc}") +if missing: + raise SystemExit("missing Python dependencies:\n " + "\n ".join(missing)) + +sam2_libs = Path(sys.argv[1]) +openpi_root = Path(sys.argv[2]) +sys.path.insert(0, str(sam2_libs)) +sys.path.insert(0, str(openpi_root / "droid")) +try: + from sam2.sam2_image_predictor import SAM2ImagePredictor # noqa: F401 +except Exception as exc: # noqa: BLE001 + raise SystemExit(f"SAM2 import failed from {sam2_libs}: {exc}") from exc +try: + from utils.cotracker_wrist_grid import load_cotracker_predictor # noqa: F401 +except Exception as exc: # noqa: BLE001 + raise SystemExit(f"CoTracker import failed from {openpi_root}: {exc}") from exc +PY + + IFS=',' read -r -a RAW_GPUS <<< "$GPU_LIST" + GPUS_NORMALIZED=() + declare -A SEEN_GPUS=() + for gpu in "${RAW_GPUS[@]}"; do + gpu="${gpu//[[:space:]]/}" + if [[ -n "$gpu" ]]; then + if [[ -n "${SEEN_GPUS[$gpu]:-}" ]]; then + echo "ERROR: duplicate GPU in --gpus: $gpu" >&2 + exit 2 + fi + SEEN_GPUS["$gpu"]=1 + GPUS_NORMALIZED+=("$gpu") + fi + done + if [[ "${#GPUS_NORMALIZED[@]}" -eq 0 ]]; then + echo "ERROR: --gpus must contain at least one GPU ID or UUID" >&2 + exit 1 + fi + if ! command -v nvidia-smi >/dev/null 2>&1; then + echo "ERROR: nvidia-smi is required for memory-aware scheduling" >&2 + exit 1 + fi + for gpu in "${GPUS_NORMALIZED[@]}"; do + gpu_total="$( + nvidia-smi -i "$gpu" \ + --query-gpu=memory.total \ + --format=csv,noheader,nounits 2>/dev/null + )" || { + echo "ERROR: cannot query GPU: $gpu" >&2 + exit 1 + } + gpu_total="${gpu_total//[[:space:]]/}" + if [[ ! "$gpu_total" =~ ^[0-9]+$ ]]; then + echo "ERROR: invalid memory.total returned for GPU $gpu: $gpu_total" >&2 + exit 1 + fi + if ((GPU_MIN_FREE_MIB > gpu_total)); then + echo "ERROR: --min-free-gpu-mib=$GPU_MIN_FREE_MIB exceeds GPU $gpu total ${gpu_total} MiB" >&2 + exit 2 + fi + done + + IFS=',' read -r -a RAW_IDLE_GPUS <<< "$GPU_REQUIRE_IDLE" + IDLE_GPUS_NORMALIZED=() + for gpu in "${RAW_IDLE_GPUS[@]}"; do + gpu="${gpu//[[:space:]]/}" + if [[ -n "$gpu" ]]; then + IDLE_GPUS_NORMALIZED+=("$gpu") + fi + done + + mkdir -p "$TRACK_CACHE" + RUN_DIR="$DATASET_ROOT/meta/trex_track_force_prepare" + mkdir -p "$RUN_DIR" + if [[ "$START_EPISODE" == "auto" ]]; then + if [[ "$FORCE_REBUILD" == "1" ]]; then + RESUME_START=0 + else + RESUME_START="$( + "$PYTHON_BIN" - "$TRACK_CACHE" "$TOTAL_EPISODES" <<'PY' +import sys +import zipfile +from pathlib import Path + +track_cache = Path(sys.argv[1]) +total_episodes = int(sys.argv[2]) +required_members = { + "tracks.npy", + "vis.npy", + "track_layout_version.npy", +} + +for episode_index in range(total_episodes): + path = track_cache / f"episode_{episode_index:06d}.npz" + if not path.is_file() or path.stat().st_size == 0: + print(episode_index) + break + try: + with zipfile.ZipFile(path) as archive: + if not required_members.issubset(archive.namelist()): + print(episode_index) + break + except (OSError, zipfile.BadZipFile): + print(episode_index) + break +else: + print(total_episodes) +PY + )" + fi + else + RESUME_START="$START_EPISODE" + fi + if ((RESUME_START > TOTAL_EPISODES)); then + echo "ERROR: start episode $RESUME_START exceeds total $TOTAL_EPISODES" >&2 + exit 2 + fi + + SESSION_ID="$(date -u +%Y%m%dT%H%M%SZ)_$$" + SESSION_DIR="$RUN_DIR/run_$SESSION_ID" + QUEUE_DIR="$SESSION_DIR/queue" + mkdir -p "$QUEUE_DIR" + + echo "Preparing the complete T-Rex v2 dataset" + echo " episodes: $TOTAL_EPISODES" + echo " output: $DATASET_ROOT" + echo " track cache: $TRACK_CACHE" + echo " GPUs: ${GPUS_NORMALIZED[*]}" + echo " force: existing observation.tactile_force [60] -> [10,6]" + echo " tracks: SAM2 + CoTracker -> [250,2] + visibility [250]" + echo " resume from: episode $RESUME_START" + echo " queue: $EPISODES_PER_JOB episodes/task" + echo " memory gate: ${GPU_MIN_FREE_MIB} MiB free" + if [[ "${#IDLE_GPUS_NORMALIZED[@]}" -gt 0 ]]; then + echo " idle gate: ${IDLE_GPUS_NORMALIZED[*]}" + fi + + for ((start = RESUME_START; start < TOTAL_EPISODES; start += EPISODES_PER_JOB)); do + end=$((start + EPISODES_PER_JOB)) + if ((end > TOTAL_EPISODES)); then + end="$TOTAL_EPISODES" + fi + printf -v task_name 'pending_%06d_%06d.task' "$start" "$end" + : > "$QUEUE_DIR/$task_name" + done + + queue_has_pending() { + compgen -G "$QUEUE_DIR/pending_*.task" >/dev/null + } + + gpu_requires_idle() { + local candidate="$1" + local idle_gpu + for idle_gpu in "${IDLE_GPUS_NORMALIZED[@]:-}"; do + if [[ "$candidate" == "$idle_gpu" ]]; then + return 0 + fi + done + return 1 + } + + wait_for_gpu() { + local gpu="$1" + local min_free_mib="$2" + local require_idle="$3" + local stop_when_queue_empty="$4" + local announced=0 + local free_mib process_output has_process + while true; do + if [[ "$stop_when_queue_empty" == "1" ]] && ! queue_has_pending; then + return 2 + fi + free_mib="$( + nvidia-smi -i "$gpu" \ + --query-gpu=memory.free \ + --format=csv,noheader,nounits 2>/dev/null + )" || { + echo "ERROR: GPU $gpu memory query failed" >&2 + return 1 + } + free_mib="${free_mib//[[:space:]]/}" + if [[ ! "$free_mib" =~ ^[0-9]+$ ]]; then + echo "ERROR: GPU $gpu returned invalid free memory: $free_mib" >&2 + return 1 + fi + process_output="$( + nvidia-smi -i "$gpu" \ + --query-compute-apps=pid \ + --format=csv,noheader,nounits 2>/dev/null + )" || { + echo "ERROR: GPU $gpu process query failed" >&2 + return 1 + } + has_process=0 + if [[ -n "${process_output//[[:space:]]/}" ]]; then + has_process=1 + fi + if ((free_mib >= min_free_mib)) \ + && { [[ "$require_idle" == "0" ]] || [[ "$has_process" == "0" ]]; }; then + if [[ "$announced" == "1" ]]; then + echo "GPU $gpu is ready (${free_mib} MiB free)" + fi + return 0 + fi + if [[ "$announced" == "0" ]]; then + if [[ "$require_idle" == "1" && "$has_process" == "1" ]]; then + echo "GPU $gpu is waiting for existing compute processes to exit" + else + echo "GPU $gpu is waiting for ${min_free_mib} MiB free (now ${free_mib} MiB)" + fi + announced=1 + fi + sleep "$GPU_POLL_SECONDS" + done + } + + claim_next_task() { + local worker_id="$1" + local candidate suffix claimed + while true; do + for candidate in "$QUEUE_DIR"/pending_*.task; do + if [[ ! -e "$candidate" ]]; then + return 1 + fi + suffix="${candidate##*/pending_}" + claimed="$QUEUE_DIR/running_${worker_id}_${suffix}" + if mv -- "$candidate" "$claimed" 2>/dev/null; then + printf '%s\n' "$claimed" + return 0 + fi + done + done + } + + run_gpu_worker() { + local worker_id="$1" + local gpu="$2" + local base_require_idle=0 + local wait_status task_path task_file task_suffix start end + local task_name task_manifest task_log attempt task_min_free + local task_require_idle command_status tee_status + local -a command pipeline_status + + if gpu_requires_idle "$gpu"; then + base_require_idle=1 + fi + + while queue_has_pending; do + set +e + wait_for_gpu "$gpu" "$GPU_MIN_FREE_MIB" "$base_require_idle" 1 + wait_status="$?" + set -e + if [[ "$wait_status" == "2" ]]; then + return 0 + elif [[ "$wait_status" != "0" ]]; then + return "$wait_status" + fi + + task_path="$(claim_next_task "$worker_id")" || continue + task_file="${task_path##*/}" + task_suffix="${task_file#running_${worker_id}_}" + task_suffix="${task_suffix%.task}" + start="${task_suffix%%_*}" + end="${task_suffix##*_}" + task_name="task_${start}_${end}" + task_manifest="$SESSION_DIR/${task_name}.json" + task_log="$SESSION_DIR/${task_name}.log" + + command=( + "$PYTHON_BIN" "$BUILDER" + --dataset-root "$DATASET_ROOT" + --track-cache "$TRACK_CACHE" + --manifest-path "$task_manifest" + --episode-range "$start" "$end" + --extract-missing + --no-update-metadata + "$FK_FLAG" + --calib-path "$CALIB_PATH" + --openpi-root "$OPENPI_ROOT" + --cotracker-checkpoint "$COTRACKER_CHECKPOINT" + --cotracker-device cuda:0 + --sam2-model "$SAM2_MODEL" + --sam2-device cuda:0 + --sam2-libs "$SAM2_LIBS" + --no-save-viz + --no-save-sam2-masks + ) + if [[ "$FORCE_REBUILD" == "1" ]]; then + command+=(--force) + fi + + attempt=1 + task_min_free="$GPU_MIN_FREE_MIB" + task_require_idle="$base_require_idle" + while true; do + if ((attempt > 1)); then + wait_for_gpu "$gpu" "$task_min_free" "$task_require_idle" 0 + fi + echo "$task_name -> GPU $gpu (attempt $attempt; log: $task_log)" \ + | tee -a "$task_log" + set +e + ( + export CUDA_VISIBLE_DEVICES="$gpu" + export CUDA_MODULE_LOADING="${CUDA_MODULE_LOADING:-LAZY}" + export PYTHONUNBUFFERED=1 + export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" + "${command[@]}" + ) 2>&1 | tee -a "$task_log" + pipeline_status=("${PIPESTATUS[@]}") + set -e + command_status="${pipeline_status[0]}" + tee_status="${pipeline_status[1]}" + + if [[ "$command_status" == "0" && "$tee_status" == "0" ]]; then + mv -- "$task_path" "$QUEUE_DIR/done_${task_suffix}.task" + break + fi + + if [[ "$attempt" == "1" ]] \ + && command -v rg >/dev/null 2>&1 \ + && rg -q 'OutOfMemoryError|CUDA out of memory' "$task_log"; then + echo "WARNING: GPU $gpu hit OOM; retrying only after it is fully idle" \ + | tee -a "$task_log" >&2 + attempt=2 + task_min_free="$GPU_RETRY_FREE_MIB" + task_require_idle=1 + continue + fi + + mv -- "$task_path" "$QUEUE_DIR/failed_${task_suffix}.task" + if [[ "$command_status" == "0" ]]; then + command_status="$tee_status" + fi + echo "ERROR: $task_name failed on GPU $gpu" >&2 + return "$command_status" + done + done + } + + PIDS=() + WORKER_NAMES=() + cleanup_workers() { + local pid + for pid in "${PIDS[@]:-}"; do + if kill -0 "$pid" 2>/dev/null; then + pkill -TERM -P "$pid" 2>/dev/null || true + kill -TERM "$pid" 2>/dev/null || true + fi + done + } + trap 'cleanup_workers; exit 130' INT TERM HUP + + for worker in "${!GPUS_NORMALIZED[@]}"; do + gpu="${GPUS_NORMALIZED[$worker]}" + printf -v worker_name 'gpu_worker_%02d_gpu_%s' "$worker" "$gpu" + echo " $worker_name started" + run_gpu_worker "$worker" "$gpu" & + PIDS+=("$!") + WORKER_NAMES+=("$worker_name") + done + + failed=0 + for index in "${!PIDS[@]}"; do + if ! wait "${PIDS[$index]}"; then + echo "ERROR: ${WORKER_NAMES[$index]} failed; see $SESSION_DIR" >&2 + failed=1 + fi + done + trap - INT TERM HUP + + shopt -s nullglob + unfinished_tasks=( + "$QUEUE_DIR"/pending_*.task + "$QUEUE_DIR"/running_*.task + "$QUEUE_DIR"/failed_*.task + ) + shopt -u nullglob + if [[ "${#unfinished_tasks[@]}" -gt 0 ]]; then + echo "Preprocessing stopped with ${#unfinished_tasks[@]} unfinished task(s)." >&2 + echo "Valid episodes and tracks are resumable; rerun the same command." >&2 + exit 1 + fi + if [[ "$failed" == "1" ]]; then + echo "WARNING: one or more GPU workers exited, but all queue tasks completed." >&2 + fi + + echo "Finalizing global metadata and canonical manifest..." + "$PYTHON_BIN" "$BUILDER" \ + --dataset-root "$DATASET_ROOT" \ + --track-cache "$TRACK_CACHE" \ + --manifest-path "$CANONICAL_MANIFEST" \ + --all \ + --no-extract-missing \ + --update-metadata \ + "$FK_FLAG" + + echo "Running final all-episode validation..." + "$PYTHON_BIN" "$BUILDER" \ + --dataset-root "$DATASET_ROOT" \ + --track-cache "$TRACK_CACHE" \ + --manifest-path "$CANONICAL_MANIFEST" \ + --all \ + --validate-only \ + "$FK_FLAG" + + echo "Prepared dataset: $DATASET_ROOT" +fi + +if [[ "$RUN_TRAINING" == "1" ]]; then + echo "Starting training with TREX_DATA_ROOT=$DATASET_ROOT" + export TREX_DATA_ROOT="$DATASET_ROOT" + export TRACK_CACHE="$TRACK_CACHE" + export PATH="$(dirname "$PYTHON_BIN"):$PATH" + exec bash "$TRAIN_SCRIPT" "${TRAIN_ARGS[@]}" +fi + +cat < str: + return datetime.now(timezone.utc).isoformat() + + +def _read_json(path: Path) -> dict: + return json.loads(path.read_text()) + + +def _read_jsonl(path: Path) -> list[dict]: + return [ + json.loads(line) + for line in path.read_text().splitlines() + if line.strip() + ] + + +def _atomic_write_json(path: Path, value: object) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fd, temporary_name = tempfile.mkstemp( + prefix=f".{path.name}.", + suffix=".tmp", + dir=path.parent, + ) + try: + with os.fdopen(fd, "w") as file: + json.dump(value, file, indent=2) + file.write("\n") + os.replace(temporary_name, path) + except Exception: + Path(temporary_name).unlink(missing_ok=True) + raise + + +def _atomic_write_jsonl(path: Path, rows: Iterable[dict]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fd, temporary_name = tempfile.mkstemp( + prefix=f".{path.name}.", + suffix=".tmp", + dir=path.parent, + ) + try: + with os.fdopen(fd, "w") as file: + for row in rows: + file.write(json.dumps(row) + "\n") + os.replace(temporary_name, path) + except Exception: + Path(temporary_name).unlink(missing_ok=True) + raise + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as file: + while chunk := file.read(8 * 1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def _video_keys(info: dict) -> list[str]: + return sorted( + key + for key, feature in info["features"].items() + if feature.get("dtype") == "video" + ) + + +def _source_episode_path(root: Path, episode_index: int, suffix: str) -> Path: + return ( + root + / suffix + / f"chunk-{episode_index // 1000:03d}" + / f"episode_{episode_index:06d}.parquet" + ) + + +def _episode_parquet_path(root: Path, episode_index: int) -> Path: + return ( + root + / "data" + / f"chunk-{episode_index // 1000:03d}" + / f"episode_{episode_index:06d}.parquet" + ) + + +def _episode_video_path(root: Path, episode_index: int, video_key: str) -> Path: + return ( + root + / "videos" + / f"chunk-{episode_index // 1000:03d}" + / video_key + / f"episode_{episode_index:06d}.mp4" + ) + + +def _track_path(root: Path, episode_index: int) -> Path: + return root / TRACK_CACHE_NAME / f"episode_{episode_index:06d}.npz" + + +def _build_registry( + episodes: Sequence[dict], + source_tasks: Sequence[dict], + *, + excluded: set[int], + source_limit: int | None, +) -> tuple[list[EpisodeRecord], list[dict]]: + tasks_by_text = { + str(row["task"]): int(row["task_index"]) + for row in source_tasks + } + selected = [ + episode + for episode in episodes + if int(episode["episode_index"]) not in excluded + and ( + source_limit is None + or int(episode["episode_index"]) < int(source_limit) + ) + ] + retained_tasks = { + str(task) + for episode in selected + for task in episode["tasks"] + } + missing_tasks = retained_tasks.difference(tasks_by_text) + if missing_tasks: + raise ValueError(f"episodes reference unknown tasks: {sorted(missing_tasks)[:3]}") + ordered_tasks = sorted(retained_tasks, key=tasks_by_text.__getitem__) + new_task_by_text = { + task: task_index + for task_index, task in enumerate(ordered_tasks) + } + task_rows = [ + {"task_index": task_index, "task": task} + for task, task_index in new_task_by_text.items() + ] + + registry: list[EpisodeRecord] = [] + global_index_start = 0 + for new_episode_index, episode in enumerate(selected): + tasks = [str(task) for task in episode["tasks"]] + if len(tasks) != 1: + raise ValueError( + f"episode {episode['episode_index']} has {len(tasks)} tasks; expected one" + ) + length = int(episode["length"]) + registry.append( + EpisodeRecord( + source_episode_index=int(episode["episode_index"]), + episode_index=new_episode_index, + length=length, + task=tasks[0], + task_index=new_task_by_text[tasks[0]], + global_index_start=global_index_start, + ) + ) + global_index_start += length + return registry, task_rows + + +def _replace_primitive_column(table, name: str, values: np.ndarray): + import pyarrow as pa + + index = table.schema.get_field_index(name) + if index < 0: + raise ValueError(f"missing required parquet column {name!r}") + field = table.schema.field(index) + return table.set_column(index, field, pa.array(values, type=field.type)) + + +def _parquet_is_ready( + path: Path, + *, + episode_index: int, + task_index: int, + global_index_start: int, + expected_rows: int, + force: bool, +) -> bool: + if not path.is_file(): + return False + try: + import pyarrow.parquet as pq + + parquet_file = pq.ParquetFile(path) + if int(parquet_file.metadata.num_rows) != int(expected_rows): + return False + metadata = parquet_file.schema_arrow.metadata or {} + if force and metadata.get(PARQUET_SCHEMA_METADATA_KEY) != ( + force_builder.SCHEMA_VERSION.encode() + ): + return False + table = pq.read_table( + path, + columns=["episode_index", "task_index", "frame_index", "index"], + ) + stored_episode = np.asarray(table["episode_index"].to_numpy()) + stored_task = np.asarray(table["task_index"].to_numpy()) + frame_index = np.asarray(table["frame_index"].to_numpy()) + global_index = np.asarray(table["index"].to_numpy()) + return bool( + np.all(stored_episode == int(episode_index)) + and np.all(stored_task == int(task_index)) + and np.array_equal(frame_index, np.arange(expected_rows)) + and np.array_equal( + global_index, + np.arange( + global_index_start, + global_index_start + expected_rows, + ), + ) + ) + except Exception: + return False + + +def _rewrite_parquet_job(job: tuple) -> tuple[int, str]: + ( + source_path_text, + destination_path_text, + episode_index, + task_index, + global_index_start, + expected_rows, + force, + ) = job + source_path = Path(source_path_text) + destination_path = Path(destination_path_text) + if _parquet_is_ready( + destination_path, + episode_index=episode_index, + task_index=task_index, + global_index_start=global_index_start, + expected_rows=expected_rows, + force=force, + ): + return episode_index, "skipped" + + import pyarrow.parquet as pq + + table = pq.read_table(source_path) + if int(table.num_rows) != int(expected_rows): + raise ValueError( + f"{source_path}: {table.num_rows} rows != expected {expected_rows}" + ) + original_metadata = table.schema.metadata + table = _replace_primitive_column( + table, + "episode_index", + np.full(expected_rows, episode_index, dtype=np.int64), + ) + table = _replace_primitive_column( + table, + "task_index", + np.full(expected_rows, task_index, dtype=np.int64), + ) + table = _replace_primitive_column( + table, + "frame_index", + np.arange(expected_rows, dtype=np.int64), + ) + table = _replace_primitive_column( + table, + "index", + np.arange( + global_index_start, + global_index_start + expected_rows, + dtype=np.int64, + ), + ) + table = table.replace_schema_metadata(original_metadata) + if force: + metadata = table.schema.metadata or {} + if metadata.get(PARQUET_SCHEMA_METADATA_KEY) != ( + force_builder.SCHEMA_VERSION.encode() + ): + raise ValueError(f"{source_path}: force schema metadata is missing") + + destination_path.parent.mkdir(parents=True, exist_ok=True) + temporary_path = destination_path.with_suffix(".tmp.parquet") + temporary_path.unlink(missing_ok=True) + try: + pq.write_table( + table, + temporary_path, + compression="zstd", + use_dictionary=True, + ) + os.replace(temporary_path, destination_path) + except Exception: + temporary_path.unlink(missing_ok=True) + raise + return episode_index, "written" + + +def _track_is_ready( + path: Path, + *, + episode_index: int, + expected_frames: int, +) -> bool: + if not path.is_file(): + return False + try: + with np.load(path, allow_pickle=False) as payload: + return bool( + int(np.asarray(payload["episode_index"]).item()) + == int(episode_index) + and int(np.asarray(payload["num_steps"]).item()) + == int(expected_frames) + ) + except Exception: + return False + + +def _rewrite_track_job(job: tuple) -> tuple[int, str, str]: + ( + source_path_text, + destination_path_text, + source_episode_index, + episode_index, + expected_frames, + ) = job + source_path = Path(source_path_text) + destination_path = Path(destination_path_text) + if _track_is_ready( + destination_path, + episode_index=episode_index, + expected_frames=expected_frames, + ): + return episode_index, _sha256(destination_path), "skipped" + + destination_path.parent.mkdir(parents=True, exist_ok=True) + if int(source_episode_index) == int(episode_index): + destination_path.unlink(missing_ok=True) + os.link(source_path, destination_path) + return episode_index, _sha256(destination_path), "linked" + + with np.load(source_path, allow_pickle=False) as archive: + payload = { + name: np.asarray(archive[name]).copy() + for name in archive.files + } + payload["episode_index"] = np.array(episode_index, dtype=np.int32) + payload["source_episode_index"] = np.array( + source_episode_index, + dtype=np.int32, + ) + temporary_path = destination_path.with_suffix(".tmp.npz") + temporary_path.unlink(missing_ok=True) + try: + with temporary_path.open("wb") as file: + np.savez_compressed(file, **payload) + os.replace(temporary_path, destination_path) + except Exception: + temporary_path.unlink(missing_ok=True) + raise + return episode_index, _sha256(destination_path), "written" + + +def _hardlink(source: Path, destination: Path) -> str: + if destination.exists(): + if os.path.samefile(source, destination): + return "skipped" + raise FileExistsError(f"destination is not the expected hard link: {destination}") + destination.parent.mkdir(parents=True, exist_ok=True) + os.link(source, destination) + return "linked" + + +def _link_job(job: tuple[str, str]) -> str: + return _hardlink(Path(job[0]), Path(job[1])) + + +def _run_jobs( + jobs: Sequence[tuple], + worker, + *, + workers: int, + process: bool, + label: str, +) -> list: + if not jobs: + return [] + executor_type = ProcessPoolExecutor if process else ThreadPoolExecutor + results = [] + with executor_type(max_workers=max(1, workers)) as executor: + futures = [executor.submit(worker, job) for job in jobs] + for completed, future in enumerate(as_completed(futures), start=1): + results.append(future.result()) + if completed % 250 == 0 or completed == len(futures): + print(f"{label}: {completed}/{len(futures)}", flush=True) + return results + + +def _write_variant_metadata_skeleton( + *, + variant_root: Path, + final_root: Path, + source_info: dict, + source_modality: dict, + source_embodiment: dict, + registry: Sequence[EpisodeRecord], + tasks: Sequence[dict], + video_keys: Sequence[str], + excluded_source_indices: Sequence[int], + force: bool, + source_dataset: Path, + blacklist_sha256: str, +) -> None: + info = copy.deepcopy(source_info) + info["total_episodes"] = len(registry) + info["total_frames"] = int(sum(record.length for record in registry)) + info["total_tasks"] = len(tasks) + info["total_videos"] = len(registry) * len(video_keys) + info["total_chunks"] = math.ceil(len(registry) / int(info["chunks_size"])) + info["splits"] = {"train": f"0:{len(registry)}"} + info.pop("discarded_episode_indices", None) + info.pop("trex_track_force", None) + info["trex_filter"] = { + "source_dataset": str(source_dataset), + "source_episode_count": int(source_info["total_episodes"]), + "source_to_new_map": "meta/source_episode_index_map.json", + "excluded_source_indices": "meta/excluded_source_episode_indices.json", + "blacklist_sha256": blacklist_sha256, + "force_variant": bool(force), + "created_at": _utc_now(), + } + metadata_dir = variant_root / "meta" + metadata_dir.mkdir(parents=True, exist_ok=True) + _atomic_write_json(metadata_dir / "info.json", info) + _atomic_write_json(metadata_dir / "modality.json", source_modality) + _atomic_write_json(metadata_dir / "embodiment.json", source_embodiment) + _atomic_write_jsonl( + metadata_dir / "episodes.jsonl", + ( + { + "episode_index": record.episode_index, + "tasks": [record.task], + "length": record.length, + } + for record in registry + ), + ) + _atomic_write_jsonl(metadata_dir / "tasks.jsonl", tasks) + _atomic_write_json( + metadata_dir / "source_episode_index_map.json", + { + str(record.source_episode_index): record.episode_index + for record in registry + }, + ) + _atomic_write_json( + metadata_dir / "episode_index_provenance.json", + [asdict(record) for record in registry], + ) + _atomic_write_json( + metadata_dir / "excluded_source_episode_indices.json", + list(excluded_source_indices), + ) + _atomic_write_json( + metadata_dir / "dataset_variant.json", + { + "name": final_root.name, + "final_root": str(final_root), + "force": bool(force), + "episodes": len(registry), + "frames": int(sum(record.length for record in registry)), + "tasks": len(tasks), + "video_keys": list(video_keys), + "created_at": _utc_now(), + }, + ) + + +def _rewrite_parquets( + *, + source_root: Path, + destination_root: Path, + registry: Sequence[EpisodeRecord], + workers: int, + force: bool, +) -> None: + jobs = [ + ( + str(_episode_parquet_path(source_root, record.source_episode_index)), + str(_episode_parquet_path(destination_root, record.episode_index)), + record.episode_index, + record.task_index, + record.global_index_start, + record.length, + force, + ) + for record in registry + ] + _run_jobs( + jobs, + _rewrite_parquet_job, + workers=workers, + process=True, + label=f"{destination_root.name} parquet", + ) + + +def _link_videos( + *, + source_root: Path, + destination_root: Path, + registry: Sequence[EpisodeRecord], + video_keys: Sequence[str], + source_uses_new_indices: bool, + workers: int, +) -> None: + jobs: list[tuple[str, str]] = [] + for record in registry: + source_episode_index = ( + record.episode_index + if source_uses_new_indices + else record.source_episode_index + ) + for video_key in video_keys: + jobs.append( + ( + str( + _episode_video_path( + source_root, + source_episode_index, + video_key, + ) + ), + str( + _episode_video_path( + destination_root, + record.episode_index, + video_key, + ) + ), + ) + ) + _run_jobs( + jobs, + _link_job, + workers=workers, + process=False, + label=f"{destination_root.name} videos", + ) + + +def _rewrite_tracks( + *, + source_root: Path, + destination_root: Path, + registry: Sequence[EpisodeRecord], + workers: int, +) -> dict[int, str]: + jobs = [ + ( + str(_track_path(source_root, record.source_episode_index)), + str(_track_path(destination_root, record.episode_index)), + record.source_episode_index, + record.episode_index, + record.length, + ) + for record in registry + ] + results = _run_jobs( + jobs, + _rewrite_track_job, + workers=workers, + process=True, + label=f"{destination_root.name} tracks", + ) + return {int(episode_index): sha256 for episode_index, sha256, _ in results} + + +def _link_subset_files( + *, + source_root: Path, + destination_root: Path, + registry: Sequence[EpisodeRecord], + video_keys: Sequence[str], + force: bool, + workers: int, +) -> dict[int, str]: + parquet_jobs = [ + ( + str(_episode_parquet_path(source_root, record.episode_index)), + str(_episode_parquet_path(destination_root, record.episode_index)), + ) + for record in registry + ] + _run_jobs( + parquet_jobs, + _link_job, + workers=workers, + process=False, + label=f"{destination_root.name} parquet links", + ) + _link_videos( + source_root=source_root, + destination_root=destination_root, + registry=registry, + video_keys=video_keys, + source_uses_new_indices=True, + workers=workers, + ) + track_hashes: dict[int, str] = {} + if force: + track_jobs = [ + ( + str(_track_path(source_root, record.episode_index)), + str(_track_path(destination_root, record.episode_index)), + ) + for record in registry + ] + _run_jobs( + track_jobs, + _link_job, + workers=workers, + process=False, + label=f"{destination_root.name} track links", + ) + track_hashes = { + record.episode_index: _sha256( + _track_path(destination_root, record.episode_index) + ) + for record in registry + } + return track_hashes + + +def _aggregate_source_manifest_entries(source_force_root: Path) -> dict[int, dict]: + manifests = sorted( + ( + source_force_root + / "meta" + / "trex_track_force_prepare" + ).glob("run_*/task_*.json"), + key=lambda path: path.stat().st_mtime_ns, + ) + entries: dict[int, dict] = {} + for path in manifests: + try: + payload = _read_json(path) + except Exception: + continue + for key, entry in payload.get("episodes", {}).items(): + if entry.get("status") == "complete": + entries[int(key)] = copy.deepcopy(entry) + expected = int( + _read_json(source_force_root / "meta" / "info.json")["total_episodes"] + ) + missing = [index for index in range(expected) if index not in entries] + if missing: + raise ValueError( + f"source task manifests do not cover every episode: {missing[:10]}" + ) + return entries + + +def _build_force_manifest( + *, + variant_root: Path, + final_root: Path, + registry: Sequence[EpisodeRecord], + source_entries: dict[int, dict], + track_hashes: dict[int, str], +) -> dict: + final_track_cache = final_root / TRACK_CACHE_NAME + manifest = force_builder._new_manifest(final_root, final_track_cache) + manifest["filter_provenance"] = { + "source_to_new_map": "meta/source_episode_index_map.json", + "excluded_source_indices": "meta/excluded_source_episode_indices.json", + } + for record in registry: + entry = copy.deepcopy(source_entries[record.source_episode_index]) + entry["episode_index"] = record.episode_index + entry["source_episode_index"] = record.source_episode_index + entry["num_frames"] = record.length + entry["parquet"] = str( + _episode_parquet_path(Path("."), record.episode_index) + ) + entry["track_npz"] = str( + final_track_cache / f"episode_{record.episode_index:06d}.npz" + ) + entry["track_sha256"] = track_hashes[record.episode_index] + entry["status"] = "complete" + entry["skipped"] = False + manifest["episodes"][f"{record.episode_index:06d}"] = entry + _atomic_write_json( + variant_root / "meta" / "trex_track_force_manifest.json", + manifest, + ) + return manifest + + +def _compute_base_stats(dataset_root: Path, registry: Sequence[EpisodeRecord]) -> dict: + import pyarrow.parquet as pq + + stats: dict[str, dict] = {} + for column in ("observation.state", "action", "timestamp"): + parts: list[np.ndarray] = [] + for record in registry: + table = pq.read_table( + _episode_parquet_path(dataset_root, record.episode_index), + columns=[column], + ) + values = force_builder._column_to_numpy(table, column) + if values.ndim == 1: + values = values[:, None] + parts.append(values) + stats[column] = force_builder._statistics(np.concatenate(parts, axis=0)) + return stats + + +def _compute_base_relative_stats( + dataset_root: Path, + registry: Sequence[EpisodeRecord], +) -> dict: + import pyarrow.parquet as pq + + action_offsets = range(24) + output: dict[str, dict] = {} + for name, selection in ( + ("left_arm", slice(0, 7)), + ("right_arm", slice(29, 36)), + ): + parts: list[np.ndarray] = [] + for record in registry: + table = pq.read_table( + _episode_parquet_path(dataset_root, record.episode_index), + columns=["observation.state", "action"], + ) + state = force_builder._column_to_numpy( + table, + "observation.state", + )[:, selection] + action = force_builder._column_to_numpy( + table, + "action", + )[:, selection] + usable_length = len(state) - max(action_offsets) + if usable_length <= 0: + continue + reference = state[:usable_length] + parts.extend( + action[offset : offset + usable_length] - reference + for offset in action_offsets + ) + if not parts: + raise ValueError(f"no relative action samples for {name}") + output[name] = force_builder._statistics( + np.concatenate(parts, axis=0) + ) + return output + + +def _finalize_metadata( + *, + variant_root: Path, + registry: Sequence[EpisodeRecord], + force: bool, +) -> None: + stats_path = variant_root / "meta" / "stats.json" + if force: + base_variant_root = variant_root.parent / variant_root.name.removesuffix( + "_force" + ) + base_stats_path = base_variant_root / "meta" / "stats.json" + if not base_stats_path.is_file(): + raise FileNotFoundError( + f"base variant stats must be finalized first: {base_stats_path}" + ) + _atomic_write_json(stats_path, _read_json(base_stats_path)) + else: + _atomic_write_json( + stats_path, + _compute_base_stats(variant_root, registry), + ) + if not force: + _atomic_write_json( + variant_root / "meta" / "relative_stats_dreamzero.json", + _compute_base_relative_stats(variant_root, registry), + ) + return + result = force_builder.update_metadata( + variant_root, + assume_all_converted=True, + ) + manifest_path = ( + variant_root / "meta" / "trex_track_force_manifest.json" + ) + manifest = _read_json(manifest_path) + manifest["metadata"] = result + manifest["updated_at"] = _utc_now() + _atomic_write_json(manifest_path, manifest) + force_builder.validate_metadata(variant_root) + for backup in (variant_root / "meta").glob("*.trex_track_force.bak"): + backup.unlink() + + +def _copy_audit(source_force_root: Path, destination_force_root: Path) -> None: + source = source_force_root / "audit" + destination = destination_force_root / "audit" / "source_dataset" + if destination.exists() or not source.exists(): + return + shutil.copytree(source, destination) + + +def _validate_variant( + *, + root: Path, + expected_registry: Sequence[EpisodeRecord], + video_keys: Sequence[str], + force: bool, +) -> dict: + import pyarrow.parquet as pq + + info = _read_json(root / "meta" / "info.json") + episodes = _read_jsonl(root / "meta" / "episodes.jsonl") + tasks = _read_jsonl(root / "meta" / "tasks.jsonl") + if int(info["total_episodes"]) != len(expected_registry): + raise ValueError(f"{root}: total_episodes is stale") + if int(info["total_frames"]) != sum(r.length for r in expected_registry): + raise ValueError(f"{root}: total_frames is stale") + if int(info["total_tasks"]) != len(tasks): + raise ValueError(f"{root}: total_tasks is stale") + if int(info["total_videos"]) != len(expected_registry) * len(video_keys): + raise ValueError(f"{root}: total_videos is stale") + if [int(row["episode_index"]) for row in episodes] != list( + range(len(expected_registry)) + ): + raise ValueError(f"{root}: episodes.jsonl is not dense") + if [int(row["task_index"]) for row in tasks] != list(range(len(tasks))): + raise ValueError(f"{root}: tasks.jsonl is not dense") + + manifest = None + if force: + force_builder.validate_metadata(root) + manifest = _read_json( + root / "meta" / "trex_track_force_manifest.json" + ) + if len(manifest.get("episodes", {})) != len(expected_registry): + raise ValueError(f"{root}: force manifest count is stale") + + for completed, record in enumerate(expected_registry, start=1): + path = _episode_parquet_path(root, record.episode_index) + if not _parquet_is_ready( + path, + episode_index=record.episode_index, + task_index=record.task_index, + global_index_start=record.global_index_start, + expected_rows=record.length, + force=force, + ): + raise ValueError(f"{root}: invalid parquet {path}") + for video_key in video_keys: + video_path = _episode_video_path( + root, + record.episode_index, + video_key, + ) + if not video_path.is_file() or video_path.stat().st_size <= 0: + raise ValueError(f"{root}: missing video {video_path}") + if force: + track_path = _track_path(root, record.episode_index) + if not _track_is_ready( + track_path, + episode_index=record.episode_index, + expected_frames=record.length, + ): + raise ValueError(f"{root}: invalid track cache {track_path}") + entry = manifest["episodes"].get( + f"{record.episode_index:06d}", + {}, + ) + if ( + entry.get("status") != "complete" + or int(entry.get("source_episode_index", -1)) + != record.source_episode_index + ): + raise ValueError( + f"{root}: invalid manifest entry {record.episode_index}" + ) + if completed % 250 == 0 or completed == len(expected_registry): + print( + f"validate {root.name}: {completed}/{len(expected_registry)}", + flush=True, + ) + + sample_indices = sorted( + { + 0, + len(expected_registry) // 2, + len(expected_registry) - 1, + } + ) + for episode_index in sample_indices: + parquet_file = pq.ParquetFile( + _episode_parquet_path(root, episode_index) + ) + if int(parquet_file.metadata.num_rows) <= 0: + raise ValueError(f"{root}: empty sample parquet {episode_index}") + if force: + force_builder.validate_episode_parquet( + _episode_parquet_path(root, episode_index), + verify_source_fk=False, + ) + result = { + "root": str(root), + "episodes": len(expected_registry), + "frames": int(sum(record.length for record in expected_registry)), + "tasks": len(tasks), + "videos": len(expected_registry) * len(video_keys), + "force": force, + "validated_at": _utc_now(), + } + _atomic_write_json(root / "meta" / "dataset_ready.json", result) + return result + + +def _stage_paths(stage_root: Path) -> dict[str, Path]: + return {name: stage_root / name for name in VARIANT_NAMES} + + +def _load_inputs(args: argparse.Namespace): + source_info = _read_json(args.base_root / "meta" / "info.json") + source_modality = _read_json(args.base_root / "meta" / "modality.json") + source_embodiment = _read_json(args.base_root / "meta" / "embodiment.json") + episodes = _read_jsonl(args.base_root / "meta" / "episodes.jsonl") + source_tasks = _read_jsonl(args.base_root / "meta" / "tasks.jsonl") + excluded = set(json.loads(args.blacklist.read_text())) + if len(excluded) != 115: + raise ValueError( + f"expected 115 excluded episodes, got {len(excluded)}" + ) + full_registry, full_tasks = _build_registry( + episodes, + source_tasks, + excluded=excluded, + source_limit=None, + ) + small_registry, small_tasks = _build_registry( + episodes, + source_tasks, + excluded=excluded, + source_limit=100, + ) + if len(full_registry) != 5349 or len(small_registry) != 100: + raise ValueError( + f"unexpected registry sizes: full={len(full_registry)}, " + f"small={len(small_registry)}" + ) + return ( + source_info, + source_modality, + source_embodiment, + episodes, + source_tasks, + excluded, + full_registry, + full_tasks, + small_registry, + small_tasks, + ) + + +def _build(args: argparse.Namespace) -> None: + ( + source_info, + source_modality, + source_embodiment, + _, + _, + excluded, + full_registry, + full_tasks, + small_registry, + small_tasks, + ) = _load_inputs(args) + paths = _stage_paths(args.stage_root) + video_keys = _video_keys(source_info) + blacklist_sha256 = _sha256(args.blacklist) + final_paths = { + name: args.data_root / name + for name in VARIANT_NAMES + } + args.stage_root.mkdir(parents=True, exist_ok=True) + + for name, registry, tasks, force, source_dataset in ( + ( + "trex_full", + full_registry, + full_tasks, + False, + args.base_root, + ), + ( + "trex_full_force", + full_registry, + full_tasks, + True, + args.force_root, + ), + ( + "trex_small", + small_registry, + small_tasks, + False, + args.base_root, + ), + ( + "trex_small_force", + small_registry, + small_tasks, + True, + args.force_root, + ), + ): + _write_variant_metadata_skeleton( + variant_root=paths[name], + final_root=final_paths[name], + source_info=source_info, + source_modality=source_modality, + source_embodiment=source_embodiment, + registry=registry, + tasks=tasks, + video_keys=video_keys, + excluded_source_indices=( + sorted(excluded) if "full" in name else [] + ), + force=force, + source_dataset=source_dataset, + blacklist_sha256=blacklist_sha256, + ) + + _rewrite_parquets( + source_root=args.base_root, + destination_root=paths["trex_full"], + registry=full_registry, + workers=args.parquet_workers, + force=False, + ) + _link_videos( + source_root=args.base_root, + destination_root=paths["trex_full"], + registry=full_registry, + video_keys=video_keys, + source_uses_new_indices=False, + workers=args.link_workers, + ) + + _rewrite_parquets( + source_root=args.force_root, + destination_root=paths["trex_full_force"], + registry=full_registry, + workers=args.parquet_workers, + force=True, + ) + _link_videos( + source_root=paths["trex_full"], + destination_root=paths["trex_full_force"], + registry=full_registry, + video_keys=video_keys, + source_uses_new_indices=True, + workers=args.link_workers, + ) + full_track_hashes = _rewrite_tracks( + source_root=args.force_root, + destination_root=paths["trex_full_force"], + registry=full_registry, + workers=args.track_workers, + ) + + source_entries = _aggregate_source_manifest_entries(args.force_root) + _build_force_manifest( + variant_root=paths["trex_full_force"], + final_root=final_paths["trex_full_force"], + registry=full_registry, + source_entries=source_entries, + track_hashes=full_track_hashes, + ) + _copy_audit(args.force_root, paths["trex_full_force"]) + + for small_record, full_record in zip( + small_registry, + full_registry[: len(small_registry)], + ): + if ( + small_record.source_episode_index + != full_record.source_episode_index + or small_record.episode_index != full_record.episode_index + or small_record.task_index != full_record.task_index + or small_record.global_index_start + != full_record.global_index_start + ): + raise ValueError("small is not an identity prefix of full") + + _link_subset_files( + source_root=paths["trex_full"], + destination_root=paths["trex_small"], + registry=small_registry, + video_keys=video_keys, + force=False, + workers=args.link_workers, + ) + small_track_hashes = _link_subset_files( + source_root=paths["trex_full_force"], + destination_root=paths["trex_small_force"], + registry=small_registry, + video_keys=video_keys, + force=True, + workers=args.link_workers, + ) + _build_force_manifest( + variant_root=paths["trex_small_force"], + final_root=final_paths["trex_small_force"], + registry=small_registry, + source_entries=source_entries, + track_hashes=small_track_hashes, + ) + _atomic_write_json( + args.stage_root / "build_complete.json", + { + "completed_at": _utc_now(), + "variants": list(VARIANT_NAMES), + }, + ) + + +def _metadata(args: argparse.Namespace) -> None: + ( + _, + _, + _, + _, + _, + _, + full_registry, + _, + small_registry, + _, + ) = _load_inputs(args) + paths = _stage_paths(args.stage_root) + selected = set(args.metadata_variants) + for name, registry, force in ( + ("trex_small", small_registry, False), + ("trex_small_force", small_registry, True), + ("trex_full", full_registry, False), + ("trex_full_force", full_registry, True), + ): + if name not in selected: + continue + print(f"Finalizing metadata: {name}", flush=True) + _finalize_metadata( + variant_root=paths[name], + registry=registry, + force=force, + ) + _atomic_write_json( + args.stage_root / "metadata_complete.json", + {"completed_at": _utc_now()}, + ) + + +def _validate(args: argparse.Namespace) -> None: + ( + source_info, + _, + _, + _, + _, + _, + full_registry, + _, + small_registry, + _, + ) = _load_inputs(args) + paths = _stage_paths(args.stage_root) + video_keys = _video_keys(source_info) + results = [] + for name, registry, force in ( + ("trex_small", small_registry, False), + ("trex_small_force", small_registry, True), + ("trex_full", full_registry, False), + ("trex_full_force", full_registry, True), + ): + results.append( + _validate_variant( + root=paths[name], + expected_registry=registry, + video_keys=video_keys, + force=force, + ) + ) + _atomic_write_json( + args.stage_root / "validation_complete.json", + { + "validated_at": _utc_now(), + "results": results, + }, + ) + + +def _install(args: argparse.Namespace) -> None: + validation_marker = args.stage_root / "validation_complete.json" + if not validation_marker.is_file(): + raise RuntimeError("staging validation marker is missing") + paths = _stage_paths(args.stage_root) + if any(not path.is_dir() for path in paths.values()): + raise RuntimeError("one or more staged variants are missing") + + backup_root = args.data_root / ( + f".trex_old_{datetime.now().strftime('%Y%m%dT%H%M%S')}" + ) + backup_root.mkdir(parents=True, exist_ok=False) + moved_old: list[tuple[Path, Path]] = [] + installed: list[tuple[Path, Path]] = [] + try: + for name in dict.fromkeys((*LEGACY_NAMES, *VARIANT_NAMES)): + current = args.data_root / name + if current.exists(): + backup = backup_root / name + os.replace(current, backup) + moved_old.append((backup, current)) + for name in VARIANT_NAMES: + staged = paths[name] + final = args.data_root / name + os.replace(staged, final) + installed.append((final, staged)) + + ( + source_info, + _, + _, + _, + _, + _, + full_registry, + _, + small_registry, + _, + ) = _load_inputs_from_backups(args, backup_root) + video_keys = _video_keys(source_info) + for name, registry, force in ( + ("trex_small", small_registry, False), + ("trex_small_force", small_registry, True), + ("trex_full", full_registry, False), + ("trex_full_force", full_registry, True), + ): + _validate_variant( + root=args.data_root / name, + expected_registry=registry, + video_keys=video_keys, + force=force, + ) + except Exception: + for final, staged in reversed(installed): + if final.exists(): + os.replace(final, staged) + for backup, current in reversed(moved_old): + if backup.exists(): + os.replace(backup, current) + backup_root.rmdir() + raise + + shutil.rmtree(backup_root) + args.stage_root.mkdir(parents=True, exist_ok=True) + for marker in args.stage_root.glob("*.json"): + marker.unlink() + args.stage_root.rmdir() + + +def _load_inputs_from_backups(args: argparse.Namespace, backup_root: Path): + backup_args = copy.copy(args) + backup_args.base_root = backup_root / "trex_datasetv2" + backup_args.force_root = backup_root / "trex_track_force_v2" + backup_args.blacklist = ( + backup_args.force_root + / "audit" + / "track_quality" + / "frozen_wrist_episode_indices.json" + ) + return _load_inputs(backup_args) + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--data-root", + type=Path, + default=_DREAMZERO_ROOT / "data", + ) + parser.add_argument( + "--base-root", + type=Path, + default=_DREAMZERO_ROOT / "data" / "trex_datasetv2", + ) + parser.add_argument( + "--force-root", + type=Path, + default=_DREAMZERO_ROOT / "data" / "trex_track_force_v2", + ) + parser.add_argument( + "--blacklist", + type=Path, + default=( + _DREAMZERO_ROOT + / "data" + / "trex_track_force_v2" + / "audit" + / "track_quality" + / "frozen_wrist_episode_indices.json" + ), + ) + parser.add_argument( + "--stage-root", + type=Path, + default=_DREAMZERO_ROOT / "data" / ".trex_variants_staging", + ) + parser.add_argument( + "--phase", + choices=("build", "metadata", "validate", "install", "all"), + default="all", + ) + parser.add_argument("--parquet-workers", type=int, default=8) + parser.add_argument("--track-workers", type=int, default=8) + parser.add_argument("--link-workers", type=int, default=32) + parser.add_argument( + "--metadata-variants", + nargs="+", + choices=VARIANT_NAMES, + default=list(VARIANT_NAMES), + help="variants to process during the metadata phase", + ) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + args = _build_parser().parse_args(argv) + for name in ("data_root", "base_root", "force_root", "blacklist", "stage_root"): + setattr(args, name, getattr(args, name).expanduser().resolve()) + if args.phase in ("build", "all"): + _build(args) + if args.phase in ("metadata", "all"): + _metadata(args) + if args.phase in ("validate", "all"): + _validate(args) + if args.phase in ("install", "all"): + _install(args) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/data/render_trex_track_samples.py b/scripts/data/render_trex_track_samples.py new file mode 100644 index 0000000000000000000000000000000000000000..28984ea556abb0daf6a8ea257e275e5c681a8c43 --- /dev/null +++ b/scripts/data/render_trex_track_samples.py @@ -0,0 +1,161 @@ +#!/usr/bin/env python3 +"""Render uniformly sampled T-Rex track-cache episodes without rerunning models.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import numpy as np + +_SCRIPT_DIR = Path(__file__).resolve().parent +_REPO_ROOT = _SCRIPT_DIR.parents[1] +_SCRIPTS_DIR = _REPO_ROOT / "scripts" +if str(_SCRIPTS_DIR) not in sys.path: + sys.path.insert(0, str(_SCRIPTS_DIR)) + +from extract_track import load_episode_videos # noqa: E402 +from trex_track.layout import VIEW_ORDER # noqa: E402 +from trex_track.trex_viz_tracks import render_three_view_combined_video # noqa: E402 + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--dataset-root", + type=Path, + default=_REPO_ROOT / "data" / "trex_full_force", + ) + parser.add_argument("--track-cache", type=Path, default=None) + parser.add_argument( + "--output-dir", + type=Path, + default=_REPO_ROOT / "outputs" / "trex_track_samples", + ) + parser.add_argument("--num-samples", type=int, default=3) + parser.add_argument("--episodes", type=int, nargs="*", default=None) + parser.add_argument("--fps", type=int, default=0) + parser.add_argument("--trail", type=int, default=15) + return parser.parse_args() + + +def _uniform_episode_indices(total_episodes: int, count: int) -> list[int]: + if count < 1: + raise ValueError("--num-samples must be positive") + if count > total_episodes: + raise ValueError("--num-samples cannot exceed total episodes") + return [ + int(index) + for index in np.rint( + np.linspace(0, total_episodes - 1, count, dtype=np.float64) + ) + ] + + +def _load_tracks(path: Path) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]: + if not path.is_file(): + raise FileNotFoundError(path) + with np.load(path, allow_pickle=False) as payload: + tracks = { + view: np.asarray(payload[f"tracks_{view}"], dtype=np.float32) + for view in VIEW_ORDER + } + visibility = { + view: np.asarray(payload[f"vis_{view}"], dtype=np.float32) + for view in VIEW_ORDER + } + for view in VIEW_ORDER: + if visibility[view].shape != tracks[view].shape[:2]: + raise ValueError( + f"{path}: {view} visibility {visibility[view].shape} " + f"does not match tracks {tracks[view].shape}" + ) + return tracks, visibility + + +def main() -> int: + args = _parse_args() + dataset_root = args.dataset_root.expanduser().resolve() + track_cache = ( + args.track_cache.expanduser().resolve() + if args.track_cache is not None + else dataset_root / "tracks_trex_track_force_v2" + ) + output_dir = args.output_dir.expanduser().resolve() + info = json.loads((dataset_root / "meta" / "info.json").read_text()) + total_episodes = int(info["total_episodes"]) + video_shape = info["features"]["observation.images.head_left"]["shape"] + out_hw = (int(video_shape[0]), int(video_shape[1])) + fps = int(args.fps) if args.fps > 0 else int(info["fps"]) + + episode_indices = ( + [int(index) for index in args.episodes] + if args.episodes + else _uniform_episode_indices(total_episodes, int(args.num_samples)) + ) + if any(index < 0 or index >= total_episodes for index in episode_indices): + raise ValueError(f"episode indices must be in [0, {total_episodes})") + + output_dir.mkdir(parents=True, exist_ok=True) + rendered: list[dict[str, object]] = [] + for episode_index in episode_indices: + cache_path = track_cache / f"episode_{episode_index:06d}.npz" + output_path = output_dir / f"episode_{episode_index:06d}_tracks.mp4" + print(f"Rendering episode {episode_index}: {output_path}", flush=True) + tracks, visibility = _load_tracks(cache_path) + videos = load_episode_videos( + dataset_root, + episode_index, + out_hw=out_hw, + ) + frame_counts = { + view: int(videos[view].shape[0]) + for view in VIEW_ORDER + } + for view in VIEW_ORDER: + if int(tracks[view].shape[0]) != frame_counts[view]: + raise ValueError( + f"episode {episode_index} {view}: " + f"{tracks[view].shape[0]} track frames != " + f"{frame_counts[view]} video frames" + ) + render_three_view_combined_video( + view_images=videos, + view_tracks=tracks, + view_vis=visibility, + out_path=output_path, + fps=fps, + draw_trail=int(args.trail), + dim_low_vis=True, + ) + rendered.append( + { + "episode_index": episode_index, + "frames": min(frame_counts.values()), + "fps": fps, + "track_cache": str(cache_path), + "video": str(output_path), + } + ) + del tracks, visibility, videos + + summary_path = output_dir / "samples.json" + summary_path.write_text( + json.dumps( + { + "dataset_root": str(dataset_root), + "sampling": "uniform endpoints and midpoint", + "episodes": rendered, + }, + indent=2, + ) + + "\n" + ) + print(f"Wrote {summary_path}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/eval/trex_track_force_inference.py b/scripts/eval/trex_track_force_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..d2d79a681e7700d75c4e852971bb9866740766d9 --- /dev/null +++ b/scripts/eval/trex_track_force_inference.py @@ -0,0 +1,146 @@ +#!/usr/bin/env python3 +"""Run one independent T-Rex Track-Force 16-step chunk from an NPZ prefix. + +Required NPZ arrays: + head_left, left_wrist, right_wrist: uint8 [9,H,W,3] + state_eef62: float [62] + track_past_xy: float [16,250,2] in [0,1] + track_past_visibility: float/bool [16,250] + tactile_force_history: float [16,10,6] (raw sensor units) + +The last force-history sample is used as current force. Output contains the +normalized model action, physical delta-base EEF62 action, absolute EEF62 +targets, and predicted future tracks. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import numpy as np +import torch +from transformers import AutoTokenizer + +from groot.vla.model.dreamzero.transform.dreamzero_cotrain import ( + basic_clean, + whitespace_clean, +) +from groot.vla.model.trex_track_force.runtime import ( + TrexRuntimeStatistics, + delta_base_to_absolute, +) +from groot.vla.model.trex_track_force.vla import TrexTrackForceVLA + + +def _grid_three_views(archive: np.lib.npyio.NpzFile) -> np.ndarray: + views = [ + np.asarray(archive[name], dtype=np.uint8) + for name in ("head_left", "left_wrist", "right_wrist") + ] + if any(view.ndim != 4 or view.shape[0] != 9 or view.shape[-1] != 3 for view in views): + raise ValueError("each RGB view must be uint8 [9,H,W,3]") + if len({view.shape for view in views}) != 1: + raise ValueError("all three RGB histories must have the same shape") + _, height, width, channels = views[0].shape + grid = np.zeros((9, 2 * height, 2 * width, channels), dtype=np.uint8) + grid[:, :height, :width] = views[0] + grid[:, height:, :width] = views[1] + grid[:, :height, width:] = views[2] + return grid + + +def _pad64(values: np.ndarray) -> np.ndarray: + return np.pad(values, ((0, 0), (0, 2)), mode="constant") + + +def run(args: argparse.Namespace) -> None: + device = torch.device(args.device) + dtype = torch.bfloat16 if args.bf16 else torch.float32 + stats = TrexRuntimeStatistics.from_dataset(args.dataset_root) + with np.load(args.input_npz, allow_pickle=False) as archive: + history_images = _grid_three_views(archive) + reference_state = np.asarray(archive["state_eef62"], dtype=np.float32) + track_xy = np.asarray(archive["track_past_xy"], dtype=np.float32) + track_visibility = np.asarray( + archive["track_past_visibility"], dtype=np.float32 + ) + force_history_raw = np.asarray( + archive["tactile_force_history"], dtype=np.float32 + ) + if reference_state.shape != (62,): + raise ValueError("state_eef62 must be [62]") + if track_xy.shape != (16, 250, 2) or track_visibility.shape != (16, 250): + raise ValueError("past tracks must be [16,250,2] and [16,250]") + if force_history_raw.shape != (16, 10, 6): + raise ValueError("tactile_force_history must be [16,10,6]") + + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path) + instruction = whitespace_clean(basic_clean(args.instruction)) + text = tokenizer( + instruction, + max_length=512, + padding="max_length", + truncation=True, + return_tensors="pt", + ) + model = TrexTrackForceVLA.load_lora(str(args.checkpoint)) + model.eval().requires_grad_(False) + model.to(device=device, dtype=dtype) + + normalized_state = _pad64(stats.normalize_state(reference_state)[None]) + normalized_force_history = stats.normalize_force(force_history_raw) + inputs = { + "history_images": torch.from_numpy(history_images[None]).to(device), + "state": torch.from_numpy(normalized_state[:, None]).to(device, dtype), + "track_past_xy": torch.from_numpy(track_xy[None, None]).to(device, dtype), + "track_past_visibility": torch.from_numpy( + track_visibility[None, None] + ).to(device, dtype), + "current_force": torch.from_numpy( + normalized_force_history[-1:][None] + ).to(device, dtype), + "tactile_force_history": torch.from_numpy( + normalized_force_history[None, None] + ).to(device, dtype), + "text": text.input_ids.to(device), + "text_attention_mask": text.attention_mask.to(device), + } + with torch.inference_mode(): + prediction = model.get_action(inputs) + normalized_action = prediction["action_pred"].float().cpu().numpy()[0] + delta_base_action = stats.denormalize_action(normalized_action) + absolute_action = delta_base_to_absolute(reference_state, delta_base_action) + output = { + "normalized_action64": normalized_action, + "delta_base_action62": delta_base_action, + "absolute_action62": absolute_action, + "track_pred": prediction["track_pred"].float().cpu().numpy()[0], + } + if "video_latents_pred" in prediction: + output["video_latents_pred"] = ( + prediction["video_latents_pred"].float().cpu().numpy()[0] + ) + args.output_npz.parent.mkdir(parents=True, exist_ok=True) + np.savez_compressed(args.output_npz, **output) + print(f"wrote {args.output_npz} with 16 actions at 20 Hz") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--input-npz", type=Path, required=True) + parser.add_argument("--output-npz", type=Path, required=True) + parser.add_argument("--dataset-root", type=Path, required=True) + parser.add_argument("--tokenizer-path", type=Path, required=True) + parser.add_argument( + "--instruction", + default="Perform the requested bimanual manipulation.", + ) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--bf16", action=argparse.BooleanOptionalAction, default=True) + return parser.parse_args() + + +if __name__ == "__main__": + run(parse_args()) diff --git a/scripts/extract_track.py b/scripts/extract_track.py new file mode 100644 index 0000000000000000000000000000000000000000..3e5c4c03db643545c05a57b85904971bd65e1b1e --- /dev/null +++ b/scripts/extract_track.py @@ -0,0 +1,862 @@ +""" +Extract T-Rex 250-point tracks for LeRobot v2 episodes with SAM2 + CoTracker. + +Frame-0 masks: one SAM2 predict per hand/view (fixed prompts, no appearance +auto-repair). Then CoTracker tracks: + - head_left: 100 pts (left hand/arm 50 + right hand/arm 50) + - each wrist: 75 pts (fixed 5×5 background grid 25 + hand 50) + +Tune prompts in ``trex_track/sam2_prompt_hands.py``. + +Example:: + + CUDA_VISIBLE_DEVICES=4 python scripts/extract_track.py \\ + --dataset-root data/trex_small --episode-index 0 \\ + --cotracker-device cuda:0 --sam2-device cuda:0 +""" + +from __future__ import annotations + +import argparse +import os +import sys +import tempfile +from pathlib import Path +from typing import NamedTuple + +import numpy as np + +_SCRIPT_DIR = Path(__file__).resolve().parent +_DREAMZERO_ROOT = _SCRIPT_DIR.parent +if str(_SCRIPT_DIR) not in sys.path: + sys.path.insert(0, str(_SCRIPT_DIR)) + +from trex_track.layout import ( # noqa: E402 + NUM_COMBINED_POINTS, + NUM_HEAD_LEFT, + NUM_HEAD_PER_HAND as NUM_HAND_POINTS, + NUM_HEAD_POINTS, + NUM_WRIST_BACKGROUND as NUM_WRIST_GRID, + NUM_WRIST_HAND, + NUM_WRIST_POINTS, + POINT_SLICES, + TRACK_LAYOUT_VERSION, + VIEW_ORDER, + identity_metadata, +) + +VIDEO_FOLDERS = { + "head_left": "observation.images.head_left", + "left_wrist": "observation.images.left_wrist", + "right_wrist": "observation.images.right_wrist", +} + +DEFAULT_OPENPI_ROOT = Path("/scratch2/home/zhicao/openpi") +DEFAULT_CALIB = _DREAMZERO_ROOT / "assets" / "trex_camera_calib.json" +DEFAULT_SAM2_MODEL = os.environ.get("SAM2_MODEL", "facebook/sam2-hiera-large") +DEFAULT_SAM2_LIBS = os.environ.get("SAM2_LIBS", "/scratch1/home/zhicao/physctrl/libs") + + +class TrackingRuntime(NamedTuple): + """Lazily-created heavy models shared across an episode batch.""" + + calib: dict + out_hw: tuple[int, int] + cotracker_model: object + cotracker_device: object + sam2_predictor: object + + +def _ensure_openpi_on_path(openpi_root: Path) -> Path: + root = openpi_root.expanduser().resolve() + droid = root / "droid" + if not droid.is_dir(): + raise FileNotFoundError(f"openpi droid package not found: {droid}") + p = str(droid) + if p not in sys.path: + sys.path.insert(0, p) + return droid + + +def _enable_cotracker_sdpa_attention(openpi_root: str | Path) -> None: + """Replace CoTracker's quadratic-memory attention with PyTorch SDPA. + + The original implementation materializes an ``[B,H,T,T]`` attention + matrix. Long T-Rex episodes can therefore require more than 80 GB even on + an otherwise empty H100. SDPA uses Flash Attention for CUDA BF16 inputs, + preserving full-sequence attention without materializing that matrix. + """ + + import torch.nn.functional as F + + cotracker_root = Path(openpi_root).expanduser().resolve() / "co-tracker" + cotracker_path = str(cotracker_root) + if not cotracker_root.is_dir(): + raise FileNotFoundError(f"CoTracker package not found: {cotracker_root}") + if cotracker_path not in sys.path: + sys.path.insert(0, cotracker_path) + + from cotracker.models.core.cotracker.blocks import Attention + + if bool(getattr(Attention, "_trex_sdpa_enabled", False)): + return + + def _sdpa_forward(self, x, context=None, attn_bias=None): + batch, query_steps, _ = x.shape + heads = int(self.heads) + query = self.to_q(x) + inner_dim = int(query.shape[-1]) + head_dim = inner_dim // heads + query = query.reshape(batch, query_steps, heads, head_dim).transpose(1, 2) + + context = x if context is None else context + key, value = self.to_kv(context).chunk(2, dim=-1) + context_steps = int(context.shape[1]) + key = key.reshape(batch, context_steps, heads, head_dim).transpose(1, 2) + value = value.reshape(batch, context_steps, heads, head_dim).transpose(1, 2) + + attended = F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=attn_bias, + dropout_p=0.0, + is_causal=False, + ) + attended = attended.transpose(1, 2).reshape(batch, query_steps, inner_dim) + return self.to_out(attended) + + Attention.forward = _sdpa_forward + Attention._trex_sdpa_enabled = True + + +def normalize_tracks_xy(tracks: np.ndarray, img_w: int, img_h: int) -> np.ndarray: + out = np.asarray(tracks, dtype=np.float32).copy() + w, h = max(float(img_w), 1.0), max(float(img_h), 1.0) + out = np.nan_to_num(out, nan=0.0, posinf=0.0, neginf=0.0) + out[..., 0] = np.clip(out[..., 0] / w, 0.0, 1.0) + out[..., 1] = np.clip(out[..., 1] / h, 0.0, 1.0) + return out.astype(np.float32, copy=False) + + +def normalize_track_result( + tracks: np.ndarray, + visibility: np.ndarray, + img_w: int, + img_h: int, +) -> tuple[np.ndarray, np.ndarray]: + """Normalize XY and clear visibility for non-finite/out-of-frame points.""" + + pixels = np.asarray(tracks, dtype=np.float32) + vis = np.asarray(visibility, dtype=np.float32) + if pixels.ndim != 3 or pixels.shape[-1] != 2: + raise ValueError(f"tracks must be (T,N,2), got {pixels.shape}") + if vis.shape != pixels.shape[:2]: + raise ValueError(f"visibility {vis.shape} does not match tracks {pixels.shape}") + finite = np.isfinite(pixels).all(axis=-1) + in_frame = ( + (pixels[..., 0] >= 0.0) + & (pixels[..., 0] < float(img_w)) + & (pixels[..., 1] >= 0.0) + & (pixels[..., 1] < float(img_h)) + ) + clean_vis = ((vis > 0.5) & finite & in_frame).astype(np.float32) + return normalize_tracks_xy(pixels, img_w, img_h), clean_vis + + +def load_episode_videos( + dataset_root: Path, + episode_index: int, + *, + out_hw: tuple[int, int], +) -> dict[str, np.ndarray]: + import cv2 + + chunk = episode_index // 1000 + frames_by_view: dict[str, list[np.ndarray]] = {k: [] for k in VIEW_ORDER} + for view in VIEW_ORDER: + rel = VIDEO_FOLDERS[view] + video_path = ( + dataset_root / "videos" / f"chunk-{chunk:03d}" / rel / f"episode_{episode_index:06d}.mp4" + ) + if not video_path.is_file(): + raise FileNotFoundError(video_path) + cap = cv2.VideoCapture(str(video_path)) + while True: + ok, bgr = cap.read() + if not ok: + break + rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) + if (rgb.shape[0], rgb.shape[1]) != out_hw: + rgb = cv2.resize(rgb, (out_hw[1], out_hw[0]), interpolation=cv2.INTER_LINEAR) + frames_by_view[view].append(rgb) + cap.release() + if not frames_by_view[view]: + raise RuntimeError(f"empty video: {video_path}") + return {v: np.stack(frames_by_view[v], axis=0) for v in VIEW_ORDER} + + +def load_episode_states(dataset_root: Path, episode_index: int) -> tuple[np.ndarray, str]: + import pandas as pd + + chunk = episode_index // 1000 + pq_path = dataset_root / "data" / f"chunk-{chunk:03d}" / f"episode_{episode_index:06d}.parquet" + df = pd.read_parquet(pq_path) + states = np.stack([np.asarray(x, dtype=np.float64) for x in df["observation.state"].values], axis=0) + task = "" + if "annotation.task" in df.columns: + task = str(df["annotation.task"].iloc[0]) + return states, task + + +def _run_cotracker_window( + model, + video_hwc: np.ndarray, + query_xy: np.ndarray, + device: object, +) -> tuple[np.ndarray, np.ndarray]: + """Run one bounded CoTracker window and immediately release its GPU tensors.""" + + import torch + + video_np = np.asarray(video_hwc, dtype=np.uint8) + queries_np = np.zeros((int(query_xy.shape[0]), 3), dtype=np.float32) + queries_np[:, 1:] = np.asarray(query_xy, dtype=np.float32) + device_type = torch.device(device).type + video_dtype = torch.bfloat16 if device_type == "cuda" else torch.float32 + video = ( + torch.from_numpy(video_np) + .permute(0, 3, 1, 2) + .unsqueeze(0) + .to(device=device, dtype=video_dtype) + ) + queries = torch.from_numpy(queries_np).unsqueeze(0).to(device) + with torch.inference_mode(), torch.autocast( + device_type=device_type, + dtype=torch.bfloat16, + enabled=device_type == "cuda", + ): + pred_tracks, pred_vis = model( + video, + queries=queries, + backward_tracking=False, + ) + tracks = pred_tracks[0].detach().cpu().numpy().astype(np.float32) + visibility = pred_vis[0].detach().cpu().numpy() + visibility = (visibility > 0.5).astype(np.float32) + del video, queries, pred_tracks, pred_vis + if device_type == "cuda": + torch.cuda.empty_cache() + return tracks, visibility + + +def _run_cotracker(model, video_hwc: np.ndarray, query_xy: np.ndarray, device: object): + """Track frame-0 queries in bounded, overlapping temporal windows. + + ``CoTrackerPredictor`` only copies backward predictions into frames before + each query timestamp. Every query here starts at frame zero, so + ``backward_tracking=True`` cannot change the result and nearly doubles the + peak memory for long T-Rex episodes. + + Even with Flash Attention, CoTracker's feature/correlation tensors grow + linearly with the frame count. Each new window is initialized from the + previous trajectory, and overlapping predictions are blended to avoid a + discontinuity at the boundary. + """ + + video_np = np.asarray(video_hwc, dtype=np.uint8) + total_frames = int(video_np.shape[0]) + num_queries = int(query_xy.shape[0]) + window_frames = int(os.environ.get("TREX_COTRACKER_WINDOW_FRAMES", "768")) + overlap_frames = int(os.environ.get("TREX_COTRACKER_WINDOW_OVERLAP", "64")) + if window_frames < 2: + raise ValueError("TREX_COTRACKER_WINDOW_FRAMES must be at least 2") + if overlap_frames < 1 or overlap_frames >= window_frames: + raise ValueError( + "TREX_COTRACKER_WINDOW_OVERLAP must be in [1, WINDOW_FRAMES)" + ) + + if total_frames <= window_frames: + return _run_cotracker_window(model, video_np, query_xy, device) + + step = window_frames - overlap_frames + num_windows = 1 + (total_frames - window_frames + step - 1) // step + print( + f" CoTracker windowing: {total_frames} frames -> {num_windows} " + f"window(s), max={window_frames}, overlap={overlap_frames}" + ) + tracks = np.empty((total_frames, num_queries, 2), dtype=np.float32) + visibility = np.empty((total_frames, num_queries), dtype=np.float32) + + start = 0 + filled_end = 0 + while start < total_frames: + end = min(start + window_frames, total_frames) + seed_xy = np.asarray(query_xy if start == 0 else tracks[start], dtype=np.float32) + window_tracks, window_visibility = _run_cotracker_window( + model, + video_np[start:end], + seed_xy, + device, + ) + + overlap_end = min(filled_end, end) + existing_frames = max(0, overlap_end - start) + if existing_frames > 0: + alpha = np.linspace( + 0.0, + 1.0, + existing_frames, + dtype=np.float32, + ) + tracks[start:overlap_end] = ( + tracks[start:overlap_end] * (1.0 - alpha[:, None, None]) + + window_tracks[:existing_frames] * alpha[:, None, None] + ) + use_new_visibility = alpha >= 0.5 + visibility[start:overlap_end] = np.where( + use_new_visibility[:, None], + window_visibility[:existing_frames], + visibility[start:overlap_end], + ) + + tracks[overlap_end:end] = window_tracks[existing_frames:] + visibility[overlap_end:end] = window_visibility[existing_frames:] + filled_end = max(filled_end, end) + if end >= total_frames: + break + start = end - overlap_frames + + return tracks, visibility + + +def _tracks_episode( + *, + view_images: dict[str, np.ndarray], + out_hw: tuple[int, int], + cotracker_model, + cotracker_device: object, + sam2_predictor=None, + sam2_seed: int | None = None, +) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray], dict[str, np.ndarray | None], dict[str, str]]: + from trex_track.sam2_cotracker_hands import head_hands_50, wrist_hand_50 + from trex_track.trex_projection import make_image_grid + + view_tracks_px: dict[str, np.ndarray] = {} + view_vis: dict[str, np.ndarray] = {} + masks: dict[str, np.ndarray | None] = {} + tags: dict[str, str] = {} + + # Head: left hand/arm 50, then right hand/arm 50. + head_q, left_m, right_m, head_tag = head_hands_50( + sam2_predictor, + view_images["head_left"][0], + n_points=NUM_HAND_POINTS, + seed=sam2_seed, + ) + if int(head_q.shape[0]) != NUM_HEAD_POINTS: + raise ValueError(f"head queries expect {NUM_HEAD_POINTS}, got {head_q.shape[0]}") + head_trk, head_vis = _run_cotracker( + cotracker_model, view_images["head_left"], head_q, cotracker_device + ) + view_tracks_px["head_left"] = head_trk + view_vis["head_left"] = head_vis + masks["head_left_hand"] = left_m + masks["head_right_hand"] = right_m + tags["head_left"] = head_tag + + # Wrists: fixed 5×5 background grid + 50 SAM2 hand points. + for view in ("left_wrist", "right_wrist"): + hand_q, mask, tag = wrist_hand_50( + sam2_predictor, + view_images[view][0], + view, + n_points=NUM_WRIST_HAND, + seed=sam2_seed, + ) + if int(hand_q.shape[0]) != NUM_WRIST_HAND: + raise ValueError(f"{view} hand queries expect {NUM_WRIST_HAND}, got {hand_q.shape[0]}") + grid_q = make_image_grid(out_hw[0], out_hw[1], grid_size=5).astype(np.float32) + q = np.concatenate([grid_q, hand_q], axis=0) + if int(q.shape[0]) != NUM_WRIST_POINTS: + raise ValueError(f"{view} queries expect {NUM_WRIST_POINTS}, got {q.shape[0]}") + trk, vis = _run_cotracker(cotracker_model, view_images[view], q, cotracker_device) + view_tracks_px[view] = trk + view_vis[view] = vis + masks[view] = mask + tags[view] = f"grid25+{tag}" + + view_tracks: dict[str, np.ndarray] = {} + for view in VIEW_ORDER: + view_tracks[view], view_vis[view] = normalize_track_result( + view_tracks_px[view], + view_vis[view], + out_hw[1], + out_hw[0], + ) + return view_tracks, view_vis, masks, tags + + +def _save_seed_overlay( + out_dir: Path, + episode_index: int, + view_images: dict[str, np.ndarray], + masks: dict[str, np.ndarray | None], + view_tracks: dict[str, np.ndarray], + out_hw: tuple[int, int], +) -> None: + """Save frame-0 overlays with mask + query points for debugging.""" + import cv2 + + out_dir.mkdir(parents=True, exist_ok=True) + h, w = out_hw + + def _draw(view: str, mask: np.ndarray | None, tracks_norm: np.ndarray, color_bgr): + rgb = view_images[view][0] + bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR) + if mask is not None and np.asarray(mask).any(): + m = np.asarray(mask, dtype=bool) + bgr[m] = (bgr[m] * 0.35 + np.asarray(color_bgr, dtype=np.float32) * 0.65).astype(np.uint8) + pts = tracks_norm[0] * np.array([w, h], dtype=np.float32) + for p in pts.astype(int): + cv2.circle(bgr, (int(p[0]), int(p[1])), 3, (0, 255, 255), -1, cv2.LINE_AA) + return bgr + + # Head: both masks + head = cv2.cvtColor(view_images["head_left"][0], cv2.COLOR_RGB2BGR) + lm, rm = masks.get("head_left_hand"), masks.get("head_right_hand") + if lm is not None and np.asarray(lm).any(): + m = np.asarray(lm, dtype=bool) + head[m] = (head[m] * 0.35 + np.array([0, 255, 0], dtype=np.float32) * 0.65).astype(np.uint8) + if rm is not None and np.asarray(rm).any(): + m = np.asarray(rm, dtype=bool) + head[m] = (head[m] * 0.35 + np.array([0, 165, 255], dtype=np.float32) * 0.65).astype(np.uint8) + pts = view_tracks["head_left"][0] * np.array([w, h], dtype=np.float32) + for i, p in enumerate(pts.astype(int)): + col = (0, 255, 255) if i < NUM_HEAD_LEFT else (255, 255, 0) + cv2.circle(head, (int(p[0]), int(p[1])), 3, col, -1, cv2.LINE_AA) + cv2.imwrite(str(out_dir / f"episode_{episode_index:06d}_head_left_seeds.png"), head) + + for view in ("left_wrist", "right_wrist"): + rgb = view_images[view][0] + bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR) + mask = masks.get(view) + if mask is not None and np.asarray(mask).any(): + m = np.asarray(mask, dtype=bool) + bgr[m] = (bgr[m] * 0.40 + np.array([0, 0, 220], dtype=np.float32) * 0.60).astype(np.uint8) + pts = view_tracks[view][0] * np.array([w, h], dtype=np.float32) + for i, p in enumerate(pts.astype(int)): + # yellow = fixed background grid; magenta = force-aligned hand + col = (0, 255, 255) if i < NUM_WRIST_GRID else (255, 0, 255) + cv2.drawMarker( + bgr, + (int(p[0]), int(p[1])), + col, + markerType=cv2.MARKER_STAR, + markerSize=8, + thickness=1, + line_type=cv2.LINE_AA, + ) + cv2.imwrite(str(out_dir / f"episode_{episode_index:06d}_{view}_seeds.png"), bgr) + + +def _save_prompt_overlays( + out_dir: Path, + episode_index: int, + view_images: dict[str, np.ndarray], + masks: dict[str, np.ndarray | None], +) -> None: + """Save SAM2 prompts (pos/neg/box) + mask overlays; reuse track masks (no 2nd SAM2).""" + import cv2 + + from trex_track.sam2_prompt_hands import ( + build_head_prompts, + build_wrist_prompts, + draw_prompt_overlay, + ) + + out_dir.mkdir(parents=True, exist_ok=True) + panels: list[np.ndarray] = [] + meta_dump: dict[str, object] = {} + + def _one( + key: str, + rgb: np.ndarray, + mask: np.ndarray | None, + coords: np.ndarray, + labels: np.ndarray, + box: np.ndarray, + mask_bgr: tuple[int, int, int], + ) -> np.ndarray: + vis = draw_prompt_overlay( + rgb, + mask=mask, + coords=coords, + labels=labels, + box=box, + mask_bgr=mask_bgr, + ) + # Legend + n_pos = int((labels == 1).sum()) + n_neg = int((labels == 0).sum()) + area = int(np.asarray(mask).sum()) if mask is not None else 0 + cv2.putText( + vis, + f"{key} +pos={n_pos} -neg={n_neg} mask_px={area}", + (6, 14), + cv2.FONT_HERSHEY_SIMPLEX, + 0.40, + (255, 255, 255), + 1, + cv2.LINE_AA, + ) + path = out_dir / f"episode_{episode_index:06d}_{key}_prompt.png" + cv2.imwrite(str(path), vis) + print(f"Wrote {path}") + meta_dump[f"{key}_coords"] = np.asarray(coords, dtype=np.float32) + meta_dump[f"{key}_labels"] = np.asarray(labels, dtype=np.int32) + meta_dump[f"{key}_box"] = np.asarray(box, dtype=np.float32) + if mask is not None: + meta_dump[f"{key}_mask"] = np.asarray(mask, dtype=bool) + return vis + + # Wrist views + for view, col in (("left_wrist", (0, 0, 220)), ("right_wrist", (0, 0, 220))): + rgb = view_images[view][0] + h, w = rgb.shape[:2] + coords, labels, box = build_wrist_prompts(h, w, view) + panels.append(_one(view, rgb, masks.get(view), coords, labels, box, col)) + + # Head: left / right hands (same RGB, two prompt sets) + rgb = view_images["head_left"][0] + h, w = rgb.shape[:2] + for side, col, mkey in ( + ("left", (0, 255, 0), "head_left_hand"), + ("right", (0, 165, 255), "head_right_hand"), + ): + coords, labels, box = build_head_prompts(h, w, side) + panels.append(_one(f"head_{side}", rgb, masks.get(mkey), coords, labels, box, col)) + + # Combined 2x2 panel for quick inspection + if len(panels) == 4: + top = np.concatenate(panels[2:4], axis=1) # head_left | head_right + bot = np.concatenate(panels[0:2], axis=1) # left_wrist | right_wrist + # Resize to same width if needed + if top.shape[1] != bot.shape[1]: + tw = max(top.shape[1], bot.shape[1]) + top = cv2.resize(top, (tw, top.shape[0])) + bot = cv2.resize(bot, (tw, bot.shape[0])) + grid = np.concatenate([top, bot], axis=0) + grid_path = out_dir / f"episode_{episode_index:06d}_all_prompts.png" + cv2.imwrite(str(grid_path), grid) + print(f"Wrote {grid_path}") + + meta_path = out_dir / f"episode_{episode_index:06d}_prompts.npz" + np.savez_compressed(meta_path, **meta_dump) + print(f"Wrote {meta_path}") + + +def _atomic_savez(path: Path, **arrays: object) -> None: + """Write an NPZ in the destination directory, then atomically replace.""" + + path.parent.mkdir(parents=True, exist_ok=True) + fd, tmp_name = tempfile.mkstemp( + prefix=f".{path.name}.", + suffix=".tmp", + dir=path.parent, + ) + try: + with os.fdopen(fd, "wb") as file: + np.savez_compressed(file, **arrays) + file.flush() + os.fsync(file.fileno()) + os.replace(tmp_name, path) + except BaseException: + try: + os.unlink(tmp_name) + except FileNotFoundError: + pass + raise + + +def process_episode( + *, + dataset_root: Path, + episode_index: int, + output_path: Path, + calib: dict | None, + out_hw: tuple[int, int], + cotracker_model, + cotracker_device: object, + save_viz: bool, + viz_out_dir: Path, + viz_fps: int, + viz_trail: int, + sam2_predictor=None, + sam2_seed: int | None = None, + save_sam2_masks_flag: bool = True, + sam2_masks_dir: Path | None = None, +) -> Path: + from trex_track.sam2_wrist_hand import save_sam2_masks + from trex_track.trex_viz_tracks import render_three_view_combined_video + + del calib # Reserved for provenance/backward-compatible callers. + states, task = load_episode_states(dataset_root, episode_index) + view_images = load_episode_videos(dataset_root, episode_index, out_hw=out_hw) + frame_counts = {"parquet": int(states.shape[0])} + frame_counts.update({view: int(view_images[view].shape[0]) for view in VIEW_ORDER}) + if len(set(frame_counts.values())) != 1: + raise ValueError( + f"episode {episode_index}: frame-count mismatch; refusing to truncate: " + f"{frame_counts}" + ) + t_len = int(states.shape[0]) + + view_tracks, view_vis, masks, tags = _tracks_episode( + view_images=view_images, + out_hw=out_hw, + cotracker_model=cotracker_model, + cotracker_device=cotracker_device, + sam2_predictor=sam2_predictor, + sam2_seed=sam2_seed, + ) + + tracks_combined = np.concatenate([view_tracks[v] for v in VIEW_ORDER], axis=1) + vis_combined = np.concatenate([view_vis[v] for v in VIEW_ORDER], axis=1) + if tracks_combined.shape != (t_len, NUM_COMBINED_POINTS, 2): + raise ValueError( + f"episode {episode_index}: combined tracks have {tracks_combined.shape}, " + f"expected {(t_len, NUM_COMBINED_POINTS, 2)}" + ) + if vis_combined.shape != (t_len, NUM_COMBINED_POINTS): + raise ValueError( + f"episode {episode_index}: combined visibility has {vis_combined.shape}" + ) + + out_npz = output_path / f"episode_{episode_index:06d}.npz" + identities = identity_metadata() + _atomic_savez( + out_npz, + tracks=tracks_combined, + vis=vis_combined, + tracks_head_left=view_tracks["head_left"], + tracks_left_wrist=view_tracks["left_wrist"], + tracks_right_wrist=view_tracks["right_wrist"], + vis_head_left=view_vis["head_left"], + vis_left_wrist=view_vis["left_wrist"], + vis_right_wrist=view_vis["right_wrist"], + language=np.array(task), + episode_index=np.array(episode_index, dtype=np.int32), + num_steps=np.array(t_len, dtype=np.int32), + point_slices=np.array(POINT_SLICES, dtype=np.int32), + point_view_ids=np.asarray(identities["view_ids"], dtype=np.int8), + point_hand_ids=np.asarray(identities["hand_ids"], dtype=np.int8), + point_role_ids=np.asarray(identities["role_ids"], dtype=np.int8), + point_local_ids=np.asarray(identities["local_ids"], dtype=np.int16), + point_global_ids=np.asarray(identities["global_ids"], dtype=np.int16), + point_names=np.asarray(identities["point_names"]), + points_per_hand=np.array(NUM_HAND_POINTS, dtype=np.int32), + wrist_grid_points=np.array(NUM_WRIST_GRID, dtype=np.int32), + wrist_hand_points=np.array(NUM_WRIST_HAND, dtype=np.int32), + head_query_source=np.array(tags.get("head_left", "sam2")), + left_wrist_query_source=np.array(tags.get("left_wrist", "sam2")), + right_wrist_query_source=np.array(tags.get("right_wrist", "sam2")), + track_source=np.array("sam2_once_prompt_cotracker"), + tracks_coord_space=np.array("normalized_div_wh"), + track_layout_version=np.array(TRACK_LAYOUT_VERSION), + proj_image_hw=np.array(out_hw, dtype=np.int32), + ) + + masks_dir = sam2_masks_dir if sam2_masks_dir is not None else output_path / "sam2_masks" + if save_sam2_masks_flag: + # Save binary masks for wrist + head hands + save_map = { + "left_wrist": masks.get("left_wrist"), + "right_wrist": masks.get("right_wrist"), + "head_left_hand": masks.get("head_left_hand"), + "head_right_hand": masks.get("head_right_hand"), + } + for p in save_sam2_masks(masks_dir, episode_index, save_map): + print(f"Wrote {p}") + _save_seed_overlay(masks_dir, episode_index, view_images, masks, view_tracks, out_hw) + _save_prompt_overlays(masks_dir, episode_index, view_images, masks) + + if save_viz: + render_three_view_combined_video( + view_images=view_images, + view_tracks=view_tracks, + view_vis=view_vis, + out_path=viz_out_dir / f"episode_{episode_index:06d}.mp4", + fps=viz_fps, + draw_trail=viz_trail, + ) + print(f"Wrote {out_npz}") + if save_viz: + print(f"Wrote {viz_out_dir / f'episode_{episode_index:06d}.mp4'}") + print(f" tags: {tags}") + return out_npz + + +def create_tracking_runtime( + *, + calib_path: str | Path = DEFAULT_CALIB, + openpi_root: str | Path = DEFAULT_OPENPI_ROOT, + cotracker_checkpoint: str | Path | None = None, + cotracker_device: str = "", + sam2_model: str = DEFAULT_SAM2_MODEL, + sam2_device: str = "", + sam2_libs: str | Path = DEFAULT_SAM2_LIBS, + image_height: int = 0, + image_width: int = 0, +) -> TrackingRuntime: + """Load CoTracker and SAM2 once; safe to call from the batch builder.""" + + import torch + + from trex_track.sam2_wrist_hand import load_sam2_predictor + from trex_track.trex_projection import load_camera_calib + + calib = load_camera_calib(calib_path) + if bool(image_height > 0) != bool(image_width > 0): + raise ValueError("image-height and image-width must be set together") + if image_height > 0: + out_hw = (int(image_height), int(image_width)) + else: + out_hw = tuple(int(x) for x in calib.get("video_hw", [180, 320])) + + _ensure_openpi_on_path(Path(openpi_root)) + from utils.cotracker_wrist_grid import ( # type: ignore + default_cotracker_checkpoint, + load_cotracker_predictor, + ) + _enable_cotracker_sdpa_attention(openpi_root) + + checkpoint_text = str(cotracker_checkpoint or "").strip() + checkpoint = ( + Path(checkpoint_text).expanduser().resolve() + if checkpoint_text + else default_cotracker_checkpoint() + ) + device_text = cotracker_device.strip() or ( + "cuda:0" if torch.cuda.is_available() else "cpu" + ) + device = torch.device(device_text) + print(f" CoTracker: {checkpoint} on {device}") + if device.type == "cuda": + print(" CoTracker attention: BF16 PyTorch SDPA (Flash-compatible)") + cotracker_model = load_cotracker_predictor(checkpoint, device) + + sam_device = sam2_device.strip() or device_text + print(f" SAM2: {sam2_model} on {sam_device}") + sam2_predictor = load_sam2_predictor( + model_id=str(sam2_model), + device=sam_device, + sam2_libs=sam2_libs, + ) + return TrackingRuntime( + calib=calib, + out_hw=out_hw, + cotracker_model=cotracker_model, + cotracker_device=device, + sam2_predictor=sam2_predictor, + ) + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Extract canonical T-Rex SAM2+CoTracker tracks (250 points)" + ) + parser.add_argument( + "--dataset-root", + type=str, + default=str(_DREAMZERO_ROOT / "data" / "trex_small"), + ) + parser.add_argument("--episode-index", type=int, default=0) + parser.add_argument( + "--output-path", + type=str, + default=str(_DREAMZERO_ROOT / "data" / "trex_small_tracks"), + ) + parser.add_argument("--calib-path", type=str, default=str(DEFAULT_CALIB)) + parser.add_argument("--openpi-root", type=str, default=str(DEFAULT_OPENPI_ROOT)) + parser.add_argument("--cotracker-checkpoint", type=str, default="") + parser.add_argument("--cotracker-device", type=str, default="") + parser.add_argument("--image-height", type=int, default=0) + parser.add_argument("--image-width", type=int, default=0) + parser.add_argument("--save-viz", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--viz-out-dir", type=str, default="") + parser.add_argument("--viz-fps", type=int, default=10) + parser.add_argument("--viz-trail", type=int, default=15) + parser.add_argument( + "--sam2-model", + type=str, + default=DEFAULT_SAM2_MODEL, + ) + parser.add_argument("--sam2-device", type=str, default="") + parser.add_argument("--sam2-seed", type=int, default=0) + parser.add_argument("--save-sam2-masks", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--sam2-masks-dir", type=str, default="") + parser.add_argument( + "--sam2-libs", + type=str, + default=DEFAULT_SAM2_LIBS, + ) + args = parser.parse_args() + + dataset_root = Path(args.dataset_root).expanduser().resolve() + output_path = Path(args.output_path).expanduser().resolve() + runtime = create_tracking_runtime( + calib_path=args.calib_path, + openpi_root=args.openpi_root, + cotracker_checkpoint=args.cotracker_checkpoint, + cotracker_device=args.cotracker_device, + sam2_model=args.sam2_model, + sam2_device=args.sam2_device, + sam2_libs=args.sam2_libs, + image_height=args.image_height, + image_width=args.image_width, + ) + + viz_out_dir = ( + Path(args.viz_out_dir).expanduser().resolve() + if args.viz_out_dir.strip() + else output_path / "viz_tracks" + ) + sam2_masks_dir = ( + Path(args.sam2_masks_dir).expanduser().resolve() + if args.sam2_masks_dir.strip() + else output_path / "sam2_masks" + ) + + print("T-Rex SAM2+CoTracker extraction") + print(f" Dataset: {dataset_root}") + print(f" Episode: {args.episode_index}") + print(f" Output: {output_path}") + print(f" Image: {runtime.out_hw[1]}x{runtime.out_hw[0]}") + print( + f" Points: head={NUM_HEAD_POINTS} (50+50), " + f"wrist={NUM_WRIST_POINTS}x2 (grid{NUM_WRIST_GRID}+hand{NUM_WRIST_HAND}), " + f"total={NUM_COMBINED_POINTS}" + ) + + process_episode( + dataset_root=dataset_root, + episode_index=int(args.episode_index), + output_path=output_path, + calib=runtime.calib, + out_hw=runtime.out_hw, + cotracker_model=runtime.cotracker_model, + cotracker_device=runtime.cotracker_device, + save_viz=bool(args.save_viz), + viz_out_dir=viz_out_dir, + viz_fps=int(args.viz_fps), + viz_trail=int(args.viz_trail), + sam2_predictor=runtime.sam2_predictor, + sam2_seed=int(args.sam2_seed), + save_sam2_masks_flag=bool(args.save_sam2_masks), + sam2_masks_dir=sam2_masks_dir, + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/inference/build_trt_engine.py b/scripts/inference/build_trt_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..227906c7f50f669323c970de1b0427d5b2809d83 --- /dev/null +++ b/scripts/inference/build_trt_engine.py @@ -0,0 +1,243 @@ +"""Build a TensorRT engine from a DreamZero checkpoint. + +Must be launched via build_trt_engine.sh (or with ENABLE_TENSORRT=true already +set) so that flash-attention compatibility mode is active before any groot model +modules are imported. + +Launched via torchrun so that RANK / WORLD_SIZE / MASTER_* env vars exist for +GrootSimPolicy's distributed initialisation. + +Calibration: + For quantized precisions (nvfp4, fp8), ModelOpt calibrates quantization + parameters by observing activation statistics during forward passes. Using + real dataset trajectories produces a significantly more accurate engine than + random dummy inputs. Pass --dataset-path to enable real calibration. +""" + +import os +import sys +import argparse +import logging +from types import SimpleNamespace + +# Verify ENABLE_TENSORRT was exported before any groot imports occur. +if os.getenv("ENABLE_TENSORRT", "").lower() != "true": + print( + "ERROR: ENABLE_TENSORRT must be 'true' before importing this script.\n" + "Use build_trt_engine.sh instead of calling this script directly.", + file=sys.stderr, + ) + sys.exit(1) + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", +) +logger = logging.getLogger(__name__) + +import numpy as np +import torch +import torch.distributed as dist +from tianshou.data import Batch +from torch.distributed.device_mesh import init_device_mesh + +from groot.vla.data.schema import EmbodimentTag +from groot.vla.model.n1_5.sim_policy import GrootSimPolicy +from groot.control.tensorrt_utils import ( + wan_trt_quantize_and_load_engine, + create_wan_test_inputs, +) + +# DreamZero-DROID uses the ar_14B_droid model type in tensorrt_utils. +_MODEL_TYPE = "ar_14B_droid" + + +def _init_single_gpu_mesh(): + """Initialise a single-GPU device mesh (launched via torchrun --nproc_per_node=1).""" + dist.init_process_group("nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + torch.cuda.set_device(rank) + mesh = init_device_mesh( + device_type="cuda", + mesh_shape=(world_size,), + mesh_dim_names=("ip",), + ) + return mesh + + +def _make_dummy_forward_loop(): + """Fallback calibration using random dummy inputs. + + Acceptable for fp16 (no quantization), but may reduce accuracy for + nvfp4/fp8 since the activation distribution differs from real data. + Prefer _make_dataset_forward_loop when a dataset is available. + """ + def forward_loop(model): + trt_forward = getattr(model, "_forward_inference_trt_droid", model.forward) + test_inputs = create_wan_test_inputs(None, device="cuda", model_type=_MODEL_TYPE) + for _ in range(16): + with torch.no_grad(): + trt_forward(*test_inputs) + + return forward_loop + + +def _make_dataset_forward_loop(policy, dataset_path: str, num_calibration_trajs: int = 2): + """Real-data calibration loop — mirrors the internal droid_video_pred.sh approach. + + Loads ``num_calibration_trajs`` trajectories from the LeRobot dataset and + runs ``policy.lazy_joint_forward_causal`` at each action-horizon step, + exercising the DiT model with realistic activation distributions. + """ + from groot.vla.data.dataset.lerobot import LeRobotSingleDataset + + def forward_loop(model): + logger.info( + "Calibration: loading dataset from %s (%d trajs)", dataset_path, num_calibration_trajs + ) + dataset = LeRobotSingleDataset( + dataset_path=dataset_path, + modality_configs=policy.modality_configs, + embodiment_tag=policy.embodiment_tag, + video_backend="torchvision_av", + video_backend_kwargs=None, + transforms=None, # policy.lazy_joint_forward_causal applies transforms + use_global_metadata=False, + ) + + action_horizon = policy.trained_model.action_head.action_horizon + num_frame_per_block = policy.trained_model.action_head.num_frame_per_block + torch._dynamo.config.recompile_limit = 500 + + for traj_id in range(min(num_calibration_trajs, len(dataset.trajectory_lengths))): + logger.info("Calibration trajectory %d / %d", traj_id + 1, num_calibration_trajs) + traj_len = int(dataset.trajectory_lengths[traj_id]) + latent_video = None + + # Step through the trajectory at action-horizon intervals (same cadence as + # real inference) for up to 5 chunks — enough to cover the KV-cache build-up + # and the cached inference path that the TRT engine will handle. + max_steps = min(traj_len, 5 * action_horizon) + for step in range(0, max_steps, action_horizon): + # Clamp delta indices to valid range for this trajectory. + indices = { + k: np.clip(v + step, 0, traj_len - 1) + for k, v in dataset.delta_indices.items() + } + data_point = dataset.get_step_data(traj_id, indices) + batch = Batch(obs=data_point) + + dist.barrier() + with torch.no_grad(): + result_batch, video_pred = policy.lazy_joint_forward_causal( + batch, latent_video=latent_video + ) + dist.barrier() + + # Feed the last generated frame back as context for the next step, + # matching autoregressive inference behaviour. + if video_pred is not None: + latent_video = video_pred[:, :, -num_frame_per_block:] + + # Reset AR state between trajectories. + policy.trained_model.action_head.current_start_frame = 0 + policy.trained_model.action_head.kv_cache1 = None + policy.trained_model.action_head.kv_cache_neg = None + policy.trained_model.action_head.crossattn_cache = None + policy.trained_model.action_head.crossattn_cache_neg = None + + return forward_loop + + +def main(): + parser = argparse.ArgumentParser( + description="Build TensorRT engine for the DreamZero DiT model." + ) + parser.add_argument("--model-path", required=True, help="Path to checkpoint directory.") + parser.add_argument( + "--tensorrt", + required=True, + choices=["nvfp4", "fp8", "fp16"], + help="TensorRT quantization / precision format.", + ) + parser.add_argument( + "--dataset-path", + default=None, + help=( + "Path to a LeRobot-format DROID dataset for real calibration. " + "Strongly recommended for nvfp4/fp8 — random dummy inputs are used as " + "fallback but may reduce quantization accuracy." + ), + ) + parser.add_argument( + "--num-calibration-trajs", + type=int, + default=2, + help="Number of dataset trajectories used for calibration (default: 2).", + ) + args = parser.parse_args() + + if args.tensorrt in ("nvfp4", "fp8") and args.dataset_path is None: + logger.warning( + "No --dataset-path provided for %s quantization. " + "Falling back to random dummy inputs — this may reduce engine accuracy. " + "Re-run with --dataset-path for best results.", + args.tensorrt, + ) + + engine_dir = os.path.join(args.model_path, "tensorrt", "wan") + engine_path = os.path.join(engine_dir, f"WanModel_{args.tensorrt}.trt") + onnx_path = os.path.join(engine_dir, f"CausalWanModel.onnx") + os.makedirs(engine_dir, exist_ok=True) + + if os.path.exists(engine_path): + logger.info("TRT engine already exists: %s", engine_path) + logger.info("Delete it first if you want to rebuild.") + return + + logger.info("Loading DreamZero policy from : %s", args.model_path) + logger.info("Target engine path : %s", engine_path) + logger.info("Quantization precision : %s", args.tensorrt) + + device_mesh = _init_single_gpu_mesh() + + policy = GrootSimPolicy( + embodiment_tag=EmbodimentTag("oxe_droid"), + model_path=args.model_path, + device="cuda" if torch.cuda.is_available() else "cpu", + device_mesh=device_mesh, + ) + + # Build calibration forward loop — prefer real data for quantized precisions. + if args.dataset_path is not None: + forward_loop = _make_dataset_forward_loop( + policy, args.dataset_path, args.num_calibration_trajs + ) + logger.info( + "Calibration: using %d real trajectories from %s", + args.num_calibration_trajs, + args.dataset_path, + ) + else: + forward_loop = _make_dummy_forward_loop() + logger.info("Calibration: using random dummy inputs (no --dataset-path given).") + + # cfg mimics the Hydra config used by the internal eval script. + cfg = SimpleNamespace(inference_mode="trt_build", quantize_dtype=args.tensorrt) + + logger.info("Building TensorRT engine (ONNX export + trtexec, may take 10-30 min) ...") + wan_trt_quantize_and_load_engine( + policy=policy, + cfg=cfg, + onnx_path=onnx_path, + engine_path=engine_path, + model_type=_MODEL_TYPE, + forward_loop=forward_loop, + ) + + logger.info("TRT engine saved to: %s", engine_path) + + +if __name__ == "__main__": + main() diff --git a/scripts/inference/build_trt_engine.sh b/scripts/inference/build_trt_engine.sh new file mode 100644 index 0000000000000000000000000000000000000000..5c943cc6bf670a18930f853d1d1a31dbf20ada4e --- /dev/null +++ b/scripts/inference/build_trt_engine.sh @@ -0,0 +1,142 @@ +#!/usr/bin/env bash +# Build a TensorRT engine from a DreamZero checkpoint. +# +# Usage (recommended — with real calibration data): +# bash scripts/inference/build_trt_engine.sh \ +# --model-path ./checkpoints/DreamZero-DROID \ +# --tensorrt nvfp4 \ +# --dataset-path ./data/droid_lerobot \ +# --cuda-device 0 +# +# Usage (without dataset — acceptable for fp16, not recommended for nvfp4/fp8): +# bash scripts/inference/build_trt_engine.sh \ +# --model-path ./checkpoints/DreamZero-DROID \ +# --tensorrt nvfp4 \ +# --cuda-device 0 +# +# The engine is saved to: +# {model_path}/tensorrt/wan/WanModel_{precision}.trt +# +# Supported precisions: nvfp4 (recommended), fp8, fp16 +# +# For quantized precisions (nvfp4, fp8), ModelOpt calibrates quantization +# parameters using real forward passes. Providing --dataset-path is strongly +# recommended — random dummy inputs are used as fallback but reduce accuracy. +# +# ENABLE_TENSORRT=true must be set before any groot modules are imported +# (it controls flash-attention compatibility mode for ONNX/TRT export). +# This script sets it and launches the Python build script via torchrun so +# that RANK / WORLD_SIZE env vars are available for GrootSimPolicy init. + +# export HF_HUB_CACHE=/mnt/aws-lfs-02/shared/ckpts +set -euo pipefail + +MODEL_PATH="" +TENSORRT_PRECISION="" +CUDA_DEVICE="0" +DATASET_PATH="" +NUM_CALIBRATION_TRAJS="2" + +while [[ $# -gt 0 ]]; do + case $1 in + --model-path) + MODEL_PATH="$2" + shift 2 + ;; + --tensorrt) + TENSORRT_PRECISION="$2" + shift 2 + ;; + --cuda-device) + CUDA_DEVICE="$2" + shift 2 + ;; + --dataset-path) + DATASET_PATH="$2" + shift 2 + ;; + --num-calibration-trajs) + NUM_CALIBRATION_TRAJS="$2" + shift 2 + ;; + -h|--help) + echo "Usage: $0 --model-path --tensorrt [OPTIONS]" + echo "" + echo "Options:" + echo " --model-path PATH Path to DreamZero checkpoint directory" + echo " --tensorrt PRECISION TRT precision: nvfp4 (recommended), fp8, fp16" + echo " --dataset-path PATH LeRobot dataset for real calibration (recommended for nvfp4/fp8)" + echo " --num-calibration-trajs N Number of calibration trajectories (default: 2)" + echo " --cuda-device ID CUDA device index (default: 0)" + exit 0 + ;; + *) + echo "Unknown argument: $1" >&2 + exit 1 + ;; + esac +done + +if [[ -z "$MODEL_PATH" ]]; then + echo "Error: --model-path is required" >&2 + exit 1 +fi + +if [[ -z "$TENSORRT_PRECISION" ]]; then + echo "Error: --tensorrt is required (e.g. nvfp4, fp8, fp16)" >&2 + exit 1 +fi + +if [[ ! -d "$MODEL_PATH" ]]; then + echo "Error: checkpoint directory not found: $MODEL_PATH" >&2 + exit 1 +fi + +ENGINE_PATH="${MODEL_PATH}/tensorrt/wan/WanModel_${TENSORRT_PRECISION}.trt" + +echo "==========================================" +echo "DreamZero TensorRT Engine Builder" +echo " Checkpoint : $MODEL_PATH" +echo " Precision : $TENSORRT_PRECISION" +echo " CUDA device : $CUDA_DEVICE" +echo " Dataset (calibrate): ${DATASET_PATH:-}" +echo " Calibration trajs : $NUM_CALIBRATION_TRAJS" +echo " Output : $ENGINE_PATH" +echo "==========================================" + +# ENABLE_TENSORRT must be set before Python imports any groot model modules +# (it activates flash-attention compatibility mode required for ONNX/TRT export). +export ENABLE_TENSORRT=true +export CUDA_VISIBLE_DEVICES="$CUDA_DEVICE" +export ATTENTION_BACKEND="TE" +export HYDRA_FULL_ERROR=1 + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" + +# Build the Python argument list. +PYTHON_ARGS=( + --model-path "$MODEL_PATH" + --tensorrt "$TENSORRT_PRECISION" + --num-calibration-trajs "$NUM_CALIBRATION_TRAJS" +) +if [[ -n "$DATASET_PATH" ]]; then + PYTHON_ARGS+=(--dataset-path "$DATASET_PATH") +fi + +# torchrun sets RANK / WORLD_SIZE / MASTER_ADDR / MASTER_PORT which are +# required by GrootSimPolicy's distributed init. +torchrun \ + --standalone \ + --nproc_per_node=1 \ + "${REPO_ROOT}/scripts/inference/build_trt_engine_droid.py" \ + "${PYTHON_ARGS[@]}" + +echo "==========================================" +echo "Engine built successfully: $ENGINE_PATH" +echo "" +echo "Run inference with:" +echo " CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.run --standalone --nproc_per_node=2 \\" +echo " socket_test_optimized_AR.py --port 5000 --enable-dit-cache \\" +echo " --model-path ${MODEL_PATH} --tensorrt ${TENSORRT_PRECISION}" +echo "==========================================" diff --git a/scripts/inference_trex.py b/scripts/inference_trex.py new file mode 100644 index 0000000000000000000000000000000000000000..8b487103c00724126b932723ae315c68b274c737 --- /dev/null +++ b/scripts/inference_trex.py @@ -0,0 +1,304 @@ +#!/usr/bin/env python3 +"""Standalone T-Rex inference (no websocket server). + +Loads your LoRA checkpoint, runs causal video+action diffusion on dataset frames, +and saves the *predicted* future video as MP4. + +Usage: + CUDA_VISIBLE_DEVICES=0 python scripts/inference_trex.py \ + --model_path /scratch1/home/zhicao/dreamzero/checkpoints/dreamzero_trex_wan22_lora/checkpoint-5000 \ + --dataset_path /scratch1/home/zhicao/dreamzero/data/trex_small \ + --episode 0 \ + --num_chunks 8 \ + --output_dir results_trex_infer + +Why not server/client? + The official server (socket_test_optimized_AR.py) + test_client_AR.py split + model hosting from the robot/eval loop for multi-GPU distributed inference. + For offline debugging you only need GrootSimPolicy + lazy_joint_forward_causal. +""" + +from __future__ import annotations + +import argparse +import glob +import os +import time + +import cv2 +import imageio +import numpy as np +import pyarrow.parquet as pq +import torch +import torch._dynamo +import torch.distributed as dist +from einops import rearrange +from tianshou.data import Batch + +torch._dynamo.config.disable = True + +from groot.vla.data.schema import EmbodimentTag +from groot.vla.data.transform import ComposedModalityTransform +from groot.vla.model.n1_5.sim_policy import GrootSimPolicy + +# Modality keys (must match transform_trex / modality_config_trex) +VIDEO_KEYS = [ + "video.head_left", + "video.left_wrist", + "video.right_wrist", +] +STATE_KEYS = { + "state.left_arm": (0, 7), + "state.left_hand": (7, 29), + "state.right_arm": (29, 36), + "state.right_hand": (36, 58), +} +VIDEO_FOLDERS = { + "video.head_left": "observation.images.head_left", + "video.left_wrist": "observation.images.left_wrist", + "video.right_wrist": "observation.images.right_wrist", +} + +# Causal chunk schedule (same idea as test_client_AR.py / DROID server) +RELATIVE_OFFSETS = [-23, -16, -8, 0] +ACTION_HORIZON = 24 + + +def get_expected_video_resolution(policy: GrootSimPolicy) -> tuple[int, int]: + """Return (height, width) that eval_transform VideoToTensor expects.""" + cfg = policy.trained_model.action_head.config + target_h = getattr(cfg, "target_video_height", None) + target_w = getattr(cfg, "target_video_width", None) + if target_h is not None and target_w is not None: + return int(target_h), int(target_w) + + eval_transform = getattr(policy, "eval_transform", None) + if isinstance(eval_transform, ComposedModalityTransform): + for t in eval_transform.transforms: + res = getattr(t, "original_resolutions", None) + if res: + w, h = next(iter(res.values())) + return int(h), int(w) + return 160, 320 + + +def resize_frames(frames: np.ndarray, target_h: int, target_w: int) -> np.ndarray: + """Resize (H,W,C) or (T,H,W,C) uint8 frames to (target_h, target_w).""" + if frames.ndim == 3: + if (frames.shape[0], frames.shape[1]) == (target_h, target_w): + return frames + return cv2.resize(frames, (target_w, target_h), interpolation=cv2.INTER_LINEAR) + return np.stack( + [cv2.resize(f, (target_w, target_h), interpolation=cv2.INTER_LINEAR) for f in frames], + axis=0, + ) + + +class TrexEpisode: + """One T-Rex episode from LeRobot v2 layout.""" + + def __init__(self, dataset_root: str, episode_index: int): + pq_path = os.path.join( + dataset_root, + "data", + f"chunk-{episode_index // 1000:03d}", + f"episode_{episode_index:06d}.parquet", + ) + if not os.path.isfile(pq_path): + raise FileNotFoundError(pq_path) + self.table = pq.read_table(pq_path) + self.length = self.table.num_rows + self.episode_index = episode_index + self.root = dataset_root + + self.video_dirs = {} + for key, folder in VIDEO_FOLDERS.items(): + pattern = os.path.join( + dataset_root, + "videos", + "**", + folder, + f"episode_{episode_index:06d}.mp4", + ) + hits = sorted(glob.glob(pattern, recursive=True)) + if not hits: + raise FileNotFoundError(f"No video for {key}: {pattern}") + self.video_dirs[key] = hits[0] + + print( + f"TrexEpisode {episode_index}: {self.length} steps, " + f"{len(self.video_dirs)} cameras" + ) + + def get_task(self, row: int) -> str: + try: + return str(self.table.column("annotation.task")[row].as_py()) + except Exception: + return "" + + def get_state(self, row: int) -> np.ndarray: + return np.array(self.table.column("observation.state")[row].as_py(), dtype=np.float64) + + def get_frame(self, row: int, video_key: str) -> np.ndarray: + cap = cv2.VideoCapture(self.video_dirs[video_key]) + cap.set(cv2.CAP_PROP_POS_FRAMES, row) + ok, frame = cap.read() + cap.release() + if not ok: + raise RuntimeError(f"Failed frame {row} from {self.video_dirs[video_key]}") + return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + + +def build_chunk_schedule(total_frames: int, num_chunks: int) -> list[list[int]]: + """First chunk: 1 frame; later chunks: 4 frames ending at anchor.""" + chunks: list[list[int]] = [] + anchor = 0 + for i in range(num_chunks): + if i == 0: + indices = [0] + else: + indices = [max(anchor + off, 0) for off in RELATIVE_OFFSETS] + if indices[-1] >= total_frames: + break + chunks.append(indices) + anchor += ACTION_HORIZON + return chunks + + +def build_obs( + episode: TrexEpisode, + frame_indices: list[int], + prompt: str, + video_height: int, + video_width: int, +) -> dict: + obs: dict = {} + anchor = frame_indices[-1] + state = episode.get_state(anchor) + + for key in VIDEO_KEYS: + frames = np.stack([episode.get_frame(i, key) for i in frame_indices], axis=0) + frames = resize_frames(frames.astype(np.uint8), video_height, video_width) + if len(frame_indices) == 1: + obs[key] = frames[0] # (H, W, 3) + else: + obs[key] = frames # (T, H, W, 3) + + for key, (s, e) in STATE_KEYS.items(): + obs[key] = state[s:e].reshape(1, -1).astype(np.float64) + + obs["annotation.task"] = prompt + return obs + + +def decode_video_latents(policy: GrootSimPolicy, video_chunks: list[torch.Tensor]) -> np.ndarray: + """Concat latent chunks along time, VAE decode -> (T, H, W, 3) uint8.""" + if not video_chunks: + raise ValueError("No video chunks to decode") + cat = torch.cat(video_chunks, dim=2) + ah = policy.trained_model.action_head + frames = ah.vae.decode( + cat, + tiled=ah.tiled, + tile_size=(ah.tile_size_height, ah.tile_size_width), + tile_stride=(ah.tile_stride_height, ah.tile_stride_width), + ) + frames = rearrange(frames, "B C T H W -> B T H W C")[0] + frames = ((frames.float() + 1) * 127.5).clip(0, 255).cpu().numpy().astype(np.uint8) + return frames + + +def save_mp4(path: str, frames: np.ndarray, fps: int = 5) -> None: + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + imageio.mimsave(path, list(frames), fps=fps, codec="libx264") + print(f"Saved {len(frames)} frames -> {path}") + + +def run(args: argparse.Namespace) -> None: + if not dist.is_initialized(): + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", "29500") + dist.init_process_group(backend="gloo", world_size=1, rank=0) + + print(f"Loading checkpoint: {args.model_path}") + policy = GrootSimPolicy( + embodiment_tag=EmbodimentTag.TREX, + model_path=args.model_path, + device=args.device, + ) + print("Model loaded.") + video_height, video_width = get_expected_video_resolution(policy) + print(f"Resizing input video to {video_height}x{video_width} (HxW) for eval_transform") + + episode = TrexEpisode(args.dataset_path, args.episode) + prompt = episode.get_task(0) if args.use_dataset_prompt else args.prompt + schedule = build_chunk_schedule(episode.length, args.num_chunks) + print(f"Prompt: {prompt!r}") + print(f"Running {len(schedule)} causal chunks: {schedule}") + + os.makedirs(args.output_dir, exist_ok=True) + video_chunks: list[torch.Tensor] = [] + times = [] + + for ci, frame_indices in enumerate(schedule): + obs = build_obs(episode, frame_indices, prompt, video_height, video_width) + t0 = time.perf_counter() + with torch.inference_mode(): + result, video_pred = policy.lazy_joint_forward_causal(Batch(obs=obs)) + elapsed = time.perf_counter() - t0 + times.append(elapsed) + video_chunks.append(video_pred) + + act = result.act + print( + f" chunk {ci:02d} frames={frame_indices} " + f"infer={elapsed:.2f}s start_frame={policy.trained_model.action_head.current_start_frame}" + ) + if ci == 0: + print(f" action keys: {[k for k in dir(act) if k.startswith('action.')]}") + + pred_frames = decode_video_latents(policy, video_chunks) + out_pred = os.path.join( + args.output_dir, + f"ep{args.episode:06d}_pred.mp4", + ) + save_mp4(out_pred, pred_frames, fps=args.fps) + + if args.save_input_clip: + # Save the conditioning frames (head camera) for reference + input_frames = [] + max_idx = min(schedule[-1][-1] + 1, episode.length) + for i in range(max_idx): + input_frames.append(episode.get_frame(i, "video.head_left")) + save_mp4( + os.path.join(args.output_dir, f"ep{args.episode:06d}_input_head_left.mp4"), + np.stack(input_frames, axis=0), + fps=args.fps, + ) + + print(f"Avg inference time per chunk: {np.mean(times):.2f}s") + print(f"Done. Output dir: {os.path.abspath(args.output_dir)}") + + +def main() -> None: + p = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) + p.add_argument("--model_path", required=True, help="checkpoint-5000 directory") + p.add_argument( + "--dataset_path", + default="/scratch1/home/zhicao/dreamzero/data/trex_small", + ) + p.add_argument("--episode", type=int, default=0) + p.add_argument("--num_chunks", type=int, default=8, + help="Number of causal chunks (more -> longer predicted video)") + p.add_argument("--prompt", default="perform the task") + p.add_argument("--use_dataset_prompt", action="store_true") + p.add_argument("--device", default="cuda:0") + p.add_argument("--output_dir", default="results_trex_infer") + p.add_argument("--fps", type=int, default=5) + p.add_argument("--save_input_clip", action="store_true", + help="Also save input head_left frames for comparison") + run(p.parse_args()) + + +if __name__ == "__main__": + main() diff --git a/scripts/open_loop_yam.py b/scripts/open_loop_yam.py new file mode 100644 index 0000000000000000000000000000000000000000..b5a5e989ea56ab77e77cb050a7189c78bb766769 --- /dev/null +++ b/scripts/open_loop_yam.py @@ -0,0 +1,348 @@ +#!/usr/bin/env python3 +"""Offline open-loop evaluation for DreamZero on YAM data. + +Loads a model checkpoint directly (no server needed), reads YAM dataset +(parquet + MP4), runs inference, and compares predicted vs ground-truth actions. + +Usage: + python scripts/open_loop_yam.py \ + --model_path /path/to/checkpoint \ + --dataset_path Dataset/YAM_play_data \ + --device cuda:0 \ + --num_samples 200 +""" + +import torch._dynamo +torch._dynamo.config.disable = True + +import argparse +import glob +import os +import time + +import cv2 +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import pyarrow.parquet as pq +import torch +import torch.distributed as dist +from tianshou.data import Batch + +from groot.vla.data.schema import EmbodimentTag +from groot.vla.model.n1_5.sim_policy import GrootSimPolicy + + +# --------------------------------------------------------------------------- +# YAM layout (from Dataset/YAM_play_data/meta/modality.json) +# --------------------------------------------------------------------------- + +VIDEO_CAMERAS = { + "video.top_camera-images-rgb": "observation.images.top_camera-images-rgb", + "video.left_camera-images-rgb": "observation.images.left_camera-images-rgb", + "video.right_camera-images-rgb": "observation.images.right_camera-images-rgb", +} + +STATE_SLICES = { + "state.left_joint_pos": (34, 40), + "state.left_gripper_pos": (32, 33), + "state.right_joint_pos": (40, 46), + "state.right_gripper_pos": (33, 34), +} + +ACTION_SLICES = { + "action.left_joint_pos": (34, 40), + "action.left_gripper_pos": (32, 33), + "action.right_joint_pos": (40, 46), + "action.right_gripper_pos": (33, 34), +} + +ACTION_KEY_ORDER = [ + "action.left_joint_pos", + "action.left_gripper_pos", + "action.right_joint_pos", + "action.right_gripper_pos", +] + + +# --------------------------------------------------------------------------- +# Dataset reader (LeRobot chunked format) +# --------------------------------------------------------------------------- + +class YAMDataset: + """Reads LeRobot-style chunked parquet + MP4.""" + + def __init__(self, dataset_path: str): + self.root = dataset_path + + data_dir = os.path.join(dataset_path, "data") + parquet_files = sorted(glob.glob(os.path.join(data_dir, "**", "episode_*.parquet"), recursive=True)) + if not parquet_files: + raise FileNotFoundError(f"No episode_*.parquet found under {data_dir}") + + self.episodes = [] + self.cum_lengths = [0] + for pf in parquet_files: + table = pq.read_table(pf) + self.episodes.append(table) + self.cum_lengths.append(self.cum_lengths[-1] + table.num_rows) + self.total_rows = self.cum_lengths[-1] + + videos_root = os.path.join(dataset_path, "videos") + self.video_dirs = {} + for server_key, folder_name in VIDEO_CAMERAS.items(): + candidates = sorted(glob.glob(os.path.join(videos_root, "**", folder_name), recursive=True)) + if candidates: + self.video_dirs[server_key] = candidates[0] + + print(f"YAMDataset: {len(self.episodes)} episodes, " + f"{self.total_rows} rows, {len(self.video_dirs)} cameras") + + def __len__(self): + return self.total_rows + + def _locate(self, idx): + for ep in range(len(self.episodes)): + if idx < self.cum_lengths[ep + 1]: + return ep, idx - self.cum_lengths[ep] + raise IndexError(f"Index {idx} out of range ({self.total_rows})") + + def get_state(self, idx) -> np.ndarray: + ep, row = self._locate(idx) + return np.array(self.episodes[ep].column("observation.state")[row].as_py(), dtype=np.float64) + + def get_action(self, idx) -> np.ndarray: + ep, row = self._locate(idx) + return np.array(self.episodes[ep].column("action")[row].as_py(), dtype=np.float64) + + def get_task(self, idx) -> str: + ep, row = self._locate(idx) + try: + return str(self.episodes[ep].column("annotation.task")[row].as_py()) + except Exception: + return "" + + def get_frame(self, idx, server_key) -> np.ndarray: + """Read one video frame → (H, W, 3) uint8 RGB.""" + ep, row = self._locate(idx) + mp4 = os.path.join(self.video_dirs[server_key], f"episode_{ep:06d}.mp4") + cap = cv2.VideoCapture(mp4) + cap.set(cv2.CAP_PROP_POS_FRAMES, row) + ret, frame = cap.read() + cap.release() + if not ret: + raise RuntimeError(f"Failed to read frame {row} from {mp4}") + return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + + +# --------------------------------------------------------------------------- +# Observation builder +# --------------------------------------------------------------------------- + +def build_obs(dataset: YAMDataset, idx: int, prompt: str) -> dict: + """Build an obs dict matching what GrootSimPolicy.forward() expects.""" + obs = {} + + for server_key in dataset.video_dirs: + frame = dataset.get_frame(idx, server_key) + obs[server_key] = frame[np.newaxis, ...].astype(np.uint8) # (1, H, W, C) + + state = dataset.get_state(idx) + for key, (start, end) in STATE_SLICES.items(): + obs[key] = state[start:end].reshape(1, -1).astype(np.float64) # (1, D) + + obs["annotation.task"] = prompt + + return obs + + +def get_gt_action_dict(dataset: YAMDataset, idx: int) -> dict: + """Split the flat GT action vector into per-key arrays.""" + action_flat = dataset.get_action(idx) + gt = {} + for key in ACTION_KEY_ORDER: + s, e = ACTION_SLICES[key] + gt[key] = action_flat[s:e] + return gt + + +# --------------------------------------------------------------------------- +# Plotting +# --------------------------------------------------------------------------- + +def save_plots(all_preds, all_gts, key_names, output_dir): + """Plot pred vs gt for each action dimension across all keys.""" + pred_flat = np.concatenate([all_preds[k] for k in key_names], axis=-1) + gt_flat = np.concatenate([all_gts[k] for k in key_names], axis=-1) + D = pred_flat.shape[1] + mse_dim = np.mean((pred_flat - gt_flat) ** 2, axis=0) + + for d in range(D): + plt.figure(figsize=(10, 4)) + plt.plot(gt_flat[:, d], label="gt", alpha=0.8) + plt.plot(pred_flat[:, d], label="pred", alpha=0.8) + plt.title(f"Action dim {d} (MSE={mse_dim[d]:.6f})") + plt.xlabel("sample index"); plt.ylabel("value") + plt.legend(); plt.grid(True, alpha=0.3); plt.tight_layout() + plt.savefig(os.path.join(output_dir, f"action_dim_{d}.png"), dpi=150) + plt.close() + + ncols = 4 + nrows = (D + ncols - 1) // ncols + fig, axes = plt.subplots(nrows, ncols, figsize=(5 * ncols, 3.5 * nrows), squeeze=False) + overall_mse = float(np.mean(mse_dim)) + fig.suptitle(f"All action dims (overall MSE={overall_mse:.6f})", fontsize=14) + for d in range(D): + ax = axes[d // ncols][d % ncols] + ax.plot(gt_flat[:, d], label="gt", alpha=0.7, lw=0.8) + ax.plot(pred_flat[:, d], label="pred", alpha=0.7, lw=0.8) + ax.set_title(f"dim {d} (MSE={mse_dim[d]:.4f})", fontsize=9) + ax.tick_params(labelsize=7); ax.grid(True, alpha=0.2) + if d == 0: ax.legend(fontsize=7) + for d in range(D, nrows * ncols): + axes[d // ncols][d % ncols].set_visible(False) + fig.tight_layout(rect=[0, 0, 1, 0.96]) + fig.savefig(os.path.join(output_dir, "all_action_dims.png"), dpi=200) + plt.close(fig) + + # Per-key summary plot + fig2, axes2 = plt.subplots(1, len(key_names), figsize=(5 * len(key_names), 4), squeeze=False) + for i, k in enumerate(key_names): + ax = axes2[0][i] + p, g = all_preds[k], all_gts[k] + for d in range(p.shape[1]): + ax.plot(g[:, d], '--', alpha=0.5, lw=0.8) + ax.plot(p[:, d], alpha=0.7, lw=0.8) + key_mse = float(np.mean((p - g) ** 2)) + ax.set_title(f"{k}\nMSE={key_mse:.6f}", fontsize=9) + ax.grid(True, alpha=0.2); ax.tick_params(labelsize=7) + fig2.suptitle("Per-key pred (solid) vs gt (dashed)", fontsize=12) + fig2.tight_layout(rect=[0, 0, 1, 0.94]) + fig2.savefig(os.path.join(output_dir, "per_key_summary.png"), dpi=200) + plt.close(fig2) + + return mse_dim, overall_mse + + +# --------------------------------------------------------------------------- +# Main evaluation +# --------------------------------------------------------------------------- + +def evaluate(args): + # Single-process distributed init (GrootSimPolicy uses dist.get_rank()) + if not dist.is_initialized(): + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", "29500") + dist.init_process_group(backend="gloo", world_size=1, rank=0) + + print(f"Loading model from {args.model_path} ...") + policy = GrootSimPolicy( + embodiment_tag=EmbodimentTag.YAM, + model_path=args.model_path, + device=args.device, + ) + print("Model loaded.") + + dataset = YAMDataset(args.dataset_path) + os.makedirs(args.output_dir, exist_ok=True) + + num = min(args.num_samples, len(dataset)) + preds_per_key = {k: [] for k in ACTION_KEY_ORDER} + gts_per_key = {k: [] for k in ACTION_KEY_ORDER} + times = [] + + print(f"\nEvaluating {num} samples (start={args.start_idx}) ...") + print("-" * 60) + + for i in range(num): + idx = args.start_idx + i + + prompt = args.prompt + if args.use_dataset_prompt: + task = dataset.get_task(idx) + if task: + prompt = task + + obs = build_obs(dataset, idx, prompt) + + t0 = time.perf_counter() + with torch.inference_mode(): + result, _ = policy.lazy_joint_forward_causal(Batch(obs=obs)) + elapsed = time.perf_counter() - t0 + times.append(elapsed) + + gt = get_gt_action_dict(dataset, idx) + + for k in ACTION_KEY_ORDER: + if k in result.act: + pred_val = result.act[k] + if isinstance(pred_val, torch.Tensor): + pred_val = pred_val.cpu().numpy() + # First timestep: squeeze may collapse (1,24,1) -> (24,) + pred_val = np.atleast_1d(pred_val[0]).flatten() + preds_per_key[k].append(pred_val) + gts_per_key[k].append(gt[k]) + + if i % args.log_every == 0: + if i == 0: + print(f" Action keys in output: {list(result.act.keys())}") + for k in ACTION_KEY_ORDER: + if k in result.act: + v = result.act[k] + shape = v.shape if hasattr(v, 'shape') else "?" + print(f" {k}: pred_shape={shape}, gt_shape={gt[k].shape}") + print(f" [{i:>5d}/{num}] idx={idx} infer={elapsed:.3f}s prompt={prompt!r:.60}") + + # Stack results + valid_keys = [k for k in ACTION_KEY_ORDER if len(preds_per_key[k]) > 0] + if not valid_keys: + print("No predictions!"); return + + stacked_preds = {k: np.stack(preds_per_key[k]) for k in valid_keys} + stacked_gts = {k: np.stack(gts_per_key[k]) for k in valid_keys} + + pred_all = np.concatenate([stacked_preds[k] for k in valid_keys], axis=-1) + gt_all = np.concatenate([stacked_gts[k] for k in valid_keys], axis=-1) + overall_mse = float(np.mean((pred_all - gt_all) ** 2)) + + print(f"\n{'='*60}") + print(f"Overall MSE: {overall_mse:.6f} | Avg inference time: {np.mean(times):.4f}s") + for k in valid_keys: + k_mse = float(np.mean((stacked_preds[k] - stacked_gts[k]) ** 2)) + print(f" {k}: MSE={k_mse:.6f}") + print(f"{'='*60}") + + mse_dim, _ = save_plots(stacked_preds, stacked_gts, valid_keys, args.output_dir) + + with open(os.path.join(args.output_dir, "mse.txt"), "w") as f: + f.write(f"overall_mse,{overall_mse}\n") + for k in valid_keys: + k_mse = float(np.mean((stacked_preds[k] - stacked_gts[k]) ** 2)) + f.write(f"{k},{k_mse}\n") + for d, v in enumerate(mse_dim): + f.write(f"dim_{d},{v}\n") + + print(f"Results saved to {os.path.abspath(args.output_dir)}/") + + +def main(): + p = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) + p.add_argument("--model_path", required=True, + help="Path to model checkpoint dir (contains config.json, model.safetensors, experiment_cfg/)") + p.add_argument("--dataset_path", required=True, + help="Root of YAM dataset (contains data/, videos/, meta/)") + p.add_argument("--device", default="cuda:0") + p.add_argument("--prompt", default="pick up the object") + p.add_argument("--use_dataset_prompt", action="store_true", + help="Read task annotation from parquet instead of --prompt") + p.add_argument("--num_samples", type=int, default=300) + p.add_argument("--start_idx", type=int, default=0) + p.add_argument("--output_dir", default="results_yam") + p.add_argument("--log_every", type=int, default=10) + main_args = p.parse_args() + evaluate(main_args) + + +if __name__ == "__main__": + main() diff --git a/scripts/train/agibot_training.sh b/scripts/train/agibot_training.sh new file mode 100644 index 0000000000000000000000000000000000000000..e8702b45e2c739ff3cce16079f1f59a547109453 --- /dev/null +++ b/scripts/train/agibot_training.sh @@ -0,0 +1,101 @@ +#!/bin/bash +# DreamZero AGIbot Training Script +# +# Usage: +# # Set your dataset path and output directory, then run: +# bash scripts/train/agibot_training.sh +# +# Prerequisites: +# - AGIbot dataset in LeRobot format at AGIBOT_DATA_ROOT (state 32, action 22, 3 views: top_head, hand_left, hand_right) +# See docs/DATASET_TO_GEAR_AND_TRAIN.md for conversion instructions +# - Wan2.1-I2V-14B-480P weights (auto-downloaded or pre-downloaded from HuggingFace) +# Download: huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir ./checkpoints/Wan2.1-I2V-14B-480P +# - umt5-xxl tokenizer (auto-downloaded or pre-downloaded from HuggingFace) +# Download: huggingface-cli download google/umt5-xxl --local-dir ./checkpoints/umt5-xxl +# - DreamZero-AgiBot pretrained checkpoint (for loading LoRA weights before fine-tuning) +# git clone https://huggingface.co/GEAR-Dreams/DreamZero-AgiBot ./checkpoints/DreamZero-AgiBot + +export HYDRA_FULL_ERROR=1 + +# ============ CHANGE THESE VARIABLES ============ +# Dataset path (AGIbot in LeRobot format: state 32, action 22, videos top_head, hand_left, hand_right) +AGIBOT_DATA_ROOT=${AGIBOT_DATA_ROOT:-"./data/agibot_lerobot"} + +# Output directory for training checkpoints +OUTPUT_DIR=${OUTPUT_DIR:-"./checkpoints/dreamzero_agibot_lora_5k"} + +# Number of GPUs to use (default: all visible GPUs, so 4-GPU machines use 4 without setting NUM_GPUS) +if [ -z "${NUM_GPUS}" ]; then + NUM_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) +fi +NUM_GPUS=${NUM_GPUS:-8} + +# Model weight paths (download from HuggingFace if not already present) +WAN_CKPT_DIR=${WAN_CKPT_DIR:-"./checkpoints/Wan2.1-I2V-14B-480P"} +TOKENIZER_DIR=${TOKENIZER_DIR:-"./checkpoints/umt5-xxl"} +# ============================================= + +# ============ AUTO-DOWNLOAD WEIGHTS ============ +if [ ! -d "$WAN_CKPT_DIR" ] || [ -z "$(ls -A "$WAN_CKPT_DIR" 2>/dev/null)" ]; then + echo "Wan2.1-I2V-14B-480P not found at $WAN_CKPT_DIR. Downloading from HuggingFace..." + huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir "$WAN_CKPT_DIR" +fi + +if [ ! -d "$TOKENIZER_DIR" ] || [ -z "$(ls -A "$TOKENIZER_DIR" 2>/dev/null)" ]; then + echo "umt5-xxl tokenizer not found at $TOKENIZER_DIR. Downloading from HuggingFace..." + huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR" +fi +# ================================================ + +# Validate dataset exists +if [ ! -d "$AGIBOT_DATA_ROOT" ]; then + echo "ERROR: AGIbot dataset not found at $AGIBOT_DATA_ROOT" + echo "Set AGIBOT_DATA_ROOT to your LeRobot-format AGIbot dataset (e.g. Dataset/3222_raw_assemble)" + exit 1 +fi + +torchrun --nproc_per_node $NUM_GPUS --standalone groot/vla/experiment/experiment.py \ + report_to=wandb \ + data=dreamzero/agibot_relative \ + wandb_project=dreamzero \ + train_architecture=lora \ + num_frames=33 \ + action_horizon=24 \ + num_views=3 \ + model=dreamzero/vla \ + model/dreamzero/action_head=wan_flow_matching_action_tf \ + model/dreamzero/transform=dreamzero_cotrain \ + num_frame_per_block=2 \ + num_action_per_block=24 \ + num_state_per_block=1 \ + seed=42 \ + training_args.learning_rate=1e-5 \ + training_args.deepspeed="groot/vla/configs/deepspeed/zero2.json" \ + save_steps=2500 \ + training_args.warmup_ratio=0.05 \ + output_dir=$OUTPUT_DIR \ + per_device_train_batch_size=1 \ + max_steps=5000 \ + weight_decay=1e-5 \ + save_total_limit=10 \ + upload_checkpoints=false \ + bf16=true \ + tf32=true \ + eval_bf16=true \ + dataloader_pin_memory=false \ + dataloader_num_workers=1 \ + image_resolution_width=320 \ + image_resolution_height=176 \ + save_lora_only=true \ + max_chunk_size=4 \ + frame_seqlen=880 \ + save_strategy=steps \ + agibot_data_root=$AGIBOT_DATA_ROOT \ + dit_version=$WAN_CKPT_DIR \ + text_encoder_pretrained_path=$WAN_CKPT_DIR/models_t5_umt5-xxl-enc-bf16.pth \ + image_encoder_pretrained_path=$WAN_CKPT_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth \ + vae_pretrained_path=$WAN_CKPT_DIR/Wan2.1_VAE.pth \ + tokenizer_path=$TOKENIZER_DIR \ + pretrained_model_path=./checkpoints/DreamZero-AgiBot \ + ++action_head_cfg.config.skip_component_loading=true \ + ++action_head_cfg.config.defer_lora_injection=true diff --git a/scripts/train/droid_training_full_finetune.sh b/scripts/train/droid_training_full_finetune.sh new file mode 100644 index 0000000000000000000000000000000000000000..2e7c2e65e08bc9ae9c7da71bd39a63184a79cdb8 --- /dev/null +++ b/scripts/train/droid_training_full_finetune.sh @@ -0,0 +1,93 @@ +#!/bin/bash +# DreamZero DROID Full Fine-Tuning Script (8x H100, ZeRO-2 + CPU Offload) +# +# Usage: +# bash scripts/train/droid_training_full_finetune.sh +# +# Prerequisites: +# - DROID dataset in LeRobot format at DROID_DATA_ROOT +# Download: huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir ./data/droid_lerobot +# Or convert from scratch: see scripts/data/convert_droid.py +# - Wan2.1-I2V-14B-480P weights (auto-downloaded or pre-downloaded from HuggingFace) +# Download: huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir ./checkpoints/Wan2.1-I2V-14B-480P +# - umt5-xxl tokenizer (auto-downloaded or pre-downloaded from HuggingFace) +# Download: huggingface-cli download google/umt5-xxl --local-dir ./checkpoints/umt5-xxl + +export HYDRA_FULL_ERROR=1 + +# ============ USER CONFIGURATION ============ +# Dataset path (DROID in LeRobot format) +DROID_DATA_ROOT=${DROID_DATA_ROOT:-"./data/droid_lerobot"} + +# Output directory for training checkpoints +OUTPUT_DIR=${OUTPUT_DIR:-"./checkpoints/dreamzero_droid_full_finetune"} + +# Number of GPUs to use (8x H100 for ZeRO-2 + CPU offload full fine-tuning) +NUM_GPUS=${NUM_GPUS:-8} + +# Model weight paths (download from HuggingFace if not already present) +WAN_CKPT_DIR=${WAN_CKPT_DIR:-"./checkpoints/Wan2.1-I2V-14B-480P"} +TOKENIZER_DIR=${TOKENIZER_DIR:-"./checkpoints/umt5-xxl"} +# ============================================= + +# ============ AUTO-DOWNLOAD WEIGHTS ============ +if [ ! -d "$WAN_CKPT_DIR" ] || [ -z "$(ls -A "$WAN_CKPT_DIR" 2>/dev/null)" ]; then + echo "Wan2.1-I2V-14B-480P not found at $WAN_CKPT_DIR. Downloading from HuggingFace..." + huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir "$WAN_CKPT_DIR" +fi + +if [ ! -d "$TOKENIZER_DIR" ] || [ -z "$(ls -A "$TOKENIZER_DIR" 2>/dev/null)" ]; then + echo "umt5-xxl tokenizer not found at $TOKENIZER_DIR. Downloading from HuggingFace..." + huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR" +fi +# ================================================ + +# Validate dataset exists +if [ ! -d "$DROID_DATA_ROOT" ]; then + echo "ERROR: DROID dataset not found at $DROID_DATA_ROOT" + echo "Download with: huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir $DROID_DATA_ROOT" + exit 1 +fi + +torchrun --nproc_per_node $NUM_GPUS --standalone groot/vla/experiment/experiment.py \ + report_to=none \ + data=dreamzero/droid_relative \ + wandb_project=dreamzero \ + train_architecture=full \ + num_frames=33 \ + action_horizon=24 \ + num_views=3 \ + model=dreamzero/vla \ + model/dreamzero/action_head=wan_flow_matching_action_tf \ + model/dreamzero/transform=dreamzero_cotrain \ + num_frame_per_block=2 \ + num_action_per_block=24 \ + num_state_per_block=1 \ + seed=42 \ + training_args.learning_rate=1e-5 \ + training_args.deepspeed="groot/vla/configs/deepspeed/zero2_offload.json" \ + save_steps=1000 \ + training_args.warmup_ratio=0.05 \ + output_dir=$OUTPUT_DIR \ + per_device_train_batch_size=1 \ + max_steps=100 \ + weight_decay=1e-5 \ + save_total_limit=10 \ + upload_checkpoints=false \ + bf16=true \ + tf32=true \ + eval_bf16=true \ + dataloader_pin_memory=false \ + dataloader_num_workers=1 \ + image_resolution_width=320 \ + image_resolution_height=176 \ + save_lora_only=false \ + max_chunk_size=4 \ + frame_seqlen=880 \ + save_strategy=no \ + droid_data_root=$DROID_DATA_ROOT \ + dit_version=$WAN_CKPT_DIR \ + text_encoder_pretrained_path=$WAN_CKPT_DIR/models_t5_umt5-xxl-enc-bf16.pth \ + image_encoder_pretrained_path=$WAN_CKPT_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth \ + vae_pretrained_path=$WAN_CKPT_DIR/Wan2.1_VAE.pth \ + tokenizer_path=$TOKENIZER_DIR diff --git a/scripts/train/droid_training_full_finetune_wan21.sh b/scripts/train/droid_training_full_finetune_wan21.sh new file mode 100644 index 0000000000000000000000000000000000000000..a43a8d8d4a323b738f6b33b5c47ed2c0563a3984 --- /dev/null +++ b/scripts/train/droid_training_full_finetune_wan21.sh @@ -0,0 +1,139 @@ +#!/bin/bash +# DreamZero DROID Full Fine-Tuning Script with Wan2.1-I2V-14B-480P backbone +# +# Usage: +# bash scripts/train/droid_training_full_finetune_wan21.sh +# +# Prerequisites: +# - DROID dataset in LeRobot format at DROID_DATA_ROOT +# Download: huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir ./data/droid_lerobot +# - Wan2.1-I2V-14B-480P weights (auto-downloaded or pre-downloaded from HuggingFace) +# huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir ./checkpoints/Wan2.1-I2V-14B-480P +# - umt5-xxl tokenizer (auto-downloaded or pre-downloaded) +# huggingface-cli download google/umt5-xxl --local-dir ./checkpoints/umt5-xxl + +export HYDRA_FULL_ERROR=1 + +# Repo root (same logic as droid_training_wan22.sh) +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SCRIPT_REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +if [ -n "$DREAMZERO_ROOT" ] && [ -d "$DREAMZERO_ROOT/groot" ]; then + : +elif [ -d "/root/yejink/dreamzero/groot" ]; then + DREAMZERO_ROOT=/root/yejink/dreamzero +elif [ -d "/root/dreamzero/groot" ]; then + DREAMZERO_ROOT=/root/dreamzero +elif [ -d "$SCRIPT_REPO_ROOT/groot" ]; then + DREAMZERO_ROOT="$SCRIPT_REPO_ROOT" +else + DREAMZERO_ROOT="${DREAMZERO_ROOT:-/root/yejink/dreamzero}" +fi +if [ ! -d "$DREAMZERO_ROOT/groot" ]; then + echo "ERROR: No groot/ under $DREAMZERO_ROOT. Set DREAMZERO_ROOT to the dreamzero repo root that contains groot/." + exit 1 +fi + +# ============ USER CONFIGURATION ============ +DROID_DATA_ROOT=${DROID_DATA_ROOT:-"$DREAMZERO_ROOT/data/droid_lerobot"} +if [ "$DROID_DATA_ROOT" = "./data/droid_lerobot" ]; then + DROID_DATA_ROOT="$DREAMZERO_ROOT/data/droid_lerobot" +fi +# Fixed output dir so Beaker restarts resume from latest checkpoint in this folder. +# For concurrent runs use: OUTPUT_DIR=.../dreamzero_droid_wan21_full_finetune_$(date +%Y%m%d_%H%M%S) +# +# Resume from 100k to 200k: use this same OUTPUT_DIR and run this script with max_steps=200000. +# If the 100k run completed (saved final model), the code will see config.json and skip training. +# To force resume: remove config.json from OUTPUT_DIR so the latest checkpoint-* is used, then run. +OUTPUT_DIR=${OUTPUT_DIR:-"$DREAMZERO_ROOT/checkpoints/dreamzero_droid_wan21_full_finetune"} + +NUM_GPUS=${NUM_GPUS:-4} +PER_DEVICE_BS=${PER_DEVICE_BS:-1} +# Global batch: default = NUM_GPUS * PER_DEVICE_BS. Override for larger effective batch, e.g. GLOBAL_BATCH_SIZE=128 +# (uses gradient accumulation; memory stays same. With 4 GPUs, 128 → grad_accum=32.) +GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE:-$((NUM_GPUS * PER_DEVICE_BS))} + +# Wan2.1-I2V-14B-480P checkpoint (includes CLIP image encoder, T5 text encoder, and VAE) +WAN21_CKPT_DIR=${WAN21_CKPT_DIR:-"$DREAMZERO_ROOT/checkpoints/Wan2.1-I2V-14B-480P"} +TOKENIZER_DIR=${TOKENIZER_DIR:-"$DREAMZERO_ROOT/checkpoints/umt5-xxl"} +# ============================================= + +# ============ AUTO-DOWNLOAD WEIGHTS ============ +if [ ! -d "$WAN21_CKPT_DIR" ] || [ -z "$(ls -A "$WAN21_CKPT_DIR" 2>/dev/null)" ]; then + echo "Wan2.1-I2V-14B-480P not found at $WAN21_CKPT_DIR. Downloading from HuggingFace..." + huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir "$WAN21_CKPT_DIR" +fi + +if [ ! -d "$TOKENIZER_DIR" ] || [ -z "$(ls -A "$TOKENIZER_DIR" 2>/dev/null)" ]; then + echo "umt5-xxl tokenizer not found at $TOKENIZER_DIR. Downloading from HuggingFace..." + huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR" +fi +# ================================================ + +if [ ! -d "$DROID_DATA_ROOT" ]; then + echo "ERROR: DROID dataset not found at $DROID_DATA_ROOT" + echo "Download with: huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir $DROID_DATA_ROOT" + exit 1 +fi + +EXPERIMENT_PY="$DREAMZERO_ROOT/groot/vla/experiment/experiment.py" +if [ ! -f "$EXPERIMENT_PY" ]; then + echo "ERROR: Not found: $EXPERIMENT_PY" + exit 1 +fi +PYTHON_311="/usr/bin/python3.11" +if [ -x "$PYTHON_311" ]; then + if [ -n "${FIX_NUMPY_IN_SCRIPT:-}" ]; then + "$PYTHON_311" -m pip install "numpy==1.26.4" --force-reinstall -q 2>/dev/null || true + fi + RUN_CMD=( "$PYTHON_311" -m torch.distributed.run --nproc_per_node "$NUM_GPUS" --standalone "$EXPERIMENT_PY" ) + echo "Using image Python 3.11: $PYTHON_311" +else + RUN_CMD=( python3 -m torch.distributed.run --nproc_per_node "$NUM_GPUS" --standalone "$EXPERIMENT_PY" ) + echo "Using: $(command -v python3)" +fi +cd "$DREAMZERO_ROOT" + +# Full fine-tune: train_architecture=full, save_lora_only=false, ZeRO-2 (+ optional CPU offload) +DEEPSPEED_CFG=${DEEPSPEED_CFG:-zero2_offload} +"${RUN_CMD[@]}" \ + report_to=wandb \ + data=dreamzero/droid_relative \ + wandb_project=dreamzero \ + train_architecture=full \ + num_frames=33 \ + action_horizon=24 \ + num_views=3 \ + model=dreamzero/vla \ + model/dreamzero/action_head=wan_flow_matching_action_tf \ + model/dreamzero/transform=dreamzero_cotrain \ + num_frame_per_block=2 \ + num_action_per_block=24 \ + num_state_per_block=1 \ + seed=42 \ + training_args.learning_rate=1e-5 \ + training_args.deepspeed="groot/vla/configs/deepspeed/${DEEPSPEED_CFG}.json" \ + save_steps=500 \ + training_args.warmup_ratio=0.05 \ + output_dir=$OUTPUT_DIR \ + per_device_train_batch_size=$PER_DEVICE_BS \ + global_batch_size=$GLOBAL_BATCH_SIZE \ + max_steps=200000 \ + weight_decay=1e-5 \ + save_total_limit=10 \ + upload_checkpoints=false \ + bf16=true \ + tf32=true \ + eval_bf16=true \ + dataloader_pin_memory=true \ + dataloader_num_workers=4 \ + image_resolution_width=320 \ + image_resolution_height=176 \ + save_lora_only=false \ + max_chunk_size=4 \ + save_strategy=steps \ + droid_data_root=$DROID_DATA_ROOT \ + dit_version=$WAN21_CKPT_DIR \ + text_encoder_pretrained_path=$WAN21_CKPT_DIR/models_t5_umt5-xxl-enc-bf16.pth \ + image_encoder_pretrained_path=$WAN21_CKPT_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth \ + vae_pretrained_path=$WAN21_CKPT_DIR/Wan2.1_VAE.pth \ + tokenizer_path=$TOKENIZER_DIR diff --git a/scripts/train/droid_training_full_finetune_wan22.sh b/scripts/train/droid_training_full_finetune_wan22.sh new file mode 100644 index 0000000000000000000000000000000000000000..fa10a25df3fa09e7c8c3222d17ad6f4e420a8ca1 --- /dev/null +++ b/scripts/train/droid_training_full_finetune_wan22.sh @@ -0,0 +1,147 @@ +#!/bin/bash +# DreamZero DROID Full Fine-Tuning Script with Wan2.2-TI2V-5B backbone +# +# Usage: +# bash scripts/train/droid_training_full_finetune_wan22.sh +# +# Prerequisites: +# - DROID dataset in LeRobot format at DROID_DATA_ROOT +# Download: huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir ./data/droid_lerobot +# - Wan2.2-TI2V-5B weights (auto-downloaded or pre-downloaded from HuggingFace) +# huggingface-cli download Wan-AI/Wan2.2-TI2V-5B --local-dir ./checkpoints/Wan2.2-TI2V-5B +# - Image encoder (CLIP) from Wan2.1 - Wan2.2-TI2V-5B does not include it +# huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir ./checkpoints/Wan2.1-I2V-14B-480P +# - umt5-xxl tokenizer (auto-downloaded or pre-downloaded) +# huggingface-cli download google/umt5-xxl --local-dir ./checkpoints/umt5-xxl + +export HYDRA_FULL_ERROR=1 + +# Repo root (same logic as droid_training_wan22.sh) +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SCRIPT_REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +if [ -n "$DREAMZERO_ROOT" ] && [ -d "$DREAMZERO_ROOT/groot" ]; then + : +elif [ -d "/root/yejink/dreamzero/groot" ]; then + DREAMZERO_ROOT=/root/yejink/dreamzero +elif [ -d "/root/dreamzero/groot" ]; then + DREAMZERO_ROOT=/root/dreamzero +elif [ -d "$SCRIPT_REPO_ROOT/groot" ]; then + DREAMZERO_ROOT="$SCRIPT_REPO_ROOT" +else + DREAMZERO_ROOT="${DREAMZERO_ROOT:-/root/yejink/dreamzero}" +fi +if [ ! -d "$DREAMZERO_ROOT/groot" ]; then + echo "ERROR: No groot/ under $DREAMZERO_ROOT. Set DREAMZERO_ROOT to the dreamzero repo root that contains groot/." + exit 1 +fi + +# ============ USER CONFIGURATION ============ +DROID_DATA_ROOT=${DROID_DATA_ROOT:-"$DREAMZERO_ROOT/data/droid_lerobot"} +if [ "$DROID_DATA_ROOT" = "./data/droid_lerobot" ]; then + DROID_DATA_ROOT="$DREAMZERO_ROOT/data/droid_lerobot" +fi +# Fixed output dir so Beaker restarts resume from latest checkpoint in this folder. +# For concurrent runs use: OUTPUT_DIR=.../dreamzero_droid_wan22_full_finetune_$(date +%Y%m%d_%H%M%S) +# +# Resume from 100k to 200k: use this same OUTPUT_DIR and run this script with max_steps=200000. +# If the 100k run completed (saved final model), the code will see config.json and skip training. +# To force resume: remove config.json from OUTPUT_DIR so the latest checkpoint-* is used, then run. +OUTPUT_DIR=${OUTPUT_DIR:-"$DREAMZERO_ROOT/checkpoints/dreamzero_droid_wan22_full_finetune"} + +NUM_GPUS=${NUM_GPUS:-4} +PER_DEVICE_BS=${PER_DEVICE_BS:-1} +# Global batch: default = NUM_GPUS * PER_DEVICE_BS. Override for larger effective batch, e.g. GLOBAL_BATCH_SIZE=128 +# (uses gradient accumulation; memory stays same. With 4 GPUs, 128 → grad_accum=32.) +GLOBAL_BATCH_SIZE=${GLOBAL_BATCH_SIZE:-$((NUM_GPUS * PER_DEVICE_BS))} + +# Wan2.2-TI2V-5B checkpoint +WAN22_CKPT_DIR=${WAN22_CKPT_DIR:-"$DREAMZERO_ROOT/checkpoints/Wan2.2-TI2V-5B"} +IMAGE_ENCODER_DIR=${IMAGE_ENCODER_DIR:-"$DREAMZERO_ROOT/checkpoints/Wan2.1-I2V-14B-480P"} +TOKENIZER_DIR=${TOKENIZER_DIR:-"$DREAMZERO_ROOT/checkpoints/umt5-xxl"} +# ============================================= + +# ============ AUTO-DOWNLOAD WEIGHTS ============ +if [ ! -d "$WAN22_CKPT_DIR" ] || [ -z "$(ls -A "$WAN22_CKPT_DIR" 2>/dev/null)" ]; then + echo "Wan2.2-TI2V-5B not found at $WAN22_CKPT_DIR. Downloading from HuggingFace..." + huggingface-cli download Wan-AI/Wan2.2-TI2V-5B --local-dir "$WAN22_CKPT_DIR" +fi + +if [ ! -d "$TOKENIZER_DIR" ] || [ -z "$(ls -A "$TOKENIZER_DIR" 2>/dev/null)" ]; then + echo "umt5-xxl tokenizer not found at $TOKENIZER_DIR. Downloading from HuggingFace..." + huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR" +fi + +if [ ! -f "$IMAGE_ENCODER_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" ]; then + echo "Image encoder not found. Downloading Wan2.1-I2V-14B-480P (for CLIP only)..." + huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir "$IMAGE_ENCODER_DIR" +fi +# ================================================ + +if [ ! -d "$DROID_DATA_ROOT" ]; then + echo "ERROR: DROID dataset not found at $DROID_DATA_ROOT" + echo "Download with: huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir $DROID_DATA_ROOT" + exit 1 +fi + +EXPERIMENT_PY="$DREAMZERO_ROOT/groot/vla/experiment/experiment.py" +if [ ! -f "$EXPERIMENT_PY" ]; then + echo "ERROR: Not found: $EXPERIMENT_PY" + exit 1 +fi +PYTHON_311="/usr/bin/python3.11" +if [ -x "$PYTHON_311" ]; then + if [ -n "${FIX_NUMPY_IN_SCRIPT:-}" ]; then + "$PYTHON_311" -m pip install "numpy==1.26.4" --force-reinstall -q 2>/dev/null || true + fi + RUN_CMD=( "$PYTHON_311" -m torch.distributed.run --nproc_per_node "$NUM_GPUS" --standalone "$EXPERIMENT_PY" ) + echo "Using image Python 3.11: $PYTHON_311" +else + RUN_CMD=( python3 -m torch.distributed.run --nproc_per_node "$NUM_GPUS" --standalone "$EXPERIMENT_PY" ) + echo "Using: $(command -v python3)" +fi +cd "$DREAMZERO_ROOT" + +# Full fine-tune: train_architecture=full, save_lora_only=false, ZeRO-2 (+ optional CPU offload) +DEEPSPEED_CFG=${DEEPSPEED_CFG:-zero2_offload} +"${RUN_CMD[@]}" \ + report_to=wandb \ + data=dreamzero/droid_relative_wan22 \ + wandb_project=dreamzero \ + train_architecture=full \ + num_frames=33 \ + action_horizon=24 \ + num_views=3 \ + model=dreamzero/vla \ + model/dreamzero/action_head=wan_flow_matching_action_tf_wan22 \ + model/dreamzero/transform=dreamzero_cotrain \ + num_frame_per_block=2 \ + num_action_per_block=24 \ + num_state_per_block=1 \ + seed=42 \ + training_args.learning_rate=1e-5 \ + training_args.deepspeed="groot/vla/configs/deepspeed/${DEEPSPEED_CFG}.json" \ + save_steps=1000 \ + training_args.warmup_ratio=0.05 \ + output_dir=$OUTPUT_DIR \ + per_device_train_batch_size=$PER_DEVICE_BS \ + global_batch_size=$GLOBAL_BATCH_SIZE \ + max_steps=200000 \ + weight_decay=1e-5 \ + save_total_limit=10 \ + upload_checkpoints=false \ + bf16=true \ + tf32=true \ + eval_bf16=true \ + dataloader_pin_memory=true \ + dataloader_num_workers=4 \ + image_resolution_width=320 \ + image_resolution_height=160 \ + save_lora_only=false \ + max_chunk_size=4 \ + save_strategy=steps \ + droid_data_root=$DROID_DATA_ROOT \ + dit_version=$WAN22_CKPT_DIR \ + text_encoder_pretrained_path=$WAN22_CKPT_DIR/models_t5_umt5-xxl-enc-bf16.pth \ + image_encoder_pretrained_path=$IMAGE_ENCODER_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth \ + vae_pretrained_path=$WAN22_CKPT_DIR/Wan2.2_VAE.pth \ + tokenizer_path=$TOKENIZER_DIR diff --git a/scripts/train/droid_training_lora.sh b/scripts/train/droid_training_lora.sh new file mode 100644 index 0000000000000000000000000000000000000000..978693bc352d899aa099f4c60c422fd36f3c1d39 --- /dev/null +++ b/scripts/train/droid_training_lora.sh @@ -0,0 +1,94 @@ +#!/bin/bash +# DreamZero DROID Training Script +# +# Usage: +# # Set your dataset path and output directory, then run: +# bash scripts/train/droid_training.sh +# +# Prerequisites: +# - DROID dataset in LeRobot format at DROID_DATA_ROOT +# Download: huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir ./data/droid_lerobot +# Or convert from scratch: see scripts/data/convert_droid.py +# - Wan2.1-I2V-14B-480P weights (auto-downloaded or pre-downloaded from HuggingFace) +# Download: huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir ./checkpoints/Wan2.1-I2V-14B-480P +# - umt5-xxl tokenizer (auto-downloaded or pre-downloaded from HuggingFace) +# Download: huggingface-cli download google/umt5-xxl --local-dir ./checkpoints/umt5-xxl + +export HYDRA_FULL_ERROR=1 + +# ============ USER CONFIGURATION ============ +# Dataset path (DROID in LeRobot format) +DROID_DATA_ROOT=${DROID_DATA_ROOT:-"./data/droid_lerobot"} + +# Output directory for training checkpoints +OUTPUT_DIR=${OUTPUT_DIR:-"./checkpoints/dreamzero_droid_lora"} + +# Number of GPUs to use +NUM_GPUS=${NUM_GPUS:-8} + +# Model weight paths (download from HuggingFace if not already present) +WAN_CKPT_DIR=${WAN_CKPT_DIR:-"./checkpoints/Wan2.1-I2V-14B-480P"} +TOKENIZER_DIR=${TOKENIZER_DIR:-"./checkpoints/umt5-xxl"} +# ============================================= + +# ============ AUTO-DOWNLOAD WEIGHTS ============ +if [ ! -d "$WAN_CKPT_DIR" ] || [ -z "$(ls -A "$WAN_CKPT_DIR" 2>/dev/null)" ]; then + echo "Wan2.1-I2V-14B-480P not found at $WAN_CKPT_DIR. Downloading from HuggingFace..." + huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir "$WAN_CKPT_DIR" +fi + +if [ ! -d "$TOKENIZER_DIR" ] || [ -z "$(ls -A "$TOKENIZER_DIR" 2>/dev/null)" ]; then + echo "umt5-xxl tokenizer not found at $TOKENIZER_DIR. Downloading from HuggingFace..." + huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR" +fi +# ================================================ + +# Validate dataset exists +if [ ! -d "$DROID_DATA_ROOT" ]; then + echo "ERROR: DROID dataset not found at $DROID_DATA_ROOT" + echo "Download with: huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir $DROID_DATA_ROOT" + exit 1 +fi + +torchrun --nproc_per_node $NUM_GPUS --standalone groot/vla/experiment/experiment.py \ + report_to=none \ + data=dreamzero/droid_relative \ + wandb_project=dreamzero \ + train_architecture=lora \ + num_frames=33 \ + action_horizon=24 \ + num_views=3 \ + model=dreamzero/vla \ + model/dreamzero/action_head=wan_flow_matching_action_tf \ + model/dreamzero/transform=dreamzero_cotrain \ + num_frame_per_block=2 \ + num_action_per_block=24 \ + num_state_per_block=1 \ + seed=42 \ + training_args.learning_rate=1e-4 \ + training_args.deepspeed="groot/vla/configs/deepspeed/zero2.json" \ + save_steps=1000 \ + training_args.warmup_ratio=0.05 \ + output_dir=$OUTPUT_DIR \ + per_device_train_batch_size=1 \ + max_steps=100 \ + weight_decay=1e-5 \ + save_total_limit=10 \ + upload_checkpoints=false \ + bf16=true \ + tf32=true \ + eval_bf16=true \ + dataloader_pin_memory=false \ + dataloader_num_workers=1 \ + image_resolution_width=320 \ + image_resolution_height=176 \ + save_lora_only=true \ + max_chunk_size=4 \ + frame_seqlen=880 \ + save_strategy=no \ + droid_data_root=$DROID_DATA_ROOT \ + dit_version=$WAN_CKPT_DIR \ + text_encoder_pretrained_path=$WAN_CKPT_DIR/models_t5_umt5-xxl-enc-bf16.pth \ + image_encoder_pretrained_path=$WAN_CKPT_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth \ + vae_pretrained_path=$WAN_CKPT_DIR/Wan2.1_VAE.pth \ + tokenizer_path=$TOKENIZER_DIR \ No newline at end of file diff --git a/scripts/train/droid_training_wan22.sh b/scripts/train/droid_training_wan22.sh new file mode 100644 index 0000000000000000000000000000000000000000..da56a65b95ac96f66e0a874445dd9c61a6c67bde --- /dev/null +++ b/scripts/train/droid_training_wan22.sh @@ -0,0 +1,141 @@ +#!/bin/bash +# DreamZero DROID Training Script with Wan2.2-TI2V-5B backbone +# +# Usage: +# bash scripts/train/droid_training_wan22.sh +# +# Prerequisites: +# - DROID dataset in LeRobot format at DROID_DATA_ROOT +# Download: huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir ./data/droid_lerobot +# - Wan2.2-TI2V-5B weights (download from HuggingFace) +# huggingface-cli download Wan-AI/Wan2.2-TI2V-5B --local-dir ./checkpoints/Wan2.2-TI2V-5B +# - Image encoder (CLIP) from Wan2.1 - Wan2.2-TI2V-5B does not include it +# Option A: huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir ./checkpoints/Wan2.1-I2V-14B-480P +# Option B: Set IMAGE_ENCODER_DIR to a path containing models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth +# - umt5-xxl tokenizer (auto-downloaded or pre-downloaded) +# huggingface-cli download google/umt5-xxl --local-dir ./checkpoints/umt5-xxl + +export HYDRA_FULL_ERROR=1 + +# Repo root: must be a directory that contains groot/ (so experiment.py can be found). +# Beaker/weka uses /root/yejink/dreamzero; image has /root/dreamzero; else script location. +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SCRIPT_REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +if [ -n "$DREAMZERO_ROOT" ] && [ -d "$DREAMZERO_ROOT/groot" ]; then + : # keep existing and valid +elif [ -d "/root/yejink/dreamzero/groot" ]; then + DREAMZERO_ROOT=/root/yejink/dreamzero +elif [ -d "/root/dreamzero/groot" ]; then + DREAMZERO_ROOT=/root/dreamzero +elif [ -d "$SCRIPT_REPO_ROOT/groot" ]; then + DREAMZERO_ROOT="$SCRIPT_REPO_ROOT" +else + DREAMZERO_ROOT="${DREAMZERO_ROOT:-/root/yejink/dreamzero}" +fi +if [ ! -d "$DREAMZERO_ROOT/groot" ]; then + echo "ERROR: No groot/ under $DREAMZERO_ROOT. Set DREAMZERO_ROOT to the dreamzero repo root that contains groot/." + exit 1 +fi + +# ============ USER CONFIGURATION ============ +NUM_GPUS=${NUM_GPUS:-8} +DROID_DATA_ROOT=${DROID_DATA_ROOT:-"$DREAMZERO_ROOT/data/droid_lerobot"} +# If env set the old relative default, resolve to repo root (e.g. Beaker image env) +if [ "$DROID_DATA_ROOT" = "./data/droid_lerobot" ]; then + DROID_DATA_ROOT="$DREAMZERO_ROOT/data/droid_lerobot" +fi +OUTPUT_DIR=${OUTPUT_DIR:-"$DREAMZERO_ROOT/checkpoints/dreamzero_droid_wan22_lora"} + +# Wan2.2-TI2V-5B checkpoint (contains: diffusion weights, T5, VAE) +WAN22_CKPT_DIR=${WAN22_CKPT_DIR:-"$DREAMZERO_ROOT/checkpoints/Wan2.2-TI2V-5B"} + +# Image encoder: Wan2.2-TI2V-5B does NOT include CLIP - use Wan2.1's or standalone +IMAGE_ENCODER_DIR=${IMAGE_ENCODER_DIR:-"$DREAMZERO_ROOT/checkpoints/Wan2.1-I2V-14B-480P"} + +TOKENIZER_DIR=${TOKENIZER_DIR:-"$DREAMZERO_ROOT/checkpoints/umt5-xxl"} +# ============================================= + +# ============ AUTO-DOWNLOAD WEIGHTS ============ +if [ ! -d "$WAN22_CKPT_DIR" ] || [ -z "$(ls -A "$WAN22_CKPT_DIR" 2>/dev/null)" ]; then + echo "Wan2.2-TI2V-5B not found at $WAN22_CKPT_DIR. Downloading from HuggingFace..." + huggingface-cli download Wan-AI/Wan2.2-TI2V-5B --local-dir "$WAN22_CKPT_DIR" +fi + +if [ ! -d "$TOKENIZER_DIR" ] || [ -z "$(ls -A "$TOKENIZER_DIR" 2>/dev/null)" ]; then + echo "umt5-xxl tokenizer not found at $TOKENIZER_DIR. Downloading from HuggingFace..." + huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR" +fi + +# Image encoder: download Wan2.1 if not present (only need CLIP from it) +if [ ! -f "$IMAGE_ENCODER_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" ]; then + echo "Image encoder not found. Downloading Wan2.1-I2V-14B-480P (for CLIP only)..." + huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir "$IMAGE_ENCODER_DIR" +fi +# ================================================ + +# Validate dataset exists +if [ ! -d "$DROID_DATA_ROOT" ]; then + echo "ERROR: DROID dataset not found at $DROID_DATA_ROOT" + echo "Download with: huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir $DROID_DATA_ROOT" + exit 1 +fi + +# Use image Python 3.11 when available (Dockerfile installs dreamzero with python3.11 -m pip). +# Absolute path so worker processes open the correct file even if their cwd differs. +EXPERIMENT_PY="$DREAMZERO_ROOT/groot/vla/experiment/experiment.py" +if [ ! -f "$EXPERIMENT_PY" ]; then + echo "ERROR: Not found: $EXPERIMENT_PY" + exit 1 +fi +PYTHON_311="/usr/bin/python3.11" +if [ -x "$PYTHON_311" ]; then + if [ -n "${FIX_NUMPY_IN_SCRIPT:-}" ]; then + "$PYTHON_311" -m pip install "numpy==1.26.4" --force-reinstall -q 2>/dev/null || true + fi + RUN_CMD=( "$PYTHON_311" -m torch.distributed.run --nproc_per_node "$NUM_GPUS" --standalone "$EXPERIMENT_PY" ) + echo "Using image Python 3.11: $PYTHON_311" +else + RUN_CMD=( python3 -m torch.distributed.run --nproc_per_node "$NUM_GPUS" --standalone "$EXPERIMENT_PY" ) + echo "Using: $(command -v python3)" +fi +cd "$DREAMZERO_ROOT" + +"${RUN_CMD[@]}" \ + report_to=wandb \ + data=dreamzero/droid_relative_wan22 \ + wandb_project=dreamzero \ + train_architecture=lora \ + num_frames=33 \ + action_horizon=24 \ + num_views=3 \ + model=dreamzero/vla \ + model/dreamzero/action_head=wan_flow_matching_action_tf_wan22 \ + model/dreamzero/transform=dreamzero_cotrain \ + num_frame_per_block=2 \ + num_action_per_block=24 \ + num_state_per_block=1 \ + seed=42 \ + training_args.learning_rate=1e-5 \ + training_args.deepspeed="groot/vla/configs/deepspeed/zero2.json" \ + save_steps=1000 \ + training_args.warmup_ratio=0.05 \ + output_dir=$OUTPUT_DIR \ + per_device_train_batch_size=1 \ + max_steps=100 \ + weight_decay=1e-5 \ + save_total_limit=10 \ + upload_checkpoints=false \ + bf16=true \ + tf32=true \ + eval_bf16=true \ + dataloader_pin_memory=false \ + dataloader_num_workers=1 \ + save_lora_only=true \ + max_chunk_size=4 \ + save_strategy=no \ + droid_data_root=$DROID_DATA_ROOT \ + dit_version=$WAN22_CKPT_DIR \ + text_encoder_pretrained_path=$WAN22_CKPT_DIR/models_t5_umt5-xxl-enc-bf16.pth \ + image_encoder_pretrained_path=$IMAGE_ENCODER_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth \ + vae_pretrained_path=$WAN22_CKPT_DIR/Wan2.2_VAE.pth \ + tokenizer_path=$TOKENIZER_DIR diff --git a/scripts/train/trex_mini_force_training_wan22.sh b/scripts/train/trex_mini_force_training_wan22.sh new file mode 100644 index 0000000000000000000000000000000000000000..5ba22e9ad0c673aea0c77774920ff9a552335d52 --- /dev/null +++ b/scripts/train/trex_mini_force_training_wan22.sh @@ -0,0 +1,62 @@ +#!/bin/bash +set -euo pipefail + +# One-demo T-Rex Track-Force overfit over the FULL mini episode. +# Uses every valid causal anchor (~537 on trex_mini_force), not a pinned clip. +# Starts from Wan2.2 foundation weights (no prior T-Rex LoRA checkpoint). + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +DREAMZERO_ROOT="${DREAMZERO_ROOT:-$REPO_ROOT}" +TREX_DATA_ROOT="${TREX_DATA_ROOT:-$DREAMZERO_ROOT/data/trex_mini_force}" +OUTPUT_DIR="${OUTPUT_DIR:-$DREAMZERO_ROOT/checkpoints/trex_mini_force_wan22_lora_track_overfit_full}" +MAX_STEPS="${MAX_STEPS:-20000}" +SAVE_STEPS="${SAVE_STEPS:-1000}" +# Viz preview length: each chunk is 0.8s of 20 Hz track. Training still uses +# all anchors; this only controls how much of the demo W&B renders. +VIZ_NUM_CHUNKS="${VIZ_NUM_CHUNKS:-16}" +VIZ_START_CHUNK="${VIZ_START_CHUNK:-0}" + +if [[ ! -f "$TREX_DATA_ROOT/meta/info.json" ]]; then + echo "ERROR: mini dataset not found: $TREX_DATA_ROOT" >&2 + echo "Build it with: python scripts/data/build_trex_mini_force.py" >&2 + exit 1 +fi +if ! [[ "$MAX_STEPS" =~ ^[1-9][0-9]*$ ]]; then + echo "ERROR: MAX_STEPS must be a positive integer: $MAX_STEPS" >&2 + exit 1 +fi +if ! [[ "$SAVE_STEPS" =~ ^[1-9][0-9]*$ ]]; then + echo "ERROR: SAVE_STEPS must be a positive integer: $SAVE_STEPS" >&2 + exit 1 +fi +if [[ -f "$OUTPUT_DIR/config.json" ]] \ + || compgen -G "$OUTPUT_DIR/checkpoint-*" >/dev/null; then + echo "ERROR: scratch output already contains trained state: $OUTPUT_DIR" >&2 + echo "Choose a new OUTPUT_DIR to start again from step 0." >&2 + exit 1 +fi + +exec env \ + DREAMZERO_ROOT="$DREAMZERO_ROOT" \ + TREX_DATA_ROOT="$TREX_DATA_ROOT" \ + OUTPUT_DIR="$OUTPUT_DIR" \ + bash "$SCRIPT_DIR/trex_track_force_training_wan22.sh" \ + wandb_project=trex_mini_force_track_overfit \ + max_steps="$MAX_STEPS" \ + save_steps="$SAVE_STEPS" \ + save_total_limit=5 \ + dataset_shard_sampling_rate=1.0 \ + max_training_anchors=null \ + pin_anchor_rank=null \ + action_head_cfg.config.track_loss_weight=20.0 \ + action_head_cfg.config.track_position_loss_weight=5.0 \ + action_head_cfg.config.track_smoothness_loss_weight=2.0 \ + action_head_cfg.config.track_motion_loss_boost=4.0 \ + action_head_cfg.config.zero_invisible_track_xy=true \ + wandb_video_reconstruction_num_chunks="$VIZ_NUM_CHUNKS" \ + wandb_video_start_chunk_index="$VIZ_START_CHUNK" \ + wandb_video_track_trail_steps=2 \ + wandb_video_reconstruction_inference_steps=6 \ + wandb_video_compare_tracks_on_gt=true \ + "$@" diff --git a/scripts/train/trex_track_force_training_wan22.sh b/scripts/train/trex_track_force_training_wan22.sh new file mode 100644 index 0000000000000000000000000000000000000000..a516b5f783faecc11d5e806f25f840e37b70c58d --- /dev/null +++ b/scripts/train/trex_track_force_training_wan22.sh @@ -0,0 +1,111 @@ +#!/bin/bash +set -euo pipefail + +# Independent T-Rex Track-Force training entry point. +# Contract: 16 actions at 20 Hz, 5 Hz force refresh (stride 4), three RGB +# views, 250 tracks, 62 physical EEF/hand DoF, and unchanged WAN AR blocks. + +export HYDRA_FULL_ERROR=1 + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +DREAMZERO_ROOT="${DREAMZERO_ROOT:-$REPO_ROOT}" +TREX_DATA_ROOT="${TREX_DATA_ROOT:-$DREAMZERO_ROOT/data/trex_full_force}" +OUTPUT_DIR="${OUTPUT_DIR:-$DREAMZERO_ROOT/checkpoints/trex_track_force_wan22_lora}" +WAN22_CKPT_DIR="${WAN22_CKPT_DIR:-$DREAMZERO_ROOT/checkpoints/Wan2.2-TI2V-5B}" +IMAGE_ENCODER_DIR="${IMAGE_ENCODER_DIR:-$DREAMZERO_ROOT/checkpoints/Wan2.1-I2V-14B-480P}" +TOKENIZER_DIR="${TOKENIZER_DIR:-$DREAMZERO_ROOT/checkpoints/umt5-xxl}" +ALLOW_DOWNLOAD="${ALLOW_DOWNLOAD:-1}" + +if [ ! -d "$DREAMZERO_ROOT/groot" ]; then + echo "ERROR: DREAMZERO_ROOT must contain groot/: $DREAMZERO_ROOT" >&2 + exit 1 +fi +if [ ! -f "$TREX_DATA_ROOT/meta/info.json" ]; then + echo "ERROR: LeRobot-v2 T-Rex dataset not found: $TREX_DATA_ROOT" >&2 + exit 1 +fi + +python "$DREAMZERO_ROOT/scripts/data/check_trex_dataset_ready.py" \ + --dataset-root "$TREX_DATA_ROOT" \ + --require-force + +if [ "$ALLOW_DOWNLOAD" = "1" ]; then + if [ ! -f "$WAN22_CKPT_DIR/Wan2.2_VAE.pth" ]; then + hf download Wan-AI/Wan2.2-TI2V-5B --local-dir "$WAN22_CKPT_DIR" + fi + if [ ! -f "$IMAGE_ENCODER_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" ]; then + hf download Wan-AI/Wan2.1-I2V-14B-480P --include \ + models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth \ + --local-dir "$IMAGE_ENCODER_DIR" + fi + if [ ! -f "$TOKENIZER_DIR/tokenizer_config.json" ]; then + hf download google/umt5-xxl --local-dir "$TOKENIZER_DIR" + fi +fi + +for required in \ + "$WAN22_CKPT_DIR/Wan2.2_VAE.pth" \ + "$WAN22_CKPT_DIR/models_t5_umt5-xxl-enc-bf16.pth" \ + "$IMAGE_ENCODER_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" \ + "$TOKENIZER_DIR/tokenizer_config.json"; do + if [ ! -f "$required" ]; then + echo "ERROR: missing checkpoint file: $required" >&2 + exit 1 + fi +done + +cd "$DREAMZERO_ROOT" +torchrun --nproc_per_node=gpu --standalone \ + "$DREAMZERO_ROOT/groot/vla/experiment/experiment.py" \ + report_to=wandb \ + wandb_project=trex_track_force \ + model=trex_track_force/vla \ + data=dreamzero/trex_track_force_wan22 \ + train_architecture=lora \ + load_pretrained_components=true \ + allow_checkpoint_download=false \ + num_frames=33 \ + action_horizon=16 \ + num_action_per_block=16 \ + num_frame_per_block=2 \ + num_state_per_block=1 \ + max_action_dim=64 \ + max_state_dim=64 \ + num_views=3 \ + seed=42 \ + training_args.learning_rate=1e-5 \ + training_args.deepspeed="groot/vla/configs/deepspeed/zero2.json" \ + training_args.warmup_ratio=0.05 \ + per_device_train_batch_size=1 \ + max_steps=100000 \ + weight_decay=1e-5 \ + save_strategy=steps \ + save_steps=8000 \ + save_total_limit=5 \ + output_dir="$OUTPUT_DIR" \ + upload_checkpoints=false \ + save_lora_only=true \ + bf16=true \ + tf32=true \ + eval_bf16=true \ + dataloader_pin_memory=false \ + dataloader_num_workers=1 \ + max_chunk_size=4 \ + enable_wandb_video_reconstruction=true \ + wandb_video_reconstruction_steps=50 \ + wandb_video_reconstruction_episode=0 \ + wandb_video_reconstruction_num_chunks=1 \ + wandb_video_reconstruction_fps=10 \ + wandb_video_overlay_tracks=true \ + wandb_video_save_local=false \ + wandb_video_save_tracks=false \ + wandb_video_track_trail_steps=8 \ + wandb_video_reconstruction_inference_steps=1 \ + trex_data_root="$TREX_DATA_ROOT" \ + dit_version="$WAN22_CKPT_DIR" \ + text_encoder_pretrained_path="$WAN22_CKPT_DIR/models_t5_umt5-xxl-enc-bf16.pth" \ + image_encoder_pretrained_path="$IMAGE_ENCODER_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" \ + vae_pretrained_path="$WAN22_CKPT_DIR/Wan2.2_VAE.pth" \ + tokenizer_path="$TOKENIZER_DIR" \ + "$@" diff --git a/scripts/train/trex_training_wan22.sh b/scripts/train/trex_training_wan22.sh new file mode 100644 index 0000000000000000000000000000000000000000..9770dc1d13c2d70050c38051d85d78113a074c58 --- /dev/null +++ b/scripts/train/trex_training_wan22.sh @@ -0,0 +1,125 @@ +#!/bin/bash +# DreamZero T-Rex Training Script with Wan2.2-TI2V-5B backbone +# +# T-Rex: Dexmate Vega-1 dual-arm + 2x Sharpa Wave dexterous hands. +# State/action are 58-dim (7 arm + 22 hand per side) -> max_action_dim=64. +# +# Usage: +# bash scripts/train/trex_training_wan22.sh +# +# Prerequisites: +# - T-Rex dataset converted to LeRobot v2 + GEAR format at TREX_DATA_ROOT +# (scripts/data/convert_trex_v3_to_v2.py then scripts/data/convert_lerobot_to_gear.py) +# - Wan2.2-TI2V-5B weights (download from HuggingFace) +# hf download Wan-AI/Wan2.2-TI2V-5B --local-dir ./checkpoints/Wan2.2-TI2V-5B +# - Image encoder (CLIP) from Wan2.1 - Wan2.2-TI2V-5B does not include it +# hf download Wan-AI/Wan2.1-I2V-14B-480P --local-dir ./checkpoints/Wan2.1-I2V-14B-480P +# - umt5-xxl tokenizer +# hf download google/umt5-xxl --local-dir ./checkpoints/umt5-xxl + +export HYDRA_FULL_ERROR=1 + +# Repo root: must be a directory that contains groot/ (so experiment.py can be found). +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SCRIPT_REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +if [ -n "$DREAMZERO_ROOT" ] && [ -d "$DREAMZERO_ROOT/groot" ]; then + : # keep existing and valid +elif [ -d "$SCRIPT_REPO_ROOT/groot" ]; then + DREAMZERO_ROOT="$SCRIPT_REPO_ROOT" +else + echo "ERROR: Set DREAMZERO_ROOT to the dreamzero repo root that contains groot/." + exit 1 +fi + +# ============ USER CONFIGURATION ============ +TREX_DATA_ROOT=${TREX_DATA_ROOT:-"$DREAMZERO_ROOT/data/trex_small"} +OUTPUT_DIR=${OUTPUT_DIR:-"$DREAMZERO_ROOT/checkpoints/dreamzero_trex_wan22_lora"} + +# Wan2.2-TI2V-5B checkpoint (contains: diffusion weights, T5, VAE) +WAN22_CKPT_DIR=${WAN22_CKPT_DIR:-"$DREAMZERO_ROOT/checkpoints/Wan2.2-TI2V-5B"} + +# Image encoder: Wan2.2-TI2V-5B does NOT include CLIP - use Wan2.1's or standalone +IMAGE_ENCODER_DIR=${IMAGE_ENCODER_DIR:-"$DREAMZERO_ROOT/checkpoints/Wan2.1-I2V-14B-480P"} + +TOKENIZER_DIR=${TOKENIZER_DIR:-"$DREAMZERO_ROOT/checkpoints/umt5-xxl"} +# ============================================= + +# ============ AUTO-DOWNLOAD WEIGHTS ============ +if [ ! -d "$WAN22_CKPT_DIR" ] || [ -z "$(ls -A "$WAN22_CKPT_DIR" 2>/dev/null)" ]; then + echo "Wan2.2-TI2V-5B not found at $WAN22_CKPT_DIR. Downloading from HuggingFace..." + hf download Wan-AI/Wan2.2-TI2V-5B --local-dir "$WAN22_CKPT_DIR" +fi + +if [ ! -d "$TOKENIZER_DIR" ] || [ -z "$(ls -A "$TOKENIZER_DIR" 2>/dev/null)" ]; then + echo "umt5-xxl tokenizer not found at $TOKENIZER_DIR. Downloading from HuggingFace..." + hf download google/umt5-xxl --local-dir "$TOKENIZER_DIR" +fi + +# Image encoder: download Wan2.1 if not present (only need CLIP from it) +if [ ! -f "$IMAGE_ENCODER_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" ]; then + echo "Image encoder not found. Downloading Wan2.1-I2V-14B-480P (for CLIP only)..." + hf download Wan-AI/Wan2.1-I2V-14B-480P --local-dir "$IMAGE_ENCODER_DIR" +fi +# ================================================ + +# Validate dataset exists and is GEAR-converted +if [ ! -d "$TREX_DATA_ROOT" ]; then + echo "ERROR: T-Rex dataset not found at $TREX_DATA_ROOT" + exit 1 +fi +if [ ! -f "$TREX_DATA_ROOT/meta/embodiment.json" ]; then + echo "ERROR: $TREX_DATA_ROOT/meta/embodiment.json missing - run convert_lerobot_to_gear.py first" + exit 1 +fi +python "$DREAMZERO_ROOT/scripts/data/check_trex_dataset_ready.py" \ + --dataset-root "$TREX_DATA_ROOT" + +EXPERIMENT_PY="$DREAMZERO_ROOT/groot/vla/experiment/experiment.py" +cd "$DREAMZERO_ROOT" + +torchrun --nproc_per_node=gpu --standalone "$EXPERIMENT_PY" \ + report_to=wandb \ + data=dreamzero/trex_relative_wan22 \ + wandb_project=dreamzero \ + train_architecture=lora \ + num_frames=33 \ + action_horizon=24 \ + max_action_dim=64 \ + ++action_head_cfg.config.diffusion_model_cfg.action_dim=64 \ + num_views=3 \ + model=dreamzero/vla \ + model/dreamzero/action_head=wan_flow_matching_action_tf_wan22 \ + model/dreamzero/transform=dreamzero_cotrain \ + num_frame_per_block=2 \ + num_action_per_block=24 \ + num_state_per_block=1 \ + seed=42 \ + training_args.learning_rate=1e-5 \ + training_args.deepspeed="groot/vla/configs/deepspeed/zero2.json" \ + save_steps=8000 \ + training_args.warmup_ratio=0.05 \ + output_dir="$OUTPUT_DIR" \ + per_device_train_batch_size=1 \ + max_steps=100000 \ + weight_decay=1e-5 \ + save_total_limit=5 \ + upload_checkpoints=false \ + bf16=true \ + tf32=true \ + eval_bf16=true \ + dataloader_pin_memory=false \ + dataloader_num_workers=1 \ + save_lora_only=true \ + max_chunk_size=4 \ + save_strategy=steps \ + enable_wandb_video_reconstruction=true \ + wandb_video_reconstruction_steps=1000 \ + wandb_video_reconstruction_episode=0 \ + wandb_video_reconstruction_num_chunks=4 \ + trex_data_root="$TREX_DATA_ROOT" \ + dit_version="$WAN22_CKPT_DIR" \ + text_encoder_pretrained_path="$WAN22_CKPT_DIR/models_t5_umt5-xxl-enc-bf16.pth" \ + image_encoder_pretrained_path="$IMAGE_ENCODER_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" \ + vae_pretrained_path="$WAN22_CKPT_DIR/Wan2.2_VAE.pth" \ + tokenizer_path="$TOKENIZER_DIR" \ + "$@" diff --git a/scripts/train/yam_training.sh b/scripts/train/yam_training.sh new file mode 100644 index 0000000000000000000000000000000000000000..49f83079a585b4b1bf71ac8dea3660a9840bae85 --- /dev/null +++ b/scripts/train/yam_training.sh @@ -0,0 +1,103 @@ +#!/bin/bash +# DreamZero YAM Training Script +# +# Usage: +# # Set your dataset path and output directory, then run: +# bash scripts/train/yam_training.sh +# +# Prerequisites: +# - YAM dataset in LeRobot format at YAM_DATA_ROOT (state 14, action 14, 3 views: top, left, right) +# meta/embodiment.json must have "embodiment_tag": "yam" +# modality: state (left_joint_pos, left_gripper_pos, right_joint_pos, right_gripper_pos), +# action (same keys), video (top, left, right), annotation.human.task_description +# - Wan2.1-I2V-14B-480P weights (auto-downloaded or pre-downloaded from HuggingFace) +# Download: huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir ./checkpoints/Wan2.1-I2V-14B-480P +# - umt5-xxl tokenizer (auto-downloaded or pre-downloaded from HuggingFace) +# Download: huggingface-cli download google/umt5-xxl --local-dir ./checkpoints/umt5-xxl +# - DreamZero-AgiBot pretrained checkpoint (for loading LoRA weights before fine-tuning) +# git clone https://huggingface.co/GEAR-Dreams/DreamZero-AgiBot ./checkpoints/DreamZero-AgiBot + +export HYDRA_FULL_ERROR=1 + +# ============ CHANGE THESE VARIABLES ============ +# Dataset path (YAM in LeRobot format: state 14, action 14, videos top, left, right) +YAM_DATA_ROOT=${YAM_DATA_ROOT:-"./data/yam_lerobot"} + +# Output directory for training checkpoints +OUTPUT_DIR=${OUTPUT_DIR:-"./checkpoints/dreamzero_yam_lora_dz_pretrained_100k_folding"} + +# Number of GPUs to use (default: all visible GPUs, so 4-GPU machines use 4 without setting NUM_GPUS) +if [ -z "${NUM_GPUS}" ]; then + NUM_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) +fi +NUM_GPUS=${NUM_GPUS:-8} + +# Model weight paths (download from HuggingFace if not already present) +WAN_CKPT_DIR=${WAN_CKPT_DIR:-"./checkpoints/Wan2.1-I2V-14B-480P"} +TOKENIZER_DIR=${TOKENIZER_DIR:-"./checkpoints/umt5-xxl"} +# ============================================= + +# ============ AUTO-DOWNLOAD WEIGHTS ============ +if [ ! -d "$WAN_CKPT_DIR" ] || [ -z "$(ls -A "$WAN_CKPT_DIR" 2>/dev/null)" ]; then + echo "Wan2.1-I2V-14B-480P not found at $WAN_CKPT_DIR. Downloading from HuggingFace..." + huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir "$WAN_CKPT_DIR" +fi + +if [ ! -d "$TOKENIZER_DIR" ] || [ -z "$(ls -A "$TOKENIZER_DIR" 2>/dev/null)" ]; then + echo "umt5-xxl tokenizer not found at $TOKENIZER_DIR. Downloading from HuggingFace..." + huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR" +fi +# ================================================ + +# Validate dataset exists +if [ ! -d "$YAM_DATA_ROOT" ]; then + echo "ERROR: YAM dataset not found at $YAM_DATA_ROOT" + echo "Set YAM_DATA_ROOT to your LeRobot-format YAM dataset (meta/embodiment.json with embodiment_tag: yam)" + exit 1 +fi + +torchrun --nproc_per_node $NUM_GPUS --standalone groot/vla/experiment/experiment.py \ + report_to=wandb \ + data=dreamzero/yam_relative \ + wandb_project=dreamzero \ + train_architecture=lora \ + num_frames=33 \ + action_horizon=24 \ + num_views=3 \ + model=dreamzero/vla \ + model/dreamzero/action_head=wan_flow_matching_action_tf \ + model/dreamzero/transform=dreamzero_cotrain \ + num_frame_per_block=2 \ + num_action_per_block=24 \ + num_state_per_block=1 \ + seed=42 \ + training_args.learning_rate=1e-5 \ + training_args.deepspeed="groot/vla/configs/deepspeed/zero2.json" \ + save_steps=10000 \ + training_args.warmup_ratio=0.05 \ + output_dir=$OUTPUT_DIR \ + per_device_train_batch_size=4 \ + max_steps=100000 \ + weight_decay=1e-5 \ + save_total_limit=10 \ + upload_checkpoints=false \ + bf16=true \ + tf32=true \ + eval_bf16=true \ + dataloader_pin_memory=false \ + dataloader_num_workers=1 \ + image_resolution_width=320 \ + image_resolution_height=176 \ + save_lora_only=true \ + max_chunk_size=4 \ + frame_seqlen=880 \ + save_strategy=steps \ + yam_data_root=$YAM_DATA_ROOT \ + dit_version=$WAN_CKPT_DIR \ + text_encoder_pretrained_path=$WAN_CKPT_DIR/models_t5_umt5-xxl-enc-bf16.pth \ + image_encoder_pretrained_path=$WAN_CKPT_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth \ + vae_pretrained_path=$WAN_CKPT_DIR/Wan2.1_VAE.pth \ + tokenizer_path=$TOKENIZER_DIR \ + pretrained_model_path=./checkpoints/DreamZero-AgiBot \ + ++action_head_cfg.config.skip_component_loading=true \ + ++action_head_cfg.config.defer_lora_injection=true diff --git a/scripts/trex_track/__init__.py b/scripts/trex_track/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/scripts/trex_track/layout.py b/scripts/trex_track/layout.py new file mode 100644 index 0000000000000000000000000000000000000000..dc244a17ee6adc505c7540691754d6b3d1eb80cd --- /dev/null +++ b/scripts/trex_track/layout.py @@ -0,0 +1,138 @@ +"""Canonical 250-point layout for T-Rex three-view tracks. + +The integer identities in this module are part of the on-disk dataset schema. +Do not reorder groups without also changing :data:`TRACK_LAYOUT_VERSION`. +""" + +from __future__ import annotations + +from typing import Final + +TRACK_LAYOUT_VERSION: Final = "trex_track_250_v1" + +VIEW_ORDER: Final = ("head_left", "left_wrist", "right_wrist") +VIEW_IDS: Final = {name: index for index, name in enumerate(VIEW_ORDER)} + +NUM_HEAD_PER_HAND: Final = 50 +NUM_HEAD_LEFT: Final = NUM_HEAD_PER_HAND +NUM_HEAD_RIGHT: Final = NUM_HEAD_PER_HAND +NUM_HEAD_POINTS: Final = NUM_HEAD_LEFT + NUM_HEAD_RIGHT + +NUM_WRIST_BACKGROUND: Final = 25 +NUM_WRIST_HAND: Final = 50 +NUM_WRIST_POINTS: Final = NUM_WRIST_BACKGROUND + NUM_WRIST_HAND + +NUM_COMBINED_POINTS: Final = NUM_HEAD_POINTS + 2 * NUM_WRIST_POINTS +POINT_SLICES: Final = ( + 0, + NUM_HEAD_POINTS, + NUM_HEAD_POINTS + NUM_WRIST_POINTS, + NUM_COMBINED_POINTS, +) + +HAND_NONE: Final = 0 +HAND_LEFT: Final = 1 +HAND_RIGHT: Final = 2 +HAND_NAMES: Final = ("none", "left", "right") + +ROLE_HEAD_HAND: Final = 0 +ROLE_WRIST_BACKGROUND: Final = 1 +ROLE_WRIST_HAND: Final = 2 +ROLE_NAMES: Final = ("head_hand", "wrist_background", "wrist_hand") + +# Half-open slices in the canonical concatenation order. +COMPONENT_SLICES: Final = { + "head_left_hand": (0, 50), + "head_right_hand": (50, 100), + "left_wrist_background": (100, 125), + "left_wrist_hand": (125, 175), + "right_wrist_background": (175, 200), + "right_wrist_hand": (200, 250), +} +VIEW_SLICES: Final = { + "head_left": (0, 100), + "left_wrist": (100, 175), + "right_wrist": (175, 250), +} +VIEW_POINT_COUNTS: Final = { + view: end - start for view, (start, end) in VIEW_SLICES.items() +} + + +def identity_metadata() -> dict[str, list[int] | list[str]]: + """Return stable per-point IDs in canonical concatenation order.""" + + view_ids: list[int] = [] + hand_ids: list[int] = [] + role_ids: list[int] = [] + local_ids: list[int] = [] + point_names: list[str] = [] + + groups = ( + ("head_left", "left_hand", NUM_HEAD_LEFT, HAND_LEFT, ROLE_HEAD_HAND), + ("head_left", "right_hand", NUM_HEAD_RIGHT, HAND_RIGHT, ROLE_HEAD_HAND), + ( + "left_wrist", + "background", + NUM_WRIST_BACKGROUND, + HAND_NONE, + ROLE_WRIST_BACKGROUND, + ), + ("left_wrist", "left_hand", NUM_WRIST_HAND, HAND_LEFT, ROLE_WRIST_HAND), + ( + "right_wrist", + "background", + NUM_WRIST_BACKGROUND, + HAND_NONE, + ROLE_WRIST_BACKGROUND, + ), + ("right_wrist", "right_hand", NUM_WRIST_HAND, HAND_RIGHT, ROLE_WRIST_HAND), + ) + for view, component, count, hand_id, role_id in groups: + view_ids.extend([VIEW_IDS[view]] * count) + hand_ids.extend([hand_id] * count) + role_ids.extend([role_id] * count) + local_ids.extend(range(count)) + point_names.extend(f"{view}.{component}.{index:03d}" for index in range(count)) + + lengths = { + len(view_ids), + len(hand_ids), + len(role_ids), + len(local_ids), + len(point_names), + } + if lengths != {NUM_COMBINED_POINTS}: + raise AssertionError(f"invalid identity lengths: {sorted(lengths)}") + return { + "view_ids": view_ids, + "hand_ids": hand_ids, + "role_ids": role_ids, + "local_ids": local_ids, + "global_ids": list(range(NUM_COMBINED_POINTS)), + "point_names": point_names, + } + + +def layout_metadata() -> dict[str, object]: + """Return the JSON-serializable schema metadata stored with the dataset.""" + + return { + "version": TRACK_LAYOUT_VERSION, + "coordinate_space": "normalized_xy_div_wh", + "point_value_order": ["x", "y", "visibility"], + "view_order": list(VIEW_ORDER), + "point_slices": list(POINT_SLICES), + "view_slices": {key: list(value) for key, value in VIEW_SLICES.items()}, + "component_slices": { + key: list(value) for key, value in COMPONENT_SLICES.items() + }, + "view_point_counts": dict(VIEW_POINT_COUNTS), + "hand_id_names": list(HAND_NAMES), + "role_id_names": list(ROLE_NAMES), + **identity_metadata(), + } + + +if NUM_COMBINED_POINTS != 250 or POINT_SLICES != (0, 100, 175, 250): + raise AssertionError("the canonical T-Rex track layout must contain 250 points") diff --git a/scripts/trex_track/sam2_cotracker_hands.py b/scripts/trex_track/sam2_cotracker_hands.py new file mode 100644 index 0000000000000000000000000000000000000000..72a54b600d28653b0c8655a3acd80a0a6b396f2a --- /dev/null +++ b/scripts/trex_track/sam2_cotracker_hands.py @@ -0,0 +1,84 @@ +"""Hand seeds for the canonical T-Rex 250-point CoTracker layout. + +No appearance auto-repair. Tune prompts in ``sam2_prompt_hands.py``. +""" + +from __future__ import annotations + +import numpy as np + +from trex_track.layout import NUM_HEAD_PER_HAND + +NUM_HAND_POINTS = NUM_HEAD_PER_HAND + +__all__ = [ + "NUM_HAND_POINTS", + "wrist_hand_50", + "head_hands_50", + # Compatibility aliases for callers of the old 25-point implementation. + "wrist_hand_25", + "head_hands_25", +] + + +def _stable_spatial_order(points_xy: np.ndarray) -> np.ndarray: + """Give sampled slots deterministic top-to-bottom, then left-to-right IDs.""" + + points = np.asarray(points_xy, dtype=np.float32).reshape(-1, 2) + order = np.lexsort((points[:, 0], points[:, 1])) + return points[order] + + +def wrist_hand_50( + predictor, + rgb_hwc: np.ndarray, + view: str, + *, + n_points: int = NUM_HAND_POINTS, + seed: int | None = 0, +) -> tuple[np.ndarray, np.ndarray, str]: + """One SAM2 pass → 50 stably ordered queries on a wrist hand mask. + + Background 5×5 grid is NOT produced here — callers prepend ``make_image_grid``. + """ + from trex_track.sam2_prompt_hands import sample_points_on_mask, sam2_wrist_mask + + mask, tag, _meta = sam2_wrist_mask(predictor, rgb_hwc, view) + rng = np.random.default_rng(seed) + pts = sample_points_on_mask(mask, n_points, rng=rng, distal_bias=True) + h, w = rgb_hwc.shape[:2] + pts[:, 0] = np.clip(pts[:, 0], 0, w - 1) + pts[:, 1] = np.clip(pts[:, 1], 0, h - 1) + return _stable_spatial_order(pts), mask.astype(bool), tag + + +def head_hands_50( + predictor, + rgb_hwc: np.ndarray, + *, + n_points: int = NUM_HAND_POINTS, + seed: int | None = 0, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, str]: + """One SAM2 pass per hand → 100 queries (left 50, then right 50).""" + from trex_track.sam2_prompt_hands import sample_points_on_mask, sam2_head_hand_mask + + rng = np.random.default_rng(seed) + left, tag_l, _ = sam2_head_hand_mask(predictor, rgb_hwc, "left") + right, tag_r, _ = sam2_head_hand_mask(predictor, rgb_hwc, "right") + pts_l = _stable_spatial_order( + sample_points_on_mask(left, n_points, rng=rng, distal_bias=True) + ) + pts_r = _stable_spatial_order( + sample_points_on_mask(right, n_points, rng=rng, distal_bias=True) + ) + h, w = rgb_hwc.shape[:2] + queries = np.concatenate([pts_l, pts_r], axis=0).astype(np.float32) + queries[:, 0] = np.clip(queries[:, 0], 0, w - 1) + queries[:, 1] = np.clip(queries[:, 1], 0, h - 1) + tag = f"{tag_l}|{tag_r}" + return queries, left.astype(bool), right.astype(bool), tag + + +# Keep old imports working while making their new 50-point behavior explicit. +wrist_hand_25 = wrist_hand_50 +head_hands_25 = head_hands_50 diff --git a/scripts/trex_track/sam2_prompt_hands.py b/scripts/trex_track/sam2_prompt_hands.py new file mode 100644 index 0000000000000000000000000000000000000000..fc656f64e0677fc3889a86cdd1cf2a573e855ce3 --- /dev/null +++ b/scripts/trex_track/sam2_prompt_hands.py @@ -0,0 +1,1687 @@ +"""One-shot SAM2 hand masks via fixed prompts (no appearance auto-repair). + +Design: + - Prompts are view-specific geometry (box + pos/neg points), tunable in one place. + - Exactly one ``predictor.predict`` call per hand/view. + - Post-process: largest connected component only (no geodesic / wood heuristics). + - Empty mask → raise (no silent fallback). +""" + +from __future__ import annotations + +import cv2 +import numpy as np + +from trex_track.sam2_wrist_hand import largest_connected_component +# wrist_hand_prompt_box / template kept available for experiments; wrist prompts +# now use explicit bottom-band geometry (see build_wrist_prompts). + +# --------------------------------------------------------------------------- +# Tunable prompt specs (fractional xy in [0,1], box as xyxy frac) +# --------------------------------------------------------------------------- + +# Absolute negative prompts (upper background / walls / person). +WRIST_NEG_FRAC = [ + (0.08, 0.04), + (0.50, 0.04), + (0.92, 0.04), + (0.50, 0.12), + (0.18, 0.10), + (0.82, 0.10), + (0.08, 0.22), + (0.92, 0.22), + (0.50, 0.22), + (0.30, 0.18), + (0.70, 0.18), +] + +# Head prompts tuned on trex_small ep0 medial-axis of white arm CCs. +HEAD_PROMPTS = { + "left": { + # Arm: bottom-left base → tip ~ (0.39, 0.34) + "box": (0.05, 0.28, 0.48, 1.00), + "pos": [ + (0.21, 0.95), + (0.20, 0.87), + (0.25, 0.78), + (0.29, 0.69), + (0.33, 0.60), + (0.35, 0.52), + (0.40, 0.42), + (0.39, 0.34), + ], + "neg": [ + (0.50, 0.12), + (0.50, 0.38), + (0.55, 0.55), + (0.72, 0.70), + (0.88, 0.85), + (0.12, 0.16), + (0.42, 0.26), + (0.28, 0.24), + (0.45, 0.45), + ], + }, + "right": { + # Arm is on the RIGHT edge; tip ~ (0.67, 0.36) — do NOT drift left onto table. + "box": (0.55, 0.28, 0.99, 1.00), + "pos": [ + (0.88, 0.95), + (0.91, 0.88), + (0.88, 0.79), + (0.81, 0.70), + (0.76, 0.60), + (0.72, 0.52), + (0.69, 0.44), + (0.67, 0.36), + ], + "neg": [ + (0.50, 0.12), + (0.50, 0.38), + (0.45, 0.55), + (0.28, 0.70), + (0.14, 0.85), + (0.90, 0.16), + (0.58, 0.28), + (0.62, 0.42), + (0.55, 0.48), + ], + }, +} + + +def _snap_right_four_tips( + tips_xy: np.ndarray, + rgb_hwc: np.ndarray, + *, + expected_x_frac: tuple[float, ...] = (0.287, 0.334, 0.388, 0.438), + expected_y_frac: float = 0.575, +) -> np.ndarray: + """Distal finger-pad tips — outward from knuckle, on bright pad tissue. + + Band ~0.545–0.600·h (~98–108 @180): covers finger pads for SAM2 without + dropping into floor/gaps that break finger completion. + """ + rgb = np.asarray(rgb_hwc, dtype=np.uint8) + gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY) + h, w = gray.shape + out = np.asarray(tips_xy, dtype=np.float32).reshape(4, 2).copy() + y_pref = int(np.clip(round(expected_y_frac * h), 0, h - 1)) + y_hard_lo = int(round(0.545 * h)) + y_hard_hi = int(round(0.600 * h)) + y_mid = int(round(0.700 * h)) + max_walk = int(round(0.160 * h)) + half = max(2, int(round(0.010 * w))) + x_slack = max(3, int(round(0.012 * w))) + gap_tol = 3 + min_gap = max(12.0, 0.035 * w) + + def _is_pad(yy: int, xx: int) -> bool: + if not (0 <= yy < h and 0 <= xx < w): + return False + return 110 <= int(gray[yy, xx]) <= 170 + + def _hand(yy: int, xc: int) -> bool: + if not _is_pad(yy, xc): + return False + x0, x1 = max(0, xc - half), min(w, xc + half + 1) + strip = gray[yy, x0:x1] + return int(((strip >= 108) & (strip <= 175)).sum()) >= max(2, (x1 - x0) // 2) + + def _local_ridge(xc: int, lo_x: int, hi_x: int) -> int: + ya, yb = int(0.56 * h), int(0.72 * h) + best = None + for xx in range(max(0, lo_x), min(w, hi_x + 1)): + strip = gray[ya:yb, xx].astype(np.float32) + face = strip[(strip >= 110) & (strip <= 170)] + if face.size < max(2, (yb - ya) // 5): + continue + med = float(np.median(face)) + cand = (-med + 0.45 * abs(xx - xc), abs(xx - xc), xx) + if best is None or cand < best: + best = cand + return int(best[2]) if best is not None else xc + + prev_x = -1e9 + for i, xf in enumerate(expected_x_frac): + xc_lane = int(np.clip(round(xf * w), 0, w - 1)) + lo_x = int(max(0, max(xc_lane - x_slack, int(round(prev_x + min_gap))))) + hi_x = int(min(w - 1, xc_lane + x_slack)) + if lo_x > hi_x: + lo_x = hi_x = int(np.clip(round(prev_x + min_gap), 0, w - 1)) + xc = _local_ridge(xc_lane, lo_x, hi_x) + xc = int(np.clip(max(xc, int(round(prev_x + min_gap))), 0, w - 1)) + + seed = None + for y_try in list(range(y_mid, min(h - 1, y_mid + 12) + 1)) + list( + range(y_mid - 1, max(y_hard_lo, y_mid - 20) - 1, -1) + ): + if _hand(y_try, xc): + seed = y_try + break + if seed is None: + out[i, 0] = float(xc) + out[i, 1] = float(np.clip(y_pref, y_hard_lo, y_hard_hi)) + prev_x = float(out[i, 0]) + continue + + y_floor = max(y_hard_lo, seed - max_walk) + tip_y = seed + yy = seed + miss = 0 + while yy >= y_floor: + if _hand(yy, xc): + tip_y = yy + miss = 0 + yy -= 1 + continue + miss += 1 + if miss > gap_tol: + break + yy -= 1 + tip_y = int(np.clip(tip_y, y_hard_lo, y_hard_hi)) + + best = None + for yy in range(y_hard_lo, min(y_hard_hi, tip_y + 2) + 1): + for xx in range(max(0, xc - 2), min(w, xc + 3)): + if xx < prev_x + min_gap or not _is_pad(yy, xx): + continue + g = int(gray[yy, xx]) + nb = gray[max(0, yy - 1) : yy + 2, max(0, xx - 1) : xx + 2] + if g < int(nb.max()) - 5: + continue + # Mild distal preference; brightness keeps us on pad not gap. + cand = (0.35 * abs(yy - tip_y) + 0.25 * max(0, yy - tip_y) - 0.04 * g + 0.35 * abs(xx - xc), xx, yy) + if best is None or cand < best: + best = cand + if best is None: + out[i, 0] = float(xc) + out[i, 1] = float(tip_y) + else: + out[i, 0] = float(best[1]) + out[i, 1] = float(best[2]) + out[i, 0] = float(np.clip(out[i, 0], xc_lane - x_slack, xc_lane + x_slack)) + if out[i, 0] < prev_x + min_gap: + out[i, 0] = float(min(w - 1, prev_x + min_gap)) + out[i, 1] = float(np.clip(out[i, 1], y_hard_lo, y_hard_hi)) + xi, yi = int(round(out[i, 0])), int(round(out[i, 1])) + if not _is_pad(yi, xi): + rescue = None + for yy in range(y_hard_lo, y_hard_hi + 1): + for xx in range(max(int(prev_x + min_gap), xc_lane - x_slack), min(w, xc_lane + x_slack + 1)): + if not _is_pad(yy, xx): + continue + g = int(gray[yy, xx]) + cand = (abs(yy - tip_y) - 0.04 * g + 0.2 * abs(xx - xc_lane), xx, yy) + if rescue is None or cand < rescue: + rescue = cand + if rescue is not None: + out[i, 0] = float(rescue[1]) + out[i, 1] = float(rescue[2]) + prev_x = float(out[i, 0]) + return out + + +def _refine_right_tips_on_base( + tips_xy: np.ndarray, + base_mask: np.ndarray, + rgb_hwc: np.ndarray, + *, + expected_x_frac: tuple[float, ...] = (0.287, 0.334, 0.388, 0.438), +) -> np.ndarray: + """Keep distal snap tips; only nudge further distal if base extends there.""" + base = np.asarray(base_mask, dtype=bool) + rgb = np.asarray(rgb_hwc, dtype=np.uint8) + gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY) + h, w = gray.shape + snapped = _snap_right_four_tips(tips_xy, rgb, expected_x_frac=expected_x_frac) + out = snapped.copy() + half = max(2, int(round(0.010 * w))) + y_hard_lo = int(round(0.545 * h)) + y_hard_hi = int(round(0.600 * h)) + + for i in range(4): + xc = int(round(float(out[i, 0]))) + snap_y = int(round(float(out[i, 1]))) + x0, x1 = max(0, xc - half), min(w, xc + half + 1) + col = base[:, x0:x1] + ys = np.where(col.any(axis=1))[0] + tip_y = snap_y + if len(ys): + y_edge = int(np.percentile(ys, 8)) + if y_edge < snap_y - 1: + tip_y = int(np.clip(y_edge + 1, y_hard_lo, snap_y)) + tip_y = int(np.clip(min(tip_y, snap_y + 1), y_hard_lo, y_hard_hi)) + # Stay on pad. + if 0 <= tip_y < h and 0 <= xc < w and 110 <= int(gray[tip_y, xc]) <= 170: + out[i, 1] = float(tip_y) + else: + out[i, 1] = float(snap_y) + return out + + + +def _snap_pos_to_hand( + pos_xy: np.ndarray, + rgb_hwc: np.ndarray, + *, + gray_lo: int = 90, + gray_hi: int = 160, + radius: int = 14, +) -> np.ndarray: + """Pull positive prompts onto nearby hand tissue (avoid dark gaps / shelf bg). + + Distal finger points must land on *bright peaks* (white finger faces), + not mid-gray inter-finger gaps. Prefer downward snaps over climbing into shelf. + """ + rgb = np.asarray(rgb_hwc, dtype=np.uint8) + gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY) + h, w = gray.shape + out = np.asarray(pos_xy, dtype=np.float32).copy() + for i, (x, y) in enumerate(out): + xi, yi = int(round(x)), int(round(y)) + xi = int(np.clip(xi, 0, w - 1)) + yi = int(np.clip(yi, 0, h - 1)) + # Palm darker; finger band needs bright white faces (gaps are ~40-90). + distal = yi < int(round(0.72 * h)) + lo = 80 if not distal else max(gray_lo, 118) + hi = gray_hi + g0 = int(gray[yi, xi]) + # Accept only if already a local brightness peak on tissue. + if lo <= g0 <= hi: + nb = gray[max(0, yi - 2) : yi + 3, max(0, xi - 2) : xi + 3] + if g0 >= int(nb.max()) - 2 and float(np.mean((nb >= lo - 15) & (nb <= hi))) >= 0.4: + continue + best = None + for dy in range(-radius, radius + 1): + for dx in range(-radius, radius + 1): + xx, yy = xi + dx, yi + dy + if not (0 <= xx < w and 0 <= yy < h): + continue + g = int(gray[yy, xx]) + lo_y = 80 if yy >= int(round(0.72 * h)) else max(gray_lo, 118) + if not (lo_y <= g <= hi): + continue + nb = gray[max(0, yy - 2) : yy + 3, max(0, xx - 2) : xx + 3] + if g < int(nb.max()) - 3: + continue # not a local peak → skip gap walls + pen = abs(dx) + abs(dy) + (4 if dy < 0 else 0) + # Brightest peak first, then nearest. + cand = (-g, pen, xx, yy) + if best is None or cand < best: + best = cand + if best is not None: + out[i, 0], out[i, 1] = float(best[2]), float(best[3]) + return out + + +def _frac_to_px(frac_xy: np.ndarray | list, h: int, w: int) -> np.ndarray: + arr = np.asarray(frac_xy, dtype=np.float32).reshape(-1, 2) + out = arr.copy() + out[:, 0] = np.clip(out[:, 0] * float(w), 0.0, float(w - 1)) + out[:, 1] = np.clip(out[:, 1] * float(h), 0.0, float(h - 1)) + return out + + +def _box_frac_to_px(box_xyxy: tuple[float, float, float, float], h: int, w: int) -> np.ndarray: + x0, y0, x1, y1 = box_xyxy + return np.array( + [ + np.clip(x0 * w, 0, w - 1), + np.clip(y0 * h, 0, h - 1), + np.clip(x1 * w, 0, w - 1), + np.clip(y1 * h, 0, h - 1), + ], + dtype=np.float32, + ) + + +def _fill_enclosed_holes(mask: np.ndarray) -> np.ndarray: + """Fill holes that are fully enclosed by the mask (open finger gaps stay).""" + mask_u8 = np.asarray(mask, dtype=np.uint8) + padded = np.pad(mask_u8, 1) + flooded = padded.copy() + flood_mask = np.zeros( + (flooded.shape[0] + 2, flooded.shape[1] + 2), dtype=np.uint8 + ) + cv2.floodFill(flooded, flood_mask, (0, 0), 1) + holes = flooded[1:-1, 1:-1] == 0 + return mask_u8.astype(bool) | holes + + +def _clean_right_wrist_mask(mask: np.ndarray) -> np.ndarray: + """Remove palm pinholes while leaving the finger/background boundary raw.""" + filled = _fill_enclosed_holes(mask) + kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5)) + closed = cv2.morphologyEx(filled.astype(np.uint8), cv2.MORPH_CLOSE, kernel) + cleaned = _fill_enclosed_holes(closed) + + # Close only the lower palm/thumb band. Distal fingers stay raw so the + # dark shelf behind them is not trapped between bridged fingertips. + h = int(np.asarray(mask).shape[0]) + split_y = int(round(0.70 * h)) + out = filled.copy() + out[split_y:] = cleaned[split_y:] + return out + + +def _edge_boundary_median(mask: np.ndarray, rgb_hwc: np.ndarray) -> float: + """Median distance from mask boundary pixels to Canny edges (lower=better). + + Image-border pixels are ignored: wrist hands often fill to the frame edge, + and those artificial cuts dominate the median without reflecting silhouette + quality. + """ + m = np.asarray(mask, dtype=bool) + if not m.any(): + return 1e9 + h, w = m.shape + gray = cv2.cvtColor(np.asarray(rgb_hwc, dtype=np.uint8), cv2.COLOR_RGB2GRAY) + edges = cv2.Canny(gray, 40, 120) + bd = np.zeros(m.shape, dtype=np.uint8) + contours, _ = cv2.findContours( + m.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE + ) + cv2.drawContours(bd, contours, -1, 1, 1) + ys, xs = np.where(bd > 0) + if len(xs) == 0: + return 1e9 + keep = ~((ys <= 1) | (ys >= h - 2) | (xs <= 1) | (xs >= w - 2)) + ys, xs = ys[keep], xs[keep] + if len(xs) == 0: + return 1e9 + dist = cv2.distanceTransform((~edges.astype(bool)).astype(np.uint8), cv2.DIST_L2, 3) + return float(np.median(dist[ys, xs])) + + +def _pos_hit_count(mask: np.ndarray, pos_xy: np.ndarray) -> int: + m = np.asarray(mask, dtype=bool) + h, w = m.shape + n = 0 + for x, y in np.asarray(pos_xy, dtype=np.float32): + xi = int(np.clip(round(x), 0, w - 1)) + yi = int(np.clip(round(y), 0, h - 1)) + n += int(bool(m[yi, xi])) + return n + + +def _refine_wrist_edges( + mask: np.ndarray, + rgb_hwc: np.ndarray, + pos_xy: np.ndarray, + *, + max_med_delta: float = 0.25, +) -> np.ndarray: + """Snap the silhouette onto Canny edges without losing prompt coverage. + + Peel boundary pixels that sit far from image edges; grow only onto nearby + Canny pixels that look like hand tissue. Accept only if hits are preserved + and median boundary→edge distance does not worsen beyond ``max_med_delta``. + """ + m0 = np.asarray(mask, dtype=bool) + if not m0.any(): + return m0 + rgb = np.asarray(rgb_hwc, dtype=np.uint8) + h, w = m0.shape + gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY) + edges = cv2.Canny(gray, 40, 120) + near_edge = cv2.dilate(edges, np.ones((3, 3), np.uint8)).astype(bool) + base_med = _edge_boundary_median(m0, rgb) + base_hits = _pos_hit_count(m0, pos_xy) + edge_dist = cv2.distanceTransform( + (~edges.astype(bool)).astype(np.uint8), cv2.DIST_L2, 3 + ) + + m = m0.copy() + k = np.ones((3, 3), np.uint8) + # Alternating peel / snap-grow so the contour walks onto Canny. + for _ in range(6): + ero = cv2.erode(m.astype(np.uint8), k).astype(bool) + bd = m & ~ero + n_nb = cv2.filter2D(m.astype(np.float32), -1, np.ones((3, 3), np.float32)) + spiky = bd & (n_nb <= 3.0) + far_edge = bd & (edge_dist > 1.75) + rem = bd & ( + (far_edge & (gray > 95)) + | (spiky & (gray > 105)) + | ((~near_edge) & (gray > 125)) + | ((gray > 165) & (np.arange(h)[:, None] < int(round(0.65 * h)))) + ) + if rem.any(): + m = m & ~rem + dil = cv2.dilate(m.astype(np.uint8), k).astype(bool) + add = ( + dil + & ~m + & near_edge + & (edge_dist <= 1.25) + & (gray >= 45) + & (gray <= 155) + ) + if add.any(): + m = m | add + if not rem.any() and not add.any(): + break + + # Restore any dropped positive seeds (tiny). + seed = m.astype(np.uint8) + for x, y in np.asarray(pos_xy, dtype=np.float32): + xi = int(np.clip(round(x), 0, w - 1)) + yi = int(np.clip(round(y), 0, h - 1)) + if seed[yi, xi] or int(gray[yi, xi]) < 80: + continue + cv2.circle(seed, (xi, yi), 2, 1, -1) + m = largest_connected_component(seed.astype(bool) | m) + + hits = _pos_hit_count(m, pos_xy) + med = _edge_boundary_median(m, rgb) + if hits >= base_hits and med <= base_med + max_med_delta: + return m + # If hits dropped, keep the better-edge mask only when almost all seeds remain. + if hits >= max(1, base_hits - 1) and med + 0.05 < base_med: + return m + return m0 + + +def _grow_mask_along_hand( + mask: np.ndarray, + rgb_hwc: np.ndarray, + pos_xy: np.ndarray, + *, + max_iters: int = 10, + max_med_edge: float = 1.8, +) -> np.ndarray: + """Expand SAM2 mask into nearby hand tissue without rewriting its silhouette. + + Growth is limited to bright, hand-like pixels near missed positive points and + stops if the boundary drifts too far from image edges. + """ + m0 = np.asarray(mask, dtype=bool) + if not m0.any(): + return m0 + rgb = np.asarray(rgb_hwc, dtype=np.uint8) + h, w = m0.shape + gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY).astype(np.float32) + hand_med = float(np.median(gray[m0])) + gx = cv2.Sobel(gray, cv2.CV_32F, 1, 0, ksize=3) + gy = cv2.Sobel(gray, cv2.CV_32F, 0, 1, ksize=3) + mag = cv2.magnitude(gx, gy) + edges = cv2.Canny(gray.astype(np.uint8), 40, 120) + near_edge = cv2.dilate(edges, np.ones((3, 3), np.uint8)).astype(bool) + palm = np.arange(h, dtype=np.int32)[:, None] >= int(round(0.68 * h)) + + missed: list[tuple[int, int]] = [] + for x, y in np.asarray(pos_xy, dtype=np.float32): + xi = int(np.clip(round(x), 0, w - 1)) + yi = int(np.clip(round(y), 0, h - 1)) + if not m0[yi, xi]: + missed.append((xi, yi)) + miss_map = np.full((h, w), 1e9, dtype=np.float32) + for x, y in missed: + yy, xx = np.ogrid[:h, :w] + miss_map = np.minimum( + miss_map, np.sqrt((xx - x) ** 2 + (yy - y) ** 2).astype(np.float32) + ) + + allow = ( + (gray >= 70.0) + & (np.abs(gray - hand_med) <= 55.0) + & (palm | near_edge | (mag > 14.0) | (miss_map <= 18.0) | (gray >= 115.0)) + ) + allow = cv2.dilate(allow.astype(np.uint8), np.ones((3, 3), np.uint8)).astype(bool) + + k = np.ones((3, 3), np.uint8) + grown = m0.copy() + best = m0.copy() + best_hits = sum( + 1 + for x, y in np.asarray(pos_xy, dtype=np.float32) + if m0[int(np.clip(round(y), 0, h - 1)), int(np.clip(round(x), 0, w - 1))] + ) + best_med = _edge_boundary_median(best, rgb) + + def _hits(m: np.ndarray) -> int: + n = 0 + for x, y in np.asarray(pos_xy, dtype=np.float32): + xi = int(np.clip(round(x), 0, w - 1)) + yi = int(np.clip(round(y), 0, h - 1)) + n += int(bool(m[yi, xi])) + return n + + for _ in range(max_iters): + dil = cv2.dilate(grown.astype(np.uint8), k).astype(bool) + add = dil & allow & ~grown & (miss_map <= 16.0) + if not add.any(): + break + grown = grown | add + med = _edge_boundary_median(grown, rgb) + hits = _hits(grown) + if med <= max_med_edge and (hits, -med) >= (best_hits, -best_med): + best = grown.copy() + best_hits, best_med = hits, med + if hits >= len(pos_xy) and med <= max_med_edge: + break + + out = _fill_enclosed_holes(best) + # Drop dark shelf pixels that growth may have briefly crossed. + # Keep the right-side outer finger (user 小拇指): dark tissue at x≳0.45w. + yy = np.arange(h, dtype=np.int32)[:, None] + xx = np.arange(w, dtype=np.int32)[None, :] + shelf = (gray < 65.0) & (yy < int(round(0.62 * h))) & (xx < int(round(0.45 * w))) + out = out & ~shelf + return largest_connected_component(out) + + +def _components_at_points( + mask: np.ndarray, + points_xy: np.ndarray, + *, + min_area: int = 24, +) -> np.ndarray: + """Keep mask components directly anchored by positive prompt points.""" + mask_u8 = np.asarray(mask, dtype=np.uint8) + out = np.zeros_like(mask_u8, dtype=bool) + if not mask_u8.any() or len(points_xy) == 0: + return out + n_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask_u8, connectivity=8) + h, w = mask_u8.shape + keep: set[int] = set() + for x, y in np.asarray(points_xy, dtype=np.float32): + xi = int(np.clip(round(x), 0, w - 1)) + yi = int(np.clip(round(y), 0, h - 1)) + label = int(labels[yi, xi]) + if label > 0 and int(stats[label, cv2.CC_STAT_AREA]) >= int(min_area): + keep.add(label) + for label in keep: + out |= labels == label + return out + + +def _right_wrist_prior_bitmap(h: int, w: int) -> np.ndarray: + """Calibrated full-resolution right-wrist hand prior (uint8 0/1).""" + prior = np.zeros((h, w), dtype=np.uint8) + offset_x = int(round(0.00625 * w)) # +2 px at 320x180 + offset_y = int(round(0.00556 * h)) # +1 px at 320x180 + + # Palm covers the outer/pinky-side dorsal surface without bridging finger gaps. + palm_frac = [ + (0.000, 0.680), + (0.040, 0.620), + (0.090, 0.575), + (0.150, 0.560), + (0.210, 0.575), + (0.260, 0.620), + (0.300, 0.690), + (0.355, 0.735), + (0.430, 0.760), + (0.500, 0.785), + (0.555, 0.825), + (0.588, 0.900), + (0.600, 0.994), + (0.000, 0.994), + ] + thumb_frac = [ + (0.531, 0.994), + (0.544, 0.839), + (0.562, 0.761), + (0.594, 0.711), + (0.628, 0.706), + (0.659, 0.761), + (0.691, 0.878), + (0.706, 0.994), + ] + for polygon_frac in (palm_frac, thumb_frac): + polygon = np.asarray( + [ + (round(x * w) + offset_x, round(y * h) + offset_y) + for x, y in polygon_frac + ], + dtype=np.int32, + ) + cv2.fillPoly(prior, [polygon], 1) + + # tip -> root along bright finger centerlines (keep thin to avoid shelf bridging). + finger_segments_frac = [ + ((0.315, 0.505), (0.344, 0.710)), # little finger tip pushed further distal + ((0.384, 0.578), (0.434, 0.661)), + ((0.469, 0.622), (0.528, 0.689)), + ((0.525, 0.650), (0.531, 0.783)), + ] + base_thickness = max(6, int(round(0.034 * w))) + for i, (tip_frac, root_frac) in enumerate(finger_segments_frac): + tip = ( + int(round(tip_frac[0] * w)) + offset_x, + int(round(tip_frac[1] * h)) + offset_y, + ) + root = ( + int(round(root_frac[0] * w)) + offset_x, + int(round(root_frac[1] * h)) + offset_y, + ) + thickness = base_thickness + (3 if i == 0 else 0) + cv2.line(prior, tip, root, 1, thickness, cv2.LINE_8) + if i == 0: + cv2.circle(prior, tip, max(6, thickness // 2 + 2), 1, -1) + return prior + + +def _right_wrist_mask_prior(h: int, w: int) -> np.ndarray: + """Build a calibrated whole-hand low-resolution logit prior.""" + prior = _right_wrist_prior_bitmap(h, w) + low_res = cv2.resize(prior, (256, 256), interpolation=cv2.INTER_NEAREST) + return np.where(low_res[None] > 0, 8.0, -8.0).astype(np.float32) + + +def _complete_right_pinky(mask: np.ndarray, rgb_hwc: np.ndarray) -> np.ndarray: + """Force-include distal fingers SAM2 truncates on right_wrist. + + Two common truncations on ep0: + 1) left-column pinky tip (x~75-125, y~85-115) + 2) right-side outer finger pad (x~155-195, y~98-130) + — the digit users identify as 小拇指 in this camera view. + + Only absorb dark finger tissue (no morphology that rewrites SAM2 edges). + """ + m = np.asarray(mask, dtype=bool) + if not m.any(): + return m + rgb = np.asarray(rgb_hwc, dtype=np.uint8) + h, w = m.shape + gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY) + yy = np.arange(h, dtype=np.int32)[:, None] + xx = np.arange(w, dtype=np.int32)[None, :] + k = np.ones((3, 3), np.uint8) + grown = m.copy() + + def _grow_roi(allow: np.ndarray, *, iters: int, capsule: np.ndarray | None = None) -> None: + nonlocal grown + seed = grown & allow + if capsule is not None: + seed = seed | (capsule.astype(bool) & allow) + if not seed.any() and not (grown & allow).any(): + dil0 = cv2.dilate(grown.astype(np.uint8), np.ones((7, 7), np.uint8)).astype(bool) + seed = dil0 & allow + if not seed.any(): + return + local = grown | seed + for _ in range(iters): + dil = cv2.dilate(local.astype(np.uint8), k).astype(bool) + add = dil & allow & ~local + if not add.any(): + break + local = local | add + grown = grown | local + + # Left distal tip: skip aggressive grow — it rewrites SAM2's silhouette + # (ep0: +1100 px raised med_edge 0.95 → 2.3). Rely on SAM2 + right absorb. + + # --- right outer finger (user-facing 小拇指) --- + # Dark ridge ~gray 35-95: tip crest ~y95 → pad ~y106 → mid ~y117 → root. + right_roi = ( + (xx >= int(round(0.47 * w))) + & (xx <= int(round(0.61 * w))) + & (yy >= int(round(0.50 * h))) + & (yy <= int(round(0.75 * h))) + ) + r_capsule = np.zeros((h, w), dtype=np.uint8) + r_crest = (int(round(0.560 * w)), int(round(0.528 * h))) # ~179,95 + r_tip = (int(round(0.553 * w)), int(round(0.589 * h))) # ~177,106 + r_mid = (int(round(0.516 * w)), int(round(0.650 * h))) # ~165,117 + r_root = (int(round(0.494 * w)), int(round(0.711 * h))) # ~158,128 + r_thick = max(14, int(round(0.050 * w))) + cv2.line(r_capsule, r_crest, r_tip, 1, r_thick, cv2.LINE_8) + cv2.line(r_capsule, r_tip, r_mid, 1, r_thick, cv2.LINE_8) + cv2.line(r_capsule, r_mid, r_root, 1, r_thick, cv2.LINE_8) + cv2.circle(r_capsule, r_crest, max(8, r_thick // 2), 1, -1) + cv2.circle(r_capsule, r_tip, max(9, r_thick // 2 + 1), 1, -1) + cap_dist = cv2.distanceTransform( + (r_capsule == 0).astype(np.uint8), cv2.DIST_L2, 3 + ) + near_mask = cv2.dilate(grown.astype(np.uint8), np.ones((9, 9), np.uint8)).astype(bool) + right_allow = ( + right_roi + & (gray >= 35) + & (gray <= 95) + & ((cap_dist <= max(14.0, 0.045 * w)) | near_mask) + ) + _grow_roi(right_allow, iters=55, capsule=r_capsule) + + # Absorb dark finger CCs that touch the mask or the capsule (bridge small gaps). + dark_finger = right_roi & (gray >= 35) & (gray <= 95) + touch = cv2.dilate(grown.astype(np.uint8), np.ones((7, 7), np.uint8)).astype(bool) + n_lab, labels, stats, _ = cv2.connectedComponentsWithStats( + dark_finger.astype(np.uint8), connectivity=8 + ) + for lab in range(1, n_lab): + comp = labels == lab + if int(stats[lab, cv2.CC_STAT_AREA]) < 15: + continue + if (comp & touch).any() or (comp & r_capsule.astype(bool)).any(): + grown = grown | comp + + # Final tip crest: dark pixels sitting just above the current finger top. + tip_zone = ( + right_roi + & (yy <= int(round(0.59 * h))) + & (gray >= 35) + & (gray <= 95) + & (cap_dist <= max(16.0, 0.050 * w)) + ) + tip_touch = cv2.dilate(grown.astype(np.uint8), np.ones((5, 5), np.uint8)).astype(bool) + grown = grown | (tip_zone & tip_touch) + + # Bridge bright gaps between palm and outer finger so LCC cannot drop the digit. + bridge_allow = ( + right_roi + & (gray >= 35) + & (gray <= 145) + & (cap_dist <= max(10.0, 0.032 * w)) + ) + local = grown.copy() + for _ in range(20): + dil = cv2.dilate(local.astype(np.uint8), k).astype(bool) + add = dil & bridge_allow & ~local + if not add.any(): + break + local = local | add + grown = local + + # Keep every CC that overlaps the original palm mask (not just the largest). + n_lab, labels, stats, _ = cv2.connectedComponentsWithStats( + grown.astype(np.uint8), connectivity=8 + ) + keep = np.zeros_like(grown) + for lab in range(1, n_lab): + comp = labels == lab + if (comp & m).any(): + keep = keep | comp + if not keep.any(): + keep = largest_connected_component(grown) + return keep + + +def _complete_right_thumb(mask: np.ndarray, rgb_hwc: np.ndarray) -> np.ndarray: + """Absorb distal thumb tissue SAM2 / soft x-trim truncates on right_wrist. + + Thumb sits bottom-right (x~185-220, y~130-175), gray ~60-120. + """ + m = np.asarray(mask, dtype=bool) + if not m.any(): + return m + rgb = np.asarray(rgb_hwc, dtype=np.uint8) + h, w = m.shape + gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY) + yy = np.arange(h, dtype=np.int32)[:, None] + xx = np.arange(w, dtype=np.int32)[None, :] + k = np.ones((3, 3), np.uint8) + grown = m.copy() + + thumb_roi = ( + (xx >= int(round(0.575 * w))) + & (xx <= int(round(0.720 * w))) + & (yy >= int(round(0.700 * h))) + & (yy <= int(round(0.980 * h))) + ) + capsule = np.zeros((h, w), dtype=np.uint8) + tip = (int(round(0.656 * w)), int(round(0.900 * h))) # ~210,162 + mid = (int(round(0.625 * w)), int(round(0.861 * h))) # ~200,155 + root = (int(round(0.594 * w)), int(round(0.806 * h))) # ~190,145 + thick = max(10, int(round(0.036 * w))) + cv2.line(capsule, tip, mid, 1, thick, cv2.LINE_8) + cv2.line(capsule, mid, root, 1, thick, cv2.LINE_8) + cv2.circle(capsule, tip, max(6, thick // 2), 1, -1) + cap_dist = cv2.distanceTransform((capsule == 0).astype(np.uint8), cv2.DIST_L2, 3) + near = cv2.dilate(grown.astype(np.uint8), np.ones((7, 7), np.uint8)).astype(bool) + allow = ( + thumb_roi + & (gray >= 70) + & (gray <= 115) + & ((cap_dist <= max(9.0, 0.030 * w)) | near) + ) + seed = (grown & allow) | (capsule.astype(bool) & allow) + if not seed.any(): + seed = near & allow + local = grown | seed + for _ in range(40): + dil = cv2.dilate(local.astype(np.uint8), k).astype(bool) + add = dil & allow & ~local + if not add.any(): + break + local = local | add + grown = local + + thumb_tissue = thumb_roi & (gray >= 70) & (gray <= 115) + touch = cv2.dilate(grown.astype(np.uint8), np.ones((7, 7), np.uint8)).astype(bool) + n_lab, labels, stats, _ = cv2.connectedComponentsWithStats( + thumb_tissue.astype(np.uint8), connectivity=8 + ) + for lab in range(1, n_lab): + comp = labels == lab + if int(stats[lab, cv2.CC_STAT_AREA]) < 20: + continue + if (comp & touch).any() or (comp & capsule.astype(bool)).any(): + grown = grown | comp + + # Bridge palm↔thumb across slightly brighter knuckle pixels. + bridge = ( + thumb_roi + & (gray >= 55) + & (gray <= 145) + & (cap_dist <= max(10.0, 0.035 * w)) + ) + local = grown.copy() + for _ in range(16): + dil = cv2.dilate(local.astype(np.uint8), k).astype(bool) + add = dil & bridge & ~local + if not add.any(): + break + local = local | add + grown = local + + n_lab, labels, _, _ = cv2.connectedComponentsWithStats( + grown.astype(np.uint8), connectivity=8 + ) + keep = np.zeros_like(grown) + for lab in range(1, n_lab): + comp = labels == lab + if (comp & m).any(): + keep = keep | comp + return keep if keep.any() else largest_connected_component(grown) + + +def _complete_right_four_fingers( + mask: np.ndarray, + rgb_hwc: np.ndarray, + pos_xy: np.ndarray, +) -> np.ndarray: + """Solid four-finger silhouette, pose-agnostic. + + Keeps palm+thumb untouched, fills tip→palm spokes, closes mid-finger gaps + (including bright floor/paper seams on ep50/75), and only lightly peels + bright bleed past the tip crown — no distal lane rebuild. + """ + base = np.asarray(mask, dtype=bool) + if not base.any(): + return base + rgb = np.asarray(rgb_hwc, dtype=np.uint8) + h, w = base.shape + gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY).astype(np.float32) + yy = np.arange(h, dtype=np.int32)[:, None] + xx = np.arange(w, dtype=np.int32)[None, :] + k3 = np.ones((3, 3), np.uint8) + + pos = np.asarray(pos_xy, dtype=np.float32).copy() + if len(pos) < 9: + return base.copy() + + # Use distal snaps only — do NOT refine on incomplete palm base (that retracts tips). + tips = _snap_right_four_tips(pos[5:9], rgb) + + thumb_x = float(np.min(pos[3:5, 0])) if len(pos) >= 5 else 0.58 * w + finger_x1 = int(round(min(0.56 * w, thumb_x - 0.03 * w))) + tip_y_min = float(np.min(tips[:, 1])) + tip_y_max = float(np.max(tips[:, 1])) + finger_y0 = int(round(max(0.38 * h, tip_y_min - 0.14 * h))) + finger_y1 = int(round(min(0.82 * h, tip_y_max + 0.20 * h))) + finger_roi = ( + (xx >= int(round(0.18 * w))) + & (xx <= finger_x1) + & (yy >= finger_y0) + & (yy <= finger_y1) + ) + protected = base & ~finger_roi + + ys_b, _ = np.where(base & finger_roi) + palm_y = float(np.percentile(ys_b, 80)) if len(ys_b) else min(0.78 * h, finger_y1) + + # Hand / metal pads; allow brighter specular (ep50/75 pads ~160–175). + tissue = finger_roi & (gray >= 50.0) & (gray <= 175.0) + # Modest distal allowance + light tip trim (not the aggressive rebuild). + tip_ceiling = tip_y_min - 0.050 * h + + grown = np.zeros((h, w), dtype=bool) + half = max(7.5, 0.026 * w) + thick = max(5, int(round(0.018 * w))) + lane_axes = [] # (tx, ty, bx, by, cap_dist) for later on-lane bright fill + for tip in tips: + tx = int(np.clip(round(float(tip[0])), 0, w - 1)) + ty = int(np.clip(round(float(tip[1])), 0, h - 1)) + lane_half = max(6, int(round(0.022 * w))) + xa, xb = max(0, tx - lane_half), min(w, tx + lane_half + 1) + col = base[:, xa:xb] + ys, xs = np.where(col) + if len(ys): + y_edge = int(np.percentile(ys, 15)) + bx = int(xa + np.median(xs[ys <= y_edge + 3])) if np.any(ys <= y_edge + 3) else tx + by = int(y_edge) + else: + bx, by = tx, int(min(h - 1, palm_y)) + rx, ry = bx, int(min(h - 1, max(by, palm_y))) + + cap = np.zeros((h, w), dtype=np.uint8) + cv2.line(cap, (tx, ty), (bx, by), 1, thick, cv2.LINE_8) + cv2.line(cap, (bx, by), (rx, ry), 1, thick, cv2.LINE_8) + cv2.circle(cap, (tx, ty), max(5, thick - 1), 1, -1) + cap_dist = cv2.distanceTransform((cap == 0).astype(np.uint8), cv2.DIST_L2, 3) + lane_axes.append((tx, ty, bx, by, cap_dist)) + + walk = ( + (cap_dist <= half) + & finger_roi + & tissue + & (yy >= tip_ceiling) + & (gray <= 175.0) + ) + spine = ( + (cap_dist <= max(3.2, 0.011 * w)) + & finger_roi + & (yy >= tip_ceiling) + & (gray >= 95.0) + & (gray <= 175.0) + ) + grown |= spine + seed = (base & (cap_dist <= half + 3.0)).astype(np.uint8) + cv2.circle(seed, (tx, ty), 2, 1, -1) + cv2.circle(seed, (bx, by), 2, 1, -1) + local = seed.astype(bool) | ((cap_dist <= 2.0) & walk) + for _ in range(80): + dil = cv2.dilate(local.astype(np.uint8), k3).astype(bool) + add = dil & walk & ~local + if not add.any(): + break + local = local | add + grown |= local & walk + + # Bright distal bleed past tips only (keep mid-finger bright pads). + beyond = finger_roi & ( + ((yy < tip_ceiling) & (gray > 145.0)) + | ((yy < tip_y_min - 0.015 * h) & (gray > 155.0)) + ) + fingers = (base & finger_roi & ~beyond) | (grown & ~beyond) + + # Tip pads. + near = cv2.dilate(fingers.astype(np.uint8), np.ones((7, 7), np.uint8)).astype(bool) + for tip in tips: + tx = int(np.clip(round(float(tip[0])), 0, w - 1)) + ty = int(np.clip(round(float(tip[1])), 0, h - 1)) + if int(gray[ty, tx]) > 152: + best = None + for dy in range(-6, 7): + for dx in range(-6, 7): + xx_, yy_ = tx + dx, ty + dy + if not (0 <= xx_ < w and 0 <= yy_ < h): + continue + g = int(gray[yy_, xx_]) + if g < 95 or g > 155 or not near[yy_, xx_]: + continue + cand = (abs(dx) + abs(dy), -g, xx_, yy_) + if best is None or cand < best: + best = cand + if best is None: + continue + tx, ty = best[2], best[3] + disk = np.zeros((h, w), dtype=np.uint8) + cv2.circle(disk, (tx, ty), 3, 1, -1) + fingers |= disk.astype(bool) & tissue & near & (yy >= tip_ceiling) & ~beyond + + filled = _fill_enclosed_holes(fingers) + fingers |= filled & finger_roi & tissue & (yy >= tip_ceiling) & ~beyond + + tip_xs = sorted(int(round(float(t[0]))) for t in tips) + x_left = max(int(round(0.18 * w)), tip_xs[0] - max(4, int(0.015 * w))) + x_right = min(finger_x1, tip_xs[-1] + max(4, int(0.015 * w))) + + # Inter-digit bridge (dark + moderately bright seams). + max_gap = max(12, int(round(0.060 * w))) + bridge = np.zeros((h, w), dtype=bool) + for y in range(max(0, finger_y0), min(h, finger_y1 + 1)): + if y < tip_ceiling: + continue + row = fingers[y, x_left : x_right + 1] + if not row.any(): + continue + xs = np.flatnonzero(row) + runs = [] + s = int(xs[0]) + prev = int(xs[0]) + for x in xs[1:]: + x = int(x) + if x == prev + 1: + prev = x + continue + runs.append((s, prev)) + s = x + prev = x + runs.append((s, prev)) + for (_, a1), (b0, _) in zip(runs[:-1], runs[1:]): + gap = b0 - a1 - 1 + if gap <= 0 or gap > max_gap: + continue + for x in range(x_left + a1 + 1, x_left + b0): + if float(gray[y, x]) > 200.0: + continue + bridge[y, x] = True + fingers |= bridge & finger_roi & ~beyond + + # ---- On-lane bright hole fill (ep50/75 finger residuals) ---- + # Bright paper/specular seams sit ON the finger axis between tip and palm. + # Fill only when sandwiched by mask along the lane (not free paper outside). + lane_r_fill = max(3.5, 0.012 * w) + bright_fill = np.zeros((h, w), dtype=bool) + for tx, ty, bx, by, cap_dist in lane_axes: + on_lane = (cap_dist <= lane_r_fill) & finger_roi & (yy >= tip_ceiling) & ~beyond + # Vertical sandwich: for each x near lane, find holes with mask above+below in y. + col_mask = fingers & on_lane + hole_cand = on_lane & ~fingers & (gray >= 140.0) & (gray <= 230.0) + if not hole_cand.any(): + continue + # Dilate existing mask along lane; close small y-gaps on-lane only. + axis = np.zeros((h, w), dtype=np.uint8) + cv2.line(axis, (tx, ty), (bx, by), 1, max(2, int(round(lane_r_fill))), cv2.LINE_8) + axis_b = axis.astype(bool) & finger_roi & (yy >= tip_ceiling) & ~beyond + # Morph-close ONLY on this thin axis (fills bright mid-finger seams). + sub = (fingers & axis_b).astype(np.uint8) + # Close with a tall thin kernel so y-gaps fill without lateral bleed. + k_v = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 11)) + closed = cv2.morphologyEx(sub, cv2.MORPH_CLOSE, k_v).astype(bool) + add = closed & axis_b & ~fingers + # Keep adds that are either tissue-gray or bright-but-sandwiched. + ys_a, xs_a = np.where(add) + for yi, xi in zip(ys_a.tolist(), xs_a.tolist()): + # Require mask within ±8px vertically on same/near column. + y0, y1 = max(0, yi - 8), min(h, yi + 9) + x0, x1 = max(0, xi - 2), min(w, xi + 3) + neigh = fingers[y0:y1, x0:x1] + if not neigh.any(): + continue + # Must see mask both proximal and distal of this hole (true seam). + above = fingers[max(0, yi - 8) : yi, x0:x1].any() + below = fingers[yi + 1 : min(h, yi + 9), x0:x1].any() + if above and below: + bright_fill[yi, xi] = True + elif float(gray[yi, xi]) <= 175.0 and (above or below): + bright_fill[yi, xi] = True + fingers |= bright_fill + + # Tip crown chord + soft pad (restore solid distal fingers). + if len(tips) == 4: + crown = np.zeros((h, w), dtype=np.uint8) + tip_pts = [ + (int(np.clip(round(float(t[0])), 0, w - 1)), + int(np.clip(round(float(t[1])), 0, h - 1))) + for t in tips + ] + chord_t = max(3, int(round(0.016 * w))) + for i in range(3): + cv2.line(crown, tip_pts[i], tip_pts[i + 1], 1, chord_t, cv2.LINE_8) + poly = np.array( + [ + tip_pts[0], + tip_pts[1], + tip_pts[2], + tip_pts[3], + (tip_pts[3][0], min(h - 1, tip_pts[3][1] + int(0.10 * h))), + (tip_pts[0][0], min(h - 1, tip_pts[0][1] + int(0.10 * h))), + ], + dtype=np.int32, + ) + pad = np.zeros((h, w), dtype=np.uint8) + cv2.fillConvexPoly(pad, poly, 1) + near_f = cv2.dilate(fingers.astype(np.uint8), np.ones((7, 7), np.uint8)).astype(bool) + crown_fill = ( + ((crown > 0) | (pad > 0)) + & finger_roi + & near_f + & (yy >= tip_ceiling) + & (yy <= tip_y_max + 0.12 * h) + & (gray <= 175.0) + & ~beyond + ) + fingers |= crown_fill + + # Close inside tip x-span. + if fingers.any() and x_right > x_left + 4: + band = np.zeros((h, w), dtype=bool) + band[:, x_left : x_right + 1] = True + near = cv2.dilate(fingers.astype(np.uint8), np.ones((7, 7), np.uint8)).astype(bool) + k_ell = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7)) + f_u8 = fingers.astype(np.uint8).copy() + sub = f_u8[:, x_left : x_right + 1] + sub = cv2.morphologyEx(sub, cv2.MORPH_CLOSE, k_ell) + f_u8[:, x_left : x_right + 1] = sub + close_ok = ( + band + & near + & finger_roi + & ~beyond + & (yy >= tip_ceiling) + & (gray <= 185.0) # allow bright mid-finger seams into close + ) + soft = cv2.GaussianBlur(f_u8.astype(np.float32), (5, 5), 0.8) + core = cv2.erode(fingers.astype(np.uint8), k3).astype(bool) + fingers = (((soft > 0.35) | core) & close_ok) | (fingers & finger_roi & ~beyond) + + # Light tip-band shave only (do not hollow mid-finger). + tip_band = (yy >= tip_ceiling) & (yy <= tip_y_min + 0.02 * h) + bright_bleed = tip_band & (gray > 165.0) + lane_dist = np.full((h, w), 1e9, dtype=np.float32) + for tip in tips: + tx = int(np.clip(round(float(tip[0])), 0, w - 1)) + lane_dist = np.minimum(lane_dist, np.abs(xx.astype(np.float32) - float(tx))) + far_lane = tip_band & (lane_dist > 0.034 * w) + fingers &= ~(bright_bleed | (far_lane & (gray > 155.0))) + + out = protected | fingers + n_lab, labels = cv2.connectedComponents(out.astype(np.uint8), connectivity=8) + keep = np.zeros_like(out) + for lab in range(1, n_lab): + comp = labels == lab + if (comp & protected).any() or (comp & base).any(): + keep |= comp + n_lab2, labels2, stats2, _ = cv2.connectedComponentsWithStats( + keep.astype(np.uint8), connectivity=8 + ) + cleaned = np.zeros_like(keep) + min_keep = max(80, int(0.0015 * h * w)) + for lab in range(1, n_lab2): + area = int(stats2[lab, cv2.CC_STAT_AREA]) + top = int(stats2[lab, cv2.CC_STAT_TOP]) + if area < 8: + continue + if area < min_keep and top < tip_y_min - 0.02 * h: + continue + cleaned |= labels2 == lab + return cleaned | (protected & cleaned) + + + +def build_wrist_prompts(h: int, w: int, view: str) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Return (point_coords Nx2, point_labels N, box xyxy) for one wrist view. + + Wrist cams: hand fills the lower part of the frame. Box is forced to the + bottom band (not the mid-frame template box, which can sit on background). + """ + # Bottom-heavy box covering the mechanical hand (ep0 palm + distal fingers). + # Wrist cameras are mirrored: fingers extend from each bottom corner toward + # the image centre, so prompts must follow that diagonal rather than forming + # a vertical line through the palm. + if view == "left_wrist": + # Palm enters from bottom-right; fingers extend up-left to x~178,y~104. + box_frac = (0.266, 0.433, 0.99, 0.99) + pos_frac = [ + # Palm + (0.781, 0.944), + (0.719, 0.850), + (0.812, 0.778), + # Thumb, separated from the other fingers by a visible gap + (0.381, 0.861), + (0.450, 0.800), + # Finger roots and distal links + (0.719, 0.722), + (0.678, 0.672), + (0.638, 0.639), + (0.594, 0.628), + (0.556, 0.661), + (0.659, 0.600), + (0.616, 0.578), + ] + else: # right_wrist + # Palm BL; four fingers fan up-center; thumb bottom-right (x~210). + # Points calibrated on white finger faces (g~100-130), NOT dark gaps / shelf. + # Keep the right wall past the thumb so SAM2 does not invent a + # vertical box-edge cut through the digit. + box_frac = (0.00, 0.340, 0.820, 0.99) + pos_frac = [ + # Palm + (0.219, 0.944), + (0.281, 0.850), + (0.184, 0.744), + # Thumb tip / pad (metal ridge, avoid saturated specular) + (0.647, 0.922), # tip ~207,166 + (0.606, 0.861), # pad ~194,155 + # Four finger tips on finger ridges (LTR = 小拇指→无名→中→食指). + # Lanes ~88/106/125/141 @320; centerline tissue required (no gap false+). + (0.287, 0.575), # T0 小拇指 — distal pad + (0.334, 0.575), # T1 无名指 + (0.388, 0.575), # T2 中指 + (0.438, 0.575), # T3 食指 + # Finger mids on same centerlines + (0.287, 0.694), # M0 + (0.334, 0.694), # M1 + (0.438, 0.680), # M3 食指 mid + ] + + box = _box_frac_to_px(box_frac, h, w) + pos = _frac_to_px(pos_frac, h, w) + if view == "right_wrist": + # Mirror-alignment residual versus left wrist: +2 px x, +1 px y. + pos += np.array([0.00625 * w, 0.00556 * h], dtype=np.float32) + pos[:, 0] = np.clip(pos[:, 0], 0, w - 1) + pos[:, 1] = np.clip(pos[:, 1], 0, h - 1) + neg = _frac_to_px(WRIST_NEG_FRAC, h, w) + if view == "left_wrist": + extra_neg = _frac_to_px([(0.04, 0.40), (0.04, 0.55), (0.04, 0.70), (0.30, 0.20)], h, w) + else: + extra_neg = _frac_to_px( + [ + (0.96, 0.40), + (0.96, 0.55), + (0.96, 0.70), + (0.78, 0.20), + # Background just distal of the four finger tips (keep clear of 食指). + # y kept high (small) so extended fingers are not clipped; runtime + # also repositions these above the snapped tip crown. + (0.250, 0.360), + (0.310, 0.360), + (0.370, 0.360), + (0.480, 0.350), + (0.560, 0.340), + ], + h, + w, + ) + neg = np.concatenate([neg, extra_neg], axis=0) + + coords = np.concatenate([pos, neg], axis=0).astype(np.float32) + labels = np.concatenate( + [np.ones(len(pos), dtype=np.int32), np.zeros(len(neg), dtype=np.int32)] + ) + return coords, labels, box + + +def build_head_prompts(h: int, w: int, side: str) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Return (point_coords, point_labels, box) for head left/right hand.""" + if side not in HEAD_PROMPTS: + raise KeyError(side) + spec = HEAD_PROMPTS[side] + box = _box_frac_to_px(spec["box"], h, w) + pos = _frac_to_px(spec["pos"], h, w) + neg = _frac_to_px(spec["neg"], h, w) + coords = np.concatenate([pos, neg], axis=0).astype(np.float32) + labels = np.concatenate( + [np.ones(len(pos), dtype=np.int32), np.zeros(len(neg), dtype=np.int32)] + ) + return coords, labels, box + + +def sam2_predict_once( + predictor, + rgb_hwc: np.ndarray, + *, + point_coords: np.ndarray, + point_labels: np.ndarray, + box: np.ndarray, + target_area_frac: float | None = None, + mask_input: np.ndarray | None = None, +) -> tuple[np.ndarray, float]: + """Single SAM2 forward. Pick multimask by prompt agreement (not raw score). + + Selection = cover positive prompts, avoid negative prompts, then score. + Post-process: largest connected component only. + """ + if predictor is None: + raise RuntimeError("SAM2 predictor is required (no appearance fallback).") + rgb = np.asarray(rgb_hwc, dtype=np.uint8) + h, w = rgb.shape[:2] + coords = np.asarray(point_coords, dtype=np.float32) + labels = np.asarray(point_labels, dtype=np.int32) + box = np.asarray(box, dtype=np.float32) + + predictor.set_image(rgb) + masks, scores, _ = predictor.predict( + point_coords=coords, + point_labels=labels, + box=box, + mask_input=mask_input, + multimask_output=mask_input is None, + # Pixel coords → let SAM2 normalize internally (False would treat px as [0,1]). + normalize_coords=True, + ) + + pos = coords[labels == 1] + neg = coords[labels == 0] + + def _hit_rate(m: np.ndarray, pts: np.ndarray) -> float: + if len(pts) == 0: + return 0.0 + hits = 0 + for x, y in pts: + xi, yi = int(np.clip(round(x), 0, w - 1)), int(np.clip(round(y), 0, h - 1)) + hits += int(bool(m[yi, xi])) + return hits / float(len(pts)) + + best_i, best_rank = -1, -1e18 + for i, (mask_i, score_i) in enumerate(zip(masks, scores)): + m = largest_connected_component(np.asarray(mask_i, dtype=bool)) + if not m.any(): + continue + # Prefer masks that stay mostly inside the prompt box. + x0, y0, x1, y1 = [int(round(v)) for v in box] + inside = m[max(0, y0) : min(h, y1 + 1), max(0, x0) : min(w, x1 + 1)].sum() + box_frac = float(inside) / max(float(m.sum()), 1.0) + pos_h = _hit_rate(m, pos) + neg_h = _hit_rate(m, neg) + area_frac = float(m.sum()) / float(h * w) + # Hard reject: misses most positives, eats negatives, or covers half the frame. + if pos_h < 0.50: + continue + if neg_h > 0.25: + continue + if box_frac < 0.55: + continue + if area_frac > 0.28: + continue + # Wrist hands are large; reject undersized multimasks that only cover palm. + if target_area_frac is not None and area_frac < 0.18: + continue + area_cost = ( + abs(area_frac - float(target_area_frac)) + if target_area_frac is not None + else area_frac + ) + area_weight = 3.0 if target_area_frac is not None else 1.5 + edge_med = _edge_boundary_median(m, rgb) + rank = ( + 4.0 * pos_h + - 3.0 * neg_h + + 1.0 * box_frac + - area_weight * area_cost + + 0.2 * float(score_i) + - 1.10 * min(edge_med, 8.0) # strongly prefer Canny-aligned edges + ) + if rank > best_rank: + best_rank = rank + best_i = i + + if best_i < 0: + # Soft fallback among multimasks: best pos coverage (still same one forward). + for i, (mask_i, score_i) in enumerate(zip(masks, scores)): + m = largest_connected_component(np.asarray(mask_i, dtype=bool)) + if not m.any(): + continue + pos_h = _hit_rate(m, pos) + neg_h = _hit_rate(m, neg) + area_frac = float(m.sum()) / float(h * w) + area_cost = ( + abs(area_frac - float(target_area_frac)) + if target_area_frac is not None + else area_frac + ) + rank = 4.0 * pos_h - 3.0 * neg_h - 1.0 * area_cost + 0.1 * float(score_i) + if rank > best_rank: + best_rank = rank + best_i = i + + if best_i < 0: + raise RuntimeError("SAM2 returned no usable multimask for prompts.") + + mask = largest_connected_component(np.asarray(masks[best_i], dtype=bool)) + # Do not OR other multimask components onto the winner: on wrist views that + # is the main source of blocky silhouette damage once 2× SAM2 already fits. + score = float(scores[best_i]) + if not mask.any(): + raise RuntimeError("SAM2 returned empty mask after largest-CC.") + return mask, score + + +def sam2_wrist_mask(predictor, rgb_hwc: np.ndarray, view: str) -> tuple[np.ndarray, str, dict]: + rgb = np.asarray(rgb_hwc, dtype=np.uint8) + h, w = rgb.shape[:2] + coords, labels, box = build_wrist_prompts(h, w, view) + + # ------------------------------------------------------------------ left + # Restore the previous native-resolution path (user: left was already good). + if view == "left_wrist": + raw, score = sam2_predict_once( + predictor, + rgb, + point_coords=coords, + point_labels=labels, + box=box, + target_area_frac=0.23, + mask_input=None, + ) + x0, y0, x1, y1 = [int(round(v)) for v in box] + pad = 10 + clipped = np.zeros_like(raw) + ya, yb = max(0, y0 - pad), min(h, y1 + pad + 1) + xa, xb = max(0, x0 - pad), min(w, x1 + pad + 1) + clipped[ya:yb, xa:xb] = raw[ya:yb, xa:xb] + candidate = clipped if clipped.any() else raw + pos = coords[labels == 1] + mask = largest_connected_component(candidate) + mask |= _components_at_points(candidate, pos) + mask = _fill_enclosed_holes(mask) + mask = largest_connected_component(mask) + if not mask.any(): + raise RuntimeError(f"SAM2 wrist {view} mask empty.") + meta = {"coords": coords, "labels": labels, "box": box, "score": score} + return mask, f"sam2_once_wrist_s{score:.2f}", meta + + # ----------------------------------------------------------------- right + # 2× SAM2 (smoother finger silhouettes like ep50/75) + solid finger fill. + # Tip projection is distance-limited so missing fingertips are not faked. + pos = _snap_pos_to_hand(coords[labels == 1], rgb) + if len(pos) >= 9: + pos = pos.copy() + pos[5:9] = _snap_right_four_tips(pos[5:9], rgb) + coords = coords.copy() + coords[labels == 1] = pos + + # Distal negatives at fixed y~0.44 sit ON extended fingertips and clip the + # mask. Push finger-band distal negatives to a clear strip above tip crown. + if len(pos) >= 9: + tip_y = float(np.min(pos[5:9, 1])) + neg_y = float(np.clip(tip_y - 0.065 * h, 2.0, h - 2)) + # Reposition the dedicated distal finger-band negatives (last 5 extras). + # Also catch any other neg that lands in the tip rectangle. + tip_x0 = float(np.min(pos[5:9, 0])) - 0.08 * w + tip_x1 = float(np.max(pos[5:9, 0])) + 0.08 * w + for i in range(len(coords)): + if int(labels[i]) != 0: + continue + x, y = float(coords[i, 0]), float(coords[i, 1]) + in_tip_band = tip_x0 <= x <= tip_x1 and (tip_y - 0.16 * h) <= y <= (tip_y + 0.12 * h) + # Broad distal finger strip (covers fixed extra_neg even if tip xs shift). + in_finger_strip = (0.18 * w) <= x <= (0.60 * w) and (0.30 * h) <= y <= (0.50 * h) + if in_tip_band or in_finger_strip: + coords[i, 1] = neg_y + box = np.asarray(box, dtype=np.float32).copy() + box[1] = min(float(box[1]), max(0.0, tip_y - 0.05 * h)) + + scale = 2 + rgb_in = cv2.resize(rgb, (w * scale, h * scale), interpolation=cv2.INTER_LINEAR) + coords_in = np.asarray(coords, dtype=np.float32) * float(scale) + box_in = np.asarray(box, dtype=np.float32) * float(scale) + raw_in, score = sam2_predict_once( + predictor, + rgb_in, + point_coords=coords_in, + point_labels=labels, + box=box_in, + target_area_frac=0.23, + mask_input=None, + ) + raw = ( + cv2.resize(raw_in.astype(np.float32), (w, h), interpolation=cv2.INTER_AREA) + > 0.50 + ) + + x0, y0, x1, y1 = [int(round(v)) for v in box] + pad = 14 + clipped = np.zeros_like(raw) + ya, yb = max(0, y0 - pad), min(h, y1 + pad + 1) + xa, xb = max(0, x0 - pad), min(w, x1 + pad + 1) + clipped[ya:yb, xa:xb] = raw[ya:yb, xa:xb] + candidate = clipped if clipped.any() else raw + pos = coords[labels == 1] + mask = largest_connected_component(candidate) + mask |= _components_at_points(candidate, pos) + mask = _fill_enclosed_holes(mask) + mask = largest_connected_component(mask) + mask = _complete_right_four_fingers(mask, rgb, pos) + if not mask.any(): + raise RuntimeError(f"SAM2 wrist {view} mask empty.") + + # Refine tip coords on mask, but never jump >5px (avoids palm fakes). + if len(pos) >= 9: + pos = np.asarray(coords[labels == 1], dtype=np.float32).copy() + # Tips: y-only refine on lanes. Mid prompts may snap locally onto mask. + tips = _refine_right_tips_on_base(pos[5:9], mask, rgb) + pos[5:9] = tips + gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY) + hh, ww = mask.shape + y_lo = int(round(0.545 * hh)) + y_hi = int(round(0.600 * hh)) + x_slack = max(3, int(round(0.012 * ww))) + for i in range(5, min(12, len(pos))): + is_tip = i < 9 + xi = int(np.clip(round(float(pos[i, 0])), 0, ww - 1)) + yi = int(np.clip(round(float(pos[i, 1])), 0, hh - 1)) + if is_tip: + # Keep distal tip on pad tissue; never land on dark gap/spine. + pos[i, 1] = float(np.clip(pos[i, 1], y_lo, y_hi)) + yi = int(np.clip(round(float(pos[i, 1])), 0, hh - 1)) + g_here = int(gray[yi, xi]) if 0 <= yi < hh and 0 <= xi < ww else 0 + if mask[yi, xi] and 102 <= g_here <= 172: + continue + best = None + for dy in range(-8, 3): # prefer distal / same row; avoid proximal + for dx in range(-x_slack, x_slack + 1): + xx, yy = xi + dx, yi + dy + if not (y_lo <= yy <= y_hi): + continue + if not (0 <= xx < ww and 0 <= yy < hh and mask[yy, xx]): + continue + g = int(gray[yy, xx]) + if not (102 <= g <= 172): + continue + retract = max(0, yy - yi) + cand = (abs(dy) + 0.5 * abs(dx) + 2.0 * retract, -g, xx, yy) + if best is None or cand < best: + best = cand + if best is not None: + pos[i, 0], pos[i, 1] = float(best[2]), float(best[3]) + continue + if mask[yi, xi]: + continue + best = None + for dy in range(-5, 6): + for dx in range(-5, 6): + xx, yy = xi + dx, yi + dy + if not (0 <= xx < ww and 0 <= yy < hh and mask[yy, xx]): + continue + cand = (abs(dx) + abs(dy), -int(gray[yy, xx]), xx, yy) + if best is None or cand < best: + best = cand + if best is not None and best[0] <= 5: + pos[i, 0], pos[i, 1] = float(best[2]), float(best[3]) + coords = coords.copy() + coords[labels == 1] = pos + + meta = {"coords": coords, "labels": labels, "box": box, "score": score} + return mask, f"sam2_once_wrist_s{score:.2f}", meta + + +def sam2_head_hand_mask( + predictor, rgb_hwc: np.ndarray, side: str +) -> tuple[np.ndarray, str, dict]: + rgb = np.asarray(rgb_hwc, dtype=np.uint8) + h, w = rgb.shape[:2] + coords, labels, box = build_head_prompts(h, w, side) + scale = 2 + rgb_in = cv2.resize(rgb, (w * scale, h * scale), interpolation=cv2.INTER_LINEAR) + coords_in = np.asarray(coords, dtype=np.float32) * float(scale) + box_in = np.asarray(box, dtype=np.float32) * float(scale) + raw_in, score = sam2_predict_once( + predictor, + rgb_in, + point_coords=coords_in, + point_labels=labels, + box=box_in, + ) + raw = ( + cv2.resize(raw_in.astype(np.float32), (w, h), interpolation=cv2.INTER_AREA) + > 0.45 + ) + # Soft clip inside a dilated prompt box (avoid hard box-edge walls). + x0, y0, x1, y1 = [int(round(v)) for v in box] + pad = 12 + ya, yb = max(0, y0 - pad), min(h, y1 + pad + 1) + xa, xb = max(0, x0 - pad), min(w, x1 + pad + 1) + inside = np.zeros((h, w), dtype=np.uint8) + inside[ya:yb, xa:xb] = 1 + inside = cv2.dilate(inside, np.ones((11, 11), np.uint8)) + clipped = raw & inside.astype(bool) + mask = largest_connected_component(clipped if clipped.any() else raw) + if not mask.any(): + raise RuntimeError(f"SAM2 head {side} mask empty.") + pos = coords[labels == 1] + mask = _refine_wrist_edges(mask, rgb, pos, max_med_delta=0.10) + meta = {"coords": coords, "labels": labels, "box": box, "score": score} + return mask, f"sam2_once_head_{side}_s{score:.2f}", meta + + +def draw_prompt_overlay( + rgb: np.ndarray, + *, + mask: np.ndarray | None, + coords: np.ndarray, + labels: np.ndarray, + box: np.ndarray, + mask_bgr: tuple[int, int, int] = (0, 0, 220), +) -> np.ndarray: + """Debug viz: mask + box + pos(green) / neg(red) prompts.""" + bgr = cv2.cvtColor(np.asarray(rgb, dtype=np.uint8), cv2.COLOR_RGB2BGR) + if mask is not None and np.asarray(mask).any(): + m = np.asarray(mask, dtype=bool) + bgr[m] = (bgr[m] * 0.40 + np.asarray(mask_bgr, dtype=np.float32) * 0.60).astype(np.uint8) + x0, y0, x1, y1 = [int(round(v)) for v in box] + cv2.rectangle(bgr, (x0, y0), (x1, y1), (255, 255, 0), 1, cv2.LINE_AA) + for (x, y), lab in zip(np.asarray(coords, dtype=np.float32), np.asarray(labels)): + col = (0, 255, 0) if int(lab) == 1 else (0, 0, 255) + cv2.drawMarker( + bgr, + (int(round(x)), int(round(y))), + col, + markerType=cv2.MARKER_CROSS if int(lab) == 1 else cv2.MARKER_TILTED_CROSS, + markerSize=8, + thickness=1, + line_type=cv2.LINE_AA, + ) + return bgr + + +def sample_points_on_mask( + mask: np.ndarray, + n: int, + *, + rng: np.random.Generator, + distal_bias: bool = True, +) -> np.ndarray: + """Uniform-ish sample on mask; optional bias to smaller-y (distal).""" + ys, xs = np.nonzero(np.asarray(mask, dtype=bool)) + if ys.size == 0: + raise RuntimeError("cannot sample from empty mask") + n = int(n) + if ys.size <= n: + idx = rng.choice(ys.size, size=n, replace=True) + return np.stack([xs[idx], ys[idx]], axis=1).astype(np.float32) + + chosen: list[int] = [] + order = np.argsort(ys) + if distal_bias: + n_dist = max(n // 2, 1) + pool = order[: max(n_dist, int(0.40 * ys.size))] + chosen.extend(rng.choice(pool, size=min(n_dist, pool.size), replace=False).tolist()) + remain = np.setdiff1d(np.arange(ys.size), np.asarray(chosen, dtype=np.int64), assume_unique=False) + n_left = n - len(chosen) + if remain.size and n_left > 0: + y_rem = ys[remain] + bins = np.array_split(np.argsort(y_rem), n_left) + for b in bins: + if len(b): + chosen.append(int(remain[b[rng.integers(0, len(b))]])) + while len(chosen) < n: + chosen.append(int(rng.choice(ys.size))) + idx = np.asarray(chosen[:n], dtype=np.int64) + return np.stack([xs[idx], ys[idx]], axis=1).astype(np.float32) diff --git a/scripts/trex_track/sam2_wrist_hand.py b/scripts/trex_track/sam2_wrist_hand.py new file mode 100644 index 0000000000000000000000000000000000000000..1b5f093535084a1d10e3638d9837fd7742b98713 --- /dev/null +++ b/scripts/trex_track/sam2_wrist_hand.py @@ -0,0 +1,477 @@ +"""Wrist-view hand segmentation for T-Rex CoTracker seeds. + +Wrist cameras see a large, often out-of-focus mechanical hand filling the lower +portion of the frame. Pure SAM2 box/template prompts fail because: + +1. Fixed finger templates often land on background (person/table). +2. SAM2 tiny prefers sharp background edges over blurry close-up plastic. + +Pipeline: + 1. Score pixels by (blur × grey × bottom-prior). + 2. Keep high-score components that touch the bottom edge (= hand seed). + 3. Morphologically close within grey∩lower to connect finger fragments. + 4. Optionally refine with SAM2, hard-clipped to the seed neighborhood. + 5. Random-sample query points from the final mask. +""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path + +import cv2 +import numpy as np + +from trex_track.trex_projection import NUM_WRIST_HAND_TEMPLATE, wrist_hand_template_pixels + +DEFAULT_SAM2_LIBS = Path(os.environ.get("SAM2_LIBS", "/scratch1/home/zhicao/physctrl/libs")) +DEFAULT_SAM2_MODEL = os.environ.get("SAM2_MODEL", "facebook/sam2-hiera-large") +DEFAULT_SAM2_CACHE_DIR = os.environ.get( + "SAM2_CACHE_DIR", + str(Path("/scratch1/home/zhicao/.cache/sam2")), +) + +# Appearance scoring +BLUR_QUANTILE = 0.80 +BOTTOM_CUTOFF = 0.38 +MAX_GROW_DIST = 58 # upward finger reach from palm seed +SAM2_MIN_SEED_OVERLAP = 0.50 +SAM2_MIN_POINT_RECALL = 0.30 + + +def ensure_sam2_on_path(sam2_libs: str | Path = DEFAULT_SAM2_LIBS) -> Path: + root = Path(sam2_libs).expanduser().resolve() + if not (root / "sam2").is_dir(): + raise FileNotFoundError(f"SAM2 libs not found: {root}") + p = str(root) + if p not in sys.path: + sys.path.insert(0, p) + return root + + +def load_sam2_predictor( + *, + model_id: str = DEFAULT_SAM2_MODEL, + device: str = "cuda:0", + cache_dir: str | Path = DEFAULT_SAM2_CACHE_DIR, + sam2_libs: str | Path = DEFAULT_SAM2_LIBS, +): + ensure_sam2_on_path(sam2_libs) + from sam2.sam2_image_predictor import SAM2ImagePredictor # type: ignore + + cache = str(Path(cache_dir).expanduser().resolve()) + Path(cache).mkdir(parents=True, exist_ok=True) + return SAM2ImagePredictor.from_pretrained(model_id, cache_dir=cache, device=device) + + +def largest_connected_component(mask: np.ndarray) -> np.ndarray: + mask_u8 = np.asarray(mask, dtype=np.uint8) + if mask_u8.ndim != 2: + raise ValueError(f"mask must be 2D, got shape {mask_u8.shape}") + if not mask_u8.any(): + return mask_u8.astype(bool) + n_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask_u8, connectivity=8) + if n_labels <= 1: + return mask_u8.astype(bool) + areas = stats[1:, cv2.CC_STAT_AREA] + best = 1 + int(np.argmax(areas)) + return labels == best + + +def _wall_strip_mask(gray: np.ndarray) -> np.ndarray: + h, w = gray.shape + gx = cv2.Sobel(gray, cv2.CV_32F, 1, 0, ksize=3) + vert = cv2.blur(np.abs(gx), (3, 15)) + xx = np.arange(w)[None, :] + return (vert > np.quantile(vert, 0.90)) & ((xx < 0.08 * w) | (xx > 0.92 * w)) + + +def _geodesic_dilate( + marker: np.ndarray, + allow: np.ndarray, + *, + max_iters: int = 40, + kernel_size: int = 5, +) -> np.ndarray: + """Grow marker under allow mask (morphological reconstruction).""" + marker_u8 = (np.asarray(marker, dtype=bool) & np.asarray(allow, dtype=bool)).astype(np.uint8) + allow_u8 = np.asarray(allow, dtype=np.uint8) + if not marker_u8.any() or not allow_u8.any(): + return marker_u8.astype(bool) + k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)) + cur = marker_u8.copy() + for _ in range(int(max_iters)): + nxt = cv2.bitwise_and(cv2.dilate(cur, k), allow_u8) + if np.array_equal(nxt, cur): + break + cur = nxt + return cur.astype(bool) + + +def detect_wrist_hand_mask( + rgb_hwc: np.ndarray, + *, + view: str | None = None, + blur_quantile: float = BLUR_QUANTILE, + bottom_cutoff: float = BOTTOM_CUTOFF, + max_grow_dist: float = MAX_GROW_DIST, +) -> tuple[np.ndarray, np.ndarray]: + """ + Detect the close-up mechanical hand (palm + fingers) in a wrist camera frame. + + 1) Blurry grey bottom-touching core (palm / near-camera mass) + 2) Optional template ROI prior (view-specific) so growth stays on the hand band + 3) Two-stage grow: loose bridge through finger joints, then strict plastic keep + """ + rgb = np.asarray(rgb_hwc, dtype=np.uint8) + if rgb.ndim != 3 or rgb.shape[2] != 3: + raise ValueError(f"expected HWC RGB, got {rgb.shape}") + h, w = rgb.shape[:2] + + hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV) + s = hsv[:, :, 1].astype(np.float32) + v = hsv[:, :, 2].astype(np.float32) + gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY).astype(np.float32) + lab = cv2.cvtColor(rgb, cv2.COLOR_RGB2LAB).astype(np.float32) + + lap = cv2.Laplacian(gray, cv2.CV_32F) + local_var = cv2.blur(lap**2, (21, 21)) + blur = 1.0 / (1.0 + local_var) + blur = (blur - blur.min()) / (blur.max() - blur.min() + 1e-6) + + grey = (s < 55) & (v > 65) & (v < 250) + yy = np.linspace(0.0, 1.0, h, dtype=np.float32)[:, None] + bottom = np.clip((yy - 0.32) / 0.55, 0.0, 1.0) ** 1.3 + score = blur * grey.astype(np.float32) * (0.1 + 0.9 * bottom) + + thr = float(np.quantile(score[int(0.45 * h) :], blur_quantile)) + seed = (score >= thr).astype(np.uint8) + seed[: int(bottom_cutoff * h), :] = 0 + wall = _wall_strip_mask(gray) + seed[wall] = 0 + + k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (9, 9)) + seed = cv2.morphologyEx(seed, cv2.MORPH_CLOSE, k, iterations=2) + + n, labels, stats, _cent = cv2.connectedComponentsWithStats(seed, 8) + seed_m = np.zeros((h, w), np.uint8) + for i in range(1, n): + _x, y, bw, bh, area = stats[i] + if area < 200: + continue + if (y + bh) < (h - 2): + continue + if bw < 0.15 * w and bh > 0.35 * h: + continue + seed_m[labels == i] = 1 + + if not seed_m.any(): + return np.zeros((h, w), dtype=bool), np.zeros((h, w), dtype=bool) + + palm_seed = seed_m.astype(bool) + mean_c = lab[palm_seed].mean(axis=0) + mean_v = float(v[palm_seed].mean()) + mean_l = float(mean_c[0]) + color_dist = np.linalg.norm(lab - mean_c[None, None, :], axis=2) + chroma = np.sqrt((lab[:, :, 1] - 128.0) ** 2 + (lab[:, :, 2] - 128.0) ** 2) + mean_ch = float(chroma[palm_seed].mean()) + + sy, sx = np.nonzero(palm_seed) + x0, x1 = int(sx.min()), int(sx.max()) + y0, y1 = int(sy.min()), int(sy.max()) + bw, bh = max(x1 - x0, 1), max(y1 - y0, 1) + roi = np.zeros((h, w), dtype=bool) + roi[ + max(0, int(y0 - 0.85 * bh)) : min(h, int(y1 + 0.04 * bh)) + 1, + max(0, int(x0 - 0.10 * bw)) : min(w, int(x1 + 0.10 * bw)) + 1, + ] = True + roi |= cv2.dilate(seed_m, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (11, 11))).astype(bool) + + near = cv2.dilate(seed_m, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (11, 11))).astype(bool) + far_ok = ( + (chroma <= mean_ch + 2.0) + & (s < 38) + & (color_dist < 20) + & (v > 55) + & (v < mean_v + 40) + & (lab[:, :, 0] < mean_l + 38) + ) + near_ok = ( + (chroma <= mean_ch + 5.0) + & (s < 65) + & (color_dist < 28) + & (v > 45) + & (v < 225) + ) + plastic = np.where(near, near_ok, far_ok) & roi & (yy > 0.10) + plastic[wall] = False + plastic[: int(0.08 * h), :] = False + + finger_u8 = np.zeros((h, w), np.uint8) + if view is not None: + try: + tmpl = wrist_hand_template_pixels((h, w), view=view).astype(np.float32) + palm_pts = tmpl[tmpl[:, 1] >= np.median(tmpl[:, 1])] + dy = float(sy.mean()) - float(palm_pts[:, 1].mean()) + dx = 0.20 * (float(sx.mean()) - float(palm_pts[:, 0].mean())) + tmpl[:, 0] = np.clip(tmpl[:, 0] + dx, 0, w - 1) + tmpl[:, 1] = np.clip(tmpl[:, 1] + dy, 0, h - 1) + for px, py in np.rint(tmpl).astype(int): + best = None + for rad in range(0, 12): + y_lo, y_hi = max(0, py - rad), min(h, py + rad + 1) + x_lo, x_hi = max(0, px - rad), min(w, px + rad + 1) + ys_l, xs_l = np.nonzero(plastic[y_lo:y_hi, x_lo:x_hi]) + if ys_l.size == 0: + continue + d2 = (xs_l + x_lo - px) ** 2 + (ys_l + y_lo - py) ** 2 + j = int(np.argmin(d2)) + best = (int(xs_l[j] + x_lo), int(ys_l[j] + y_lo)) + break + if best is not None: + cv2.circle(finger_u8, best, 10, 1, -1) + except Exception: + pass + finger_disk = finger_u8.astype(bool) & plastic + + markers = (palm_seed | finger_disk).astype(np.uint8) + bridge = plastic.astype(np.uint8) + grown = _geodesic_dilate(markers, bridge, max_iters=40, kernel_size=5) + grown = cv2.morphologyEx( + grown.astype(np.uint8), + cv2.MORPH_CLOSE, + cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 7)), + iterations=1, + ) + grown = _geodesic_dilate(grown, bridge, max_iters=15, kernel_size=3) + + keep = (grown.astype(bool) & plastic) | palm_seed | finger_disk + dist = cv2.distanceTransform((~palm_seed).astype(np.uint8), cv2.DIST_L2, 5) + seed_cy = float(sy.mean()) + Y = np.arange(h)[:, None] + cap = np.where(Y < seed_cy, float(max_grow_dist) + 10.0, float(max_grow_dist) * 0.45) + keep = (keep & (dist <= cap) & roi) | palm_seed + + dil_palm = cv2.dilate( + palm_seed.astype(np.uint8), cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (37, 37)) + ) + n, labels, stats, _ = cv2.connectedComponentsWithStats(keep.astype(np.uint8), 8) + out = np.zeros((h, w), dtype=bool) + for i in range(1, n): + if stats[i, -1] < 40: + continue + comp = labels == i + if (comp & palm_seed).any() or ( + (comp & finger_disk).any() and (comp & dil_palm.astype(bool)).any() + ): + out |= comp + out |= finger_disk & dil_palm.astype(bool) + if not out.any(): + out = palm_seed + if int(out.sum()) > 0.38 * h * w: + out = palm_seed | (finger_disk & dil_palm.astype(bool)) + + return out.astype(bool), palm_seed.astype(bool) + + + + + + + + +def _sample_seeds_from_mask(mask: np.ndarray, n: int, rng: np.random.Generator) -> np.ndarray: + ys, xs = np.nonzero(mask) + if ys.size == 0: + return np.zeros((0, 2), dtype=np.float32) + # Mix bottom-biased and spatially spread samples so fingers get prompts too. + n = min(int(n), ys.size) + n_bottom = max(n // 2, 1) + n_spread = n - n_bottom + wts = (ys.astype(np.float64) + 1.0) ** 1.5 + wts /= wts.sum() + idx_b = rng.choice(ys.size, size=n_bottom, replace=False, p=wts) + remain = np.setdiff1d(np.arange(ys.size), idx_b, assume_unique=False) + if n_spread > 0 and remain.size: + # Farthest-point style: pick remaining by y-spread bins. + y_rem = ys[remain] + bins = np.linspace(y_rem.min(), y_rem.max() + 1e-3, num=min(n_spread, remain.size) + 1) + chosen = [] + for i in range(len(bins) - 1): + in_bin = np.where((y_rem >= bins[i]) & (y_rem < bins[i + 1]))[0] + if in_bin.size: + chosen.append(remain[rng.choice(in_bin)]) + if len(chosen) < n_spread: + extra = rng.choice(remain, size=n_spread - len(chosen), replace=False) + chosen.extend(extra.tolist()) + idx = np.concatenate([idx_b, np.asarray(chosen[:n_spread], dtype=np.int64)]) + else: + idx = idx_b + return np.stack([xs[idx], ys[idx]], axis=1).astype(np.float32) + + +def sample_points_from_mask( + mask: np.ndarray, + n_points: int = NUM_WRIST_HAND_TEMPLATE, + *, + rng: np.random.Generator | None = None, +) -> np.ndarray | None: + mask_bool = np.asarray(mask, dtype=bool) + ys, xs = np.nonzero(mask_bool) + if ys.size == 0: + return None + gen = rng if rng is not None else np.random.default_rng() + n_points = int(n_points) + # Spread across y so distal finger regions are represented, not only palm mass. + if ys.size <= n_points: + idx = gen.choice(ys.size, size=n_points, replace=True) + else: + order = np.argsort(ys) + bins = np.array_split(order, min(n_points, ys.size)) + idx = np.array([gen.choice(b) for b in bins if len(b)], dtype=np.int64) + if idx.size < n_points: + extra = gen.choice(ys.size, size=n_points - idx.size, replace=False) + idx = np.concatenate([idx, extra]) + idx = idx[:n_points] + return np.stack([xs[idx], ys[idx]], axis=1).astype(np.float32) + + +def refine_mask_with_sam2( + predictor, + rgb_hwc: np.ndarray, + prop: np.ndarray, + *, + rng: np.random.Generator | None = None, +) -> tuple[np.ndarray, str]: + """Optional SAM2 refine; falls back to prop if SAM drifts off the hand.""" + if predictor is None or not np.asarray(prop).any(): + return np.asarray(prop, dtype=bool), "appearance" + + rgb = np.asarray(rgb_hwc, dtype=np.uint8) + prop = np.asarray(prop, dtype=bool) + h, w = rgb.shape[:2] + rng = rng or np.random.default_rng(0) + + seeds = _sample_seeds_from_mask(prop, 14, rng) + if len(seeds) < 4: + return prop, "appearance" + + ys, xs = np.nonzero(prop) + box = np.array( + [xs.min() - 4, ys.min() - 4, xs.max() + 4, ys.max() + 4], + dtype=np.float32, + ) + box[0] = max(0.0, box[0]) + box[1] = max(0.0, box[1]) + box[2] = min(float(w - 1), box[2]) + box[3] = min(float(h - 1), box[3]) + + neg = np.array( + [ + [0.08 * w, 0.06 * h], + [0.50 * w, 0.06 * h], + [0.92 * w, 0.06 * h], + [0.50 * w, max(0.0, float(box[1] - 20))], + ], + dtype=np.float32, + ) + coords = np.concatenate([seeds, neg], axis=0) + labels = np.concatenate( + [np.ones(len(seeds), dtype=np.int32), np.zeros(len(neg), dtype=np.int32)] + ) + + dil_prop = cv2.dilate(prop.astype(np.uint8), np.ones((15, 15), np.uint8)).astype(bool) + seed_img = np.zeros((h, w), np.uint8) + for px, py in seeds.astype(int): + seed_img[int(np.clip(py, 0, h - 1)), int(np.clip(px, 0, w - 1))] = 255 + dist = cv2.distanceTransform((seed_img == 0).astype(np.uint8), cv2.DIST_L2, 5) + + predictor.set_image(rgb) + masks, scores, _ = predictor.predict( + point_coords=coords, + point_labels=labels, + box=box, + multimask_output=True, + normalize_coords=False, + ) + + best = None + best_rank = -np.inf + for mask, score in zip(masks, scores): + m = largest_connected_component(np.asarray(mask, dtype=bool) & dil_prop & (dist <= 50)) + area = int(m.sum()) + if area < 300: + continue + ov = float((m & prop).sum()) / max(area, 1) + hit = sum( + 1 + for px, py in seeds.astype(int) + if m[int(np.clip(py, 0, h - 1)), int(np.clip(px, 0, w - 1))] + ) + pr = hit / max(len(seeds), 1) + if ov < SAM2_MIN_SEED_OVERLAP or pr < SAM2_MIN_POINT_RECALL: + continue + rank = 2.5 * pr + 2.0 * ov + 0.3 * float(score) + if rank > best_rank: + best_rank = rank + best = m + + if best is None: + return prop, "appearance" + return best, "appearance_sam2" + + +def sam2_wrist_hand_queries( + predictor, + frame0_rgb: np.ndarray, + image_hw: tuple[int, int], + view: str, + *, + n_points: int = NUM_WRIST_HAND_TEMPLATE, + seed: int | None = None, +) -> tuple[np.ndarray, str, np.ndarray | None]: + """ + Segment wrist-view mechanical hand, then sample CoTracker query points. + + Returns (queries_xy, source_tag, mask_or_none). + """ + template = wrist_hand_template_pixels(image_hw, view=view).astype(np.float32) + rgb = np.asarray(frame0_rgb, dtype=np.uint8) + if (rgb.shape[0], rgb.shape[1]) != tuple(image_hw): + rgb = cv2.resize(rgb, (image_hw[1], image_hw[0]), interpolation=cv2.INTER_LINEAR) + + prop, _seed = detect_wrist_hand_mask(rgb, view=view) + if not prop.any(): + return template, "template_fallback_no_hand", None + + rng = np.random.default_rng(seed) + mask, tag = refine_mask_with_sam2(predictor, rgb, prop, rng=rng) + if not mask.any(): + return template, "template_fallback_empty_mask", None + + pts = sample_points_from_mask(mask, n_points=n_points, rng=rng) + if pts is None: + return template, "template_fallback_sample", mask + + pts[..., 0] = np.clip(pts[..., 0], 0.0, float(image_hw[1] - 1)) + pts[..., 1] = np.clip(pts[..., 1], 0.0, float(image_hw[0] - 1)) + return pts, f"{tag}_random22", mask + + +def save_sam2_masks( + out_dir: str | Path, + episode_index: int, + masks_by_view: dict[str, np.ndarray | None], +) -> list[Path]: + root = Path(out_dir).expanduser().resolve() + root.mkdir(parents=True, exist_ok=True) + written: list[Path] = [] + for view, mask in masks_by_view.items(): + out_path = root / f"episode_{episode_index:06d}_{view}.png" + if mask is None: + continue + mask_u8 = (np.asarray(mask, dtype=bool) * 255).astype(np.uint8) + cv2.imwrite(str(out_path), mask_u8) + written.append(out_path) + return written diff --git a/scripts/trex_track/trex_fk.py b/scripts/trex_track/trex_fk.py new file mode 100644 index 0000000000000000000000000000000000000000..d48788008b28d2ebdc7bd9ba1c5217542fd540bf --- /dev/null +++ b/scripts/trex_track/trex_fk.py @@ -0,0 +1,187 @@ +"""T-Rex Pinocchio FK helpers (no dexmate_urdf / coal install required).""" + +from __future__ import annotations + +import sys +import types +from functools import lru_cache +from pathlib import Path +from typing import Callable + +import numpy as np +import pinocchio as pin +from pinocchio.robot_wrapper import RobotWrapper + +_TREX_ROOT = Path("/scratch1/home/zhicao/T-Rex") +_QS_SRC = _TREX_ROOT / "dataset_quickstart" / "src" +_DEXMATE_PKG = _TREX_ROOT / "hardware_code" / "third_party" / "dexmate-urdf" / "src" +_VEGA_DIR = _DEXMATE_PKG / "dexmate_urdf" / "robots" / "humanoid" / "vega_1" +VEGA_URDF = str(_VEGA_DIR / "vega_1.urdf") +VEGA_PKG = str(_VEGA_DIR) +VEGA_SRDF = str(_VEGA_DIR / "vega_1.srdf") + +DEFAULT_TORSO = np.array([0.9, 1.57, 0.1]) +DEFAULT_HEAD = np.array([0.28, 0.0, 0.0]) + +SHARPA_LEFT_HAND_JOINT_ORDER = [ + "left_thumb_CMC_FE", + "left_thumb_CMC_AA", + "left_thumb_MCP_FE", + "left_thumb_MCP_AA", + "left_thumb_IP", + "left_index_MCP_FE", + "left_index_MCP_AA", + "left_index_PIP", + "left_index_DIP", + "left_middle_MCP_FE", + "left_middle_MCP_AA", + "left_middle_PIP", + "left_middle_DIP", + "left_ring_MCP_FE", + "left_ring_MCP_AA", + "left_ring_PIP", + "left_ring_DIP", + "left_pinky_CMC", + "left_pinky_MCP_FE", + "left_pinky_MCP_AA", + "left_pinky_PIP", + "left_pinky_DIP", +] +SHARPA_RIGHT_HAND_JOINT_ORDER = [ + "right_thumb_CMC_FE", + "right_thumb_CMC_AA", + "right_thumb_MCP_FE", + "right_thumb_MCP_AA", + "right_thumb_IP", + "right_index_MCP_FE", + "right_index_MCP_AA", + "right_index_PIP", + "right_index_DIP", + "right_middle_MCP_FE", + "right_middle_MCP_AA", + "right_middle_PIP", + "right_middle_DIP", + "right_ring_MCP_FE", + "right_ring_MCP_AA", + "right_ring_PIP", + "right_ring_DIP", + "right_pinky_CMC", + "right_pinky_MCP_FE", + "right_pinky_MCP_AA", + "right_pinky_PIP", + "right_pinky_DIP", +] + + +def _bootstrap_trex_imports() -> None: + if "coal" not in sys.modules: + coal = types.ModuleType("coal") + coal.Box = lambda *args, **kwargs: None + sys.modules["coal"] = coal + + vega_1 = types.SimpleNamespace( + urdf=VEGA_URDF, + _parent_dir=VEGA_PKG, + srdf=VEGA_SRDF, + ) + vega_1_mod = types.ModuleType("dexmate_urdf.robots.humanoid.vega_1") + vega_1_mod.vega_1 = vega_1 + humanoid = types.ModuleType("dexmate_urdf.robots.humanoid") + humanoid.vega_1 = vega_1_mod + robots = types.ModuleType("dexmate_urdf.robots") + robots.humanoid = humanoid + dexmate_urdf = types.ModuleType("dexmate_urdf") + dexmate_urdf.robots = robots + for name, mod in [ + ("dexmate_urdf", dexmate_urdf), + ("dexmate_urdf.robots", robots), + ("dexmate_urdf.robots.humanoid", humanoid), + ("dexmate_urdf.robots.humanoid.vega_1", vega_1_mod), + ]: + sys.modules.setdefault(name, mod) + + qs = str(_QS_SRC) + if qs not in sys.path: + sys.path.insert(0, qs) + + +@lru_cache(maxsize=1) +def _load_robot_modules(): + _bootstrap_trex_imports() + from trex_dataset_quickstart.robot import ( # noqa: WPS433 + build_full_robot, + build_reduced_bimanual_robot, + forward_kinematics, + ) + + return build_full_robot, build_reduced_bimanual_robot, forward_kinematics + + +@lru_cache(maxsize=1) +def get_full_robot() -> tuple[RobotWrapper, Callable, Callable]: + build_full_robot, _, _ = _load_robot_modules() + return build_full_robot({"head": DEFAULT_HEAD, "torso": DEFAULT_TORSO}) + + +@lru_cache(maxsize=1) +def get_bimanual_robot() -> tuple[RobotWrapper, Callable, Callable]: + _, build_reduced_bimanual_robot, _ = _load_robot_modules() + return build_reduced_bimanual_robot({"head": DEFAULT_HEAD, "torso": DEFAULT_TORSO}) + + +def state_to_components(state58: np.ndarray) -> dict[str, np.ndarray]: + s = np.asarray(state58, dtype=np.float64).reshape(-1) + if s.shape[0] != 58: + raise ValueError(f"expected 58-dim state, got {s.shape[0]}") + return { + "left_arm": s[0:7], + "left_hand": s[7:29], + "right_arm": s[29:36], + "right_hand": s[36:58], + } + + +def hand_joint_positions_3d( + robot: RobotWrapper, + qpos: np.ndarray, + joint_names: list[str], +) -> np.ndarray: + pin.forwardKinematics(robot.model, robot.data, qpos) + pin.updateFramePlacements(robot.model, robot.data) + pts = np.zeros((len(joint_names), 3), dtype=np.float64) + for i, name in enumerate(joint_names): + fid = robot.model.getFrameId(name, pin.FrameType.JOINT) + pts[i] = robot.data.oMf[fid].translation + return pts + + +def frame_pose_matrix(robot: RobotWrapper, qpos: np.ndarray, frame_name: str) -> np.ndarray: + pin.forwardKinematics(robot.model, robot.data, qpos) + pin.updateFramePlacements(robot.model, robot.data) + fid = robot.model.getFrameId(frame_name) + se3 = robot.data.oMf[fid] + T = np.eye(4, dtype=np.float64) + T[:3, :3] = se3.rotation + T[:3, 3] = se3.translation + return T + + +def hand_keypoints_world( + state58: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Return left/right hand (22,3) keypoints and L_ee/R_ee cam-to-world 4x4.""" + comps = state_to_components(state58) + full_robot, assemble_full, _ = get_full_robot() + bimanual_robot, assemble_bimanual, _ = get_bimanual_robot() + + q_full = assemble_full(comps) + q_bi = assemble_bimanual( + {"left_arm": comps["left_arm"], "right_arm": comps["right_arm"]}, + ) + + left_pts = hand_joint_positions_3d(full_robot, q_full, SHARPA_LEFT_HAND_JOINT_ORDER) + right_pts = hand_joint_positions_3d(full_robot, q_full, SHARPA_RIGHT_HAND_JOINT_ORDER) + T_left_ee = frame_pose_matrix(full_robot, q_full, "L_ee") + T_right_ee = frame_pose_matrix(full_robot, q_full, "R_ee") + T_head_cam = frame_pose_matrix(bimanual_robot, q_bi, "zed_left_camera") + return left_pts, right_pts, T_head_cam, (T_left_ee, T_right_ee) diff --git a/scripts/trex_track/trex_projection.py b/scripts/trex_track/trex_projection.py new file mode 100644 index 0000000000000000000000000000000000000000..512fd6f6651025bd68f13b540207f0fe200d55d7 --- /dev/null +++ b/scripts/trex_track/trex_projection.py @@ -0,0 +1,296 @@ +"""Project T-Rex hand keypoints with camera intrinsics / extrinsics.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import cv2 +import numpy as np +from scipy.spatial.transform import Rotation + +EULER_CONVENTION = "xyz" +TREX_NATIVE_HW = (360, 640) # (H, W) — robot stream resolution before v2 downscale + + +def pose6_to_matrix(pose6: np.ndarray, *, euler: str = EULER_CONVENTION) -> np.ndarray: + x = np.asarray(pose6, dtype=np.float64).reshape(-1) + rot = Rotation.from_euler(euler, x[3:6]).as_matrix() + T = np.eye(4, dtype=np.float64) + T[:3, :3] = rot + T[:3, 3] = x[:3] + return T + + +def matrix_to_pose6(T: np.ndarray, *, euler: str = EULER_CONVENTION) -> np.ndarray: + T = np.asarray(T, dtype=np.float64).reshape(4, 4) + xyz = T[:3, 3] + rpy = Rotation.from_matrix(T[:3, :3]).as_euler(euler) + return np.concatenate([xyz, rpy], axis=0) + + +def scale_K_to_image_hw( + K: np.ndarray, + native_hw: tuple[int, int], + image_hw: tuple[int, int], +) -> np.ndarray: + href, wref = native_hw + hnew, wnew = image_hw + K = np.asarray(K, dtype=np.float64).reshape(3, 3).copy() + K[0, 0] *= wnew / float(wref) + K[0, 2] *= wnew / float(wref) + K[1, 1] *= hnew / float(href) + K[1, 2] *= hnew / float(href) + return K + + +def scale_pixels_between_hw( + pixels: np.ndarray, + src_hw: tuple[int, int], + dst_hw: tuple[int, int], +) -> np.ndarray: + out = np.asarray(pixels, dtype=np.float64).copy() + sh, sw = src_hw + dh, dw = dst_hw + out[..., 0] *= dw / float(sw) + out[..., 1] *= dh / float(sh) + return out + + +def make_image_grid(img_h: int, img_w: int, grid_size: int = 5) -> np.ndarray: + xs = np.linspace(0.0, float(img_w - 1), grid_size, dtype=np.float64) + ys = np.linspace(0.0, float(img_h - 1), grid_size, dtype=np.float64) + return np.stack(np.meshgrid(xs, ys, indexing="xy"), axis=-1).reshape(-1, 2) + + +# Canonical 22-joint hand layout in a unit box (x,y in [0,1], palm down, fingers up). +_HAND_TEMPLATE_UNIT = np.array( + [ + [0.18, 0.82], + [0.22, 0.72], + [0.26, 0.62], + [0.30, 0.52], + [0.34, 0.42], + [0.42, 0.88], + [0.46, 0.74], + [0.50, 0.60], + [0.54, 0.46], + [0.50, 0.90], + [0.54, 0.76], + [0.58, 0.62], + [0.62, 0.48], + [0.58, 0.88], + [0.62, 0.74], + [0.66, 0.60], + [0.70, 0.46], + [0.72, 0.86], + [0.76, 0.72], + [0.80, 0.58], + [0.84, 0.44], + [0.88, 0.34], + ], + dtype=np.float64, +) + +# T-Rex wrist cameras: hand sits in the lower band; left/right are mirrored in x. +# Anchors tuned on trex_small episode_000000 frame-0 hand blob. +_WRIST_HAND_VIEW_ANCHORS = { + "left_wrist": {"cx": 0.60, "cy": 0.66, "sx": 0.36, "sy": 0.44, "mirror_x": False}, + "right_wrist": {"cx": 0.50, "cy": 0.66, "sx": 0.38, "sy": 0.46, "mirror_x": True}, +} + + +def _build_wrist_hand_template_frac(view: str) -> np.ndarray: + if view not in _WRIST_HAND_VIEW_ANCHORS: + raise KeyError(f"unknown wrist view {view!r}") + cfg = _WRIST_HAND_VIEW_ANCHORS[view] + pts = _HAND_TEMPLATE_UNIT.copy() + if cfg["mirror_x"]: + pts[:, 0] = 1.0 - pts[:, 0] + pts[:, 0] = (pts[:, 0] - 0.5) * cfg["sx"] + cfg["cx"] + pts[:, 1] = (pts[:, 1] - 0.5) * cfg["sy"] + cfg["cy"] + pts[:, 0] = np.clip(pts[:, 0], 0.02, 0.98) + pts[:, 1] = np.clip(pts[:, 1], 0.02, 0.98) + return pts + + +WRIST_HAND_TEMPLATE_FRAC_BY_VIEW = { + view: _build_wrist_hand_template_frac(view) for view in _WRIST_HAND_VIEW_ANCHORS +} +# Back-compat alias (left wrist). +WRIST_HAND_TEMPLATE_FRAC = WRIST_HAND_TEMPLATE_FRAC_BY_VIEW["left_wrist"] +NUM_WRIST_HAND_TEMPLATE = int(_HAND_TEMPLATE_UNIT.shape[0]) + +# FK seeds are only trusted when enough points land in the lower hand band. +WRIST_HAND_FK_MIN_VALID = 8 +WRIST_HAND_FK_MIN_MEAN_Y_FRAC = 0.45 + + +def wrist_hand_template_pixels(image_hw: tuple[int, int], view: str = "left_wrist") -> np.ndarray: + h, w = image_hw + frac = WRIST_HAND_TEMPLATE_FRAC_BY_VIEW[view] + return (frac * np.array([float(w), float(h)], dtype=np.float64)).astype(np.float64) + + +def wrist_hand_prompt_box(image_hw: tuple[int, int], view: str = "left_wrist") -> np.ndarray: + """XYXY box prompt for SAM2, derived from per-view hand anchors.""" + if view not in _WRIST_HAND_VIEW_ANCHORS: + raise KeyError(f"unknown wrist view {view!r}") + h, w = image_hw + cfg = _WRIST_HAND_VIEW_ANCHORS[view] + cx = float(cfg["cx"]) * float(w) + cy = float(cfg["cy"]) * float(h) + half_w = float(cfg["sx"]) * float(w) * 0.5 + half_h = float(cfg["sy"]) * float(h) * 0.5 + x0 = np.clip(cx - half_w, 0.0, float(w - 1)) + y0 = np.clip(cy - half_h, 0.0, float(h - 1)) + x1 = np.clip(cx + half_w, 0.0, float(w - 1)) + y1 = np.clip(cy + half_h, 0.0, float(h - 1)) + if x1 <= x0: + x1 = min(float(w - 1), x0 + 1.0) + if y1 <= y0: + y1 = min(float(h - 1), y0 + 1.0) + return np.array([x0, y0, x1, y1], dtype=np.float32) + + +def fk_hand_queries_trusted( + queries: np.ndarray, + valid: np.ndarray, + image_hw: tuple[int, int], +) -> bool: + """Reject wrist FK seeds that sit in the background instead of on the hand.""" + valid = np.asarray(valid, dtype=bool).reshape(-1) + if int(valid.sum()) < WRIST_HAND_FK_MIN_VALID: + return False + h = max(float(image_hw[0]), 1.0) + mean_y_frac = float(np.asarray(queries, dtype=np.float32)[valid, 1].mean() / h) + return mean_y_frac >= WRIST_HAND_FK_MIN_MEAN_Y_FRAC + + +def queries_in_image( + pixels: np.ndarray, + vis: np.ndarray, + image_hw: tuple[int, int], +) -> np.ndarray: + h, w = image_hw + pts = np.asarray(pixels, dtype=np.float64) + mask = np.asarray(vis, dtype=bool).reshape(-1) + mask &= np.isfinite(pts).all(axis=1) + mask &= (pts[:, 0] >= 0.0) & (pts[:, 0] < float(w)) + mask &= (pts[:, 1] >= 0.0) & (pts[:, 1] < float(h)) + return mask + + +def align_hand_queries_to_template( + queries: np.ndarray, + valid: np.ndarray, + image_hw: tuple[int, int], + *, + view: str = "left_wrist", +) -> np.ndarray: + """Keep valid frame-0 projections; place invalid joints from a hand template.""" + out = wrist_hand_template_pixels(image_hw, view=view).astype(np.float32, copy=True) + valid = np.asarray(valid, dtype=bool).reshape(-1) + q = np.asarray(queries, dtype=np.float32) + if valid.any(): + vpts = q[valid] + vmin, vmax = vpts.min(axis=0), vpts.max(axis=0) + tmpl = out.copy() + tmin, tmax = tmpl.min(axis=0), tmpl.max(axis=0) + scale = (vmax - vmin) / np.maximum(tmax - tmin, 1e-3) + out = (tmpl - tmin) * scale + vmin + out[valid] = q[valid] + out[..., 0] = np.clip(out[..., 0], 0.0, float(image_hw[1] - 1)) + out[..., 1] = np.clip(out[..., 1], 0.0, float(image_hw[0] - 1)) + return out + + +def project_world_points( + world_pts: np.ndarray, + cam_to_world: np.ndarray, + *, + K: np.ndarray, + dist: np.ndarray | None = None, + native_hw: tuple[int, int] = TREX_NATIVE_HW, + out_hw: tuple[int, int], +) -> tuple[np.ndarray, np.ndarray]: + """Project (N,3) world points -> (N,2) pixels and visibility mask.""" + world = np.asarray(world_pts, dtype=np.float64).reshape(-1, 3) + T_c2w = np.asarray(cam_to_world, dtype=np.float64).reshape(4, 4) + w2c = np.linalg.inv(T_c2w) + R = w2c[:3, :3] + t = w2c[:3, 3].reshape(3, 1) + + K_native = np.asarray(K, dtype=np.float64).reshape(3, 3) + dist_use = None + if dist is not None: + dist_use = np.asarray(dist, dtype=np.float64).reshape(-1) + if dist_use.size >= 5 and np.allclose(dist_use[:5], 0.0): + dist_use = None + + rvec, _ = cv2.Rodrigues(R) + if dist_use is not None: + pts2d, _ = cv2.projectPoints( + world.reshape(-1, 1, 3), + rvec, + t, + K_native, + dist_use[:5], + ) + else: + pts2d, _ = cv2.projectPoints(world.reshape(-1, 1, 3), rvec, t, K_native, np.zeros(5)) + pixels_native = pts2d.reshape(-1, 2).astype(np.float64) + + cam_z = (R @ world.T + t).T[:, 2] + vis = cam_z > 0.01 + + pixels_out = scale_pixels_between_hw(pixels_native, native_hw, out_hw) + href, wref = native_hw + dh, dw = out_hw + vis &= ( + (pixels_native[:, 0] >= 0.0) + & (pixels_native[:, 0] < float(wref)) + & (pixels_native[:, 1] >= 0.0) + & (pixels_native[:, 1] < float(href)) + ) + vis &= ( + (pixels_out[:, 0] >= 0.0) + & (pixels_out[:, 0] < float(dw)) + & (pixels_out[:, 1] >= 0.0) + & (pixels_out[:, 1] < float(dh)) + ) + return pixels_out.astype(np.float32), vis.astype(np.float32) + + +def wrist_cam_to_world(T_ee: np.ndarray, T_ee_to_cam: np.ndarray) -> np.ndarray: + return np.asarray(T_ee, dtype=np.float64) @ np.asarray(T_ee_to_cam, dtype=np.float64) + + +def load_camera_calib(path: str | Path) -> dict: + p = Path(path).expanduser().resolve() + with p.open("r", encoding="utf-8") as f: + return json.load(f) + + +def view_intrinsics( + calib: dict, + view: str, + out_hw: tuple[int, int], +) -> tuple[np.ndarray, np.ndarray | None]: + native_hw = tuple(calib.get("native_hw", TREX_NATIVE_HW)) + view_cfg = calib["views"][view] + K = np.asarray(view_cfg["K_native"], dtype=np.float64).reshape(3, 3) + dist = view_cfg.get("dist") + if dist is not None: + dist = np.asarray(dist, dtype=np.float64) + K_out = scale_K_to_image_hw(K, native_hw, out_hw) + return K_out, dist + + +def view_ee_to_cam(calib: dict, view: str) -> np.ndarray | None: + view_cfg = calib["views"][view] + if "T_ee_to_cam" in view_cfg: + return np.asarray(view_cfg["T_ee_to_cam"], dtype=np.float64).reshape(4, 4) + if "ee_to_cam_pose6" in view_cfg: + return pose6_to_matrix(np.asarray(view_cfg["ee_to_cam_pose6"], dtype=np.float64)) + return None diff --git a/scripts/trex_track/trex_viz_tracks.py b/scripts/trex_track/trex_viz_tracks.py new file mode 100644 index 0000000000000000000000000000000000000000..a58bd9678a31bc93fe24289951485fcfe3ce73d9 --- /dev/null +++ b/scripts/trex_track/trex_viz_tracks.py @@ -0,0 +1,121 @@ +"""Draw T-Rex three-view tracks (head_left | left_wrist | right_wrist).""" + +from __future__ import annotations + +from pathlib import Path + +import cv2 +import numpy as np + +from trex_track.layout import ( + NUM_HEAD_LEFT, + NUM_WRIST_BACKGROUND as NUM_WRIST_GRID, + VIEW_ORDER, +) + +_COLOR_LEFT_HAND = (0, 255, 0) # green +_COLOR_RIGHT_HAND = (255, 128, 0) # blue-orange +_COLOR_WRIST_GRID = (0, 255, 255) # yellow +_COLOR_WRIST_HAND = (255, 0, 255) # purple + + +def _clamp_tracks_pixels(tracks: np.ndarray, img_w: int, img_h: int, *, normalized: bool) -> np.ndarray: + trk = np.asarray(tracks, dtype=np.float32).copy() + if normalized: + trk[..., 0] *= float(img_w) + trk[..., 1] *= float(img_h) + trk[..., 0] = np.clip(trk[..., 0], 0.0, float(img_w - 1)) + trk[..., 1] = np.clip(trk[..., 1], 0.0, float(img_h - 1)) + return trk + + +def _bgr_for_point(view: str, i: int, vis_val: float | None, dim_low_vis: bool) -> tuple[int, int, int]: + if view == "head_left": + color = _COLOR_LEFT_HAND if i < NUM_HEAD_LEFT else _COLOR_RIGHT_HAND + elif view in ("left_wrist", "right_wrist"): + color = _COLOR_WRIST_GRID if i < NUM_WRIST_GRID else _COLOR_WRIST_HAND + else: + color = (255, 255, 255) + if dim_low_vis and vis_val is not None and float(vis_val) < 0.5: + return tuple(int(c * 0.55) for c in color) + return color + + +def _draw_tracks_on_frame_bgr( + frame_bgr: np.ndarray, + *, + tracks_pixels: np.ndarray, + vis: np.ndarray | None, + t: int, + t0: int, + view: str, + dim_low_vis: bool, +) -> None: + draw_trk = tracks_pixels + for i in range(int(draw_trk.shape[1])): + vv = float(vis[t, i]) if vis is not None else None + color = _bgr_for_point(view, i, vv, dim_low_vis) + x, y = float(draw_trk[t, i, 0]), float(draw_trk[t, i, 1]) + if not np.isfinite(x) or not np.isfinite(y): + continue + if t > t0: + pts = [ + (int(round(float(draw_trk[tt, i, 0]))), int(round(float(draw_trk[tt, i, 1])))) + for tt in range(t0, t + 1) + if np.isfinite(draw_trk[tt, i]).all() + ] + if len(pts) >= 2: + cv2.polylines(frame_bgr, [np.array(pts, dtype=np.int32)], False, color, 2, cv2.LINE_AA) + cv2.circle(frame_bgr, (int(round(x)), int(round(y))), 4, color, -1, cv2.LINE_AA) + + +def render_three_view_combined_video( + *, + view_images: dict[str, np.ndarray], + view_tracks: dict[str, np.ndarray], + view_vis: dict[str, np.ndarray], + out_path: Path, + fps: int, + draw_trail: int, + dim_low_vis: bool = False, +) -> None: + for v in VIEW_ORDER: + if v not in view_images or v not in view_tracks: + raise KeyError(f"missing view {v!r}") + + t_end = min(int(view_images[v].shape[0]) for v in VIEW_ORDER) + panel_h, panel_w = view_images[VIEW_ORDER[0]].shape[1:3] + out_w, out_h = panel_w * 3, panel_h + trail_len = max(0, int(draw_trail)) + tracks_pixels = { + view: _clamp_tracks_pixels( + view_tracks[view], + panel_w, + panel_h, + normalized=True, + ) + for view in VIEW_ORDER + } + + fourcc = cv2.VideoWriter_fourcc(*"mp4v") + out_path.parent.mkdir(parents=True, exist_ok=True) + writer = cv2.VideoWriter(str(out_path), fourcc, float(fps), (out_w, out_h)) + if not writer.isOpened(): + raise RuntimeError(f"VideoWriter failed: {out_path}") + + for t in range(t_end): + panels = [] + for view in VIEW_ORDER: + frame_bgr = cv2.cvtColor(np.ascontiguousarray(view_images[view][t]), cv2.COLOR_RGB2BGR) + _draw_tracks_on_frame_bgr( + frame_bgr, + tracks_pixels=tracks_pixels[view], + vis=view_vis.get(view), + t=t, + t0=max(0, t - trail_len), + view=view, + dim_low_vis=dim_low_vis, + ) + panels.append(frame_bgr) + writer.write(np.concatenate(panels, axis=1)) + writer.release() diff --git a/socket_test_optimized_AR.py b/socket_test_optimized_AR.py new file mode 100644 index 0000000000000000000000000000000000000000..5b93e05e9a8f7f873b81541ed8eddf4cb47f83ad --- /dev/null +++ b/socket_test_optimized_AR.py @@ -0,0 +1,836 @@ +import dataclasses +import logging +import socket +import asyncio +import os +import http +import logging +import time +import traceback +import torch +import tyro +from einops import rearrange +import datetime + +from groot.vla.model.n1_5.sim_policy import GrootSimPolicy +from groot.vla.data.schema import EmbodimentTag +import imageio +import numpy as np + +from openpi_client import base_policy as _base_policy +from openpi_client import msgpack_numpy +import websockets.asyncio.server as _server +import websockets.frames +from tianshou.data import Batch +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh, init_device_mesh + +# Use roboarena policy server interface +from eval_utils.policy_server import WebsocketPolicyServer as RoboarenaServer +from eval_utils.policy_server import PolicyServerConfig + +logger = logging.getLogger(__name__) + +@dataclasses.dataclass +class Args: + port: int = 8000 + timeout_seconds: int = 50000 # 10 hours default, configurable + model_path: str = "./checkpoints/dreamzero" + enable_dit_cache: bool = False + index: int = 0 + max_chunk_size: int | None = None # If None, use config value. Otherwise override max_chunk_size for inference. + + +class ARDroidRoboarenaPolicy: + """Wrapper policy that implements roboarena.policy.BasePolicy interface for AR_droid. + + Handles: + - Observation format conversion (roboarena -> AR_droid format) + - Frame accumulation across calls (roboarena sends single frames, AR_droid expects multi-frame video) + - Action format conversion (AR_droid dict -> roboarena array format) + - Distributed inference coordination + """ + + # Number of frames to accumulate after the first call + FRAMES_PER_CHUNK = 4 + + def __init__( + self, + groot_policy: GrootSimPolicy, + signal_group: dist.ProcessGroup, + output_dir: str | None = None, + ) -> None: + self._policy = groot_policy + self._signal_group = signal_group + self._output_dir = output_dir + + # Frame buffers for accumulation (per camera view) + self._frame_buffers: dict[str, list[np.ndarray]] = { + "video.exterior_image_1_left": [], + "video.exterior_image_2_left": [], + "video.wrist_image_left": [], + } + self._call_count = 0 + self._is_first_call = True + + # Session tracking - reset state when new session starts + self._current_session_id: str | None = None + + # Video across time for saving (similar to original server) + self.video_across_time = [] + self._msg_index = 0 + + # Create output directory if specified + if self._output_dir: + os.makedirs(self._output_dir, exist_ok=True) + + def _convert_observation(self, obs: dict) -> dict: + """Convert roboarena observation format to AR_droid format. + + Roboarena format: + - observation/exterior_image_0_left: (H, W, 3) single frame + - observation/exterior_image_1_left: (H, W, 3) single frame + - observation/wrist_image_left: (H, W, 3) single frame + - observation/joint_position: (7,) + - observation/gripper_position: (1,) + - prompt: str + + AR_droid format: + - video.exterior_image_1_left: (T, H, W, 3) multi-frame + - video.exterior_image_2_left: (T, H, W, 3) multi-frame + - video.wrist_image_left: (T, H, W, 3) multi-frame + - state.joint_position: (1, 7) + - state.gripper_position: (1, 1) + - annotation.language.action_text: str + """ + converted = {} + + # Map image keys (roboarena uses 0-indexed, AR_droid uses 1-indexed) + image_key_mapping = { + "observation/exterior_image_0_left": "video.exterior_image_1_left", + "observation/exterior_image_1_left": "video.exterior_image_2_left", + "observation/wrist_image_left": "video.wrist_image_left", + } + + # Accumulate frames for each camera view + for roboarena_key, droid_key in image_key_mapping.items(): + if roboarena_key in obs: + data = obs[roboarena_key] + if isinstance(data, np.ndarray): + if data.ndim == 4: + # Multiple frames (T, H, W, 3) + self._frame_buffers[droid_key].extend(list(data)) + else: + # Single frame (H, W, 3) + self._frame_buffers[droid_key].append(data) + + # Determine how many frames to use + if self._is_first_call: + # First call: use only 1 frame + num_frames = 1 + else: + # Subsequent calls: use exactly FRAMES_PER_CHUNK frames + num_frames = self.FRAMES_PER_CHUNK + + # Build video tensors from accumulated frames + for droid_key, buffer in self._frame_buffers.items(): + if len(buffer) > 0: + if len(buffer) >= num_frames: + # Take the last num_frames frames + frames_to_use = buffer[-num_frames:] + else: + # Pad by repeating the first frame to reach num_frames + frames_to_use = buffer.copy() + while len(frames_to_use) < num_frames: + # Prepend the first frame to pad + frames_to_use.insert(0, buffer[0]) + # Stack to (T, H, W, C) + video = np.stack(frames_to_use, axis=0) + converted[droid_key] = video + + # Convert state observations + if "observation/joint_position" in obs: + joint_pos = obs["observation/joint_position"] + # Reshape to (1, 7) if needed + if joint_pos.ndim == 1: + joint_pos = joint_pos.reshape(1, -1) + converted["state.joint_position"] = joint_pos.astype(np.float64) + else: + converted["state.joint_position"] = np.zeros((1, 7), dtype=np.float64) + + if "observation/gripper_position" in obs: + gripper_pos = obs["observation/gripper_position"] + # Reshape to (1, 1) if needed + if gripper_pos.ndim == 1: + gripper_pos = gripper_pos.reshape(1, -1) + converted["state.gripper_position"] = gripper_pos.astype(np.float64) + else: + converted["state.gripper_position"] = np.zeros((1, 1), dtype=np.float64) + + # Convert prompt + if "prompt" in obs: + converted["annotation.language.action_text"] = obs["prompt"] + else: + converted["annotation.language.action_text"] = "" + + return converted + + def _convert_action(self, action_dict: dict) -> np.ndarray: + """Convert AR_droid action dict to roboarena action array. + + AR_droid format: + - action.joint_position: (N, 7) + - action.gripper_position: (N,) or (N, 1) + + Roboarena format: + - action: (N, 8) - 7 joint positions + 1 gripper + """ + joint_action = None + gripper_action = None + + # Extract actions from dict + for key, value in action_dict.items(): + if "joint_position" in key: + joint_action = value + elif "gripper_position" in key or "gripper" in key: + gripper_action = value + + if joint_action is None: + # Fallback: return zeros + return np.zeros((1, 8), dtype=np.float32) + + # Convert to numpy if tensor + if isinstance(joint_action, torch.Tensor): + joint_action = joint_action.cpu().numpy() + + # Ensure 2D shape (N, 7) + if joint_action.ndim == 1: + joint_action = joint_action.reshape(1, -1) + + N = joint_action.shape[0] + + # Handle gripper action + if gripper_action is not None: + if isinstance(gripper_action, torch.Tensor): + gripper_action = gripper_action.cpu().numpy() + # Reshape to (N, 1) if needed + if gripper_action.ndim == 1: + gripper_action = gripper_action.reshape(-1, 1) + elif gripper_action.ndim == 0: + gripper_action = gripper_action.reshape(1, 1) + else: + gripper_action = np.zeros((N, 1), dtype=np.float32) + + # Concatenate: (N, 7) + (N, 1) -> (N, 8) + action = np.concatenate([joint_action, gripper_action], axis=-1).astype(np.float32) + + return action + + def _broadcast_batch_to_workers(self, obs: dict) -> None: + """Broadcast batch data from rank 0 to all other ranks.""" + import pickle + + # Serialize the obs + serialized = pickle.dumps(obs) + data_size = len(serialized) + + # Broadcast size first + size_tensor = torch.tensor([data_size], dtype=torch.int64, device='cuda') + dist.broadcast(size_tensor, src=0) + + # Broadcast data + data_tensor = torch.frombuffer(serialized, dtype=torch.uint8).cuda() + dist.broadcast(data_tensor, src=0) + + def infer(self, obs: dict) -> np.ndarray: + """Infer actions from observations. + + Args: + obs: Observation dict in roboarena format + + Returns: + action: (N, 8) action array + """ + # Check for session change - reset state if new session + session_id = obs.get("session_id", None) + if session_id is not None and session_id != self._current_session_id: + if self._current_session_id is not None: + logger.info(f"Session changed from '{self._current_session_id}' to '{session_id}', resetting state") + # Reset state for new session + self._reset_state() + else: + logger.info(f"New session started: '{session_id}'") + self._current_session_id = session_id + + self._msg_index += 1 + self._call_count += 1 + + # Convert observation format + converted_obs = self._convert_observation(obs) + + # Signal workers to continue (0 = continue) + signal_tensor = torch.zeros(1, dtype=torch.int32, device='cpu') + dist.broadcast(signal_tensor, src=0, group=self._signal_group) + + # Broadcast obs to workers + self._broadcast_batch_to_workers(converted_obs) + + # Create batch for policy + batch = Batch(obs=converted_obs) + + # Distributed forward pass + dist.barrier() + with torch.no_grad(): + result_batch, video_pred = self._policy.lazy_joint_forward_causal(batch) + dist.barrier() + + # Store video predictions for potential saving + self.video_across_time.append(video_pred) + + # Extract and convert action + action_chunk_dict = result_batch.act + + # Convert Batch to dict + action_dict = {} + for k in dir(action_chunk_dict): + if k.startswith("action."): + action_dict[k] = getattr(action_chunk_dict, k) + + action = self._convert_action(action_dict) + + # Update first call flag + if self._is_first_call: + self._is_first_call = False + + return action + + def _reset_state(self, save_video: bool = True) -> None: + """Internal method to reset policy state. + + Args: + save_video: Whether to save accumulated video before reset. + """ + # Optionally save accumulated video before reset + if save_video and len(self.video_across_time) > 0 and self._output_dir: + try: + frame_list = [] + video_across_time_cat = torch.cat(self.video_across_time, dim=2) + frames = self._policy.trained_model.action_head.vae.decode( + video_across_time_cat, + tiled=self._policy.trained_model.action_head.tiled, + tile_size=(self._policy.trained_model.action_head.tile_size_height, self._policy.trained_model.action_head.tile_size_width), + tile_stride=(self._policy.trained_model.action_head.tile_stride_height, self._policy.trained_model.action_head.tile_stride_width), + ) + frames = rearrange(frames, "B C T H W -> B T H W C") + frames = frames[0] + frames = ((frames.float() + 1) * 127.5).clip(0, 255).cpu().numpy().astype(np.uint8) + for frame in frames: + frame_list.append(frame) + + if len(frame_list) > 0: + sample_frame = frame_list[0] + if len(sample_frame.shape) == 3 and sample_frame.shape[2] in [1, 3, 4]: + save_dir = self._output_dir + os.makedirs(save_dir, exist_ok=True) + all_mp4_files = [f for f in os.listdir(save_dir) if f.endswith(".mp4")] + timestamp = datetime.datetime.now().strftime("%m_%d_%H_%M_%S") + num_frames = len(frame_list) + n = (num_frames - 1) // 8 + output_path = os.path.join(save_dir, f'{len(all_mp4_files):06}_{timestamp}_n{n}.mp4') + imageio.mimsave(output_path, frame_list, fps=5, codec='libx264') + logger.info(f"Saved video on reset to: {output_path}") + except Exception as e: + logger.warning(f"Failed to save video on reset: {e}") + + # Clear frame buffers + for key in self._frame_buffers: + self._frame_buffers[key] = [] + + self._call_count = 0 + self._is_first_call = True + self.video_across_time = [] + + def reset(self, reset_info: dict) -> None: + """Reset the policy state for a new episode. + + Clears frame buffers and resets call count. + """ + self._reset_state(save_video=True) + + +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: _base_policy.BasePolicy, + host: str = "0.0.0.0", + port: int | None = None, + metadata: dict | None = None, + output_dir: str | None = None, + signal_group: dist.ProcessGroup | None = None, + ) -> None: + self._policy = policy + self._host = host + self._port = port + self._metadata = metadata or {} + self._output_dir = output_dir + logging.getLogger("websockets.server").setLevel(logging.INFO) + self.video_across_time = [] + self._msg_index = 0 + self._signal_group = signal_group + # Create output directory if specified + if self._output_dir: + os.makedirs(self._output_dir, exist_ok=True) + os.makedirs(os.path.join(self._output_dir, "inputs"), exist_ok=True) + + def _save_input_obs(self, obs: dict) -> None: + """Save incoming observation images per message. + + Expected format: THWC (Time, Height, Width, Channel) with 4 frames. + Saves each frame as a separate PNG image: HWC format (uint8). + + Directory structure: + output_dir/inputs/{msg_index:06d}_{timestamp}/{obs_key}/f{frame_idx:02d}.png + """ + if not self._output_dir: + return + timestamp = datetime.datetime.now().strftime("%m_%d_%H_%M_%S") + base_dir = os.path.join(self._output_dir, "inputs", f"{self._msg_index:06d}_{timestamp}") + try: + os.makedirs(base_dir, exist_ok=True) + except Exception: + return + + for key in ("video.exterior_image_1_left", "video.exterior_image_2_left", "video.wrist_image_left"): + if key not in obs: + continue + value = obs[key] + try: + # Convert to numpy if tensor + if isinstance(value, torch.Tensor): + arr = value.detach().cpu().numpy() + else: + arr = np.asarray(value) + + # Expected format: THWC (Time, Height, Width, Channel) + if arr.ndim != 4: + logger.warning(f"obs key '{key}' has shape {arr.shape}, expected 4D (T,H,W,C)") + continue + + # arr is (T, H, W, C) + T, H, W, C = arr.shape + + # Normalize to uint8 + if arr.dtype == np.uint8: + frames_u8 = arr + else: + f = arr.astype(np.float32) + # Common conventions: [-1,1] or [0,1] + min_val = float(np.nanmin(f)) + max_val = float(np.nanmax(f)) + if min_val >= -1.1 and max_val <= 1.1: + # Assume [-1,1] range + frames_u8 = ((f + 1.0) * 127.5).clip(0, 255).astype(np.uint8) + else: + # Min-max scaling + denom = (max_val - min_val) if (max_val - min_val) > 1e-6 else 1.0 + frames_u8 = ((f - min_val) / denom * 255.0).clip(0, 255).astype(np.uint8) + + # Save each frame: frames_u8[i] is (H, W, C) + key_dir = os.path.join(base_dir, key.replace("/", "_")) + os.makedirs(key_dir, exist_ok=True) + for frame_idx in range(T): + frame = frames_u8[frame_idx] # (H, W, C) + # Handle grayscale (H, W) -> (H, W, 1) + if frame.ndim == 2: + frame = np.expand_dims(frame, axis=-1) + imageio.imwrite(os.path.join(key_dir, f"f{frame_idx:02d}.png"), frame) + + except Exception as e: + logger.warning(f"Failed to save obs key '{key}': {e}") + continue + + + + def serve_forever(self, rank: int = 0) -> None: + asyncio.run(self.run(rank)) + + async def run(self, rank: int = 0): + if rank == 0: + async with _server.serve( + self._handler, + self._host, + self._port, + compression=None, + max_size=None, + process_request=_health_check, + ping_interval=None, + ) as server: + await server.serve_forever() + else: + # Non-rank-0 processes run a worker loop + await self._worker_loop() + + async def _worker_loop(self): + """Worker loop for non-rank-0 processes to participate in distributed inference.""" + logger.info(f"Worker loop started for rank {dist.get_rank()}") + signal_tensor = torch.zeros(1, dtype=torch.int32, device='cpu') + while True: + try: + # Wait for obs broadcast from rank 0 + # Create a dummy obs dict structure - will be filled by broadcast + # obs = {} + + dist.broadcast(signal_tensor, src=0, group=self._signal_group) + + signal = signal_tensor.item() + if signal == 1: + logger.info(f"Rank {dist.get_rank()} received shutdown signal") + break + + # --- ADD THIS ELIF BLOCK --- + elif signal == 2: + logger.info(f"Rank {dist.get_rank()} received idle signal. Waiting for next client.") + # Loop back to the top and wait for the next signal + continue + + # Receive the batch data via broadcast/gather mechanism + # This is a simplified version - the actual obs structure needs to be broadcasted + batch = self._receive_batch_from_rank0() + # Participate in distributed forward pass + dist.barrier() + with torch.no_grad(): + result_batch, video_pred = self._policy.lazy_joint_forward_causal(batch) + dist.barrier() + + except Exception as e: + logger.error(f"Worker loop error on rank {dist.get_rank()}: {e}") + traceback.print_exc() + break + + def _receive_batch_from_rank0(self): + """Receive batch data from rank 0 using torch.distributed primitives.""" + import pickle + + # Receive the size of the pickled data first + size_tensor = torch.zeros(1, dtype=torch.int64, device='cuda') + dist.broadcast(size_tensor, src=0) + data_size = size_tensor.item() + + # Receive the actual data + data_tensor = torch.zeros(data_size, dtype=torch.uint8, device='cuda') + dist.broadcast(data_tensor, src=0) + + # Deserialize + obs = pickle.loads(data_tensor.cpu().numpy().tobytes()) + return Batch(obs=obs) + + def _broadcast_batch_to_workers(self, obs): + """Broadcast batch data from rank 0 to all other ranks.""" + import pickle + + # Serialize the obs + serialized = pickle.dumps(obs) + data_size = len(serialized) + + # Broadcast size first + size_tensor = torch.tensor([data_size], dtype=torch.int64, device='cuda') + dist.broadcast(size_tensor, src=0) + + # Broadcast data + data_tensor = torch.frombuffer(serialized, dtype=torch.uint8).cuda() + dist.broadcast(data_tensor, src=0) + + async def _handler(self, websocket: _server.ServerConnection): + logger.info(f"Connection from {websocket.remote_address} opened") + packer = msgpack_numpy.Packer() + + await websocket.send(packer.pack(self._metadata)) + + prev_total_time = None + signal_tensor = torch.zeros(1, dtype=torch.int32, device='cpu') + + try: + while True: + try: + start_time = time.perf_counter() + data = await websocket.recv() + recv_done = time.perf_counter() + obs = msgpack_numpy.unpackb(data) + print(f"Wait Time: {recv_done - start_time:.2f} seconds") + self._msg_index += 1 + + infer_start_time = time.perf_counter() + + # Signal other ranks to continue (0 = continue) + signal_tensor.zero_() + dist.broadcast(signal_tensor, src=0, group=self._signal_group) # <-- USE GLOO GROUP + + # Broadcast the obs to all ranks for distributed inference + self._broadcast_batch_to_workers(obs) + batch = Batch(obs=obs) + + # All ranks need to participate in the forward pass + dist.barrier() + forward_start_time = time.perf_counter() + with torch.no_grad(): + result_batch, video_pred = self._policy.lazy_joint_forward_causal(batch) + dist.barrier() + print(f"Forward Time: {time.perf_counter() - forward_start_time:.2f} seconds") + + action_chunk_dict = result_batch.act + video_chunk = video_pred + + print(f"Inference Time: {time.perf_counter() - infer_start_time:.2f} seconds") + + self.video_across_time.append(video_chunk) + + if len(self.video_across_time) > 10: + frame_list = [] + video_across_time_cat = torch.cat(self.video_across_time, dim=2) + frames = self._policy.trained_model.action_head.vae.decode( + video_across_time_cat, + tiled=self._policy.trained_model.action_head.tiled, + tile_size=(self._policy.trained_model.action_head.tile_size_height, self._policy.trained_model.action_head.tile_size_width), + tile_stride=(self._policy.trained_model.action_head.tile_stride_height, self._policy.trained_model.action_head.tile_stride_width), + ) + frames = rearrange(frames, "B C T H W -> B T H W C") + frames = frames[0] + frames = ((frames.float() + 1) * 127.5).clip(0, 255).cpu().numpy().astype(np.uint8) + # Add each frame individually to the list + for frame in frames: + frame_list.append(frame) + + sample_frame = frame_list[0] + if len(sample_frame.shape) == 3 and sample_frame.shape[2] in [1, 3, 4]: + # Save all frames as a single MP4 file + save_dir = self._output_dir if self._output_dir else "." + os.makedirs(save_dir, exist_ok=True) + all_mp4_files = [f for f in os.listdir(save_dir) if f.endswith(".mp4")] + timestamp = datetime.datetime.now().strftime("%m_%d_%H_%M_%S") + num_frames = len(frame_list) + n = (num_frames - 1) // 8 # num_frames = 8n+1, so n = (num_frames-1)/8 + output_path = os.path.join(save_dir, f'{len(all_mp4_files):06}_{timestamp}_n{n}.mp4') + imageio.mimsave(output_path, frame_list, fps=5, codec='libx264') + print(f"Saved video to: {output_path}") + else: + print(f"Warning: Invalid frame shape {sample_frame.shape}. Expected (H, W, C) with C in [1, 3, 4]. Skipping video save.") + + self.video_across_time = [] + elif self._policy.trained_model.action_head.current_start_frame == 1 + self._policy.trained_model.action_head.num_frame_per_block and len(self.video_across_time) > 1: + print("current_start_frame == 1 + num_frame_per_block and len(self.video_across_time) > 1") + frame_list = [] + video_across_time_cat = torch.cat(self.video_across_time[:-1], dim=2) + frames = self._policy.trained_model.action_head.vae.decode( + video_across_time_cat, + tiled=self._policy.trained_model.action_head.tiled, + tile_size=(self._policy.trained_model.action_head.tile_size_height, self._policy.trained_model.action_head.tile_size_width), + tile_stride=(self._policy.trained_model.action_head.tile_stride_height, self._policy.trained_model.action_head.tile_stride_width), + ) + frames = rearrange(frames, "B C T H W -> B T H W C") + frames = frames[0] + frames = ((frames.float() + 1) * 127.5).clip(0, 255).cpu().numpy().astype(np.uint8) + # Add each frame individually to the list + for frame in frames: + frame_list.append(frame) + sample_frame = frame_list[0] + if len(sample_frame.shape) == 3 and sample_frame.shape[2] in [1, 3, 4]: + # Save all frames as a single MP4 file + save_dir = self._output_dir if self._output_dir else "." + os.makedirs(save_dir, exist_ok=True) + all_mp4_files = [f for f in os.listdir(save_dir) if f.endswith(".mp4")] + timestamp = datetime.datetime.now().strftime("%m_%d_%H_%M_%S") + num_frames = len(frame_list) + n = (num_frames - 1) // 8 # num_frames = 8n+1, so n = (num_frames-1)/8 + output_path = os.path.join(save_dir, f'{len(all_mp4_files):06}_{timestamp}_n{n}.mp4') + imageio.mimsave(output_path, frame_list, fps=5, codec='libx264') + print(f"Saved video to: {output_path}") + self.video_across_time = [video_chunk] + + + def batch_to_dict(batch): + out = {} + for k in dir(batch): + if not k.startswith("action."): + continue + out[k] = getattr(batch, k) + return out + action_chunk_dict = batch_to_dict(action_chunk_dict) + await websocket.send(packer.pack(action_chunk_dict)) + + except websockets.ConnectionClosed: + logger.info(f"Connection from {websocket.remote_address} closed") + if len(self.video_across_time) > 0: + frame_list = [] + video_across_time_cat = torch.cat(self.video_across_time, dim=2) + frames = self._policy.trained_model.action_head.vae.decode( + video_across_time_cat, + tiled=self._policy.trained_model.action_head.tiled, + tile_size=(self._policy.trained_model.action_head.tile_size_height, self._policy.trained_model.action_head.tile_size_width), + tile_stride=(self._policy.trained_model.action_head.tile_stride_height, self._policy.trained_model.action_head.tile_stride_width), + ) + frames = rearrange(frames, "B C T H W -> B T H W C") + frames = frames[0] + frames = ((frames.float() + 1) * 127.5).clip(0, 255).cpu().numpy().astype(np.uint8) + # Add each frame individually to the list + for frame in frames: + frame_list.append(frame) + + sample_frame = frame_list[0] + if len(sample_frame.shape) == 3 and sample_frame.shape[2] in [1, 3, 4]: + # Save all frames as a single MP4 file + save_dir = self._output_dir if self._output_dir else "." + os.makedirs(save_dir, exist_ok=True) + all_mp4_files = [f for f in os.listdir(save_dir) if f.endswith(".mp4")] + timestamp = datetime.datetime.now().strftime("%m_%d_%H_%M_%S") + num_frames = len(frame_list) + n = (num_frames - 1) // 8 # num_frames = 8n+1, so n = (num_frames-1)/8 + output_path = os.path.join(save_dir, f'{len(all_mp4_files):06}_{timestamp}_n{n}.mp4') + imageio.mimsave(output_path, frame_list, fps=5, codec='libx264') + print(f"Saved video to: {output_path}") + else: + print(f"Warning: Invalid frame shape {sample_frame.shape}. Expected (H, W, C) with C in [1, 3, 4]. Skipping video save.") + + self.video_across_time = [] + 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 + finally: + logger.info(f"Rank 0: Client session ended. Sending idle signal (2) to workers.") + signal_tensor.fill_(2) # Set tensor value to 2 + dist.broadcast(signal_tensor, src=0, group=self._signal_group) + # When connection closes, signal other ranks to continue waiting for next connection + # (or implement proper shutdown if needed) + + +def init_mesh() -> DeviceMesh: + # env vars set by torchrun + dist.init_process_group("nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + print(f"Rank {rank}/{world_size} (PID: {os.getpid()}) setting device to {rank}") + + torch.cuda.set_device(rank) + device = torch.device(f"cuda:{rank}") + + mesh = init_device_mesh( + device_type="cuda", + mesh_shape=(world_size, ), + mesh_dim_names=("ip", ), + ) + print(f"Rank {rank}/{world_size} (PID: {os.getpid()}) using device {device}") + + return mesh + +def _health_check(connection: _server.ServerConnection, request: _server.Request) -> _server.Response | None: + if request.path == "/healthz": + return connection.respond(http.HTTPStatus.OK, "OK\n") + # Continue with the normal request handling. + return None + + +def main(args: Args) -> None: + # Set environment variable for DIT cache. + os.environ["ENABLE_DIT_CACHE"] = "true" if args.enable_dit_cache else "false" + + # Use TE cuDNN backend for attention. + os.environ["ATTENTION_BACKEND"] = "TE" + + # Increase the recompile limit to 100 for inference due + # to autoregressive nature of the model (several possible shapes). + torch._dynamo.config.recompile_limit = 800 + + embodiment_tag = "oxe_droid" + model_path = args.model_path + policy_metadata = { + "embodiment": embodiment_tag, + "model_name": "dreamzero", + "model_path": model_path, + } + + device_mesh = init_mesh() + rank = dist.get_rank() + + timeout_delta = datetime.timedelta(seconds=args.timeout_seconds) + signal_group = dist.new_group(backend="gloo", timeout=timeout_delta) + logger.info(f"Rank {rank} initialized signal_group (gloo)") + + policy = GrootSimPolicy( + embodiment_tag=EmbodimentTag(embodiment_tag), + model_path=model_path, + device="cuda" if torch.cuda.is_available() else "cpu", + device_mesh=device_mesh, + ) + + # Create server for all ranks - rank 0 handles websocket, others run worker loop + hostname = socket.gethostname() + local_ip = socket.gethostbyname(hostname) + + if rank == 0: + logging.info("Creating server (host: %s, ip: %s)", hostname, local_ip) + # Create output directory for videos + # Extract parent directory and checkpoint name from model_path + parent_dir = os.path.dirname(model_path) + date_suffix = datetime.datetime.now().strftime("%Y%m%d") + checkpoint_name = os.path.basename(model_path) + output_dir = os.path.join(parent_dir, f"real_world_eval_gen_{date_suffix}_{args.index}", checkpoint_name) + os.makedirs(output_dir, exist_ok=True) + logging.info("Videos will be saved to: %s", output_dir) + else: + output_dir = None + logging.info(f"Rank {rank} starting as worker for distributed inference...") + + # Create wrapper policy that converts between roboarena and AR_droid formats + wrapper_policy = ARDroidRoboarenaPolicy( + groot_policy=policy, + signal_group=signal_group, + output_dir=output_dir, + ) + + # Configure server for AR_droid (2 external cameras, wrist camera, joint position actions) + server_config = PolicyServerConfig( + image_resolution=(180, 320), # AR_droid expects 180x320 images + needs_wrist_camera=True, + n_external_cameras=2, + needs_stereo_camera=False, + needs_session_id=True, # Track session to reset state for new clients + action_space="joint_position", + ) + + if rank == 0: + logging.info("Using roboarena policy server interface") + logging.info(f"Server config: {server_config}") + roboarena_server = RoboarenaServer( + policy=wrapper_policy, + server_config=server_config, + host="0.0.0.0", + port=args.port, + ) + roboarena_server.serve_forever() + else: + # Non-rank-0 processes need to run worker loop for distributed inference + # We'll use the existing WebsocketPolicyServer's worker loop mechanism + server = WebsocketPolicyServer( + policy=policy, + host="0.0.0.0", + port=args.port, + metadata=policy_metadata, + output_dir=output_dir, + signal_group=signal_group, + ) + asyncio.run(server._worker_loop()) + + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO, force=True) + args = tyro.cli(Args) + main(args) \ No newline at end of file diff --git a/test_client_AR.py b/test_client_AR.py new file mode 100644 index 0000000000000000000000000000000000000000..e52370c21483ff5c1ce92db112206d9a83bff7bd --- /dev/null +++ b/test_client_AR.py @@ -0,0 +1,331 @@ +#!/usr/bin/env python3 +"""Test client for AR_droid policy server using roboarena interface. + +Sends real video frames from debug_image/ directory instead of zero dummy images. + +Frame schedule (matching debug_inference.py): + - Step 0 (initial): send frame [0] (1 frame, H W 3) + - Step 1: send frames [0, 7, 15, 23] (4 frames, 4 H W 3) + - Step 2: send frames [24, 31, 39, 47] (4 frames) + - Step 3: send frames [48, 55, 63, 71] (4 frames) + - ... + +Expected server configuration: + - image_resolution: (180, 320) + - n_external_cameras: 2 + - needs_wrist_camera: True + - action_space: "joint_position" + +Usage: + # Start server with roboarena interface: + torchrun --nproc_per_node=8 socket_test_optimized_AR.py --port 8000 + + # Run this test: + python test_client_AR.py --host --port 8000 + + # Use zero images instead of real video (old behavior): + python test_client_AR.py --host --port 8000 --use-zero-images +""" + +import argparse +import logging +import os +import time +import uuid + +import cv2 +import numpy as np + +import eval_utils.policy_server as policy_server +from eval_utils.policy_client import WebsocketClientPolicy + +VIDEO_DIR = os.path.join(os.path.dirname(__file__), "debug_image") + +# roboarena key -> video filename +CAMERA_FILES = { + "observation/exterior_image_0_left": "exterior_image_1_left.mp4", + "observation/exterior_image_1_left": "exterior_image_2_left.mp4", + "observation/wrist_image_left": "wrist_image_left.mp4", +} + +# Frame schedule constants (matching debug_inference.py) +RELATIVE_OFFSETS = [-23, -16, -8, 0] +ACTION_HORIZON = 24 + + +def load_all_frames(video_path: str) -> np.ndarray: + """Load all frames from a video file. Returns (N, H, W, 3) uint8 array (RGB).""" + cap = cv2.VideoCapture(video_path) + frames = [] + while True: + ret, frame = cap.read() + if not ret: + break + frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) + cap.release() + if not frames: + raise RuntimeError(f"No frames loaded from {video_path}") + return np.stack(frames, axis=0) + + +def load_camera_frames() -> dict[str, np.ndarray]: + """Load all video frames for each camera from the debug_image/ directory. + + Returns: + Dict mapping roboarena camera keys to (N, H, W, 3) uint8 arrays. + """ + camera_frames: dict[str, np.ndarray] = {} + for cam_key, fname in CAMERA_FILES.items(): + path = os.path.join(VIDEO_DIR, fname) + camera_frames[cam_key] = load_all_frames(path) + logging.info(f"Loaded {cam_key}: {camera_frames[cam_key].shape}") + return camera_frames + + +def build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]: + """Build the frame index schedule for multi-frame chunks. + + Returns a list of frame-index lists. Each inner list has 4 indices. + """ + chunks: list[list[int]] = [] + current_frame = 23 # first anchor frame + for _ in range(num_chunks): + indices = [max(current_frame + off, 0) for off in RELATIVE_OFFSETS] + if indices[-1] >= total_frames: + logging.info( + f"Frame {indices[-1]} >= {total_frames}, stopping at {len(chunks)} chunks" + ) + break + chunks.append(indices) + current_frame += ACTION_HORIZON + return chunks + + +def _make_obs_from_video( + camera_frames: dict[str, np.ndarray], + frame_indices: list[int], + prompt: str, + session_id: str, +) -> dict: + """Build an observation dict from real video frames. + + For 1 frame: each image key is (H, W, 3). + For 4 frames: each image key is (4, H, W, 3). + """ + obs: dict = {} + for cam_key, all_frames in camera_frames.items(): + selected = all_frames[frame_indices] # (T, H, W, 3) + if len(frame_indices) == 1: + selected = selected[0] # (H, W, 3) + obs[cam_key] = selected + + obs["observation/joint_position"] = np.zeros(7, dtype=np.float32) + obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32) + obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32) + obs["prompt"] = prompt + obs["session_id"] = session_id + return obs + + +def _make_zero_observation( + server_config: policy_server.PolicyServerConfig, + prompt: str = "pick up the object", + session_id: str | None = None, +) -> dict: + """Create a dummy observation matching AR_droid expectations. + + AR_droid expects: + - 2 external cameras (exterior_image_0_left, exterior_image_1_left) + - 1 wrist camera (wrist_image_left) + - Image resolution: 180x320 (H x W) + - joint_position: 7 DoF + - gripper_position: 1 DoF + """ + obs = {} + + # Determine image resolution + if server_config.image_resolution is not None: + h, w = server_config.image_resolution + else: + # Default for AR_droid + h, w = 180, 320 + + # External cameras (0-indexed in roboarena) + for i in range(server_config.n_external_cameras): + obs[f"observation/exterior_image_{i}_left"] = np.zeros((h, w, 3), dtype=np.uint8) + if server_config.needs_stereo_camera: + obs[f"observation/exterior_image_{i}_right"] = np.zeros((h, w, 3), dtype=np.uint8) + + # Wrist camera + if server_config.needs_wrist_camera: + obs["observation/wrist_image_left"] = np.zeros((h, w, 3), dtype=np.uint8) + if server_config.needs_stereo_camera: + obs["observation/wrist_image_right"] = np.zeros((h, w, 3), dtype=np.uint8) + + # Session ID - should be passed in to ensure consistency within a session + if server_config.needs_session_id: + import uuid + # Generate unique session ID if not provided + obs["session_id"] = session_id if session_id else str(uuid.uuid4()) + + # State observations (AR_droid: 7 DoF arm + 1 gripper) + obs["observation/joint_position"] = np.zeros(7, dtype=np.float32) + obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32) + obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32) + + # Language prompt + obs["prompt"] = prompt + + return obs + + +def test_ar_droid_policy_server( + host: str = "localhost", + port: int = 8000, + num_chunks: int = 15, + prompt: str = "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan", + use_zero_images: bool = False, +): + """Test the AR_droid policy server with roboarena interface. + + When use_zero_images is False (default), loads real video frames from + debug_image/ and follows the frame schedule from debug_inference.py. + """ + logging.info(f"Connecting to AR_droid server at {host}:{port}...") + + client = WebsocketClientPolicy(host=host, port=port) + + # Validate server metadata + metadata = client.get_server_metadata() + logging.info(f"Server metadata: {metadata}") + assert isinstance(metadata, dict), "Metadata should be a dict" + + try: + server_config = policy_server.PolicyServerConfig(**metadata) + except Exception as e: + logging.error(f"Error parsing metadata: {e}") + raise e + + # Validate expected AR_droid configuration + logging.info(f"Server config: {server_config}") + assert server_config.n_external_cameras == 2, f"Expected 2 external cameras, got {server_config.n_external_cameras}" + assert server_config.needs_wrist_camera, "Expected wrist camera to be enabled" + assert server_config.action_space == "joint_position", f"Expected joint_position action space, got {server_config.action_space}" + + logging.info("Server configuration validated for AR_droid") + + # Generate unique session ID for this test run + import uuid + session_id = str(uuid.uuid4()) + logging.info(f"Session ID: {session_id}") + + # ── Zero-image fallback mode ────────────────────────────────────── + if use_zero_images: + logging.info("Using ZERO dummy images (legacy mode)") + for i in range(num_chunks): + obs = _make_zero_observation(server_config, prompt=prompt, session_id=session_id) + logging.info(f"Inference {i + 1}/{num_chunks}: prompt='{prompt}'") + t0 = time.time() + actions = client.infer(obs) + dt = time.time() - t0 + _log_action(actions, dt) + + logging.info("Sending reset...") + client.reset({}) + logging.info("Done (zero-image mode).") + return + + # ── Real video frame mode ───────────────────────────────────────── + logging.info("Loading real video frames from debug_image/ directory") + camera_frames = load_camera_frames() + + total_frames = min(v.shape[0] for v in camera_frames.values()) + logging.info(f"Total frames available: {total_frames}") + + # Build frame schedule + chunks = build_frame_schedule(total_frames, num_chunks) + + logging.info("Frame schedule:") + logging.info(" Initial: [0]") + for i, indices in enumerate(chunks): + logging.info(f" Chunk {i}: {indices}") + + # Step 0: initial single frame + logging.info("=== Initial: frame [0] ===") + obs = _make_obs_from_video(camera_frames, [0], prompt, session_id) + t0 = time.time() + actions = client.infer(obs) + dt = time.time() - t0 + _log_action(actions, dt) + + # Subsequent chunks: send 4 frames at a time + for chunk_idx, frame_indices in enumerate(chunks): + logging.info(f"=== Chunk {chunk_idx}: frames {frame_indices} ===") + obs = _make_obs_from_video(camera_frames, frame_indices, prompt, session_id) + t0 = time.time() + actions = client.infer(obs) + dt = time.time() - t0 + _log_action(actions, dt) + + # Reset triggers video save on the server + logging.info("Sending reset to save video...") + client.reset({}) + + logging.info("Done.") + + +def _log_action(actions: np.ndarray, dt: float) -> None: + """Pretty-print action shape, range, and timing.""" + assert isinstance(actions, np.ndarray), f"Expected numpy array, got {type(actions)}" + assert actions.ndim == 2, f"Expected 2D array, got shape {actions.shape}" + assert actions.shape[-1] == 8, ( + f"Expected 8 action dims (7 joints + 1 gripper), got {actions.shape[-1]}" + ) + logging.info( + f" Action shape: {actions.shape}, " + f"range: [{actions.min():.4f}, {actions.max():.4f}], " + f"time: {dt:.2f}s" + ) + + +def main(): + parser = argparse.ArgumentParser( + description="Test AR_droid policy server with real video frames from debug_image/" + ) + parser.add_argument("--host", default="localhost", help="Server hostname") + parser.add_argument("--port", type=int, default=8000, help="Server port") + parser.add_argument( + "--num-chunks", + type=int, + default=15, + help="Number of 4-frame chunks to send after the initial frame (default: 15)", + ) + parser.add_argument( + "--prompt", + default="Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan", + help="Language prompt for the policy", + ) + parser.add_argument( + "--use-zero-images", + action="store_true", + help="Use zero dummy images instead of real video frames (legacy mode)", + ) + + args = parser.parse_args() + + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(message)s", + ) + + test_ar_droid_policy_server( + host=args.host, + port=args.port, + num_chunks=args.num_chunks, + prompt=args.prompt, + use_zero_images=args.use_zero_images, + ) + + +if __name__ == "__main__": + main() diff --git a/tests/trex_track_force/test_attention.py b/tests/trex_track_force/test_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..e58b191beb125f147e0007301f6adc1b4353aedf --- /dev/null +++ b/tests/trex_track_force/test_attention.py @@ -0,0 +1,114 @@ +import torch + +from groot.vla.model.trex_track_force.attention import ( + BlockTokenLayout, + CausalWanSelfAttention, + TokenType, + build_dense_allow_matrix, +) + + +def test_dense_allow_matrix_exact_information_flow(): + layout = BlockTokenLayout( + num_blocks=2, + conditioning_obs_tokens=2, + obs_tokens_per_block=3, + action_tokens_per_block=2, + state_tokens_per_block=1, + track_points=3, + ) + allow = build_dense_allow_matrix(layout) + assert allow.dtype == torch.bool + assert allow.shape == (layout.sequence_length, layout.sequence_length) + + block0 = layout.slices(0) + block1 = layout.slices(1) + + # Same-block observation/action are bidirectional and both can read track. + assert allow[block0[TokenType.OBS], block0[TokenType.ACTION]].all() + assert allow[block0[TokenType.ACTION], block0[TokenType.OBS]].all() + assert allow[block0[TokenType.OBS], block0[TokenType.TRACK_PAST]].all() + assert allow[block0[TokenType.ACTION], block0[TokenType.TRACK_FUTURE]].all() + + # Track queries never read observation, action, state, or conditioning obs. + track_rows = layout.indices((TokenType.TRACK_PAST, TokenType.TRACK_FUTURE)) + forbidden_columns = layout.indices( + ( + TokenType.CONDITIONING_OBS, + TokenType.OBS, + TokenType.ACTION, + TokenType.STATE, + ) + ) + assert not allow[track_rows][:, forbidden_columns].any() + + # Clean-past track does not consume its own noisy future target. + assert not allow[ + block0[TokenType.TRACK_PAST], block0[TokenType.TRACK_FUTURE] + ].any() + assert allow[ + block0[TokenType.TRACK_FUTURE], block0[TokenType.TRACK_PAST] + ].all() + + # Cross-block attention is causal. + assert allow[block1[TokenType.OBS], block0[TokenType.OBS]].all() + assert allow[ + block1[TokenType.TRACK_FUTURE], block0[TokenType.TRACK_FUTURE] + ].all() + assert not allow[block0[TokenType.OBS], block1[TokenType.OBS]].any() + assert not allow[ + block0[TokenType.TRACK_FUTURE], block1[TokenType.TRACK_PAST] + ].any() + + # State cannot carry information into other tokens as a query. + state_rows = torch.arange( + block0[TokenType.STATE].start, block0[TokenType.STATE].stop + ) + assert allow[state_rows].sum().item() == state_rows.numel() + + +def test_track_attention_output_is_invariant_to_obs_action_values(): + torch.manual_seed(7) + layout = BlockTokenLayout( + num_blocks=1, + obs_tokens_per_block=2, + action_tokens_per_block=2, + state_tokens_per_block=1, + track_points=3, + ) + layer = CausalWanSelfAttention(dim=16, num_heads=4).eval() + hidden = torch.randn(2, layout.sequence_length, 16) + perturbed = hidden.clone() + obs_action = layout.indices((TokenType.OBS, TokenType.ACTION, TokenType.STATE)) + perturbed[:, obs_action] = torch.randn_like(perturbed[:, obs_action]) * 100 + + output = layer(hidden, layout) + perturbed_output = layer(perturbed, layout) + track = layout.indices((TokenType.TRACK_PAST, TokenType.TRACK_FUTURE)) + torch.testing.assert_close(output[:, track], perturbed_output[:, track]) + + +def test_autoregressive_attention_keeps_configured_chunk_window(): + layout = BlockTokenLayout( + num_blocks=6, + conditioning_obs_tokens=1, + obs_tokens_per_block=2, + action_tokens_per_block=2, + state_tokens_per_block=1, + track_points=3, + ) + allow = build_dense_allow_matrix(layout, max_previous_blocks=2) + block2 = layout.slices(2) + block3 = layout.slices(3) + block5 = layout.slices(5) + assert allow[block5[TokenType.ACTION], block3[TokenType.OBS]].all() + assert not allow[block5[TokenType.ACTION], block2[TokenType.OBS]].any() + assert allow[ + block5[TokenType.TRACK_FUTURE], block3[TokenType.TRACK_PAST] + ].all() + assert not allow[ + block5[TokenType.TRACK_FUTURE], block2[TokenType.TRACK_PAST] + ].any() + assert allow[ + block5[TokenType.OBS], layout.conditioning_slice + ].all() diff --git a/tests/trex_track_force/test_dataset_loader.py b/tests/trex_track_force/test_dataset_loader.py new file mode 100644 index 0000000000000000000000000000000000000000..73501b9a83ec7fce403096fb7a3918c9c03cdd3e --- /dev/null +++ b/tests/trex_track_force/test_dataset_loader.py @@ -0,0 +1,216 @@ +from types import SimpleNamespace + +import numpy as np +import pandas as pd + +from groot.vla.data.dataset.lerobot_sharded import ( + ShardedLeRobotSubLangSingleActionChunkDatasetDROID, +) +from groot.vla.model.trex_track_force.dataset import ( + AR_BLOCKS, + FORCE_COLUMN, + TRACK_VISIBILITY_COLUMN, + TRACK_XY_COLUMN, + TrexTrackForceShardedDataset, + eef62_delta_base, + nearest_timestamp_indices, + uniform_target_times, +) +from groot.vla.model.trex_track_force.runtime import ( + TrexRuntimeStatistics, + delta_base_to_absolute, +) + + +def _pose9(translation: tuple[float, float, float], rotation: np.ndarray) -> np.ndarray: + return np.concatenate((np.asarray(translation), rotation[:, 0], rotation[:, 1])) + + +def test_nearest_timestamp_indices_builds_16_step_30_to_20_grid() -> None: + source = np.arange(61, dtype=np.float64) / 30.0 + target = uniform_target_times(0.0, range(16), 20.0) + selection = nearest_timestamp_indices(source, target) + expected = np.floor(np.arange(16, dtype=np.float64) * 1.5).astype(np.int64) + np.testing.assert_array_equal(selection.indices, expected) + assert not selection.padding_mask.any() + assert target[-1] == 0.75 + + +def test_nearest_timestamp_indices_marks_history_edge_padding() -> None: + source = np.arange(31, dtype=np.float64) / 30.0 + target = uniform_target_times(0.0, range(-15, 1), 5.0) + selection = nearest_timestamp_indices(source, target) + assert selection.padding_mask[:-1].all() + assert not selection.padding_mask[-1] + assert np.all(selection.indices[:-1] == 0) + + +def test_step_filter_rejects_float_roundoff_future_padding(monkeypatch) -> None: + timestamps = np.arange(1126, dtype=np.float64) / 30.0 + dataset = object.__new__(TrexTrackForceShardedDataset) + dataset.action_rate_hz = 20.0 + dataset.video_rate_hz = 10.0 + dataset.num_ar_blocks = AR_BLOCKS + dataset._trajectory_ids = np.array([0], dtype=np.int64) + dataset._logical_anchor_times = {} + + monkeypatch.setattr( + ShardedLeRobotSubLangSingleActionChunkDatasetDROID, + "_get_step_filter", + lambda self: {0: np.arange(timestamps.size, dtype=np.int64)}, + ) + monkeypatch.setattr(dataset, "get_parquet_path", lambda trajectory_id: "unused") + monkeypatch.setattr( + pd, + "read_parquet", + lambda *args, **kwargs: pd.DataFrame({"timestamp": timestamps}), + ) + + filtered = dataset._get_step_filter() + + assert 1027 in filtered[0] + assert 1029 not in filtered[0] + assert 1029 not in dataset._logical_anchor_times[0] + + +def test_column_array_collapses_pyarrow_nested_track_objects() -> None: + row = np.empty(2, dtype=object) + row[:] = ( + np.asarray([0.1, 0.2], dtype=np.float32), + np.asarray([0.3, 0.4], dtype=np.float32), + ) + frame = pd.DataFrame({TRACK_XY_COLUMN: [row, row.copy()]}) + + values = TrexTrackForceShardedDataset._column_array(frame, TRACK_XY_COLUMN) + + assert values.shape == (2, 2, 2) + assert values.dtype == np.float32 + np.testing.assert_allclose(values[0, 1], [0.3, 0.4]) + + +def test_loader_preserves_four_ar_blocks_with_16_actions_each() -> None: + frames = 121 + timestamps = np.arange(frames, dtype=np.float64) / 30.0 + identity_rot6d = np.array([1, 0, 0, 0, 1, 0], dtype=np.float32) + eef = np.zeros((frames, 62), dtype=np.float32) + eef[:, 3:9] = identity_rot6d + eef[:, 34:40] = identity_rot6d + eef[:, 0] = np.arange(frames, dtype=np.float32) / 30.0 + tracks = np.zeros((frames, 250, 2), dtype=np.float32) + tracks[..., 0] = np.arange(frames, dtype=np.float32)[:, None] / frames + frame = pd.DataFrame( + { + "timestamp": timestamps, + "observation.state_eef62": list(eef), + "action.eef62_absolute": list(eef), + TRACK_XY_COLUMN: list(tracks), + TRACK_VISIBILITY_COLUMN: list( + np.ones((frames, 250), dtype=np.float32) + ), + FORCE_COLUMN: list( + np.zeros((frames, 60), dtype=np.float32) + ), + } + ) + + dataset = object.__new__(TrexTrackForceShardedDataset) + dataset.action_rate_hz = 20.0 + dataset.tactile_rate_hz = 5.0 + dataset.video_rate_hz = 10.0 + dataset.num_ar_blocks = AR_BLOCKS + dataset._logical_anchor_times = {0: {0: 0.0}} + dataset._force_low = np.zeros((10, 6), dtype=np.float32) + dataset._force_high = np.ones((10, 6), dtype=np.float32) + dataset._lerobot_modality_meta = SimpleNamespace( + state={ + "eef62": SimpleNamespace( + original_key="observation.state_eef62", start=0, end=62 + ) + }, + action={ + "eef62": SimpleNamespace( + original_key="action.eef62_absolute", start=0, end=62 + ) + }, + ) + dataset.cached_shard = { + key: np.zeros((frames, 4, 6, 3), dtype=np.uint8) + for key in ("video.head_left", "video.left_wrist", "video.right_wrist") + } + dataset.shard_start_indices = {0: 0} + dataset._modality_keys = { + "video": list(dataset.cached_shard), + "language": [], + } + dataset.get_trajectory_data = lambda trajectory_id: frame + + sample = dataset.get_step_data(0, {"action.eef62": np.array([0])}) + assert sample["action.eef62"].shape == (4 * 16, 62) + assert sample["state.eef62"].shape == (4, 62) + assert sample["track_past_xy"].shape == (4, 16, 250, 2) + assert sample["track_future_xy"].shape == (4, 16, 250, 2) + np.testing.assert_allclose( + sample["track_future_xy"][:, 0], sample["track_past_xy"][:, -1] + ) + assert not sample["track_past_visibility"][0, :-1].any() + assert sample["track_past_visibility"][0, -1].all() + assert sample["current_force"].shape == (4, 4, 10, 6) + assert sample["force_history"].shape == (4, 4, 16, 10, 6) + assert np.count_nonzero(sample["force_history"][0, 0, :-1]) == 0 + np.testing.assert_array_equal(sample["current_force"][0, 0], -1.0) + assert sample["force_history_padding_mask"].shape == (4, 4, 16) + for key in dataset.modality_keys["video"]: + assert sample[key].shape == (33, 4, 6, 3) + + +def test_training_and_runtime_force_normalization_are_identical() -> None: + low = np.zeros((10, 6), dtype=np.float32) + high = np.ones((10, 6), dtype=np.float32) + high[0, 0] = 0 + force = np.linspace(-0.5, 1.5, 120, dtype=np.float32).reshape(2, 10, 6) + dataset = object.__new__(TrexTrackForceShardedDataset) + dataset._force_low = low + dataset._force_high = high + runtime = TrexRuntimeStatistics( + action_q01=np.zeros(62), + action_q99=np.ones(62), + state_q01=np.zeros(62), + state_q99=np.ones(62), + force_q01=low, + force_q99=high, + ) + np.testing.assert_allclose( + dataset._normalize_force(force), runtime.normalize_force(force) + ) + + +def test_eef62_delta_base_matches_trex_chunk_start_frame_math() -> None: + identity = np.eye(3) + angle = np.pi / 2 + reference_rotation = np.array( + [ + [np.cos(angle), -np.sin(angle), 0.0], + [np.sin(angle), np.cos(angle), 0.0], + [0.0, 0.0, 1.0], + ] + ) + reference = np.zeros(62, dtype=np.float32) + reference[:9] = _pose9((0.0, 0.0, 0.0), reference_rotation) + reference[31:40] = _pose9((0.0, 0.0, 0.0), identity) + + targets = np.zeros((2, 62), dtype=np.float32) + targets[:, :9] = _pose9((1.0, 0.0, 0.0), reference_rotation) + targets[:, 31:40] = _pose9((0.0, 2.0, 0.0), identity) + targets[:, 9:31] = 0.25 + targets[:, 40:62] = -0.5 + + delta = eef62_delta_base(reference, targets) + np.testing.assert_allclose(delta[:, :3], [[0.0, -1.0, 0.0]] * 2, atol=1e-6) + np.testing.assert_allclose( + delta[:, 3:9], [[1.0, 0.0, 0.0, 0.0, 1.0, 0.0]] * 2, atol=1e-6 + ) + np.testing.assert_allclose(delta[:, 31:34], [[0.0, 2.0, 0.0]] * 2) + np.testing.assert_allclose(delta[:, 9:31], 0.25) + np.testing.assert_allclose(delta[:, 40:62], -0.5) + reconstructed = delta_base_to_absolute(reference, delta) + np.testing.assert_allclose(reconstructed, targets, atol=1e-6) diff --git a/tests/trex_track_force/test_dataset_schema.py b/tests/trex_track_force/test_dataset_schema.py new file mode 100644 index 0000000000000000000000000000000000000000..b29a19e72297da8dd924e4d794b8a62b3e4ac796 --- /dev/null +++ b/tests/trex_track_force/test_dataset_schema.py @@ -0,0 +1,539 @@ +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +SCRIPTS = REPO_ROOT / "scripts" +DATA_SCRIPTS = SCRIPTS / "data" +for path in (str(SCRIPTS), str(DATA_SCRIPTS)): + if path not in sys.path: + sys.path.insert(0, path) + +import build_trex_track_force_v2 as builder # noqa: E402 +import extract_track # noqa: E402 +from trex_track.layout import ( # noqa: E402 + NUM_COMBINED_POINTS, + POINT_SLICES, + TRACK_LAYOUT_VERSION, + VIEW_ORDER, + VIEW_POINT_COUNTS, + identity_metadata, +) + + +def _fake_eef_converter(joints: np.ndarray) -> np.ndarray: + source = np.asarray(joints, dtype=np.float32) + output = np.zeros((source.shape[0], 62), dtype=np.float32) + identity_rot6d = np.array([1, 0, 0, 0, 1, 0], dtype=np.float32) + output[:, 0:3] = source[:, 0:3] + output[:, 3:9] = identity_rot6d + output[:, 9:31] = source[:, 7:29] + output[:, 31:34] = source[:, 29:32] + output[:, 34:40] = identity_rot6d + output[:, 40:62] = source[:, 36:58] + return output + + +def _write_fixture_dataset(root: Path, *, frames: int = 4) -> tuple[np.ndarray, np.ndarray]: + (root / "meta").mkdir(parents=True) + parquet_path = root / "data" / "chunk-000" / "episode_000000.parquet" + parquet_path.parent.mkdir(parents=True) + + rng = np.random.default_rng(4) + state = rng.normal(size=(frames, 58)).astype(np.float32) + action = rng.normal(size=(frames, 58)).astype(np.float32) + tactile_force = rng.normal(size=(frames, 60)).astype(np.float32) + table = pa.table( + { + "observation.state": pa.array(state.tolist(), type=pa.list_(pa.float32())), + "action": pa.array(action.tolist(), type=pa.list_(pa.float32())), + "observation.tactile_force": pa.array( + tactile_force.tolist(), type=pa.list_(pa.float32()) + ), + "timestamp": pa.array(np.arange(frames, dtype=np.float32) / 30.0), + "episode_index": pa.array(np.zeros(frames, dtype=np.int64)), + } + ) + pq.write_table(table, parquet_path) + + info = { + "codebase_version": "v2.1", + "total_episodes": 1, + "total_frames": frames, + "chunks_size": 1000, + "data_path": "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet", + "features": { + "observation.state": {"dtype": "float32", "shape": [58], "names": None}, + "action": {"dtype": "float32", "shape": [58], "names": None}, + "observation.tactile_force": { + "dtype": "float32", + "shape": [60], + "names": None, + }, + }, + } + modality = { + "state": { + "joint": { + "original_key": "observation.state", + "start": 0, + "end": 58, + "rotation_type": None, + "absolute": True, + "dtype": "float32", + "range": None, + } + }, + "action": { + "joint": { + "original_key": "action", + "start": 0, + "end": 58, + "rotation_type": None, + "absolute": True, + "dtype": "float32", + "range": None, + } + }, + "video": {}, + "annotation": {}, + } + stats = { + "timestamp": { + key: [value] + for key, value in { + "mean": 0.05, + "std": 0.03, + "min": 0.0, + "max": 0.1, + "q01": 0.0, + "q99": 0.1, + }.items() + } + } + for name, value in (("info.json", info), ("modality.json", modality), ("stats.json", stats)): + (root / "meta" / name).write_text(json.dumps(value), encoding="utf-8") + return state, action + + +def _write_track_npz(path: Path, *, frames: int = 4, out_of_range: bool = False) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + identities = identity_metadata() + per_view_tracks: dict[str, np.ndarray] = {} + per_view_vis: dict[str, np.ndarray] = {} + for view_index, view in enumerate(VIEW_ORDER): + count = VIEW_POINT_COUNTS[view] + track = np.empty((frames, count, 2), dtype=np.float32) + track[..., 0] = np.linspace(0.05, 0.95, count, dtype=np.float32) + track[..., 1] = 0.2 + view_index * 0.2 + per_view_tracks[view] = track + per_view_vis[view] = np.ones((frames, count), dtype=np.float32) + if out_of_range: + per_view_tracks["head_left"][0, 0, 0] = 1.5 + tracks = np.concatenate([per_view_tracks[view] for view in VIEW_ORDER], axis=1) + visibility = np.concatenate([per_view_vis[view] for view in VIEW_ORDER], axis=1) + with path.open("wb") as file: + np.savez_compressed( + file, + tracks=tracks, + vis=visibility, + **{f"tracks_{view}": per_view_tracks[view] for view in VIEW_ORDER}, + **{f"vis_{view}": per_view_vis[view] for view in VIEW_ORDER}, + episode_index=np.array(0, dtype=np.int32), + num_steps=np.array(frames, dtype=np.int32), + point_slices=np.asarray(POINT_SLICES, dtype=np.int32), + track_layout_version=np.array(TRACK_LAYOUT_VERSION), + point_view_ids=np.asarray(identities["view_ids"], dtype=np.int8), + point_hand_ids=np.asarray(identities["hand_ids"], dtype=np.int8), + point_role_ids=np.asarray(identities["role_ids"], dtype=np.int8), + point_local_ids=np.asarray(identities["local_ids"], dtype=np.int16), + point_global_ids=np.asarray(identities["global_ids"], dtype=np.int16), + point_names=np.asarray(identities["point_names"]), + ) + + +def test_layout_and_episode_selection_are_stable() -> None: + identities = identity_metadata() + assert NUM_COMBINED_POINTS == 250 + assert POINT_SLICES == (0, 100, 175, 250) + assert all(len(values) == 250 for key, values in identities.items() if key != "point_names") + assert len(identities["point_names"]) == 250 + assert identities["hand_ids"][0] == 1 + assert identities["hand_ids"][50] == 2 + assert identities["hand_ids"][100] == 0 + assert identities["role_ids"][0] == 0 + assert identities["role_ids"][100] == 1 + assert identities["role_ids"][125] == 2 + assert builder.select_episode_indices(5, episode_index=2) == [2] + assert builder.select_episode_indices(5, episode_range=(1, 4)) == [1, 2, 3] + assert builder.select_episode_indices(3, all_episodes=True) == [0, 1, 2] + with pytest.raises(ValueError): + builder.select_episode_indices(3, episode_range=(2, 2)) + + +def test_sample_timestamps_nearest_30_to_20_index_pattern() -> None: + source = np.arange(61, dtype=np.float64) / 30.0 + result = builder.sample_timestamps_nearest( + source, + target_rate_hz=20.0, + anchor_index=0, + offsets=np.arange(16), + ) + # Midpoint ties deterministically use the earlier 30 Hz frame. + expected_indices = np.floor(np.arange(16, dtype=np.float64) * 1.5).astype( + np.int64 + ) + np.testing.assert_array_equal(result["indices"], expected_indices) + assert not np.asarray(result["padding_mask"]).any() + assert np.asarray(result["alignment_errors"]).max() <= 1.0 / 60.0 + 1e-6 + + +def test_timestamp_alignment_accepts_encoder_jitter_but_rejects_dropped_frames() -> None: + deltas = np.full(60, 1.0 / 30.0) + deltas[31] += 2e-6 + deltas[32] -= 2e-6 + jittered = np.concatenate(([0.0], np.cumsum(deltas))) + result = builder.sample_timestamps_nearest(jittered, target_rate_hz=20.0) + assert result["max_source_period_seconds"] > result["source_period_seconds"] + + dropped = np.delete(np.arange(61, dtype=np.float64) / 30.0, 30) + with pytest.raises(builder.DatasetSchemaError, match="dropped-frame gap"): + builder.sample_timestamps_nearest(dropped, target_rate_hz=20.0) + + +def test_frame_zero_cotracker_queries_skip_equivalent_backward_pass() -> None: + import torch + + class FakePredictor: + def __init__(self) -> None: + self.backward_tracking: bool | None = None + + def __call__(self, video, *, queries, backward_tracking): + self.backward_tracking = backward_tracking + batch, time = video.shape[:2] + points = queries.shape[1] + tracks = queries[:, None, :, 1:].expand(batch, time, points, 2).clone() + visibility = torch.ones(batch, time, points, dtype=torch.bool) + return tracks, visibility + + predictor = FakePredictor() + video = np.zeros((5, 8, 12, 3), dtype=np.uint8) + queries = np.array([[2.0, 3.0], [7.0, 5.0]], dtype=np.float32) + tracks, visibility = extract_track._run_cotracker( + predictor, video, queries, torch.device("cpu") + ) + assert predictor.backward_tracking is False + assert tracks.shape == (5, 2, 2) + assert visibility.shape == (5, 2) + np.testing.assert_allclose(tracks[0], queries) + np.testing.assert_array_equal(visibility, 1.0) + + +def test_16_step_horizon_duration_and_boundary_padding_mask() -> None: + source = np.arange(61, dtype=np.float64) / 30.0 + summary = builder.summarize_timestamp_sampling(source) + assert summary["target_rate_hz"] == 20.0 + assert summary["action_chunk_steps"] == 16 + assert summary["action_chunk_duration_seconds"] == pytest.approx(0.8) + assert summary["action_chunk_timestamp_span_seconds"] == pytest.approx(0.75) + assert summary["action_chunk_fully_covered"] is True + + padded = builder.sample_timestamps_nearest( + source[:10], + anchor_index=0, + offsets=[-2, -1, 0, 1, 2], + ) + np.testing.assert_array_equal( + padded["padding_mask"], [True, True, False, False, False] + ) + assert padded["indices"][0] == padded["indices"][1] == padded["indices"][2] == 0 + + +@pytest.mark.parametrize( + "timestamps", + [ + np.array([0.0, 0.0, 0.1]), + np.array([0.0, 0.2, 0.1]), + np.array([0.0, np.nan, 0.1]), + ], +) +def test_sample_timestamps_nearest_rejects_invalid_timestamps( + timestamps: np.ndarray, +) -> None: + with pytest.raises(builder.DatasetSchemaError): + builder.sample_timestamps_nearest(timestamps) + + +def test_real_trex_timestamps_validate_when_available() -> None: + dataset = builder.DEFAULT_DATASET_ROOT + info_path = dataset / "meta" / "info.json" + if not info_path.is_file(): + pytest.skip(f"real fixture unavailable: {info_path}") + info = json.loads(info_path.read_text()) + parquet_path = builder.episode_parquet_path(dataset, 0, info) + if not parquet_path.is_file(): + pytest.skip(f"real fixture unavailable: {parquet_path}") + table = pq.read_table(parquet_path, columns=["timestamp", builder.FORCE_COLUMN]) + timestamps = np.asarray(table["timestamp"].to_pylist(), dtype=np.float64) + summary = builder.summarize_timestamp_sampling(timestamps) + assert summary["source_rate_hz"] == pytest.approx(30.0, rel=5e-3) + assert summary["target_rate_hz"] == 20.0 + assert summary["action_chunk_duration_seconds"] == pytest.approx(0.8) + assert builder.validate_tactile_force(table)["reshape"] == [10, 6] + + +def test_build_episode_preserves_columns_and_uses_fixed_size_lists(tmp_path: Path) -> None: + dataset = tmp_path / "dataset" + state, action = _write_fixture_dataset(dataset) + track_path = tmp_path / "tracks" / "episode_000000.npz" + _write_track_npz(track_path) + original = pq.read_table(builder.episode_parquet_path(dataset, 0)) + original_names = original.column_names + + summary = builder.build_episode( + dataset_root=dataset, + episode_index=0, + track_path=track_path, + converter=_fake_eef_converter, + ) + + parquet_path = builder.episode_parquet_path(dataset, 0) + output = pq.read_table(parquet_path) + assert summary["num_frames"] == 4 + assert summary["sampling_20hz"]["source_rate_hz"] == pytest.approx(30.0) + assert summary["sampling_20hz"]["action_chunk_duration_seconds"] == pytest.approx(0.8) + assert summary["force_only"]["history_frames"] == 16 + assert ( + output.schema.metadata[builder.PARQUET_SCHEMA_METADATA_KEY].decode() + == builder.SCHEMA_VERSION + ) + assert output.column_names[: len(original_names)] == original_names + assert output["timestamp"].to_pylist() == original["timestamp"].to_pylist() + assert ( + output[builder.FORCE_COLUMN].to_pylist() + == original[builder.FORCE_COLUMN].to_pylist() + ) + np.testing.assert_array_equal( + np.asarray(output["observation.state"].to_pylist(), dtype=np.float32), + state, + ) + np.testing.assert_array_equal( + np.asarray(output["action"].to_pylist(), dtype=np.float32), + action, + ) + assert builder.backup_path(parquet_path).is_file() + for view, column in builder.TRACK_COLUMNS.items(): + field_type = output.schema.field(column).type + assert pa.types.is_fixed_size_list(field_type) + assert field_type.list_size == VIEW_POINT_COUNTS[view] + assert pa.types.is_fixed_size_list(field_type.value_type) + assert field_type.value_type.list_size == 3 + track_xy_type = output.schema.field(builder.TRACK_XY_COLUMN).type + assert track_xy_type.list_size == 250 + assert track_xy_type.value_type.list_size == 2 + assert output.schema.field(builder.TRACK_VISIBILITY_COLUMN).type.list_size == 250 + for column in (builder.STATE_EEF_COLUMN, builder.ACTION_EEF_COLUMN): + assert output.schema.field(column).type.list_size == 62 + pandas_frame = pd.read_parquet(parquet_path) + pandas_track_xy = np.stack( + [ + np.asarray(row.tolist(), dtype=np.float32) + for row in pandas_frame[builder.TRACK_XY_COLUMN] + ] + ) + assert pandas_track_xy.shape == (4, 250, 2) + assert np.stack(pandas_frame[builder.STATE_EEF_COLUMN]).shape == (4, 62) + builder.validate_episode_parquet( + parquet_path, + expected_frames=4, + verify_source_fk=True, + converter=_fake_eef_converter, + ) + + +def test_metadata_stats_and_backups_are_atomic(tmp_path: Path) -> None: + dataset = tmp_path / "dataset" + _write_fixture_dataset(dataset) + track_path = tmp_path / "tracks" / "episode_000000.npz" + _write_track_npz(track_path) + builder.build_episode( + dataset_root=dataset, + episode_index=0, + track_path=track_path, + converter=_fake_eef_converter, + ) + + result = builder.update_metadata(dataset) + assert result["complete"] is True + for name in ("info.json", "modality.json", "stats.json"): + assert builder.backup_path(dataset / "meta" / name).is_file() + info = json.loads((dataset / "meta" / "info.json").read_text()) + modality = json.loads((dataset / "meta" / "modality.json").read_text()) + stats = json.loads((dataset / "meta" / "stats.json").read_text()) + assert info["features"]["observation.state"]["shape"] == [58] + assert info["trex_track_force"]["track_layout"]["point_slices"] == [0, 100, 175, 250] + assert info["trex_track_force"]["sampling_20hz"]["target_rate_hz"] == 20.0 + assert info["trex_track_force"]["sampling_20hz"]["source_data_overwritten"] is False + assert info["trex_track_force"]["autoregressive_training"] == { + "blocks": 4, + "action_steps_per_block": 16, + "action_steps_per_sample": 64, + "video_conditioning_frames": 1, + "video_frames_per_block": 8, + "video_frames_per_sample": 33, + } + assert info["trex_track_force"]["force_only"] == builder._force_only_metadata() + assert set(modality["track"]["views"]) == set(VIEW_ORDER) + assert modality["track"]["xy"]["original_key"] == builder.TRACK_XY_COLUMN + assert ( + modality["track"]["visibility"]["original_key"] + == builder.TRACK_VISIBILITY_COLUMN + ) + assert "joint" in modality["state"] and "joint" in modality["action"] + assert modality["state"]["eef62"]["original_key"] == builder.STATE_EEF_COLUMN + assert ( + modality["action"]["eef62_absolute"]["original_key"] + == builder.ACTION_EEF_COLUMN + ) + assert modality["action"]["eef62"]["absolute"] is False + assert modality["force"]["current"]["reshape"] == [10, 6] + assert modality["force"]["history"]["frames"] == 16 + assert modality["force"]["history"]["encoding"] == "online_model_encoder" + assert modality["force"]["history"]["vq_codes_on_disk"] is False + assert "observation.force_history_vq" not in info["features"] + assert "timestamp" in stats + assert all(column in stats for column in builder.NEW_COLUMNS) + assert len(stats[builder.FORCE_COLUMN]["q99"]) == 60 + relative_stats = json.loads( + (dataset / "meta" / builder.RELATIVE_ACTION_STATS_FILENAME).read_text() + ) + assert len(relative_stats["eef62"]["q99"]) == 62 + builder.validate_metadata(dataset) + + +def test_track_validation_rejects_range_and_rgb_payloads(tmp_path: Path) -> None: + invalid_range = tmp_path / "invalid_range.npz" + _write_track_npz(invalid_range, out_of_range=True) + with pytest.raises(builder.DatasetSchemaError, match=r"\[0,1\]"): + builder.load_track_payload(invalid_range, expected_frames=4, episode_index=0) + + valid = tmp_path / "valid.npz" + _write_track_npz(valid) + with np.load(valid, allow_pickle=False) as archive: + payload = {key: np.asarray(archive[key]) for key in archive.files} + payload["images_head_left"] = np.zeros((4, 2, 2, 3), dtype=np.uint8) + forbidden = tmp_path / "forbidden_rgb.npz" + with forbidden.open("wb") as file: + np.savez_compressed(file, **payload) + with pytest.raises(builder.DatasetSchemaError, match="full RGB"): + builder.load_track_payload(forbidden) + + +def test_dry_run_does_not_mutate_dataset(tmp_path: Path) -> None: + dataset = tmp_path / "dataset" + _write_fixture_dataset(dataset) + before = (dataset / "meta" / "info.json").read_bytes() + exit_code = builder.main( + [ + "--dataset-root", + str(dataset), + "--episode-index", + "0", + "--dry-run", + "--no-extract-missing", + ] + ) + assert exit_code == 0 + assert (dataset / "meta" / "info.json").read_bytes() == before + assert not (dataset / "meta" / "trex_track_force_manifest.json").exists() + assert not builder.backup_path(builder.episode_parquet_path(dataset, 0)).exists() + + +def test_main_writes_manifest_and_skips_valid_output( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + dataset = tmp_path / "dataset" + _write_fixture_dataset(dataset) + track_cache = tmp_path / "tracks" + _write_track_npz(track_cache / "episode_000000.npz") + monkeypatch.setattr(builder, "joint58_to_eef62_batch", _fake_eef_converter) + arguments = [ + "--dataset-root", + str(dataset), + "--track-cache", + str(track_cache), + "--episode-index", + "0", + "--no-extract-missing", + ] + assert builder.main(arguments) == 0 + parquet_path = builder.episode_parquet_path(dataset, 0) + first_mtime = parquet_path.stat().st_mtime_ns + + assert builder.main(arguments) == 0 + assert parquet_path.stat().st_mtime_ns == first_mtime + manifest = json.loads( + (dataset / "meta" / "trex_track_force_manifest.json").read_text() + ) + assert manifest["schema_version"] == "trex_track_force_v2.3" + assert manifest["episodes"]["000000"]["status"] == "complete" + assert manifest["episodes"]["000000"]["skipped"] is True + sampling = manifest["episodes"]["000000"]["sampling_20hz"] + assert sampling["source_rate_hz"] == pytest.approx(30.0) + assert sampling["target_rate_hz"] == 20.0 + assert sampling["action_chunk_steps"] == 16 + assert sampling["action_chunk_duration_seconds"] == pytest.approx(0.8) + assert sampling["action_chunk_timestamp_span_seconds"] == pytest.approx(0.75) + assert manifest["episodes"]["000000"]["force_only"]["reshape"] == [10, 6] + assert manifest["track_layout"]["point_slices"] == [0, 100, 175, 250] + assert manifest["autoregressive_training"]["action_steps_per_sample"] == 64 + assert manifest["autoregressive_training"]["video_frames_per_sample"] == 33 + assert builder.main([*arguments, "--validate-only"]) == 0 + validate_output = capsys.readouterr().out + assert "Hz target=20.0Hz" in validate_output + assert "chunk=16 steps/0.8s (timestamp span 0.75s)" in validate_output + + +def test_stale_manifest_resets_breakpoint_instead_of_skipping(tmp_path: Path) -> None: + manifest_path = tmp_path / "manifest.json" + manifest_path.write_text( + json.dumps( + { + "schema_version": "trex_track_force_v2.1", + "episodes": {"000000": {"status": "complete"}}, + } + ) + ) + fresh = builder.load_manifest( + manifest_path, + dataset_root=tmp_path / "dataset", + track_cache=tmp_path / "tracks", + ) + assert fresh["schema_version"] == builder.SCHEMA_VERSION + assert fresh["episodes"] == {} + assert fresh["supersedes_schema_version"] == "trex_track_force_v2.1" + assert fresh["stale_episode_entries_discarded"] == 1 + + +def test_tactile_force_validation_rejects_nonfinite_values() -> None: + force = np.zeros((2, 60), dtype=np.float32) + force[1, 5] = np.nan + table = pa.table( + { + builder.FORCE_COLUMN: pa.array( + force.tolist(), type=pa.list_(pa.float32()) + ) + } + ) + with pytest.raises(builder.DatasetSchemaError, match="NaN/Inf"): + builder.validate_tactile_force(table) diff --git a/tests/trex_track_force/test_force_schedule.py b/tests/trex_track_force/test_force_schedule.py new file mode 100644 index 0000000000000000000000000000000000000000..32a54af521c2c2cb7acaf68cebde738776404b24 --- /dev/null +++ b/tests/trex_track_force/test_force_schedule.py @@ -0,0 +1,497 @@ +import pytest +import torch +from transformers.feature_extraction_utils import BatchFeature + +from groot.vla.model.trex_track_force.force import ( + FORCE_OFFSETS, + ForceOnlyOutput, + ForceOnlyTransformer, + TwoStageFlowSchedule, + force_offsets, + pad_action_62_to_64, + run_force_stage, + sample_trex_timestep, +) +from groot.vla.model.trex_track_force.lora import LoRALinear +from groot.vla.model.trex_track_force.policy import ( + WANTrackForcePolicyHead, + WANTrackForcePolicyHeadConfig, +) +from groot.vla.model.trex_track_force.tactile_vq import EMAQuantizer + + +def _small_policy_config(**kwargs): + return WANTrackForcePolicyHeadConfig( + train_architecture="full", + diffusion_model_cfg={ + "dim": 16, + "ffn_dim": 32, + "freq_dim": 8, + "num_heads": 4, + "num_layers": 1, + "obs_token_dim": 8, + "max_track_history": 2, + "max_blocks": 1, + }, + force_model_cfg={ + "dim": 16, + "num_heads": 4, + "num_layers": 1, + "ffn_dim": 32, + "vq_codebook_size": 64, + "vq_code_dim": 8, + "coarse_memory_dim": 16, + }, + **kwargs, + ) + + +def test_three_view_grid_is_resized_to_wan22_token_geometry(): + policy = WANTrackForcePolicyHead(_small_policy_config()) + + class RecordingVAE(torch.nn.Module): + def __init__(self): + super().__init__() + self.seen_shape = None + + def encode(self, images, tiled=False): + self.seen_shape = tuple(images.shape) + return images[:, :1] + + vae = RecordingVAE() + policy.vae = vae + result = policy.encode_video(torch.zeros(2, 3, 9, 320, 640)) + assert vae.seen_shape == (2, 3, 9, 160, 320) + assert result.shape == (2, 1, 9, 160, 320) + + +def test_component_coverage_is_measured_against_target_not_extra_source_keys( + tmp_path, +): + target = torch.nn.Linear(3, 2) + source = { + "weight": torch.randn_like(target.weight), + "bias": torch.randn_like(target.bias), + "unused_text_branch": torch.randn(100), + } + checkpoint = tmp_path / "component.pt" + torch.save(source, checkpoint) + report = WANTrackForcePolicyHead._load_component_state( + target, str(checkpoint), minimum_numel_coverage=1.0 + ) + assert report["numel_coverage"] == 1.0 + + torch.save({"weight": source["weight"]}, checkpoint) + with pytest.raises(ValueError, match="checkpoint coverage"): + WANTrackForcePolicyHead._load_component_state( + target, str(checkpoint), minimum_numel_coverage=1.0 + ) + + +def test_two_stage_schedule_and_stride_offsets_are_exact(): + schedule = TwoStageFlowSchedule() + coarse = list(schedule.iter_steps("coarse")) + force = list(schedule.iter_steps("force")) + assert schedule.num_steps == 10 + assert len(coarse) == 6 + assert len(force) == 4 + assert coarse[0].tau == 1.0 + assert coarse[-1].tau_next == pytest.approx(0.4) + assert force[0].tau == pytest.approx(0.4) + assert force[-1].tau_next == 0.0 + assert schedule.offsets == FORCE_OFFSETS == force_offsets(16, 4) + + +def test_trex_training_timesteps_use_beta_distribution_and_split_scaling(): + torch.manual_seed(19) + action_tau = sample_trex_timestep( + (20_000,), device="cpu", dtype=torch.float32 + ) + tactile_tau = sample_trex_timestep( + (20_000,), device="cpu", dtype=torch.float32, upper=0.4 + ) + + assert 0 < action_tau.min() <= action_tau.max() <= 1.0 + assert 0 < tactile_tau.min() <= tactile_tau.max() <= 0.4 + assert action_tau.mean().item() == pytest.approx(0.6, abs=0.01) + assert tactile_tau.mean().item() == pytest.approx(0.24, abs=0.01) + + +def test_ema_quantizer_accepts_bf16_encoder_with_fp32_codebook(): + quantizer = EMAQuantizer(codebook_size=64, embed_dim=8).train() + encoded = torch.randn(20, 8, dtype=torch.bfloat16, requires_grad=True) + + output = quantizer(encoded) + + assert quantizer.embedding.dtype == torch.float32 + assert output.quantized.dtype == torch.bfloat16 + assert output.commitment_loss.dtype == torch.float32 + (output.quantized.float().square().mean() + output.commitment_loss).backward() + assert encoded.grad is not None + assert torch.isfinite(quantizer.embedding).all() + + +def test_force_transformer_bf16_path_keeps_ema_buffers_fp32(): + model = ForceOnlyTransformer( + dim=16, + num_heads=4, + num_layers=1, + ffn_dim=32, + vq_code_dim=8, + ).eval().to(torch.bfloat16) + model.tactile_vq.quantizer.float() + + with torch.no_grad(): + output = model( + torch.randn(2, 16, 62, dtype=torch.bfloat16), + torch.tensor([0.4, 0.2], dtype=torch.bfloat16), + torch.randn(2, 10, 6, dtype=torch.bfloat16), + torch.randn(2, 16, 10, 6, dtype=torch.bfloat16), + return_details=True, + ) + + assert isinstance(output, ForceOnlyOutput) + assert output.flow.dtype == torch.bfloat16 + assert output.commitment_loss.dtype == torch.float32 + + +def test_force_transformer_raw_history_vq_and_precomputed_codes(): + torch.manual_seed(5) + model = ForceOnlyTransformer( + dim=32, + num_heads=4, + num_layers=1, + ffn_dim=64, + vq_codebook_size=64, + vq_code_dim=12, + coarse_memory_dim=24, + ) + action62 = torch.randn(2, 16, 62, requires_grad=True) + current_force = torch.randn(2, 10, 6) + raw_history = torch.randn(2, 16, 10, 6) + coarse_memory = torch.randn(2, 7, 24, requires_grad=True) + details = model( + action62, + torch.tensor([0.4, 0.2]), + current_force, + raw_history, + coarse_memory=coarse_memory, + return_details=True, + ) + assert isinstance(details, ForceOnlyOutput) + assert details.flow.shape == (2, 16, 64) + assert details.tactile_codes.shape == (2, 10) + assert details.vq_loss.ndim == details.commitment_loss.ndim == 0 + ( + details.flow.square().mean() + + details.vq_loss + + 0.25 * details.commitment_loss + ).backward() + assert action62.grad is not None + assert model.tactile_vq.encoder.stem[0].weight.grad is not None + assert model.tactile_vq.decoder.head.weight.grad is not None + assert details.tactile_reconstruction.shape == raw_history.shape + assert details.tactile_perplexity.ndim == 0 + assert details.tactile_active_codes.ndim == 0 + assert not model.tactile_vq.quantizer.embedding.requires_grad + assert coarse_memory.grad is None + + codes = details.tactile_codes.detach() + code_flow = model(action62.detach(), 0.25, current_force, codes) + assert code_flow.shape == (2, 16, 64) + slot_flow = model( + action62.detach(), + 0.25, + current_force, + codes, + update_offset=torch.tensor([4, 12]), + ) + assert slot_flow.shape == (2, 16, 64) + with pytest.raises(ValueError, match="one of"): + model(action62.detach(), 0.25, current_force, codes, update_offset=3) + with pytest.raises(ValueError, match=r"\[B,10\]"): + model(action62.detach(), 0.25, current_force, torch.zeros(2, 16).long()) + + +def test_force_vq_reconstruction_ignores_padded_history_steps(): + torch.manual_seed(11) + model = ForceOnlyTransformer( + dim=16, + num_heads=4, + num_layers=1, + ffn_dim=32, + vq_code_dim=8, + ).eval() + history = torch.randn(2, 16, 10, 6) + valid = torch.ones(2, 16, dtype=torch.bool) + valid[0, :8] = False + valid[1, :4] = False + details = model( + torch.randn(2, 16, 62), + 0.25, + torch.randn(2, 10, 6), + history, + tactile_history_valid_mask=valid, + return_details=True, + ) + assert isinstance(details, ForceOnlyOutput) + weight = valid[:, :, None, None] + expected = ( + (details.tactile_reconstruction - history).square() * weight + ).sum() / (weight.sum() * 10 * 6) + torch.testing.assert_close(details.vq_loss, expected) + + +def test_tactile_dropout_removes_both_fresh_and_history_conditions(): + torch.manual_seed(7) + model = ForceOnlyTransformer( + dim=16, + num_heads=4, + num_layers=1, + ffn_dim=32, + vq_code_dim=8, + ).eval() + action = torch.randn(2, 16, 62) + keep_none = torch.zeros(2, dtype=torch.bool) + first = model( + action, + 0.2, + torch.randn(2, 10, 6), + torch.randn(2, 16, 10, 6), + tactile_keep_mask=keep_none, + ) + second = model( + action, + 0.2, + torch.randn(2, 10, 6) * 100, + torch.randn(2, 16, 10, 6) * 100, + tactile_keep_mask=keep_none, + ) + torch.testing.assert_close(first, second) + + +def test_force_stage_updates_only_action_tensor(): + action = torch.zeros(1, 16, 62) + video = torch.randn(1, 3, 4) + track = torch.randn(1, 250, 16, 2) + video_before, track_before = video.clone(), track.clone() + refined = run_force_stage( + action, lambda padded_action, tau: torch.ones_like(padded_action) + ) + assert refined.shape == (1, 16, 64) + torch.testing.assert_close(refined, -0.4 * torch.ones_like(refined)) + torch.testing.assert_close(video, video_before) + torch.testing.assert_close(track, track_before) + assert pad_action_62_to_64(action).shape[-1] == 64 + + +def test_lora_mode_keeps_track_force_and_adapters_trainable(): + config = WANTrackForcePolicyHeadConfig( + train_architecture="lora", + diffusion_model_cfg={ + "dim": 32, + "ffn_dim": 64, + "freq_dim": 16, + "num_heads": 4, + "num_layers": 1, + "obs_token_dim": 12, + "max_track_history": 4, + "max_blocks": 1, + }, + force_model_cfg={ + "dim": 32, + "num_heads": 4, + "num_layers": 1, + "ffn_dim": 64, + "vq_codebook_size": 64, + "vq_code_dim": 8, + "coarse_memory_dim": 32, + }, + ) + policy = WANTrackForcePolicyHead(config) + trainable = { + name for name, parameter in policy.named_parameters() if parameter.requires_grad + } + assert any(name.startswith("model.track_encoder") for name in trainable) + assert any(name.startswith("model.track_decoder") for name in trainable) + assert any(name.startswith("force_transformer") for name in trainable) + assert any("lora_A" in name or "lora_B" in name for name in trainable) + assert any("cross_attn.k_img.lora_" in name for name in trainable) + assert any(name.startswith("model.img_emb") for name in trainable) + assert not any( + name.startswith("model.obs_input_projection") + or name.startswith("model.obs_token_decoder") + for name in trainable + ) + assert any(isinstance(module, LoRALinear) for module in policy.model.modules()) + assert policy.model.merge_and_unload() is policy.model + assert not any(isinstance(module, LoRALinear) for module in policy.model.modules()) + + +def test_policy_training_uses_extra_detached_tau_half_memory_and_raw_vq(): + torch.manual_seed(19) + policy = WANTrackForcePolicyHead( + _small_policy_config( + require_force_inputs=True, + tactile_dropout=0.25, + vq_loss_weight=0.5, + ) + ) + observed_tau = [] + observed_offsets = [] + observed_history_valid = [] + original_forward = policy.model.forward + original_force_forward = policy.force_transformer.forward + + def record_tau(*args, **kwargs): + observed_tau.append(torch.as_tensor(kwargs["tau"]).detach().clone()) + return original_forward(*args, **kwargs) + + policy.model.forward = record_tau + + def record_force_offset(*args, **kwargs): + observed_offsets.append( + torch.as_tensor(kwargs["update_offset"]).detach().cpu().clone() + ) + observed_history_valid.append( + kwargs["tactile_history_valid_mask"].detach().cpu().clone() + ) + return original_force_forward(*args, **kwargs) + + policy.force_transformer.forward = record_force_offset + output = policy.forward_core( + { + "obs_tokens": torch.randn(1, 1, 3, 8), + "action": torch.randn(1, 16, 62).clamp(-1, 1), + "action_mask": torch.ones(1, 16, 62, dtype=torch.bool), + "state": torch.randn(1, 1, 62), + "track_past_xy": torch.rand(1, 1, 2, 250, 2), + "track_past_visibility": torch.ones(1, 1, 2, 250), + "track_future_xy": torch.rand(1, 1, 16, 250, 2), + "track_future_visibility": torch.ones(1, 1, 16, 250), + "current_force": torch.randn(1, 1, 4, 10, 6), + "tactile_force_history": torch.randn(1, 1, 4, 16, 10, 6), + "tactile_force_history_padding_mask": torch.tensor( + [[[[True] * 3 + [False] * 13] * 4]] + ), + }, + tau=torch.tensor([[0.8]]), + ) + assert len(observed_tau) == 2 + assert len(observed_offsets) == 1 + torch.testing.assert_close( + observed_offsets[0], torch.tensor([0, 4, 8, 12]) + ) + assert observed_history_valid[0].shape == (4, 16) + assert not observed_history_valid[0][:, :3].any() + assert observed_history_valid[0][:, 3:].all() + torch.testing.assert_close(observed_tau[0], torch.tensor([[0.8]])) + torch.testing.assert_close(observed_tau[1], torch.tensor([[0.4]])) + assert output["loss"].ndim == 0 + assert output["action_flow"].shape == (1, 1, 16, 64) + assert output["track_flow"].shape == (1, 1, 250, 16, 2) + assert output["vq_loss"].item() > 0 + + +def test_policy_two_stage_sampling_returns_16_steps_with_ten_codes(): + policy = WANTrackForcePolicyHead(_small_policy_config()).eval() + past_track = torch.rand(1, 1, 2, 250, 2) + result = policy.get_action( + BatchFeature(data={"backbone_marker": torch.tensor(1)}), + BatchFeature( + data={ + "obs_tokens": torch.zeros(1, 1, 2, 8), + "state": torch.zeros(1, 1, 62), + "track_past_xy": past_track, + "track_past_visibility": torch.ones(1, 1, 2, 250), + "current_force": torch.zeros(1, 10, 6), + "force_history_vq": torch.zeros(1, 10, dtype=torch.long), + } + ), + ) + assert isinstance(result, BatchFeature) + assert result["action_pred"].shape == (1, 16, 64) + assert result["track_pred"].shape == (1, 250, 16, 2) + torch.testing.assert_close(result["track_pred"][:, :, 0], past_track[:, 0, -1]) + assert result["obs_pred"].shape == (1, 1, 2, 8) + + +def test_track_target_is_absolute_xy_and_starts_at_current_gt_anchor(): + policy = WANTrackForcePolicyHead(_small_policy_config()).eval() + past = torch.zeros(1, 1, 2, 250, 2) + past[:, :, -1] = 0.5 + future = torch.full((1, 1, 16, 250, 2), 0.5) + future[:, :, 0] = 0.9 # The contract must replace this with current GT. + future[:, :, -1, :, 0] = 0.51 + + target, anchor = policy._prepare_absolute_track_target(future, past) + prediction = target.clone() + prediction[:, :, 0] = 0.1 + prediction = policy._anchor_track_prediction(prediction, anchor) + + torch.testing.assert_close(target[:, :, 0], anchor) + torch.testing.assert_close(target[:, :, -1], future[:, :, -1]) + torch.testing.assert_close(prediction[:, :, 0], anchor) + torch.testing.assert_close(prediction[:, :, 1:], target[:, :, 1:]) + + +def test_visualization_sampling_uses_coarse_and_clean_boundary_without_force(): + policy = WANTrackForcePolicyHead(_small_policy_config()).eval() + model_calls = 0 + observation_steps = [] + original_model_forward = policy.model.forward + original_advance = policy._advance_sampling_observation + + def record_model_forward(*args, **kwargs): + nonlocal model_calls + model_calls += 1 + return original_model_forward(*args, **kwargs) + + def reject_force_forward(*args, **kwargs): + raise AssertionError("visualization must not run tactile refinement") + + def record_observation_step(observation, output, tau, tau_next): + observation_steps.append((float(tau), float(tau_next))) + return original_advance(observation, output, tau, tau_next) + + policy.model.forward = record_model_forward + policy.force_transformer.forward = reject_force_forward + policy._advance_sampling_observation = record_observation_step + past_track = torch.rand(1, 1, 2, 250, 2) + result = policy.get_visualization( + BatchFeature(data={"backbone_marker": torch.tensor(1)}), + BatchFeature( + data={ + "obs_tokens": torch.zeros(1, 1, 2, 8), + "state": torch.zeros(1, 1, 62), + "track_past_xy": past_track, + "track_past_visibility": torch.ones(1, 1, 2, 250), + } + ), + coarse_steps=1, + seed=42, + ) + + assert model_calls == 2 + assert observation_steps[0] == pytest.approx((1.0, 0.4)) + assert observation_steps[1] == pytest.approx((0.4, 0.0)) + assert result["action_pred"].shape == (1, 16, 64) + assert result["track_pred"].shape == (1, 250, 16, 2) + torch.testing.assert_close(result["track_pred"][:, :, 0], past_track[:, 0, -1]) + assert result["obs_pred"].shape == (1, 1, 2, 8) + + +def test_5hz_refresh_redenoises_only_unexecuted_stride4_suffix(): + policy = WANTrackForcePolicyHead(_small_policy_config()).eval() + coarse = torch.randn(1, 16, 64) + memory = torch.randn(1, 5, 16) + refined = policy.refine_action_suffix( + coarse, + memory, + torch.randn(1, 10, 6), + torch.zeros(1, 10, dtype=torch.long), + update_offset=4, + ) + torch.testing.assert_close(refined[:, :4, :62], coarse[:, :4, :62]) + assert torch.count_nonzero(refined[..., 62:]) == 0 + assert not torch.allclose(refined[:, 4:], coarse[:, 4:]) diff --git a/tests/trex_track_force/test_rebuild_dataset_variants.py b/tests/trex_track_force/test_rebuild_dataset_variants.py new file mode 100644 index 0000000000000000000000000000000000000000..21bf0fe367972bcc0794561f9beff7e8ad6b29c8 --- /dev/null +++ b/tests/trex_track_force/test_rebuild_dataset_variants.py @@ -0,0 +1,197 @@ +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq + +_ROOT = Path(__file__).resolve().parents[2] +_SCRIPT_ROOT = _ROOT / "scripts" / "data" +if str(_SCRIPT_ROOT) not in sys.path: + sys.path.insert(0, str(_SCRIPT_ROOT)) + +import rebuild_trex_dataset_variants as rebuild # noqa: E402 +import check_trex_dataset_ready as ready_check # noqa: E402 + + +def test_build_registry_filters_and_reindexes_densely() -> None: + episodes = [ + {"episode_index": 0, "tasks": ["zero"], "length": 2}, + {"episode_index": 1, "tasks": ["removed"], "length": 3}, + {"episode_index": 2, "tasks": ["two"], "length": 4}, + ] + tasks = [ + {"task_index": 0, "task": "zero"}, + {"task_index": 1, "task": "removed"}, + {"task_index": 2, "task": "two"}, + ] + + registry, task_rows = rebuild._build_registry( + episodes, + tasks, + excluded={1}, + source_limit=None, + ) + + assert [record.source_episode_index for record in registry] == [0, 2] + assert [record.episode_index for record in registry] == [0, 1] + assert [record.task_index for record in registry] == [0, 1] + assert [record.global_index_start for record in registry] == [0, 2] + assert task_rows == [ + {"task_index": 0, "task": "zero"}, + {"task_index": 1, "task": "two"}, + ] + + +def test_rewrite_parquet_updates_all_index_columns(tmp_path: Path) -> None: + source = tmp_path / "source.parquet" + destination = tmp_path / "out" / "episode_000001.parquet" + table = pa.table( + { + "episode_index": pa.array([2, 2, 2], type=pa.int64()), + "task_index": pa.array([5, 5, 5], type=pa.int64()), + "frame_index": pa.array([0, 1, 2], type=pa.int64()), + "index": pa.array([10, 11, 12], type=pa.int64()), + "value": pa.array([1.0, 2.0, 3.0], type=pa.float32()), + } + ).replace_schema_metadata( + { + rebuild.PARQUET_SCHEMA_METADATA_KEY: ( + rebuild.force_builder.SCHEMA_VERSION.encode() + ) + } + ) + pq.write_table(table, source) + + episode_index, status = rebuild._rewrite_parquet_job( + ( + str(source), + str(destination), + 1, + 3, + 7, + 3, + True, + ) + ) + + assert episode_index == 1 + assert status == "written" + result = pq.read_table(destination) + assert result["episode_index"].to_pylist() == [1, 1, 1] + assert result["task_index"].to_pylist() == [3, 3, 3] + assert result["frame_index"].to_pylist() == [0, 1, 2] + assert result["index"].to_pylist() == [7, 8, 9] + assert ( + result.schema.metadata[rebuild.PARQUET_SCHEMA_METADATA_KEY] + == rebuild.force_builder.SCHEMA_VERSION.encode() + ) + + +def test_rewrite_track_changes_only_episode_identity(tmp_path: Path) -> None: + source = tmp_path / "episode_000002.npz" + destination = tmp_path / "tracks" / "episode_000001.npz" + tracks = np.arange(12, dtype=np.float32).reshape(3, 2, 2) + np.savez_compressed( + source, + tracks=tracks, + vis=np.ones((3, 2), dtype=np.float32), + episode_index=np.array(2, dtype=np.int32), + num_steps=np.array(3, dtype=np.int32), + ) + + episode_index, digest, status = rebuild._rewrite_track_job( + (str(source), str(destination), 2, 1, 3) + ) + + assert episode_index == 1 + assert len(digest) == 64 + assert status == "written" + with np.load(source, allow_pickle=False) as payload: + assert int(payload["episode_index"]) == 2 + with np.load(destination, allow_pickle=False) as payload: + assert int(payload["episode_index"]) == 1 + assert int(payload["source_episode_index"]) == 2 + np.testing.assert_array_equal(payload["tracks"], tracks) + + +def test_hardlink_is_resumable(tmp_path: Path) -> None: + source = tmp_path / "source.bin" + destination = tmp_path / "nested" / "destination.bin" + source.write_bytes(b"trex") + + assert rebuild._hardlink(source, destination) == "linked" + assert rebuild._hardlink(source, destination) == "skipped" + assert source.stat().st_ino == destination.stat().st_ino + + +def test_fast_ready_check_validates_force_samples(tmp_path: Path) -> None: + root = tmp_path / "trex_small_force" + (root / "meta").mkdir(parents=True) + info = { + "total_episodes": 3, + "total_frames": 3, + "total_tasks": 1, + "total_videos": 0, + } + (root / "meta" / "info.json").write_text(json.dumps(info)) + (root / "meta" / "dataset_ready.json").write_text( + json.dumps( + { + "episodes": 3, + "frames": 3, + "force": True, + "validated_at": "now", + } + ) + ) + (root / "meta" / "episodes.jsonl").write_text( + "".join( + json.dumps( + {"episode_index": index, "tasks": ["task"], "length": 1} + ) + + "\n" + for index in range(3) + ) + ) + (root / "meta" / "tasks.jsonl").write_text( + json.dumps({"task_index": 0, "task": "task"}) + "\n" + ) + for name in ( + "stats.json", + "relative_stats_dreamzero.json", + "source_episode_index_map.json", + ): + (root / "meta" / name).write_text("{}") + manifest = { + "episodes": { + f"{index:06d}": {"status": "complete"} + for index in range(3) + } + } + (root / "meta" / "trex_track_force_manifest.json").write_text( + json.dumps(manifest) + ) + + for index in range(3): + parquet_path = rebuild._episode_parquet_path(root, index) + parquet_path.parent.mkdir(parents=True, exist_ok=True) + table = pa.table({"value": [index]}).replace_schema_metadata( + {ready_check.FORCE_SCHEMA_METADATA_KEY: b"test"} + ) + pq.write_table(table, parquet_path) + track_path = rebuild._track_path(root, index) + track_path.parent.mkdir(parents=True, exist_ok=True) + np.savez_compressed( + track_path, + episode_index=np.array(index, dtype=np.int32), + num_steps=np.array(1, dtype=np.int32), + ) + + result = ready_check.check_dataset(root, require_force=True) + + assert result["episodes"] == 3 + assert result["force"] is True diff --git a/tests/trex_track_force/test_track.py b/tests/trex_track_force/test_track.py new file mode 100644 index 0000000000000000000000000000000000000000..5b1d3636945200040370fec28b439ff25c89b7b2 --- /dev/null +++ b/tests/trex_track_force/test_track.py @@ -0,0 +1,425 @@ +import json +import copy +from pathlib import Path + +import pytest +import torch +from safetensors import safe_open + +from groot.vla.model.dreamzero.modules.wan2_1_submodule import rope_params +from groot.vla.model.dreamzero.modules.wan_video_dit_action_casual_chunk import ( + CausalWanModel as DreamZeroCausalWanModel, +) +from groot.vla.model.trex_track_force.blocks import apply_multimodal_rope +from groot.vla.model.trex_track_force.model import CausalWanTrackForceModel +from groot.vla.model.trex_track_force.track import ( + NUM_TRACK_POINTS, + TRACK_SEGMENTS, + TrackDecoder, + TrackEncoder, + canonical_point_metadata, + track_motion_weights, + visibility_masked_track_loss, + visibility_masked_track_smoothness_loss, + zero_invisible_track_xy, +) + + +def test_canonical_250_point_identity_layout(): + assert sum(segment.size for segment in TRACK_SEGMENTS) == 250 + assert [(segment.start, segment.stop) for segment in TRACK_SEGMENTS] == [ + (0, 50), + (50, 100), + (100, 125), + (125, 175), + (175, 200), + (200, 250), + ] + metadata = canonical_point_metadata() + assert all(value.shape == (NUM_TRACK_POINTS,) for value in metadata.values()) + assert metadata["point"].unique().numel() == NUM_TRACK_POINTS + + +def test_track_embedding_decoder_and_visibility_loss_shapes(): + torch.manual_seed(3) + batch, blocks, history, dim = 2, 2, 5, 32 + past_xy = torch.rand(batch, blocks, history, 250, 2) + past_vis = torch.randint(0, 2, (batch, blocks, history, 250)).float() + future_xy = torch.rand(batch, blocks, 16, 250, 2) + future_vis = torch.randint(0, 2, (batch, blocks, 16, 250)).float() + + encoder = TrackEncoder(dim, max_history_frames=history, max_blocks=blocks) + decoder = TrackDecoder(dim) + assert encoder.trajectory_projection[0].in_features == 16 * 5 + embedded = encoder(past_xy, past_vis, future_xy, future_vis) + assert embedded.past_tokens.shape == (batch, blocks, 250, dim) + assert embedded.future_tokens.shape == (batch, blocks, 250, dim) + flow = decoder(embedded.future_tokens) + assert flow.shape == (batch, blocks, 250, 16, 2) + + target = torch.zeros_like(flow) + prediction = target.clone() + prediction[:, :, 0] = 100.0 + visibility = torch.ones(batch, blocks, 250, 16) + visibility[:, :, 0] = 0 + assert visibility_masked_track_loss(prediction, target, visibility).item() == 0 + visibility[:, :, 0] = 1 + assert visibility_masked_track_loss(prediction, target, visibility).item() > 0 + + +def test_zero_invisible_and_motion_weighted_losses(): + torch.manual_seed(5) + batch, blocks, points, horizon = 1, 1, NUM_TRACK_POINTS, 16 + clean = torch.zeros(batch, blocks, points, horizon, 2) + clean[:, :, 0, :, 0] = torch.linspace(0, 1, horizon) + visibility = torch.ones(batch, blocks, points, horizon) + visibility[:, :, 1, :] = 0 + polluted = clean.clone() + polluted[:, :, 1] = 9.0 + cleaned = zero_invisible_track_xy(polluted, visibility) + assert cleaned[:, :, 1].abs().sum().item() == 0 + assert torch.allclose(cleaned[:, :, 0], clean[:, :, 0]) + + weights = track_motion_weights(clean, visibility, boost=3.0) + assert weights.shape == visibility.shape + assert weights[0, 0, 0, 0].item() > weights[0, 0, 2, 0].item() + + pred = clean.clone() + pred[:, :, 0, 1:, 0] += 0.1 + smooth = visibility_masked_track_smoothness_loss(pred, clean, visibility) + assert smooth.item() > 0 + + +def test_track_embedding_preserves_which_future_timestep_moved(): + torch.manual_seed(7) + encoder = TrackEncoder(32, max_history_frames=16, max_blocks=1).eval() + past = torch.zeros(1, 1, 16, 250, 2) + visibility = torch.ones(1, 1, 16, 250) + early = torch.zeros_like(past) + late = torch.zeros_like(past) + early[:, :, 2, 0, 0] = 0.5 + late[:, :, 8, 0, 0] = 0.5 + + early_token = encoder(past, visibility, early, visibility).future_tokens + late_token = encoder(past, visibility, late, visibility).future_tokens + + # The old visible-time average produced identical tokens here because both + # coordinate and first-difference multisets were the same. + assert not torch.allclose(early_token[:, :, 0], late_token[:, :, 0]) + + +def test_causal_wan_core_has_runnable_cpu_forward(): + torch.manual_seed(11) + model = CausalWanTrackForceModel( + dim=32, + ffn_dim=64, + freq_dim=16, + num_heads=4, + num_layers=1, + obs_token_dim=12, + max_track_history=4, + max_blocks=1, + gradient_checkpointing=False, + ) + future_noisy = torch.rand(1, 1, 16, 250, 2) + output = model( + obs_tokens=torch.randn(1, 1, 3, 12), + action=torch.randn(1, 1, 16, 62), + state=torch.randn(1, 1, 62), + past_track_xy=torch.rand(1, 1, 4, 250, 2), + past_track_visibility=torch.ones(1, 1, 4, 250), + future_track_xy=future_noisy, + future_track_visibility=torch.ones(1, 1, 16, 250), + tau=torch.tensor([0.7]), + ) + assert output.action_flow.shape == (1, 1, 16, 64) + assert output.track_flow.shape == (1, 1, 250, 16, 2) + assert output.obs_flow_tokens.shape == (1, 1, 3, 12) + assert output.coarse_memory.shape[-1] == 32 + + +def test_four_block_coarse_memory_is_causally_padded_per_block(): + torch.manual_seed(13) + blocks = 4 + model = CausalWanTrackForceModel( + dim=8, + ffn_dim=16, + freq_dim=8, + num_heads=2, + num_layers=1, + obs_token_dim=8, + max_track_history=2, + max_blocks=blocks, + max_chunk_size=blocks, + ).eval() + with torch.no_grad(): + output = model( + obs_tokens=torch.randn(1, blocks, 2, 8), + action=torch.randn(1, blocks, 16, 62), + state=torch.randn(1, blocks, 62), + past_track_xy=torch.rand(1, blocks, 2, 250, 2), + past_track_visibility=torch.ones(1, blocks, 2, 250), + future_track_xy=torch.rand(1, blocks, 16, 250, 2), + future_track_visibility=torch.ones(1, blocks, 16, 250), + tau=torch.full((1, blocks), 0.7), + ) + memory_per_block = 2 + 16 + 250 + 250 + assert output.coarse_memory_by_block.shape == ( + 1, + blocks, + blocks * memory_per_block, + 8, + ) + torch.testing.assert_close( + output.coarse_memory_mask_by_block.sum(dim=1), + torch.arange(1, blocks + 1) * memory_per_block, + ) + assert ( + torch.count_nonzero( + output.coarse_memory_by_block[:, 0, memory_per_block:] + ) + == 0 + ) + + +def test_real_wan_inheritance_and_shared_state_dict_compatibility(): + common = { + "model_type": "ti2v", + "patch_size": (1, 2, 2), + "frame_seqlen": 3, + "text_len": 8, + "in_dim": 8, + "dim": 32, + "ffn_dim": 64, + "freq_dim": 16, + "text_dim": 24, + "out_dim": 8, + "num_heads": 4, + "num_layers": 1, + "max_chunk_size": -1, + "sink_size": 0, + "qk_norm": True, + "cross_attn_norm": True, + "eps": 1e-6, + "num_frame_per_block": 1, + "action_dim": 64, + "max_state_dim": 64, + "hidden_size": 32, + "num_state_per_block": 1, + "concat_first_frame_latent": False, + } + original = DreamZeroCausalWanModel( + **common, + num_action_per_block=16, + ) + model = CausalWanTrackForceModel( + **common, + action_horizon=16, + obs_token_dim=12, + max_track_history=2, + max_blocks=1, + ) + assert isinstance(model, DreamZeroCausalWanModel) + + original_state = original.state_dict() + model_state = model.state_dict() + assert all(key in model_state for key in original_state) + assert all( + model_state[key].shape == value.shape + for key, value in original_state.items() + ) + report = model.pretrained_compatibility_report(original_state) + assert report["key_coverage"] == 1.0 + assert report["numel_coverage"] == 1.0 + assert report["shape_mismatches"] == {} + model.load_pretrained_wan_state_dict(original_state) + + +def test_local_wan22_ti2v_5b_checkpoint_has_exact_shared_shape_coverage(): + checkpoint = ( + Path(__file__).resolve().parents[2] / "checkpoints" / "Wan2.2-TI2V-5B" + ) + index_path = checkpoint / "diffusion_pytorch_model.safetensors.index.json" + if not index_path.is_file(): + pytest.skip("local Wan2.2-TI2V-5B checkpoint is unavailable") + + index = json.loads(index_path.read_text()) + source = {} + for shard_name in sorted(set(index["weight_map"].values())): + with safe_open( + checkpoint / shard_name, framework="pt", device="cpu" + ) as shard: + for key in shard.keys(): + source[key] = torch.empty( + tuple(shard.get_slice(key).get_shape()), device="meta" + ) + + with torch.device("meta"): + model = CausalWanTrackForceModel( + dim=3072, + ffn_dim=14336, + freq_dim=256, + num_heads=24, + num_layers=30, + frame_seqlen=50, + in_dim=48, + out_dim=48, + text_dim=4096, + action_dim=64, + action_horizon=16, + track_points=250, + track_horizon=16, + num_frame_per_block=2, + ) + report = model.pretrained_compatibility_report(source) + assert report["matched_keys"] == report["wan_checkpoint_keys"] == 825 + assert report["numel_coverage"] == 1.0 + assert report["shape_mismatches"] == {} + + +def test_multimodal_rope_rotates_nonzero_positions(): + values = torch.zeros(1, 3, 2, 8) + values[..., 0::2] = 1 + frequencies = rope_params(3, 8) + rotated = apply_multimodal_rope(values, frequencies) + torch.testing.assert_close(rotated[:, 0], values[:, 0]) + assert not torch.allclose(rotated[:, 1], values[:, 1]) + + +def test_inherited_wan_patch_and_head_latent_forward(): + model = CausalWanTrackForceModel( + patch_size=(1, 2, 2), + frame_seqlen=1, + in_dim=4, + out_dim=4, + dim=16, + ffn_dim=32, + freq_dim=8, + num_heads=4, + num_layers=1, + num_frame_per_block=2, + max_track_history=2, + max_blocks=1, + gradient_checkpointing=False, + ).eval() + output = model( + video_latents=torch.randn(1, 4, 3, 2, 2), + action=torch.randn(1, 1, 16, 62), + state=torch.randn(1, 1, 62), + past_track_xy=torch.rand(1, 1, 2, 250, 2), + past_track_visibility=torch.ones(1, 1, 2, 250), + future_track_xy=torch.rand(1, 1, 16, 250, 2), + future_track_visibility=torch.ones(1, 1, 16, 250), + tau=0.7, + ) + assert output.video_flow.shape == (1, 4, 2, 2, 2) + + +def test_separate_nine_frame_history_and_future_latents_keep_autoregressive_block(): + model = CausalWanTrackForceModel( + patch_size=(1, 2, 2), + frame_seqlen=1, + in_dim=4, + out_dim=4, + dim=16, + ffn_dim=32, + freq_dim=8, + num_heads=4, + num_layers=1, + num_frame_per_block=2, + max_track_history=2, + max_blocks=1, + gradient_checkpointing=False, + ).eval() + output = model( + conditioning_video_latents=torch.randn(1, 4, 3, 2, 2), + video_latents=torch.randn(1, 4, 2, 2, 2), + action=torch.randn(1, 1, 16, 62), + state=torch.randn(1, 1, 62), + past_track_xy=torch.rand(1, 1, 2, 250, 2), + past_track_visibility=torch.ones(1, 1, 2, 250), + future_track_xy=torch.rand(1, 1, 16, 250, 2), + future_track_visibility=torch.ones(1, 1, 16, 250), + tau=0.7, + ) + assert output.layout.num_blocks == 1 + assert output.layout.conditioning_obs_tokens == 3 + assert output.video_flow.shape == (1, 4, 2, 2, 2) + # clean history + future obs + action + past/future tracks + assert output.coarse_memory.shape[1] == 3 + 2 + 16 + 250 + 250 + + +def test_track_output_has_no_full_block_obs_action_state_leakage(): + torch.manual_seed(29) + model = CausalWanTrackForceModel( + dim=16, + ffn_dim=32, + freq_dim=8, + num_heads=4, + num_layers=1, + obs_token_dim=8, + max_track_history=2, + max_blocks=1, + gradient_checkpointing=False, + ).eval() + torch.nn.init.normal_(model.track_decoder.projection.weight, std=0.02) + tracks = { + "past_track_xy": torch.rand(1, 1, 2, 250, 2), + "past_track_visibility": torch.ones(1, 1, 2, 250), + "future_track_xy": torch.rand(1, 1, 16, 250, 2), + "future_track_visibility": torch.ones(1, 1, 16, 250), + } + first = model( + obs_tokens=torch.zeros(1, 1, 2, 8), + action=torch.zeros(1, 1, 16, 62), + state=torch.zeros(1, 1, 62), + tau=0.4, + **tracks, + ) + second = model( + obs_tokens=torch.randn(1, 1, 2, 8) * 100, + action=torch.randn(1, 1, 16, 62) * 100, + state=torch.randn(1, 1, 62) * 100, + tau=0.4, + **tracks, + ) + torch.testing.assert_close(first.track_flow, second.track_flow) + + +def test_gradient_checkpointing_recomputes_each_original_wan_block(): + torch.manual_seed(41) + plain = CausalWanTrackForceModel( + dim=16, + ffn_dim=32, + freq_dim=8, + num_heads=4, + num_layers=2, + obs_token_dim=8, + max_track_history=2, + max_blocks=1, + gradient_checkpointing=False, + ).train() + checkpointed = copy.deepcopy(plain) + checkpointed.gradient_checkpointing = True + inputs = { + "obs_tokens": torch.randn(1, 1, 2, 8), + "action": torch.randn(1, 1, 16, 62), + "state": torch.randn(1, 1, 62), + "past_track_xy": torch.rand(1, 1, 2, 250, 2), + "past_track_visibility": torch.ones(1, 1, 2, 250), + "future_track_xy": torch.rand(1, 1, 16, 250, 2), + "future_track_visibility": torch.ones(1, 1, 16, 250), + "tau": 0.7, + } + plain_output = plain(**inputs) + checkpointed_output = checkpointed(**inputs) + plain_loss = plain_output.action_flow.square().mean() + checked_loss = checkpointed_output.action_flow.square().mean() + plain_loss.backward() + checked_loss.backward() + torch.testing.assert_close(plain_output.action_flow, checkpointed_output.action_flow) + for block_index in range(2): + torch.testing.assert_close( + plain.blocks[block_index].self_attn.q.weight.grad, + checkpointed.blocks[block_index].self_attn.q.weight.grad, + ) diff --git a/tests/trex_track_force/test_training_logging.py b/tests/trex_track_force/test_training_logging.py new file mode 100644 index 0000000000000000000000000000000000000000..e49525d3031c875c329178890740c2954c05e03d --- /dev/null +++ b/tests/trex_track_force/test_training_logging.py @@ -0,0 +1,39 @@ +import json +from types import SimpleNamespace + +from groot.vla.experiment.base import LossLoggerCallback + + +def test_loss_logger_keeps_track_force_and_tactile_metrics(tmp_path): + output_path = tmp_path / "loss_log.jsonl" + callback = LossLoggerCallback(str(output_path)) + state = SimpleNamespace(is_world_process_zero=True, global_step=500) + + callback.on_log( + args=None, + state=state, + control=None, + logs={ + "loss": 1.0, + "dynamics_loss_avg": 0.1, + "action_loss_avg": 0.2, + "track_loss_avg": 0.3, + "force_loss_avg": 0.4, + "vq_loss_avg": 0.5, + "commitment_loss_avg": 0.6, + "tactile_perplexity_avg": 12.0, + "tactile_active_codes_avg": 8.0, + "learning_rate": 1e-5, + "ignored_metric": 99.0, + }, + ) + + record = json.loads(output_path.read_text().strip()) + assert record["step"] == 500 + assert record["track_loss_avg"] == 0.3 + assert record["force_loss_avg"] == 0.4 + assert record["vq_loss_avg"] == 0.5 + assert record["commitment_loss_avg"] == 0.6 + assert record["tactile_perplexity_avg"] == 12.0 + assert record["tactile_active_codes_avg"] == 8.0 + assert "ignored_metric" not in record diff --git a/tests/trex_track_force/test_visualization.py b/tests/trex_track_force/test_visualization.py new file mode 100644 index 0000000000000000000000000000000000000000..636b1d6767de3d6d4742e09a51ea6429520c8d4e --- /dev/null +++ b/tests/trex_track_force/test_visualization.py @@ -0,0 +1,243 @@ +from pathlib import Path +from types import ModuleType, SimpleNamespace +import sys + +import numpy as np +import torch + +from groot.vla.experiment import trex_wandb_video_callback as callback_module +from groot.vla.experiment import trex_eval_utils +from groot.vla.experiment.trex_track_force_eval_utils import ( + TrexTrackForceVisualization, + _decode_future_video, + overlay_track_motion, +) +from groot.vla.experiment.trex_wandb_video_callback import TrexWandbVideoCallback + + +def test_track_motion_overlay_outputs_only_three_view_panels(): + frames = np.zeros((8, 100, 200, 3), dtype=np.uint8) + tracks = np.full((1, 16, 250, 2), 0.5, dtype=np.float32) + tracks[0, :, :, 0] += np.linspace(-0.2, 0.2, 16)[:, None] + visualization = TrexTrackForceVisualization( + video_frames=frames, + predicted_track_xy=tracks, + target_track_xy=tracks.copy(), + target_track_visibility=np.ones((1, 16, 250), dtype=np.float32), + anchor_timestamps=np.array([0.0], dtype=np.float64), + frame_track_indices=np.array( + [2, 4, 6, 8, 10, 12, 14, 15], dtype=np.int64 + ), + ) + + overlaid = overlay_track_motion(visualization, trail_steps=8) + + assert overlaid.shape == (8, 50, 300, 3) + assert np.count_nonzero(overlaid[:, :, :100]) > 0 + assert np.count_nonzero(overlaid[:, :, 100:200]) > 0 + assert np.count_nonzero(overlaid[:, :, 200:]) > 0 + + +def test_track_motion_overlay_rejects_mismatched_video_indices(): + visualization = TrexTrackForceVisualization( + video_frames=np.zeros((2, 100, 200, 3), dtype=np.uint8), + predicted_track_xy=np.zeros((1, 16, 250, 2), dtype=np.float32), + target_track_xy=np.zeros((1, 16, 250, 2), dtype=np.float32), + target_track_visibility=np.ones((1, 16, 250), dtype=np.float32), + anchor_timestamps=np.array([0.0], dtype=np.float64), + frame_track_indices=np.array([1], dtype=np.int64), + ) + + try: + overlay_track_motion(visualization) + except ValueError as exc: + assert "video and frame-track index counts" in str(exc) + else: + raise AssertionError("mismatched frame-track indices must fail") + + +def test_track_video_starts_with_model_gt_anchor(): + conditioning = np.full((250, 2), 0.5, dtype=np.float32) + anchored = np.repeat(conditioning[None], 16, axis=0) + anchored[:, :, 0] += np.linspace(0.0, 0.1, 16)[:, None] + np.testing.assert_allclose(anchored[0], conditioning) + + visualization = TrexTrackForceVisualization( + video_frames=np.zeros((9, 100, 200, 3), dtype=np.uint8), + predicted_track_xy=anchored[None], + target_track_xy=anchored[None], + target_track_visibility=np.ones((1, 16, 250), dtype=np.float32), + anchor_timestamps=np.array([0.0], dtype=np.float64), + frame_track_indices=np.array( + [-1, 2, 4, 6, 8, 10, 12, 14, 15], dtype=np.int64 + ), + conditioning_frame=np.zeros((100, 200, 3), dtype=np.uint8), + conditioning_track_xy=conditioning[None], + conditioning_track_visibility=np.ones( + (1, 250), dtype=np.float32 + ), + ) + overlaid = overlay_track_motion(visualization) + assert overlaid.shape == (9, 50, 300, 3) + assert np.count_nonzero(overlaid[0]) > 0 + + +def test_callback_saves_plain_and_track_overlay_videos_separately( + tmp_path, monkeypatch +): + visualization = TrexTrackForceVisualization( + video_frames=np.zeros((8, 100, 200, 3), dtype=np.uint8), + predicted_track_xy=np.zeros((1, 16, 250, 2), dtype=np.float32), + target_track_xy=np.zeros((1, 16, 250, 2), dtype=np.float32), + target_track_visibility=np.ones((1, 16, 250), dtype=np.float32), + anchor_timestamps=np.array([0.0], dtype=np.float64), + frame_track_indices=np.array( + [2, 4, 6, 8, 10, 12, 14, 15], dtype=np.int64 + ), + ) + monkeypatch.setattr( + callback_module, + "run_trex_track_force_prediction", + lambda *args, **kwargs: visualization, + ) + monkeypatch.setattr( + callback_module, + "overlay_track_motion", + lambda *args, **kwargs: np.full_like( + visualization.video_frames, 255 + ), + ) + + written: dict[Path, np.ndarray] = {} + + def fake_mimsave(path, frames, **kwargs): + written[Path(path)] = np.stack(frames) + + monkeypatch.setattr(callback_module.imageio, "mimsave", fake_mimsave) + + logged: list[dict] = [] + fake_wandb = ModuleType("wandb") + fake_wandb.run = object() + fake_wandb.define_metric = lambda *args, **kwargs: None + fake_wandb.Video = lambda path, **kwargs: {"path": path, **kwargs} + fake_wandb.log = logged.append + monkeypatch.setitem(sys.modules, "wandb", fake_wandb) + + callback = TrexWandbVideoCallback.__new__(TrexWandbVideoCallback) + callback.episode_index = 0 + callback.num_chunks = 1 + callback.every_n_steps = 500 + callback.fps = 10 + callback.eval_bf16 = False + callback.use_dataset_prompt = False + callback.prompt = "test task" + callback.overlay_tracks = True + callback.save_tracks = True + callback.track_trail_steps = 8 + callback.reconstruction_inference_steps = 1 + callback.eval_video_dir = tmp_path / "eval_videos" + callback.eval_track_video_dir = tmp_path / "eval_track_videos" + callback.eval_track_dir = tmp_path / "eval_tracks" + callback._episode = SimpleNamespace( + episode_index=0, + get_task=lambda row: "dataset task", + ) + callback._eval_transform = object() + + class FakeModel: + training = True + + def eval(self): + self.training = False + + def train(self): + self.training = True + + model = FakeModel() + state = SimpleNamespace(is_world_process_zero=True, global_step=500) + callback.on_train_begin(None, state, None) + callback.on_step_end(None, state, None, model=model) + + plain_path = callback.eval_video_dir / "train_step_000500.mp4" + track_path = callback.eval_track_video_dir / "train_step_000500.mp4" + assert np.count_nonzero(written[plain_path]) == 0 + assert np.all(written[track_path] == 255) + assert (callback.eval_track_dir / "train_step_000500.npz").is_file() + assert TrexWandbVideoCallback.VIDEO_METRIC in logged[0] + assert TrexWandbVideoCallback.TRACK_VIDEO_METRIC in logged[0] + assert model.training + + +def test_episode_frame_reads_are_cached_across_visualization_chunks(monkeypatch): + opens = 0 + + class FakeCapture: + def __init__(self, path): + nonlocal opens + opens += 1 + self.position = 0 + + def isOpened(self): + return True + + def set(self, prop, value): + self.position = int(value) + + def read(self): + value = self.position + self.position += 1 + return True, np.full((2, 3, 3), value, dtype=np.uint8) + + def release(self): + pass + + monkeypatch.setattr(trex_eval_utils.cv2, "VideoCapture", FakeCapture) + episode = trex_eval_utils.TrexEpisode.__new__( + trex_eval_utils.TrexEpisode + ) + episode.length = 10 + episode.video_dirs = {"video.head_left": "fake.mp4"} + episode._frame_cache = {"video.head_left": {}} + + first = episode.get_frames([0, 1, 1, 2], "video.head_left") + second = episode.get_frames([1, 2], "video.head_left") + + assert opens == 1 + assert first[:, 0, 0, 0].tolist() == [0, 1, 1, 2] + assert second[:, 0, 0, 0].tolist() == [1, 2] + + +def test_track_force_video_decode_reads_tiling_from_policy_config(): + seen = {} + + class FakeVAE: + def decode(self, latents, **kwargs): + seen["shape"] = tuple(latents.shape) + seen.update(kwargs) + return torch.zeros(1, 3, 9, 4, 8) + + action_head = SimpleNamespace( + config=SimpleNamespace( + tiled=False, + tile_size_height=20, + tile_size_width=21, + tile_stride_height=10, + tile_stride_width=11, + ), + vae=FakeVAE(), + ) + model = SimpleNamespace(action_head=action_head) + + frames = _decode_future_video( + model, + torch.zeros(1, 4, 3, 2, 2), + torch.zeros(1, 4, 2, 2, 2), + ) + + assert frames.shape == (8, 4, 8, 3) + assert seen == { + "shape": (1, 4, 3, 2, 2), + "tiled": False, + "tile_size": (20, 21), + "tile_stride": (10, 11), + }