Upload dreamzero source code (no model checkpoints)

.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

.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*/

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.

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.

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<n<10M
+configs:
+- config_name: frames
+  data_files:
+  - split: train
+    path: data/chunk-*/file-*.parquet
+  default: true
+- config_name: episodes
+  data_files:
+  - split: train
+    path: episodes_preview.parquet
+---
+
+# T-Rex Dataset
+
+A large-scale, tactile-reactive bimanual manipulation dataset, collected via teleoperation on a
+Dexmate Vega-1 robot with two Sharpa Wave dexterous hands. Stored as a
+[LeRobotDataset v3.0](https://github.com/huggingface/lerobot).
+
+[🌐 Project Page](https://tactile-rex.github.io/) · [✍️ Paper (arXiv)](https://arxiv.org/abs/2606.17055) · [💻 Code (T-Rex)](https://github.com/ZhuoyangLiu2005/T-Rex) · [🚀 Dataset Quickstart](https://github.com/ZhuoyangLiu2005/T-Rex/tree/main/dataset_quickstart) · [📓 Colab notebook](https://colab.research.google.com/github/ZhuoyangLiu2005/T-Rex/blob/main/dataset_quickstart/quickstart.ipynb)
+
+<p align="center">
+  <img src="assets/dataset_overview.png" width="100%">
+  <br>
+  <em>One episode from each of 20 motor primitives (head-camera view, cropped to the workspace), each with a different object.</em>
+</p>
+
+<p align="center">
+  <img src="assets/hardware_setup.png" width="100%">
+  <br>
+  <em>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.</em>
+</p>
+
+## 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
+
+<p align="center">
+  <img src="assets/dataset_analysis.png" width="100%">
+  <br>
+  <em>Object categories, episodes per motor primitive, and per-object episode counts.</em>
+</p>
+
+## 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/<key>/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`.

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]
+    }
+  }
+}

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 `<EMBODIMENT>` 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 <EMBODIMENT> \
+    --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": "<EMBODIMENT>"}` |
+| `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 = "<EMBODIMENT>"
+```
+
+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 = [
+    ...,
+    "<EMBODIMENT>",
+]
+```
+
+---
+
+## 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.<name>` | `state.joint_pos` |
+| Action | `action.<name>` | `action.joint_pos` |
+| Video | `video.<name>` | `video.cam0` |
+| Language | `annotation.<name>` | `annotation.task` |
+
+If the YAML keys don't match `modality.json`, training will fail with missing-key errors.
+
+### 3b. Add `modality_config_<EMBODIMENT>`
+
+```yaml
+modality_config_<EMBODIMENT>:
+  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_<EMBODIMENT>`
+
+```yaml
+transform_<EMBODIMENT>:
+  _target_: groot.vla.data.transform.ComposedModalityTransform
+  transforms:
+    # Video
+    - <<: *totensor_cfg
+      apply_to: ${modality_config_<EMBODIMENT>.video.modality_keys}
+    - <<: *crop_cfg
+      apply_to: ${modality_config_<EMBODIMENT>.video.modality_keys}
+    - <<: *resize_cfg
+      apply_to: ${modality_config_<EMBODIMENT>.video.modality_keys}
+    - <<: *color_jitter_cfg
+      apply_to: ${modality_config_<EMBODIMENT>.video.modality_keys}
+    - <<: *to_numpy_cfg
+      apply_to: ${modality_config_<EMBODIMENT>.video.modality_keys}
+
+    # State
+    - _target_: groot.vla.data.transform.StateActionToTensor
+      apply_to: ${modality_config_<EMBODIMENT>.state.modality_keys}
+    - _target_: groot.vla.data.transform.StateActionTransform
+      apply_to: ${modality_config_<EMBODIMENT>.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_<EMBODIMENT>.action.modality_keys}
+    - _target_: groot.vla.data.transform.StateActionTransform
+      apply_to: ${modality_config_<EMBODIMENT>.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_<EMBODIMENT>.video.modality_keys}
+      state_concat_order: ${modality_config_<EMBODIMENT>.state.modality_keys}
+      action_concat_order: ${modality_config_<EMBODIMENT>.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:
+  ...
+  <EMBODIMENT>: ${modality_config_<EMBODIMENT>}
+
+transforms:
+  ...
+  <EMBODIMENT>: ${transform_<EMBODIMENT>}
+
+metadata_versions:
+  ...
+  <EMBODIMENT>: '0221'
+
+fps:
+  ...
+  <EMBODIMENT>: 30          # set to your dataset's FPS
+```
+
+---
+
+## Step 4: Create a Dataset YAML
+
+Create `groot/vla/configs/data/dreamzero/<EMBODIMENT>_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
+
+<EMBODIMENT>_data_root: ???     # set via CLI or env var
+
+train_dataset:
+  _target_: ${mixture_dataset_cls}
+  _convert_: object
+  mixture_spec:
+    - dataset_path:
+        <EMBODIMENT>:            # must match key in modality_configs/transforms
+          - ${<EMBODIMENT>_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:
+
+- `<EMBODIMENT>` 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/<EMBODIMENT>_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_<EMBODIMENT>_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/<EMBODIMENT>_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 \
+    <EMBODIMENT>_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/<EMBODIMENT>_training.sh
+
+# With overrides:
+DATA_ROOT=/path/to/your_dataset OUTPUT_DIR=./checkpoints/run1 NUM_GPUS=4 \
+    bash scripts/train/<EMBODIMENT>_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.

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
+```

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`。

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.

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({})

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()
+        

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)

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)

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)
+
+

groot/vla/common/__init__.py

@@ -0,0 +1 @@
+

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

groot/vla/common/utils/data_structure/__init__.py

@@ -0,0 +1,2 @@
+from .shape_utils import *  # noqa: F403
+from .tree_utils import *  # noqa: F403

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)

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)

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

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)

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

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)

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)

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': <traceback str>}
+    """
+    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))

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 <konstantin.lepa@gmail.com>
+# 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",
+    )

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

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.")

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

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)

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

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

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

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: .

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}

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

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}

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}

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}

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}

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}

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
+    }
+}

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
+    }
+}