From 141796747f55f1e1d137e870e5ecbc65c414b401 Mon Sep 17 00:00:00 2001 From: d-finite Date: Sun, 25 Oct 2026 19:42:58 +0800 Subject: [PATCH] Refactor four-source pretraining around cleaned LeRobot v3 datasets --- README.md | 18 + configs/dataloader/agibotworld.yaml | 29 + configs/dataloader/interndata_a1.yaml | 30 + .../{pretrain_data => }/mixture.yaml | 0 configs/dataloader/oxe_droid.yaml | 30 + .../dataloader/pretrain_data/agibotworld.yaml | 69 -- .../pretrain_data/interndata_a1.yaml | 63 - .../dataloader/pretrain_data/muka_franka.yaml | 37 - .../dataloader/pretrain_data/oxe_droid.yaml | 53 - .../dataloader/pretrain_data/robocoin.yaml | 67 - configs/dataloader/robocoin.yaml | 30 + configs/train.yaml | 2 +- openwam/dataloader/agibotworld.py | 329 +---- openwam/dataloader/bases/lerobot_v3_reader.py | 148 +-- openwam/dataloader/interndata_a1.py | 1054 +--------------- openwam/dataloader/mixture.py | 58 +- openwam/dataloader/muka_franka.py | 346 ------ openwam/dataloader/oxe_droid.py | 550 +-------- openwam/dataloader/registry.py | 4 +- openwam/dataloader/robocoin.py | 1087 +---------------- openwam/dataloader/utils/lerobotv3.py | 453 +------ openwam/dataloader/utils/normalization.py | 56 +- openwam/dataloader/utils/oxe_schema.py | 110 +- .../utils/stats_computation/accumulator.py | 126 ++ .../agibotworld_stats_computation.py | 375 ------ .../ebench_stats_computation.py | 2 +- .../interndata_a1_stats_computation.py | 640 ---------- .../muka_franka_stats_computation.py | 145 --- .../oxe_stats_computation.py | 323 ----- .../robocasa_gr1_stats_computation.py | 2 +- .../robocoin_stats_computation.py | 556 --------- .../robodojo_stats_computation.py | 2 +- tests/dataloader/test_agibotworld.py | 506 ++------ tests/dataloader/test_interndata_a1.py | 644 ++-------- .../test_interndata_a1_stats_computation.py | 946 -------------- tests/dataloader/test_lerobotv3_utils.py | 27 +- tests/dataloader/test_mixture.py | 31 + tests/dataloader/test_mixture_from_config.py | 9 + tests/dataloader/test_muka_franka.py | 253 ---- tests/dataloader/test_normalization_utils.py | 102 ++ tests/dataloader/test_oxe_droid.py | 170 +++ .../test_oxe_mixture_integration.py | 6 +- tests/dataloader/test_oxe_schema.py | 132 +- .../dataloader/test_oxe_stats_computation.py | 304 ----- .../dataloader/test_robocoin_compute_stats.py | 773 ------------ tests/dataloader/test_robocoin_final.py | 298 +++++ tests/dataloader/test_robocoin_trim.py | 715 ----------- tests/dataloader/test_stats_accumulator.py | 71 ++ tests/test_public_config_defaults.py | 56 +- 49 files changed, 1656 insertions(+), 10181 deletions(-) create mode 100644 configs/dataloader/agibotworld.yaml create mode 100644 configs/dataloader/interndata_a1.yaml rename configs/dataloader/{pretrain_data => }/mixture.yaml (100%) create mode 100644 configs/dataloader/oxe_droid.yaml delete mode 100644 configs/dataloader/pretrain_data/agibotworld.yaml delete mode 100644 configs/dataloader/pretrain_data/interndata_a1.yaml delete mode 100644 configs/dataloader/pretrain_data/muka_franka.yaml delete mode 100644 configs/dataloader/pretrain_data/oxe_droid.yaml delete mode 100644 configs/dataloader/pretrain_data/robocoin.yaml create mode 100644 configs/dataloader/robocoin.yaml delete mode 100644 openwam/dataloader/muka_franka.py create mode 100644 openwam/dataloader/utils/stats_computation/accumulator.py delete mode 100644 openwam/dataloader/utils/stats_computation/agibotworld_stats_computation.py delete mode 100644 openwam/dataloader/utils/stats_computation/interndata_a1_stats_computation.py delete mode 100644 openwam/dataloader/utils/stats_computation/muka_franka_stats_computation.py delete mode 100644 openwam/dataloader/utils/stats_computation/oxe_stats_computation.py delete mode 100644 openwam/dataloader/utils/stats_computation/robocoin_stats_computation.py delete mode 100644 tests/dataloader/test_interndata_a1_stats_computation.py delete mode 100644 tests/dataloader/test_muka_franka.py create mode 100644 tests/dataloader/test_oxe_droid.py delete mode 100644 tests/dataloader/test_oxe_stats_computation.py delete mode 100644 tests/dataloader/test_robocoin_compute_stats.py create mode 100644 tests/dataloader/test_robocoin_final.py delete mode 100644 tests/dataloader/test_robocoin_trim.py create mode 100644 tests/dataloader/test_stats_accumulator.py diff --git a/README.md b/README.md index 633d5ac97..f47f3dd71 100644 --- a/README.md +++ b/README.md @@ -88,6 +88,8 @@ All architectures are selected via `configs/model/.yaml` with `archit --> ## News +- **[2026/10/25]** 🤗 We release all cleaned robot datasets used for OpenWAM pretraining: [AgiBotWorld-Beta](https://huggingface.co/datasets/OpenWAM/AgiBotWorld-Beta), [RoboCOIN](https://huggingface.co/datasets/wayrise/RoboCOIN), [InternData-A1](https://huggingface.co/datasets/OpenWAM/InternData-A1), and [DROID](https://huggingface.co/datasets/OpenWAM/Droid). +- **[2026/10/10]** 🌟 OpenWAM reaches **1K GitHub stars**! Thank you for your support! - **[2026/09/09]** 📄 OpenWAM [Paper](https://arxiv.org/pdf/2609.07398) is released on arXiv. - **[2026/09/06]** 🤖 OpenWAM is integrated into [XPolicyLab](https://github.com/XPolicyLab/XPolicyLab). - **[2026/09/06]** 🤗 We release all pretrained and finetuned models on [huggingface](https://huggingface.co/OpenWAM). @@ -355,6 +357,22 @@ Wan2.2-TI2V-5B video backbone, and the Mutual attention mask. reports the returned action dimension and latency. Stop the deployment process after the checks. +### Four-Source Pretraining + +Use the four cleaned LeRobot v3 datasets with the shared mixture configuration: + +~~~bash +bash scripts/train.sh dataloader=mixture +~~~ + +[`configs/dataloader/mixture.yaml`](configs/dataloader/mixture.yaml) enables +AgiBotWorld-Beta, RoboCOIN, DROID, and InternData-A1 with proportional sampling +by default. Their standalone configurations are in `configs/dataloader/` and +point to `/mnt/data/wangyuran/pretrain_data_final`. Override each source's +`dataset_dir` when moving the datasets, keeping the supplied `meta` statistics +with them. Readers use the existing statistics directly and require no trim +CSV, exclusion list, or statistics generation step. + ### OpenWAM-α Fine-Tuning 1. Download LIBERO as shown above. diff --git a/configs/dataloader/agibotworld.yaml b/configs/dataloader/agibotworld.yaml new file mode 100644 index 000000000..5b8ef7327 --- /dev/null +++ b/configs/dataloader/agibotworld.yaml @@ -0,0 +1,29 @@ +type: agibotworld + +# Family root with native LeRobot v3 buckets and meta/stats_g2a.json. +# The reader retains separate action/state statistics and native hand/mobile masks. +dataset_dir: /mnt/data/wangyuran/pretrain_data_final/AgiBotWorld-Beta + +# Windows preserve episode boundaries and align video with action/state samples. +num_frames: 33 +video_stride: 4 +window_stride: 1 +height: 384 +width: 320 +multiview: true +split: train + +# Stored normalization statistics are read from the dataset's meta directory. +normalize_mode: quantile +enable_action_supervision: true +unify_action: true + +# One consistent augmentation per training clip. +color_jitter: + brightness: 0.2 + contrast: 0.2 + saturation: 0.2 + hue: 0.0 + +# Optional reproducible footage budget; null loads the complete dataset. +total_hours: null diff --git a/configs/dataloader/interndata_a1.yaml b/configs/dataloader/interndata_a1.yaml new file mode 100644 index 000000000..9ba913f76 --- /dev/null +++ b/configs/dataloader/interndata_a1.yaml @@ -0,0 +1,30 @@ +type: interndata_a1 + +# Family root with native LeRobot v3 buckets and meta/stats_.json. +# Raw arm poses use *_ee_to_robot_pose, independently of the gripper motion. +dataset_dir: /mnt/data/wangyuran/pretrain_data_final/InternData-A1 + +# Windows preserve episode boundaries and align video with action/state samples. +num_frames: 33 +video_stride: 4 +window_stride: 1 +height: 384 +width: 320 +multiview: true +split: train + +# Stored normalization statistics are read from the dataset's meta directory. +normalize_mode: quantile +enable_action_supervision: true +unify_action: true +unify_action_map: ["0-9", "34-43"] + +# One consistent augmentation per training clip. +color_jitter: + brightness: 0.2 + contrast: 0.2 + saturation: 0.2 + hue: 0.0 + +# Optional reproducible footage budget; null loads the complete dataset. +total_hours: null diff --git a/configs/dataloader/pretrain_data/mixture.yaml b/configs/dataloader/mixture.yaml similarity index 100% rename from configs/dataloader/pretrain_data/mixture.yaml rename to configs/dataloader/mixture.yaml diff --git a/configs/dataloader/oxe_droid.yaml b/configs/dataloader/oxe_droid.yaml new file mode 100644 index 000000000..65e3a0d39 --- /dev/null +++ b/configs/dataloader/oxe_droid.yaml @@ -0,0 +1,30 @@ +type: oxe_droid + +# Single LeRobot v3 dataset with meta/eef_stats.json and resolved tasks.parquet. +# Wrist/mount poses and inverted closedness retain the canonical DROID EEF contract. +dataset_dir: /mnt/data/wangyuran/pretrain_data_final/Droid + +# Windows preserve episode boundaries and align video with action/state samples. +num_frames: 33 +video_stride: 4 +window_stride: 1 +height: 384 +width: 320 +multiview: true +split: train + +# Stored normalization statistics are read from the dataset's meta directory. +normalize_mode: quantile +enable_action_supervision: true +unify_action: true +unify_action_map: ["0-9", "34-43"] + +# One consistent augmentation per training clip. +color_jitter: + brightness: 0.2 + contrast: 0.2 + saturation: 0.2 + hue: 0.0 + +# Optional reproducible footage budget; null loads the complete dataset. +total_hours: null diff --git a/configs/dataloader/pretrain_data/agibotworld.yaml b/configs/dataloader/pretrain_data/agibotworld.yaml deleted file mode 100644 index 19d8111ac..000000000 --- a/configs/dataloader/pretrain_data/agibotworld.yaml +++ /dev/null @@ -1,69 +0,0 @@ -type: agibotworld - -# Root containing AgiBotWorld-Beta LeRobot v3 task buckets. Each bucket has -# meta/info.json, data/, videos/, meta/tasks.parquet, and meta/stats.json. -dataset_dir: /path/to/pretrain_dataset/AgiBotWorld-Beta - -# 33 raw frames at stride 4 produce 9 video frames. The latter must satisfy the -# Wan VAE temporal constraint: (num_video_frames - 1) % 4 == 0. -num_frames: 33 -video_stride: 4 -window_stride: 1 - -# Apply the public segment_flag/segment_delta annotations: keep, trim the -# leading or trailing static region, or drop a no-motion segment. -use_segment_annotations: true - -# Drop an episode if segment trimming removes at least this fraction. Set null -# to retain every trimmed episode. -segment_max_trim_ratio: 0.7 - -# multiview=true builds a 384x320 L-shaped canvas: 256x320 head camera above -# two 128x160 wrist cameras. These dimensions are required in multiview mode. -height: 384 -width: 320 - -# Camera fields are observation.images.head, hand_left, and hand_right. -multiview: true - -# Uses info.json["splits"][split] when available; train otherwise uses the full -# episode table. -split: train - -# action.ee_base and observation.state.ee_base are robot-base poses of the -# axis-7 arm flange, not a gripper-center/task TCP. No universal flange-to-TCP -# calibration exists across the release's gripper and dex-hand embodiments. - -# Load-time color jitter uses the same sampled factors for every frame in a -# clip and is active only for the training split. -color_jitter: - brightness: 0.2 - contrast: 0.2 - saturation: 0.2 - hue: 0.0 - -# Loads pooled stats from /meta/stats_g2a.json. Action and state -# use separate stats; rot6d dimensions are identity-pinned. Before -# normalization, gripper aperture is canonicalized to 0=closed, 1=open. -# Supported modes: null/none, min-max, z-score, and quantile. -# Generate stats with: -# python -m openwam.dataloader.utils.stats_computation.agibotworld_stats_computation \ -# --dataset_dir /path/to/pretrain_dataset/AgiBotWorld-Beta \ -# --segment-max-trim-ratio 0.7 -normalize_mode: quantile - -# false keeps values available but masks action and proprio supervision. -enable_action_supervision: true - -# Scatter each bucket into the shared 80-D action/proprio layout: -# 0:9 left EEF pose, 9 left grip, 10:34 left hand, -# 34:43 right EEF pose, 43 right grip, 44:68 right hand, -# 68 move-x, 69 move-y, 70 move-yaw, 71:80 reserved. -# Gripper buckets map the two 10-D EEF arms into 0:10 and 34:44. Dex-hand -# buckets map pose and available finger joints while masking unused slots. Base -# commands occupy 68/70; differential drive has no lateral-velocity slot. -unify_action: true - -# null uses the complete filtered dataset. A positive float caps effective -# footage in hours by reproducibly subsampling episodes with project.seed. -total_hours: null diff --git a/configs/dataloader/pretrain_data/interndata_a1.yaml b/configs/dataloader/pretrain_data/interndata_a1.yaml deleted file mode 100644 index 46f3b1552..000000000 --- a/configs/dataloader/pretrain_data/interndata_a1.yaml +++ /dev/null @@ -1,63 +0,0 @@ -type: interndata_a1 - -# Public InternData-A1 LeRobot v3 root. Override with INTERNDATA_A1_ROOT. -dataset_dir: ${oc.env:INTERNDATA_A1_ROOT,/path/to/InternData-A1} - -# Per-embodiment stats normally live below dataset_dir. Use an alternate root -# when the dataset mount is read-only. -stats_root: null - -# Optional CSV containing per-episode head/tail trimming metadata. Stats encode -# whether trimming was active and are rejected if their population differs. -trim_csv: null - -# 33 raw frames at stride 4 produce 9 video frames. The latter must satisfy the -# Wan VAE temporal constraint: (num_video_frames - 1) % 4 == 0. -num_frames: 33 -video_stride: 4 -window_stride: 1 - -# multiview=true builds a 384x320 L-shaped canvas: 256x320 head camera above -# two 128x160 wrist cameras. -height: 384 -width: 320 -multiview: true - -# Public A1 buckets normally expose only a train split. Stats are train-derived. -split: train - -# action/state use *_ee_to_robot_pose: a rigid terminal-arm attachment in the -# robot frame shared by both arms. It may be a hand base, final arm link, or -# gripper mount rather than a literal flange. Native *_tcp_to_robot_pose is not -# used because the shared contract requires a gripper-independent endpoint. - -# The same color-jitter factors are applied to every frame in one training clip. -color_jitter: - brightness: 0.2 - contrast: 0.2 - saturation: 0.2 - hue: 0.0 - -# Stats are per embodiment because gripper units differ. Source gripper -# positions are converted to the shared raw aperture convention 0=closed, -# 1=open before normalization; rot6d stats are identity-pinned. Supported -# modes are null/none, min-max, z-score, and quantile. -# Generate matching stats with: -# python -m openwam.dataloader.utils.stats_computation.interndata_a1_stats_computation \ -# --dataset_dir /path/to/InternData-A1 --stats_root /path/to/InternData-A1 \ -# [--trim_csv /path/to/trim.csv] --split train -normalize_mode: quantile - -# false keeps values available but masks action and proprio supervision. -enable_action_supervision: true - -# Raw 20-D EEF layout: -# [L xyz3, L rot6d6, L grip1, R xyz3, R rot6d6, R grip1]. -# It is scattered into left slots 0:10 and right slots 34:44 of the shared -# 80-D head. Every enabled mixture source must use the same action width. -unify_action: true -unify_action_map: ["0-9", "34-43"] - -# null uses the complete filtered dataset. A positive float caps effective -# footage in hours by reproducibly subsampling episodes with project.seed. -total_hours: null diff --git a/configs/dataloader/pretrain_data/muka_franka.yaml b/configs/dataloader/pretrain_data/muka_franka.yaml deleted file mode 100644 index f89828aeb..000000000 --- a/configs/dataloader/pretrain_data/muka_franka.yaml +++ /dev/null @@ -1,37 +0,0 @@ -type: muka_franka - -# Hugging Face: ewykric/muka_franka_lerobot_v3 (LeRobot v3, 15 fps). -# Raw state/action: [xyz3, extrinsic XYZ Euler3, gripper_closedness1]. -# The reader converts both to [xyz3, rot6d6, gripper_open_scale1], with the -# pretrained convention -1=closed / +1=open. -dataset_dir: /path/to/muka_franka_lerobot_v3 -action_mode: eef -gripper_convention: minus1_closed_plus1_open - -# Normalize raw EEF10, then scatter to the shared 80-D left-arm block: -# [0:3)=xyz, [3:9)=rot6d, [9]=gripper. All other slots remain zero/masked. -unify_action: true -unify_action_map: ["0-9"] -normalize_mode: min-max - -head_camera_priority: - - observation.images.left -wrist_camera_priority: - - observation.images.left_wrist - -num_frames: 33 -video_stride: 4 -window_stride: 1 -height: 384 -width: 320 -multiview: true - -# Follow the published train split; no additional holdout override is applied. -split: train -enable_action_supervision: true - -# color_jitter: -# brightness: 0.2 -# contrast: 0.2 -# saturation: 0.2 -# hue: 0.0 diff --git a/configs/dataloader/pretrain_data/oxe_droid.yaml b/configs/dataloader/pretrain_data/oxe_droid.yaml deleted file mode 100644 index 66c922d53..000000000 --- a/configs/dataloader/pretrain_data/oxe_droid.yaml +++ /dev/null @@ -1,53 +0,0 @@ -type: oxe_droid - -# Root of the DROID LeRobot v3 conversion. The reader requires the deterministic -# no-prompt exclusion manifest (meta/excluded_episodes.json) to be present. -dataset_dir: /path/to/pretrain_dataset/Droid - -# DROID is 10 fps. A 33-frame window spans 3.2 seconds; stride 4 yields 9 video -# frames and satisfies the Wan VAE temporal constraint. -num_frames: 33 -video_stride: 4 -window_stride: 1 - -# Source cameras are 180x320. multiview=true builds a 384x320 L-shaped canvas: -# 256x320 primary view, 128x160 wrist view, and a blank right-wrist tile. -height: 384 -width: 320 - -multiview: true - -split: train - -# DROID exposes separate moving task-TCP and rigid wrist/mount streams. The -# reader intentionally uses action_wrist_pose and -# observation_gripper_pose6d + state[6], in the robot-base frame. Raw gripper -# values are inverted to the shared aperture convention 0=closed, 1=open. - -# Load-time color jitter uses the same sampled factors for every frame in a -# training clip. -color_jitter: - brightness: 0.2 - contrast: 0.2 - saturation: 0.2 - hue: 0.0 - -# Loads /meta/eef_stats.json. Its strict contract records source -# columns and the exclusion population, rejecting stale TCP-based statistics. -# rot6d stats are identity-pinned. Supported modes are null/none, min-max, -# z-score, and quantile. Generate matching stats with: -# python -m openwam.dataloader.utils.stats_computation.oxe_stats_computation \ -# --dataset DROID --dataset-dir /path/to/Droid -normalize_mode: quantile - -# false keeps values available but masks action and proprio supervision. -enable_action_supervision: true - -# Scatter the single active 10-D arm into the left EEF slots 0:10 of the shared -# 80-D head. The zero-padded right half maps to 34:44 and remains masked. -unify_action: true -unify_action_map: ["0-9", "34-43"] - -# null uses the complete post-exclusion dataset. A positive float caps effective -# footage in hours by reproducibly subsampling episodes with project.seed. -total_hours: null diff --git a/configs/dataloader/pretrain_data/robocoin.yaml b/configs/dataloader/pretrain_data/robocoin.yaml deleted file mode 100644 index a60986092..000000000 --- a/configs/dataloader/pretrain_data/robocoin.yaml +++ /dev/null @@ -1,67 +0,0 @@ -type: robocoin - -# Root containing RoboCOIN LeRobot v3 task buckets and -# meta/stats_.json. The reader discovers buckets recursively and -# dispatches samples through MultiRobotCOINDataset. -dataset_dir: /path/to/pretrain_dataset/RoboCOIN_clean - -# eef_sim_pose_* uses RoboCOIN's unified base axes, but the physical endpoint -# varies by embodiment. Treat it as a rigid, gripper-independent terminal-arm -# frame; do not apply one blanket task-TCP offset across robot types. - -# 33 raw frames at stride 4 produce 9 video frames. The latter must satisfy the -# Wan VAE temporal constraint: (num_video_frames - 1) % 4 == 0. -num_frames: 33 -video_stride: 4 -window_stride: 1 - -# Optional per-episode head/tail trim manifest. Leave null to load the full -# public dataset. Episode indices are dataset-version-specific; after changing -# the manifest, regenerate normalization stats from that exact population. -trim_csv: null - -# multiview=true builds a 384x320 L-shaped canvas: 256x320 head camera above -# two 128x160 wrist cameras. Missing wrist cameras become blank tiles. -height: 384 -width: 320 - -# Head and wrist keys are resolved from the public per-embodiment camera names. -multiview: true - -# Uses info.json["splits"][split] when available; train otherwise uses the full -# episode table. -split: train - -# Load-time color jitter uses the same sampled factors for every frame in a -# training clip. -color_jitter: - brightness: 0.2 - contrast: 0.2 - saturation: 0.2 - hue: 0.0 - -# Loads per-robot-type stats from -# /meta/stats_.json. Scalar grippers are canonicalized -# to aperture 0=closed, 1=open; dexterous-hand joints retain their joint-angle -# semantics. Stats are bound to the same split, exclusions, and trim manifest; -# rot6d dimensions are identity-pinned. Supported modes are null/none, min-max, -# z-score, and quantile. Generate matching stats with: -# python -m openwam.dataloader.utils.stats_computation.robocoin_stats_computation \ -# --dataset_dir /path/to/RoboCOIN [--trim_csv /path/to/trim.csv] --split train -normalize_mode: quantile - -# false keeps values available but masks action and proprio supervision. -enable_action_supervision: true - -# Shared 80-D layout: -# 0:9 left EEF pose, 9 left grip, 10:34 left hand, -# 34:43 right EEF pose, 43 right grip, 44:68 right hand, 68:80 reserved. -# Gripper embodiments scatter two 10-D EEF arms into 0:10 and 34:44. -# Dexterous-hand embodiments scatter pose and available finger joints while -# masking scalar-gripper and unused hand slots. -unify_action: true -unify_action_map: ["0-9", "34-43"] - -# null uses the complete filtered dataset. A positive float caps effective -# footage in hours by reproducibly water-filling episode samples across buckets. -total_hours: null diff --git a/configs/dataloader/robocoin.yaml b/configs/dataloader/robocoin.yaml new file mode 100644 index 000000000..4b74957de --- /dev/null +++ b/configs/dataloader/robocoin.yaml @@ -0,0 +1,30 @@ +type: robocoin + +# Family root with native LeRobot v3 buckets and meta/stats_.json. +# Scalar grippers use aperture 0=closed, 1=open; dexterous hands retain joint angles. +dataset_dir: /mnt/data/wangyuran/pretrain_data_final/RoboCOIN + +# Windows preserve episode boundaries and align video with action/state samples. +num_frames: 33 +video_stride: 4 +window_stride: 1 +height: 384 +width: 320 +multiview: true +split: train + +# Stored normalization statistics are read from the dataset's meta directory. +normalize_mode: quantile +enable_action_supervision: true +unify_action: true +unify_action_map: ["0-9", "34-43"] + +# One consistent augmentation per training clip. +color_jitter: + brightness: 0.2 + contrast: 0.2 + saturation: 0.2 + hue: 0.0 + +# Optional reproducible footage budget; null loads the complete dataset. +total_hours: null diff --git a/configs/train.yaml b/configs/train.yaml index c466625d0..c2cbffd4e 100644 --- a/configs/train.yaml +++ b/configs/train.yaml @@ -1,6 +1,6 @@ defaults: - model: dual_system # [Optional]: dual_system(default), single_system, tri_system - - dataloader: robotwin # [Optional]: robotwin(default), ebench, robocasa365, robocasa_gr1, robodojo, vlabench, libero + - dataloader: robotwin # [Optional]: robotwin(default), mixture (four-source pretraining), ebench, robocasa365, robocasa_gr1, robodojo, vlabench, libero - _self_ training: diff --git a/openwam/dataloader/agibotworld.py b/openwam/dataloader/agibotworld.py index 7e402b397..f83f4201d 100644 --- a/openwam/dataloader/agibotworld.py +++ b/openwam/dataloader/agibotworld.py @@ -20,25 +20,20 @@ from __future__ import annotations -import hashlib import json import logging from pathlib import Path from typing import List import numpy as np -import pandas as pd from openwam.dataloader.bases import LeRobotV3Reader, MultiLeRobotV3Reader from openwam.dataloader.robocoin import GRIP_EXCLUDED_DIM_MASK, _build_dex_unify_map -from openwam.dataloader.utils.lerobotv3 import ( - DataContractError, - digest_lerobot_v3_data_population, - resolve_lerobot_v3_data_population, -) +from openwam.dataloader.utils.lerobotv3 import DataContractError from openwam.dataloader.utils.normalization import ( ROT6D_DIMS_EEF20, apply_normalization, + find_dataset_stats, materialize_eef_stats, pin_rot6d_identity, ) @@ -78,7 +73,6 @@ "closed_endpoint_m": _GRIPPER_STATE_CLOSED_POSITION_M, "output_semantics": "0_closed_1_open_aperture_fraction", } -_STATS_SCHEMA_VERSION = 3 # ``info.features`` declares both signal families, so bucket flavor is explicit. @@ -142,99 +136,6 @@ _BUCKET_STATS_FILENAME = "stats.json" -_SEGMENT_FLAG_COL = "segment_flag" -_SEGMENT_DELTA_COL = "segment_delta" - - -def _validate_trim_ratio(value) -> float | None: - """Normalize ``segment_max_trim_ratio`` to a float in (0, 1] or None. - - Raises rather than clamping: a config typo like ``70`` (percent instead of a - fraction) would silently disable the filter, and ``0`` would drop every - episode including the untrimmed ones. - """ - - if value is None: - return None - try: - ratio = float(value) - except (TypeError, ValueError) as e: - raise ValueError(f"segment_max_trim_ratio must be a float in (0, 1] or null; got {value!r}") from e - if not np.isfinite(ratio) or not 0.0 < ratio <= 1.0: - raise ValueError( - f"segment_max_trim_ratio must be a fraction in (0, 1] or null; got {ratio!r}. " - "Use 0.7 for '70% or more of the episode trimmed', not 70." - ) - return ratio - - -def _apply_segment_annotations( - eps_df: pd.DataFrame, - *, - use_segment_annotations: bool, - segment_max_trim_ratio: float | None, -) -> tuple[pd.DataFrame, dict]: - """Return the exact episode population admitted by the segment cleanup. - - This is deliberately shared by the training reader and the normalization- - stats producer. Keeping the policy in one function prevents a stats rerun - from accidentally pooling the leading/trailing frames that the reader can - never serve. - """ - - ratio = _validate_trim_ratio(segment_max_trim_ratio) - summary = { - "missing_columns": False, - "dropped_static": 0, - "dropped_empty": 0, - "dropped_over_trim": 0, - } - if not use_segment_annotations: - return eps_df.reset_index(drop=True), summary - if _SEGMENT_FLAG_COL not in eps_df.columns or _SEGMENT_DELTA_COL not in eps_df.columns: - summary["missing_columns"] = True - return eps_df.reset_index(drop=True), summary - - out = eps_df.copy() - flag = out[_SEGMENT_FLAG_COL].fillna(0).astype(np.int64).to_numpy() - delta = np.maximum(0, out[_SEGMENT_DELTA_COL].fillna(0).astype(np.int64).to_numpy()) - length = out["length"].astype(np.int64).to_numpy() - - valid_start = np.zeros(len(out), dtype=np.int64) - valid_end = length.copy() - start_mask = flag == 1 - end_mask = flag == 2 - valid_start[start_mask] = np.minimum(delta[start_mask], length[start_mask]) - valid_end[end_mask] = np.maximum(0, length[end_mask] - delta[end_mask]) - - keep = (flag != 3) & (valid_end > valid_start) - summary["dropped_static"] = int(np.sum(flag == 3)) - summary["dropped_empty"] = int(np.sum((flag != 3) & (valid_end <= valid_start))) - - if ratio is not None: - trimmed = np.maximum(0, length - np.maximum(0, valid_end - valid_start)) - with np.errstate(invalid="ignore", divide="ignore"): - trim_ratio = np.where(length > 0, trimmed / np.maximum(length, 1), 0.0) - over = (flag != 3) & keep & (trim_ratio >= ratio) - summary["dropped_over_trim"] = int(np.sum(over)) - keep = keep & ~over - - out["_valid_start"] = valid_start - out["_valid_end"] = valid_end - return out.loc[keep].reset_index(drop=True), summary - - -def _effective_segment_population_digest(eps_df: pd.DataFrame) -> str: - """Hash the kept ``(episode, start, end)`` ranges in reader order.""" - hasher = hashlib.sha256(b"openwam:agibotworld-segment-population:v1\0") - ordered = eps_df.sort_values("episode_index", kind="stable") - for _, row in ordered.iterrows(): - start = int(row.get("_valid_start", 0)) - end = int(row.get("_valid_end", row["length"])) - for value in (int(row["episode_index"]), start, end): - hasher.update(value.to_bytes(8, "little", signed=False)) - return hasher.hexdigest() - def _bucket_base_motion_flags(dataset_dir) -> tuple[bool, bool]: """Return independent ``(action_moves, state_moves)`` flags for a bucket. @@ -245,29 +146,18 @@ def _bucket_base_motion_flags(dataset_dir) -> tuple[bool, bool]: mobile buckets: action slots 68/70 remain supervised while all-zero state slots are absent from the proprio mask. - A missing / unreadable stats file returns ``(False, False)`` and logs a - warning, so no fabricated constant-zero movement is supervised. + Missing or unreadable bucket stats fail immediately: silently treating a + mobile bucket as stationary would remove real supervision. """ p = Path(dataset_dir) / "meta" / _BUCKET_STATS_FILENAME - if not p.exists(): - logger.warning( - "AgiBotWorld %s: no %s — treating base as stationary (move slots unmapped).", - dataset_dir, - _BUCKET_STATS_FILENAME, - ) - return False, False try: - with open(p) as f: - st = json.load(f) - except (OSError, ValueError) as e: - logger.warning( - "AgiBotWorld %s: could not read %s (%s) — treating base as stationary (move slots unmapped).", - dataset_dir, - p.name, - e, - ) - return False, False + with p.open(encoding="utf-8") as stream: + st = json.load(stream) + except (OSError, ValueError) as exc: + raise DataContractError(f"Cannot load AgiBotWorld movement statistics from {p}: {exc}") from exc + if not isinstance(st, dict): + raise DataContractError(f"AgiBotWorld movement statistics {p} must contain a mapping.") flags = [] for key in ("action.robot_velocity", "observation.state.robot_velocity"): @@ -275,9 +165,14 @@ def _bucket_base_motion_flags(dataset_dir) -> tuple[bool, bool]: if not blk: flags.append(False) continue - lo = np.abs(np.asarray(blk.get("min", [0.0]), dtype=np.float64)).max() - hi = np.abs(np.asarray(blk.get("max", [0.0]), dtype=np.float64)).max() - flags.append(bool(max(lo, hi) > _MOVE_EPS)) + if not isinstance(blk, dict): + raise DataContractError(f"{p}: {key} must contain movement bounds.") + bounds = [np.asarray(blk.get(field), dtype=np.float64) for field in ("min", "max")] + if any(values.shape != (3,) or not np.isfinite(values).all() for values in bounds): + raise DataContractError(f"{p}: {key} min/max must contain three finite values.") + if np.any(bounds[0] > bounds[1]): + raise DataContractError(f"{p}: {key} min exceeds max.") + flags.append(bool(max(np.abs(values).max() for values in bounds) > _MOVE_EPS)) return flags[0], flags[1] @@ -324,6 +219,7 @@ class AgiBotWorldDataset(LeRobotV3Reader): AgiBotWorld-specific bits are overridden below. """ + STRICT_BUCKET_LOADING = True DATASET_NAME = "AgiBotWorld" HEAD_CAMERA = HEAD_CAMERA LEFT_WRIST_CAMERA = LEFT_WRIST_CAMERA @@ -334,7 +230,6 @@ class AgiBotWorldDataset(LeRobotV3Reader): DEFAULT_NORMALIZE_MODE = None WRIST_DECODE_TOLERATED = (Exception,) - CONFIG_KEYS = LeRobotV3Reader.CONFIG_KEYS + ("use_segment_annotations", "segment_max_trim_ratio") def __init__( self, @@ -342,8 +237,6 @@ def __init__( *, unify_action: bool = False, unify_action_map=None, - use_segment_annotations: bool = True, - segment_max_trim_ratio: float | None = None, **kwargs, ): """Detect the bucket flavor + base motion and wire the finger / move scatter. @@ -363,14 +256,6 @@ def __init__( Under ``unify_action=False`` no movement is emitted: grippered → 20-D EEF, dex → 20-D pose with grip slots masked. - Args: - use_segment_annotations: honor ``segment_flag`` / ``segment_delta`` - (see :meth:`_filter_episodes`). False → full segments, the - pre-annotation behavior. - segment_max_trim_ratio: drop an episode outright when the segment trim - would remove at least this fraction of it (``0.7`` → an episode - keeping under 30% of its frames is discarded). None disables the - rule. Only meaningful with ``use_segment_annotations=True``. """ self._is_dex = Path(dataset_dir).name in _DEX_BUCKET_IDS @@ -382,15 +267,6 @@ def __init__( self.ACTION_DIM_MASK = None self.PROPRIO_DIM_MASK = None - self._use_segment_annotations = bool(use_segment_annotations) - self._segment_max_trim_ratio = _validate_trim_ratio(segment_max_trim_ratio) - if self._segment_max_trim_ratio is not None and not self._use_segment_annotations: - logger.warning( - "AgiBotWorld %s: segment_max_trim_ratio=%s is ignored because " - "use_segment_annotations=False (there is no trim to measure).", - dataset_dir, - self._segment_max_trim_ratio, - ) if unify_action: move_dim = _MOVE_DIM if self._has_move else 0 move_slots = list(_MOVE_SLOTS) if self._has_move else [] @@ -409,47 +285,9 @@ def __init__( self.PROPRIO_DIM_MASK = proprio_mask super().__init__(dataset_dir, unify_action=unify_action, unify_action_map=unify_action_map, **kwargs) - def _filter_episodes(self, eps_df: pd.DataFrame) -> pd.DataFrame: - """Apply segment-boundary annotations to the trainable episode span. - - ``segment_flag`` selects an unchanged segment, a leading trim, a - trailing trim, or a fully static segment to drop. ``segment_delta`` is - interpreted as an integer frame count. The resulting - ``_valid_start``/``_valid_end`` bounds are consumed by the shared - window builder, keeping parquet and video slices aligned. - - When configured, ``segment_max_trim_ratio`` drops a segment whose - boundary trim would remove at least that fraction. Mid-segment frames - are never spliced or reordered. - """ - - out, summary = _apply_segment_annotations( - eps_df, - use_segment_annotations=self._use_segment_annotations, - segment_max_trim_ratio=self._segment_max_trim_ratio, - ) - if summary["missing_columns"]: - logger.warning( - "AgiBotWorld(%s): segment annotation columns %s/%s missing; using full segments.", - self._dataset_id, - _SEGMENT_FLAG_COL, - _SEGMENT_DELTA_COL, - ) - dropped_static = summary["dropped_static"] - dropped_empty = summary["dropped_empty"] - dropped_over_trim = summary["dropped_over_trim"] - if dropped_static or dropped_empty or dropped_over_trim: - logger.info( - "AgiBotWorld(%s): segment annotations dropped %d static, %d empty-after-trim and " - "%d over-trimmed (>=%s of the episode) episodes; kept %d episodes.", - self._dataset_id, - dropped_static, - dropped_empty, - dropped_over_trim, - "n/a" if self._segment_max_trim_ratio is None else f"{self._segment_max_trim_ratio:.0%}", - len(out), - ) - return out + def _load_excluded_episode_indices(self) -> set[int]: + """The exported episodes already contain exactly the retained frames.""" + return set() def _resolve_cameras(self, info: dict): """Return the fixed head/left/right cameras, record robot_type, and pick @@ -475,123 +313,48 @@ def _resolve_cameras(self, info: dict): return self.HEAD_CAMERA, self.LEFT_WRIST_CAMERA, self.RIGHT_WRIST_CAMERA def _load_stats(self, info: dict): - """Assemble SEPARATE action / state normalization stats from the unified stats file. - - Returns None (raw pass-through) unless a config opts into normalization. - Action and state are DISTINCT distributions — e.g. ``gripper`` is a binary - command in ``action.*`` but a continuous physical opening in - ``observation.state.*`` — so each stream gets its OWN stats, built from its - own ``{prefix}.ee_base`` / ``.gripper`` | ``.dex`` / ``.robot_velocity`` blocks - (assembled to this bucket's raw layout, rot6d dims identity-pinned since the - stats file is not pinned). The two blocks are stashed on the instance - (``_action_norm_stats`` / ``_proprio_norm_stats``) and consumed by - ``_action_20d`` / ``_proprio_20d`` respectively. - - Stats source: the UNIFIED ``/meta/stats_g2a.json`` produced by - agibotworld_stats_computation — one flavor-aware stats set pooled across all - buckets (AgiBotWorld is a single embodiment, so the same physical action should - normalize identically everywhere; this mirrors RoboCOIN's per-robot-type stats). - Each column was pooled only over the buckets where it carries signal (ee_base: - all; gripper: grippered buckets; dex: dex buckets; robot_velocity: the - independently moving buckets for that action/state stream). - mean/std/min/max are exact; q01/q99 are an approximate pooled estimate (merged - reservoirs). Every mode is valid; missing file → raise (run the stats script). - Under root mode ``_dataset_dir`` is a bucket, so the file lives at - ``_dataset_dir.parent / meta``. + """Read the shipped family-root stats and assemble each signal layout. + + Action/state and mobile-base distributions remain separate. No frame + selection or statistics generation is performed by the reader. """ self._action_norm_stats = None self._proprio_norm_stats = None if not self._normalize_mode or self._normalize_mode in ("none", "null"): return None - stats_path = self._dataset_dir.parent / "meta" / _UNIFIED_STATS_FILENAME - if not stats_path.exists(): - raise FileNotFoundError( - f"normalize_mode={self._normalize_mode!r} but the unified stats file {stats_path} is missing. " - f"Run `python -m openwam.dataloader.utils.stats_computation.agibotworld_stats_computation " - f"--dataset_dir {self._dataset_dir.parent}` to generate it, or set normalize_mode=null to disable." - ) - with open(stats_path) as f: + stats_path = find_dataset_stats(self._dataset_dir, _UNIFIED_STATS_FILENAME, family_root=self._stats_root) + with open(stats_path, encoding="utf-8") as f: raw = json.load(f) + if not isinstance(raw, dict): + raise DataContractError(f"{stats_path} must contain a statistics mapping.") if raw.get("gripper_contract") != _GRIPPER_CONTRACT: raise DataContractError( f"AgiBotWorld stats {stats_path} use gripper_contract=" f"{raw.get('gripper_contract')!r}; expected {_GRIPPER_CONTRACT!r}. " - "Regenerate stats with the current canonical direction and endpoint calibration." - ) - - population = raw.get("population") - if not isinstance(population, dict) or population.get("schema_version") != _STATS_SCHEMA_VERSION: - raise DataContractError( - f"AgiBotWorld stats {stats_path} use stale population schema " - f"{None if not isinstance(population, dict) else population.get('schema_version')!r}; " - f"expected {_STATS_SCHEMA_VERSION}. Regenerate it to bind the canonical gripper " - "contract, independent action/state velocity populations, and terminal-action exclusion." - ) - recorded_buckets = population.get("buckets") - if not isinstance(recorded_buckets, dict) or not all( - isinstance(name, str) and name for name in recorded_buckets - ): - raise DataContractError( - f"AgiBotWorld stats {stats_path} has a malformed pooled contributor map; regenerate stats." - ) - - current_buckets = { - path.name - for path in self._dataset_dir.parent.iterdir() - if path.is_dir() and (path / "meta" / "info.json").is_file() - } - recorded_bucket_names = set(recorded_buckets) - if recorded_bucket_names != current_buckets: - missing = sorted(current_buckets - recorded_bucket_names) - extra = sorted(recorded_bucket_names - current_buckets) - raise DataContractError( - f"AgiBotWorld stats {stats_path} pooled contributor set no longer matches " - f"the current dataset root (missing={missing}, extra={extra}); regenerate stats." - ) - if population.get("split") != "train": - raise DataContractError( - f"AgiBotWorld stats {stats_path} were generated from split={population.get('split')!r}; " - "normalization statistics must be train-derived. Regenerate with --split train." - ) - expected_annotations = bool(population.get("use_segment_annotations")) - expected_ratio = _validate_trim_ratio(population.get("segment_max_trim_ratio")) - if expected_annotations != self._use_segment_annotations or expected_ratio != self._segment_max_trim_ratio: - raise DataContractError( - f"AgiBotWorld stats {stats_path} were generated with " - f"use_segment_annotations={expected_annotations}, segment_max_trim_ratio={expected_ratio}, " - f"but bucket {self._dataset_id} loads with " - f"use_segment_annotations={self._use_segment_annotations}, " - f"segment_max_trim_ratio={self._segment_max_trim_ratio}. Regenerate stats with matching settings." + "Use the stats shipped with the cleaned AgiBotWorld export." ) - bucket_population = recorded_buckets.get(self._dataset_id) - if not isinstance(bucket_population, dict): - raise DataContractError( - f"AgiBotWorld bucket {self._dataset_id} has no population record in {stats_path}; regenerate stats." - ) - disk_population = resolve_lerobot_v3_data_population(self._dataset_dir, info=info) - disk_digest = digest_lerobot_v3_data_population(disk_population) - if bucket_population.get("data_population_digest") != disk_digest: - raise DataContractError( - f"AgiBotWorld bucket {self._dataset_id} manifest/data mapping changed after {stats_path} was generated; " - "regenerate normalization stats." - ) - - if self._split == population.get("split") and self._max_hours is None: - effective_digest = _effective_segment_population_digest(self._eps_df) - if bucket_population.get("effective_population_digest") != effective_digest: - raise DataContractError( - f"AgiBotWorld bucket {self._dataset_id} segment/exclusion population changed after " - f"{stats_path} was generated; regenerate normalization stats." - ) def _mat(key: str, dim: int): if key not in raw: raise KeyError(f"AgiBotWorld bucket {self._dataset_id}: '{key}' missing from {stats_path}.") - return materialize_eef_stats( - raw[key], self._normalize_mode, dim=dim, strict_minmax=False, source_hint=f"{stats_path}: {key}" + block = raw[key] + if not isinstance(block, dict): + raise DataContractError(f"{stats_path}: {key} must contain a statistics mapping.") + for field in _STAT_FIELDS: + if field in block: + values = np.asarray(block[field], dtype=np.float32) + if values.shape != (dim,) or not np.isfinite(values).all(): + raise DataContractError(f"{stats_path}: {key}.{field} must contain {dim} finite values.") + stats = materialize_eef_stats( + block, self._normalize_mode, dim=dim, strict_minmax=False, source_hint=f"{stats_path}: {key}" ) + if np.any(stats["std"] < 0) or any( + np.any(stats[low] > stats[high]) for low, high in (("min", "max"), ("q01", "q99")) + ): + raise DataContractError(f"{stats_path}: {key} has invalid normalization bounds or standard deviation.") + return stats def _build(prefix: str): """Assemble the width-matched stats block for one stream ('action' or diff --git a/openwam/dataloader/bases/lerobot_v3_reader.py b/openwam/dataloader/bases/lerobot_v3_reader.py index 4fd0f9389..da341edfb 100644 --- a/openwam/dataloader/bases/lerobot_v3_reader.py +++ b/openwam/dataloader/bases/lerobot_v3_reader.py @@ -1,8 +1,8 @@ """Unified single-bucket LeRobot v3 reader base. -One base class for every LeRobot v3 single-bucket reader in this package — -the 4 OXE readers (BC-Z / Bridge / Fractal / DROID), RoboCOIN, and EgoDex are -all **direct, equal subclasses** of :class:`LeRobotV3Reader`. There is no +One base class for LeRobot v3 single-bucket readers, including the four +pretraining data sources and benchmark readers. They are +direct subclasses of :class:`LeRobotV3Reader`. There is no per-family intermediate base; the differences between readers are expressed through a small set of hook methods + class attributes, and the genuinely shared building blocks (single-arm EEF assembly, stats materialization, @@ -83,9 +83,7 @@ compute_file_local_offsets, effective_episode_frames, load_episodes_parquet, - load_tasks_annotated, parse_info_json, - resolve_prompt_by_episode, subsample_episodes_by_hours, ) from openwam.dataloader.utils.normalization import materialize_eef_stats @@ -132,6 +130,7 @@ class LeRobotV3Reader(BaseDataset): # --- subclass-overridable class attributes --- DATASET_NAME: ClassVar[str] = "LeRobotV3" + STRICT_BUCKET_LOADING: ClassVar[bool] = False NEEDED_COLS: ClassVar[Tuple[str, ...]] = () HEAD_CAMERA: ClassVar[Optional[str]] = None LEFT_WRIST_CAMERA: ClassVar[Optional[str]] = None @@ -143,9 +142,6 @@ class LeRobotV3Reader(BaseDataset): # ``_resolve_cameras``. PROPRIO_DIM_MASK: ClassVar[Optional[np.ndarray]] = None ACTION_DIM: ClassVar[int] = EEF_DIM - # Prompt convention: "task_index" (tasks.parquet, index=text — RoboCOIN / - # EgoDex) or "episode_annotated" (tasks_annotated.parquet by episode_index — OXE). - PROMPT_SOURCE: ClassVar[str] = "task_index" PROMPT_FILE_REQUIRED: ClassVar[bool] = True # In-reader normalization stats: when STATS_FILENAME is set, the default # ``_load_stats`` loads ``meta/`` (a flat stats dict) with @@ -163,9 +159,8 @@ class LeRobotV3Reader(BaseDataset): # so the artifact is authored in RAW space (NOT scattered). None → no deploy # stats written (default). DEPLOY_ACTION_MODE: ClassVar[Optional[str]] = None - # Default normalize_mode when the caller doesn't pass one. OXE readers - # override to "quantile" (their historical default); RoboCOIN/EgoDex keep - # None (no in-reader normalization unless a config opts in). + # Default normalize_mode when the caller does not pass one. A config can + # select the dataset's precomputed statistics or disable normalization. DEFAULT_NORMALIZE_MODE: ClassVar[Optional[str]] = None # Exception types a wrist (auxiliary-camera) decode failure is tolerated for # (missing / corrupt mp4 → black slot rather than a fatal sample). OXE used @@ -185,6 +180,8 @@ def __init__( multiview: bool = False, normalize_mode: Any = _NORMALIZE_MODE_UNSET, dataset_id: Optional[str] = None, + # Internal root-mode boundary for supplied family normalization stats. + stats_root: Optional[str] = None, target_camera: Optional[str] = None, camera_layout: Optional[List[str]] = None, # Octo-style head-view sampling: a list of >=2 camera keys (must include @@ -212,6 +209,7 @@ def __init__( **_unused: Any, ): self._dataset_dir = Path(dataset_dir) + self._stats_root = Path(stats_root) if stats_root is not None else None self._dataset_id = dataset_id or self._dataset_dir.name self._num_frames = int(num_frames) self._video_stride = max(1, int(video_stride)) @@ -342,12 +340,6 @@ def __init__( # LeRobot v2.1: per-episode parquet + meta/episodes.jsonl) can supply # the same eps DataFrame contract without reimplementing __init__. self._eps_df = self._build_episode_index(info) - # Stable universe for external scan tools' over-exclusion guardrails. Capture - # it before this reader applies the mutable exclusion artifact; using a - # post-exclusion length plus a later artifact snapshot races concurrent - # exclusion writers and can overstate the denominator. - self._n_episodes_before_exclusions = int(len(self._eps_df)) - # ── per-bucket episode exclusion (data-quality blacklist) ───────── # meta/excluded_episodes.json holds episode_index values that must not # be sampled (e.g. episodes whose frames fall in a truncated video @@ -368,16 +360,7 @@ def __init__( n_before, ) - # ── subclass episode filter hook (default identity) ─────────────── - # Runs AFTER info-split + excluded_episodes.json filtering and BEFORE - # the offset arrays are extracted, so dropped rows stay alignment-safe - # (the per-row _data_row_offset / _video_frame_offset columns ride - # along). Ego4D overrides this to drop episodes whose bilingual prompt - # has no usable English half; every other reader keeps the identity - # default. - self._eps_df = self._filter_episodes(self._eps_df) - - # ── optional episode-level subsample to fit a per-bucket hour budget ── + # Optional whole-episode subsampling to fit a per-bucket hour budget. if self._max_hours is not None: self._select_episodes_by_hours(float(self._max_hours), self._subsample_seed) @@ -455,7 +438,7 @@ def _select_episodes_by_hours(self, target_hours: float, seed: int) -> None: This helper only changes ``_eps_df``. The caller must rebuild the derived episode/window arrays afterwards. Keeping selection separate lets root-mode multi-bucket construction first build every leaf against - its complete split/exclusion/trim population, water-fill using those + its complete split/exclusion population, water-fill using those *effective* capacities, and then apply each allocation in place. """ n_before = len(self._eps_df) @@ -489,17 +472,7 @@ def _select_episodes_by_hours(self, target_hours: float, seed: int) -> None: def _rebuild_episode_arrays(self) -> None: """Rebuild every array derived from the current ``_eps_df`` rows.""" - effective_length = effective_episode_frames(self._eps_df) - self._ep_valid_start = ( - self._eps_df["_valid_start"].to_numpy().astype(np.int64) - if "_valid_start" in self._eps_df.columns - else np.zeros(len(self._eps_df), dtype=np.int64) - ) - self._ep_valid_end = ( - self._eps_df["_valid_end"].to_numpy().astype(np.int64) - if "_valid_end" in self._eps_df.columns - else self._eps_df["length"].to_numpy().astype(np.int64) - ) + episode_lengths = effective_episode_frames(self._eps_df) self._ep_data_row_offset = self._eps_df["_data_row_offset"].to_numpy().astype(np.int64) self._ep_video_frame_offsets = {} for cam in self._video_cameras(): @@ -509,8 +482,8 @@ def _rebuild_episode_arrays(self) -> None: min_window_len = self._num_frames if self._split == "val" else self._train_min_window_len() n_starts = np.where( - effective_length >= min_window_len, - (effective_length - min_window_len) // self._window_stride + 1, + episode_lengths >= min_window_len, + (episode_lengths - min_window_len) // self._window_stride + 1, 0, ).astype(np.int64) self._cum_n_starts = np.concatenate([[0], np.cumsum(n_starts)]).astype(np.int64) @@ -526,7 +499,7 @@ def _apply_effective_hour_budget(self, target_hours: float, seed: int) -> None: """Apply a root-mode water-fill allocation to an initialized leaf. Multi-bucket construction deliberately initializes leaves uncapped so - their split/exclusion/trim filters and data-contract checks run against + their split/exclusion filters and data-contract checks run against the complete population. It then calls this method once with the effective-hours allocation. Prompt tables, normalization stats and static action masks do not depend on the selected episode subset; only @@ -538,7 +511,7 @@ def _apply_effective_hour_budget(self, target_hours: float, seed: int) -> None: @property def effective_hours(self) -> float: - """Post-split/exclusion/trim footage represented by this leaf.""" + """Sampleable post-split/exclusion footage represented by this leaf.""" frames = self._sampleable_episode_frames(self._eps_df) return float(frames.sum()) / self._fps / 3600.0 @@ -547,9 +520,8 @@ def effective_hours(self) -> float: def _resolve_cameras(self, info: dict) -> Tuple[Optional[str], Optional[str], Optional[str]]: """Resolve (head, left_wrist, right_wrist) camera keys. - Default: a single ``target_camera`` if set (EgoDex-style), else the - class-level ``HEAD_CAMERA`` / ``LEFT_WRIST_CAMERA`` / ``RIGHT_WRIST_CAMERA`` - (OXE-style). RoboCOIN overrides to resolve by info.features priority. + Use ``target_camera`` if set, otherwise the class-level camera keys. + RoboCOIN overrides this to resolve cameras by info.features priority. """ if self._target_camera is not None: return (self._target_camera, None, None) @@ -597,46 +569,8 @@ def _train_min_window_len(self) -> int: """Min episode length to yield a train window. 1 = any single labeled step.""" return 1 - def _filter_episodes(self, eps_df: pd.DataFrame) -> pd.DataFrame: - """Optional hook to drop episodes after split / excluded_episodes filtering. - - Default: identity (no filtering). Called in ``__init__`` right after the - ``meta/excluded_episodes.json`` blacklist is applied and before the - offset arrays / window index are built, so a subclass can remove - episodes on a data-quality criterion computed from ``eps_df`` columns - (e.g. Ego4D drops rows whose bilingual ``tasks`` string has no usable - English half). Implementations MUST ``reset_index(drop=True)`` on the - returned frame. The per-row ``_data_row_offset`` / ``_video_frame_offset`` - columns are preserved across row filtering, so alignment stays correct. - - An implementation may additionally return two optional columns to trim - each episode instead of dropping it whole: - - ``_valid_start`` / ``_valid_end`` - Half-open ``[start, end)`` row range of the episode that may be - sampled. Absent → ``[0, length)`` (the whole episode). - ``__init__`` builds the window index from ``end - start`` and - ``_getitem_impl`` offsets BOTH the parquet slice and the video decode - by ``start``, so the two stay aligned and no window can reach a - trimmed row or frame. AgiBotWorld uses this for the segment-boundary - cleanup (see :meth:`~openwam.dataloader.agibotworld.AgiBotWorldDataset._filter_episodes`). - - Contract: emit them as a PAIR, with ``0 <= start < end <= length``. They - are consumed as given — an out-of-range value would silently read into - the neighbouring episode, since LeRobot v3 packs many episodes per - parquet shard and per mp4. Drop the episode rather than returning an - empty or inverted range. - """ - return eps_df - def _load_prompts(self) -> None: - """Populate the prompt lookup according to ``PROMPT_SOURCE``.""" - if self.PROMPT_SOURCE == "episode_annotated": - self._episode_idx_to_text = load_tasks_annotated( - self._dataset_dir, source_name=f"{self.DATASET_NAME}({self._dataset_id})" - ) - return - # "task_index": meta/tasks.parquet, index=text, "task_index" column=int. + """Populate the task-index lookup from ``meta/tasks.parquet``.""" tasks_path = self._dataset_dir / "meta" / "tasks.parquet" if not self.PROMPT_FILE_REQUIRED and not tasks_path.exists(): self._task_idx_to_text: Dict[int, str] = {} @@ -647,13 +581,11 @@ def _load_prompts(self) -> None: self._task_idx_to_text = dict(zip(task_idx.tolist(), task_str.tolist())) def _resolve_prompt(self, row, win: pd.DataFrame) -> str: - """Resolve this window's prompt text according to ``PROMPT_SOURCE``. + """Resolve this window's prompt text through its ``task_index``. Raises on a missing entry / empty string (a data bug should not slip into training as a fabricated prompt). """ - if self.PROMPT_SOURCE == "episode_annotated": - return resolve_prompt_by_episode(self._episode_idx_to_text, int(row["episode_index"]), self.DATASET_NAME) task_idx = int(win["task_index"].iloc[0]) if task_idx not in self._task_idx_to_text: raise KeyError( @@ -682,7 +614,7 @@ def _load_stats(self, info: dict) -> Optional[dict]: if not stats_path.exists(): raise FileNotFoundError( f"normalize_mode={self._normalize_mode!r} but {stats_path} is missing. " - f"Run the matching compute_stats script for {self.DATASET_NAME}, or set normalize_mode=null." + "Provide the dataset's precomputed statistics, or set normalize_mode=null." ) if stats_path.suffix == ".npy": raw = np.load(stats_path, allow_pickle=True).item() @@ -815,8 +747,36 @@ def _add_episode_offsets(self, eps: pd.DataFrame) -> None: for cam in self._video_cameras(): chunk_col = f"videos/{cam}/chunk_index" file_col = f"videos/{cam}/file_index" - if chunk_col in eps.columns and file_col in eps.columns: + if chunk_col not in eps.columns or file_col not in eps.columns: + continue + timestamp_col = f"videos/{cam}/from_timestamp" + if timestamp_col not in eps.columns: + # Legacy benchmark manifests omit timestamps and pack episodes + # contiguously within each video shard. eps[self._video_offset_col(cam)] = compute_file_local_offsets(eps, chunk_col, file_col) + continue + timestamps = eps[timestamp_col].to_numpy(dtype=np.float64) + frame_offsets = timestamps * self._fps + if np.any(~np.isfinite(frame_offsets) | (frame_offsets < 0) | (frame_offsets > 2**53)): + raise ValueError( + f"{self.DATASET_NAME}({self._dataset_id}): camera {cam!r} has invalid from_timestamp values" + ) + # Allow floating-point timestamp noise, but never silently move an + # episode to a different physical video frame. + if np.any(np.abs(frame_offsets - np.rint(frame_offsets)) > 0.25): + raise ValueError( + f"{self.DATASET_NAME}({self._dataset_id}): camera {cam!r} timestamps " + f"are not on frame boundaries at fps={self._fps}" + ) + end_col = f"videos/{cam}/to_timestamp" + if end_col in eps.columns: + spans = (eps[end_col].to_numpy(dtype=np.float64) - timestamps) * self._fps + if np.any(~np.isfinite(spans) | (np.abs(spans - eps["length"].to_numpy()) > 0.5)): + raise ValueError( + f"{self.DATASET_NAME}({self._dataset_id}): camera {cam!r} from/to_timestamp " + "span does not match the declared episode length" + ) + eps[self._video_offset_col(cam)] = np.rint(frame_offsets).astype(np.int64) @staticmethod def _video_offset_col(camera: str) -> str: @@ -857,11 +817,9 @@ def _safe_get(self, idx: int) -> dict: def _getitem_impl(self, idx: int) -> dict: ep_local = int(np.searchsorted(self._cum_n_starts, idx, side="right") - 1) - valid_start = int(self._ep_valid_start[ep_local]) - valid_end = int(self._ep_valid_end[ep_local]) - offset = valid_start + (idx - int(self._cum_n_starts[ep_local])) * self._window_stride + offset = (idx - int(self._cum_n_starts[ep_local])) * self._window_stride row = self._eps_df.iloc[ep_local] - actual_raw_len = min(self._num_frames, valid_end - offset) + actual_raw_len = min(self._num_frames, int(row["length"]) - offset) # 1) parquet window rows local_start = int(self._ep_data_row_offset[ep_local]) + offset @@ -1204,6 +1162,7 @@ def from_config(cls, config, split: str = "train") -> "BaseDataset": logger.info("%s.from_config: root mode, %d buckets under %s", cls.__name__, len(sub_dirs), root) base_seed = int(_get(config, "seed", 42)) total_hours = _get(config, "total_hours") + common["stats_root"] = str(root) return build_multibucket( cls, sub_dirs, @@ -1212,6 +1171,7 @@ def from_config(cls, config, split: str = "train") -> "BaseDataset": total_hours=total_hours, wrapper_cls=wrapper, source_name=cls.__name__, + strict=cls.STRICT_BUCKET_LOADING, ) diff --git a/openwam/dataloader/interndata_a1.py b/openwam/dataloader/interndata_a1.py index 192ef426a..7610521dc 100644 --- a/openwam/dataloader/interndata_a1.py +++ b/openwam/dataloader/interndata_a1.py @@ -19,16 +19,11 @@ from __future__ import annotations -import csv -import hashlib -import io import json import logging import os -import re -from dataclasses import dataclass from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Tuple import numpy as np @@ -41,8 +36,8 @@ assert_unit_quaternion, quat_wxyz_to_rot6d, ) -from openwam.dataloader.utils.lerobotv3 import DataContractError, build_multibucket -from openwam.dataloader.utils.normalization import apply_normalization, materialize_eef_stats +from openwam.dataloader.utils.lerobotv3 import build_multibucket +from openwam.dataloader.utils.normalization import apply_normalization, find_dataset_stats, materialize_eef_stats logger = logging.getLogger(__name__) @@ -153,9 +148,8 @@ def resolve_gripper_scale(bucket_dir: Path, embodiment: str, grip_col: str) -> f if observed_max / scale > _GRIPPER_SANE_MAX: logger.warning( - "InternData-A1 %s: raw source stats %s max %.4f is %.2fx the assumed " - "full-open stroke %.4f for %r — the source stats may include excluded/trimmed " - "rows; otherwise this is an out-of-range outlier or GRIPPER_FULL_OPEN needs updating.", + "InternData-A1 %s: bucket stats %s max %.4f is %.2fx the assumed " + "full-open stroke %.4f for %r; verify the gripper range or update GRIPPER_FULL_OPEN.", bucket_dir, grip_col, observed_max, @@ -235,8 +229,8 @@ def embodiment_key(robot_type: str, arm_layout: str) -> str: Unknown robot types fall back to a slug of the raw string so a newly added embodiment gets its own stats bucket instead of silently borrowing another - one's scale. The fallback is logged because it means the stats-computation - script must be re-run to produce the matching file. + one's scale. A matching precomputed stats file must exist in the dataset + tree before this embodiment can be normalized. """ if robot_type in ROBOT_TYPE_TO_EMBODIMENT: @@ -244,7 +238,7 @@ def embodiment_key(robot_type: str, arm_layout: str) -> str: slug = "".join(c if c.isalnum() else "_" for c in str(robot_type).lower()).strip("_") or "unknown" logger.warning( "InternData-A1: unrecognized robot_type %r (%s layout) -> stats key %r. " - "Add it to ROBOT_TYPE_TO_EMBODIMENT and re-run interndata_a1_stats_computation.", + "Provide matching precomputed normalization statistics for this embodiment.", robot_type, arm_layout, slug, @@ -255,34 +249,10 @@ def embodiment_key(robot_type: str, arm_layout: str) -> str: def discover_a1_buckets(root: Path) -> List[Path]: """Recursively find every LeRobot bucket (a dir holding ``meta/info.json``). - The A1 tree nests buckets at two different depths (see module docstring), so - a fixed-depth glob misses half the dataset. Each branch is pruned as soon as - a bucket is found — buckets never nest — which keeps the walk off the - ``data/`` and ``videos/`` subtrees where essentially all the inodes live. - - Symlinks ARE followed. Carving a subset out of a large tree by symlinking - buckets (rather than copying them) is the realistic usage, and the base - reader's root mode follows symlinks too — ``os.walk``'s ``followlinks=False`` - default would silently skip every such bucket and report the tree as empty. - - A visited ``(st_dev, st_ino)`` set then makes each directory yield at most - one bucket path. Unguarded, a symlink cycle does NOT hang: each lap adds one - symlink component until resolution hits ``MAXSYMLINKS`` (40 on Linux) and the - walk stops on its own — but it emits many aliases of every bucket reachable - through the cycle, i.e. silently duplicated training data. Which alias - survives is decided by traversal order, so ``dirnames`` is sorted below: raw - ``readdir`` order is a filesystem-instance property (ext4 htree hashing is - seeded per mkfs), and without the sort the same tree on two machines can keep - different aliases — shifting every later bucket's index, hence the per-bucket - subsample seeds in ``build_multibucket`` and the stats merge order. With it, - the lexicographically-first path always wins. - - Dot-prefixed directories are skipped: archive extraction stages - every archive through ``//.partial_/`` and logs into - ``.extract_logs/``. A SIGKILL / OOM / node preemption bypasses the extractor's - cleanup entirely, so a half-extracted staging tree whose ``meta/`` was already - written would otherwise be discovered as a complete bucket and then fail at - ``__getitem__`` mid-training. + Buckets occur at variable depth, so prune each branch when its metadata is + found. Follow symlinks while deduplicating physical directories to avoid + duplicate training samples and cycles. Sorted traversal keeps bucket IDs + and subsampling reproducible. Hidden staging directories are skipped. """ out: List[Path] = [] @@ -307,473 +277,20 @@ def discover_a1_buckets(root: Path) -> List[Path]: return sorted(out) -_CHUNK_DIR_RE = re.compile(r"chunk-(\d+)") -_SHARD_FILE_RE = re.compile(r"file-(\d+)\.parquet") - - -def parse_shard_path(path) -> Optional[Tuple[int, int]]: - """``(chunk, file)`` for a CANONICAL LeRobot shard path, else ``None``. - - ``fullmatch`` on both components: ``file-000.backup.parquet`` matches the - ``file-*.parquet`` glob but is not a shard. A copy left beside the original - doubled the stats row count while the reader ignored it. - - Canonical means the name is exactly what ``info.json``'s ``data_path`` - template emits — ``chunk-{chunk_index:03d}/file-{file_index:03d}.parquet``. - Sampling rebuilds the path from that template, never from the name found on - disk, so accepting a name the template cannot produce enumerates a file that - can never be read: an earlier revision took ``file-0.parquet``, scanned its - rows into the statistics and let the reader construct windows over them, - and only the first real data load failed — on a *different*, non-existent - padded path. The round-trip also collapses aliases: ``file-00.parquet`` and - ``file-000.parquet`` would otherwise both claim shard 0. - """ - - p = Path(path) - chunk = _CHUNK_DIR_RE.fullmatch(p.parent.name) - shard = _SHARD_FILE_RE.fullmatch(p.name) - if chunk is None or shard is None: - return None - ci, fi = int(chunk.group(1)), int(shard.group(1)) - if p.parent.name != f"chunk-{ci:03d}" or p.name != f"file-{fi:03d}.parquet": - return None - return ci, fi - - -def iter_data_shards(bucket) -> List[Tuple[int, int, Path]]: - """Every data shard of ``bucket`` as ``(chunk, file, path)``, in index order. - - Ordered by the PARSED indices, not by path: row offsets accumulate in this - order, so under a lexical sort an unpadded ``file-10`` would follow - ``file-1`` and shift every subsequent shard's boundary. A1 ships zero-padded - names where the two agree, which is exactly why sorting on the parsed value - costs nothing and removes the dependency on that padding. - """ - - out: List[Tuple[int, int, Path]] = [] - for pth in (Path(bucket) / "data").glob("chunk-*/file-*.parquet"): - ids = parse_shard_path(pth) - if ids is not None: - out.append((ids[0], ids[1], pth)) - return sorted(out, key=lambda t: (t[0], t[1])) - - -def load_excluded_episodes(bucket) -> set: - """``meta/excluded_episodes.json`` as a set of ints; empty set when absent. - - Strict by design. The base reader keeps the JSON values verbatim and tests - them with ``isin`` against an integer ``episode_index`` column, so a file - holding ``["0"]`` excludes NOTHING there — while an ``int()`` cast in the - stats generator excluded episode 0. Same file, two populations, no error. - - Rejecting non-integers is the one resolution that cannot drift: it gives - both sides the same answer by construction rather than by keeping two casts - in agreement. Every file this repo emits contains plain ints, so nothing - valid is refused. ``bool`` is excluded explicitly because it passes - ``isinstance(x, int)`` and would silently become episode 0/1. - """ - - path = Path(bucket) / "meta" / "excluded_episodes.json" - if not path.is_file(): - return set() - with open(path) as fh: - raw = json.load(fh)["episode_indices"] - out = set() - for x in raw: - if isinstance(x, bool) or not isinstance(x, int): - raise ValueError( - f"{path}: episode_indices must be JSON integers, got {x!r} " - f"({type(x).__name__}). The reader matches these against an integer " - "episode_index column verbatim, so a quoted or float value excludes " - "nothing there while excluding here — the two sides would then " - "normalise over different populations with no error anywhere." - ) - out.add(int(x)) - return out - - def validate_manifest_ranges(from_idx, to_idx, lengths, episode_idx, who: str) -> None: - """The manifest's ``[from, to)`` ranges must tile ``[0, total)`` exactly. - - A per-episode capacity check and an equal grand total are BOTH satisfied by - ranges that overlap: ``[0,4) [2,6) [8,12)`` ends at 12 and every range fits - inside a 12-row shard, yet episode 1 reads two of episode 0's rows and two - of its own. The ranges are the only statement of which rows belong to whom, - so they have to be checked as a partition — individually valid ranges say - nothing about whether they carve the file up consistently. - - ``length`` is checked against the range width for the same reason: the - reader uses ``to - from`` to size the window and ``length`` to bound it, so - the two disagreeing means one of them is describing a different episode. - """ - - from_idx = np.asarray(from_idx, dtype=np.int64) - to_idx = np.asarray(to_idx, dtype=np.int64) - lengths = np.asarray(lengths, dtype=np.int64) - episode_idx = np.asarray(episode_idx, dtype=np.int64) - - if (from_idx < 0).any() or (to_idx < from_idx).any(): - i = int(np.flatnonzero((from_idx < 0) | (to_idx < from_idx))[0]) - raise ValueError( - f"{who}: episode {int(episode_idx[i])} has a malformed manifest range " - f"[{int(from_idx[i])}, {int(to_idx[i])})." - ) - bad = np.flatnonzero(to_idx - from_idx != lengths) - if bad.size: - i = int(bad[0]) - raise ValueError( - f"{who}: episode {int(episode_idx[i])} spans " - f"[{int(from_idx[i])}, {int(to_idx[i])}) = {int(to_idx[i] - from_idx[i])} rows but " - f"declares length={int(lengths[i])}. The reader sizes windows from one and " - "bounds them with the other, so they must describe the same episode." - ) - - order = np.argsort(from_idx, kind="stable") - fs, ts = from_idx[order], to_idx[order] - if fs.size > 1: - overlap = np.flatnonzero(fs[1:] < ts[:-1]) - if overlap.size: - k = int(overlap[0]) - raise ValueError( - f"{who}: episode {int(episode_idx[order][k + 1])} starts at row " - f"{int(fs[k + 1])} but episode {int(episode_idx[order][k])} runs to " - f"{int(ts[k])}. Overlapping manifest ranges make two episodes read the " - "same rows — the totals and the per-episode capacity check both pass " - "while the data is paired with the wrong episode." - ) - - -def _shard_episode_bounds(pf) -> Optional[Tuple[int, int]]: - """Return the min/max ``episode_index`` advertised by an open shard. - - Row-group metadata is preferred. If statistics are absent, the function - reads the run-length-encoded ``episode_index`` column instead of treating - missing evidence as success. ``None`` means the column itself is absent. - The result is an envelope check, not proof of row order within the shard. - """ - - try: - col = pf.schema_arrow.names.index("episode_index") - except ValueError: - return None - lo = hi = None - meta = pf.metadata - for g in range(meta.num_row_groups): - st = meta.row_group(g).column(col).statistics - if st is None or not st.has_min_max: - vals = pf.read(columns=["episode_index"]).column("episode_index") - if len(vals) == 0: - return None - return int(min(vals.to_pylist())), int(max(vals.to_pylist())) - lo = st.min if lo is None else min(lo, st.min) - hi = st.max if hi is None else max(hi, st.max) - if lo is None or hi is None: - return None - return int(lo), int(hi) - - -@dataclass(frozen=True) -class _A1TrimSnapshot: - """One immutable parse of the exact trim CSV bytes used for construction.""" - - path: str - sha256: str - spec: Dict[str, Dict[int, Tuple[int, Optional[int], Optional[int]]]] - - @property - def provenance(self) -> dict: - return {"schema_version": 1, "sha256": self.sha256} - - -_TRIM_SPEC_CACHE: Dict[str, Tuple[Tuple[int, int, int, int], _A1TrimSnapshot]] = {} - - -def _trim_file_signature(path: Path) -> Tuple[int, int, int, int]: - stat = path.stat() - return (int(stat.st_dev), int(stat.st_ino), int(stat.st_size), int(stat.st_mtime_ns)) - - -def _load_trim_snapshot(path) -> _A1TrimSnapshot: - """Read, hash and parse one trim artifact from the same byte snapshot. - - Root mode constructs many readers, so a successful parse is cached while the - file's stat signature is unchanged. The SHA-256 is persisted in normalization - stats; the reader compares it before accepting those stats. - """ - - key = str(path) - trim_path = Path(key) - try: - signature_before = _trim_file_signature(trim_path) - cached = _TRIM_SPEC_CACHE.get(key) - if cached is not None and cached[0] == signature_before: - return cached[1] - raw = trim_path.read_bytes() - signature_after = _trim_file_signature(trim_path) - if signature_before != signature_after: - raise DataContractError( - f"InternDataA1 trim_csv {key} changed while it was being read; retry with an immutable file." - ) - - spec: Dict[str, Dict[int, Tuple[int, Optional[int], Optional[int]]]] = {} - seen = set() - n = 0 - reader = csv.DictReader(io.StringIO(raw.decode("utf-8"), newline="")) - required = {"dataset", "episode_index", "total_frames", "trim_head_to", "trim_tail_from"} - if reader.fieldnames is None or not required.issubset(reader.fieldnames): - missing = sorted(required.difference(reader.fieldnames or ())) - raise ValueError(f"missing required column(s): {missing}") - for row in reader: - - def _int(name): - value = (row.get(name) or "").strip() - return int(value) if value else None - - dataset = (row.get("dataset") or "").strip() - if not dataset: - raise ValueError("empty dataset key") - episode_index = _int("episode_index") - if episode_index is None or episode_index < 0: - raise ValueError(f"invalid episode_index={episode_index!r}") - entry_key = (dataset, episode_index) - if entry_key in seen: - raise ValueError(f"duplicate entry for dataset={dataset!r}, episode_index={episode_index}") - seen.add(entry_key) - - head = _int("trim_head_to") or 0 - tail = _int("trim_tail_from") - total = _int("total_frames") - if head < 0 or (tail is not None and tail < 0) or (total is not None and total <= 0): - raise ValueError( - f"invalid bounds for dataset={dataset!r}, episode_index={episode_index}: " - f"head={head}, tail={tail}, total={total}" - ) - if not head and tail is None: - continue - spec.setdefault(dataset, {})[episode_index] = (head, tail, total) - n += 1 - except (OSError, UnicodeError, csv.Error, KeyError, ValueError) as e: + """Require the cleaned manifest to describe consecutive physical rows.""" + starts = np.asarray(from_idx, dtype=np.int64) + ends = np.asarray(to_idx, dtype=np.int64) + sizes = np.asarray(lengths, dtype=np.int64) + ids = np.asarray(episode_idx, dtype=np.int64) + malformed = np.flatnonzero((starts < 0) | (sizes <= 0) | (ends - starts != sizes)) + if malformed.size: + i = int(malformed[0]) raise ValueError( - f"InternDataA1: trim_csv {key} could not be read ({e}). Fix the path or set " - "trim_csv=null to run untrimmed — it will not be skipped silently." - ) from e - - snapshot = _A1TrimSnapshot( - path=key, - sha256=hashlib.sha256(raw).hexdigest(), - spec=spec, - ) - _TRIM_SPEC_CACHE[key] = (signature_after, snapshot) - logger.info("InternDataA1: loaded %d trim entries across %d buckets from %s", n, len(spec), key) - return snapshot - - -def _load_trim_spec(path) -> Dict[str, Dict[int, Tuple[int, Optional[int], Optional[int]]]]: - """Compatibility wrapper returning the parsed mapping from a pinned snapshot.""" - return _load_trim_snapshot(path).spec - - -def _assert_trim_snapshot_current(snapshot: _A1TrimSnapshot, *, context: str) -> None: - """Fail if the configured trim bytes changed after the snapshot was pinned.""" - try: - actual_sha256 = hashlib.sha256(Path(snapshot.path).read_bytes()).hexdigest() - except OSError as e: - raise DataContractError(f"InternDataA1 trim_csv {snapshot.path} disappeared during {context}: {e}") from e - if actual_sha256 != snapshot.sha256: - raise DataContractError( - f"InternDataA1 trim_csv {snapshot.path} changed during {context}; " - f"expected sha256={snapshot.sha256}, got {actual_sha256}." - ) - - -class AmbiguousBucketKey(LookupError): - """A bare bucket name matches several keys, so the bucket cannot be identified.""" - - -def resolve_bucket_key(keys, dataset_id: str, bucket_dir, *, what: str, source: str) -> Optional[str]: - """Resolve a bucket onto a key in a mapping keyed by bucket path. - - **Shared by the reader and the stats generator on purpose** — they used to - resolve differently, so the same trim CSV could be applied by one and skipped - by the other while digest provenance saw nothing wrong. - - Matching is by EXACT key only, tried against progressively longer tails of - this bucket's real path. A bare-suffix match is deliberately not used: these - maps are frequently sparse (a trim CSV lists only trimmed buckets; a coverage - map only successfully scanned ones), and in a sparse map a unique suffix - proves nothing about identity. For example, with `good/…/apple` - scanned and `bad/…/apple` skipped, a reader for `bad` matched the sole - `apple` suffix and accepted statistics computed entirely from `good`. - - So an exact miss for an already-qualified id stays a miss, and an unqualified - id is resolved from the directory tree rather than from the map's contents. - - Raises :class:`AmbiguousBucketKey` only when the path itself cannot - disambiguate — the caller must decide, because silently picking one would - apply another bucket's numbers and regenerating produces the same names - (``///`` is a documented depth, so repeating leaf - names is expected). - """ - - if dataset_id in keys: - return dataset_id - - if "/" in dataset_id: - return None - - parts = Path(bucket_dir).resolve().parts - for n in range(min(len(parts), 6), 0, -1): - cand = "/".join(parts[-n:]) - if cand in keys: - return cand - - leaf = Path(bucket_dir).name - cands = sorted(k for k in keys if k == leaf or k.endswith("/" + leaf)) - if len(cands) > 1: - raise AmbiguousBucketKey( - f"{source}: {what} has {len(cands)} buckets whose path ends in {leaf!r} " - f"({', '.join(cands[:4])}{' ...' if len(cands) > 4 else ''}), and none of them " - "matches this bucket's own path. Bucket leaf names repeat across tasks, so this " - "one cannot be identified from its directory name. Pass dataset_id (or point " - "--dataset_dir at the corpus root) so buckets are keyed by their path relative " - "to the root, which is what these files key on." + f"{who}: episode {int(ids[i])} spans [{int(starts[i])}, {int(ends[i])}) but declares length={int(sizes[i])}" ) - return None - - -def resolve_trim_bounds(entry, length: int, min_len: int) -> Optional[Tuple[int, int]]: - """Decide one episode's kept span, or ``None`` to leave it whole. - - Shared by the reader and the stats script **on purpose**. They must make the - identical keep/skip call: if the stats path trimmed an episode the reader - leaves whole, the normalizer would describe a distribution the reader never - emits. Two copies of these four conditions would drift silently — nothing - downstream compares them. - - Returns ``None`` when the entry should not be applied at all: - * stale — recorded ``total_frames`` disagrees with the manifest length - (the signature of a CSV built against a different corpus version); - * a no-op (head 0, tail == length); - * the trim would leave fewer than ``min_len`` frames, i.e. not even one - usable window — better a whole episode than a degenerate one. - """ - - head, tail_from, total = entry - if total is not None and int(total) != int(length): - return None - tail = int(length) if tail_from is None else min(int(tail_from), int(length)) - head = max(0, min(int(head), tail)) - if head == 0 and tail == int(length): - return None - if tail - head < min_len: - return None - return head, tail - - -def effective_a1_population_provenance(eps_df, trim_spec: Optional[Dict[int, Tuple]], min_len: int) -> dict: - """Digest the exact post-split/post-exclusion episode spans used by A1. - - ``eps_df`` is the reader population immediately before trim application. - Hashing the global manifest start, raw length and effective relative - ``[head, tail)`` span detects changes to the manifest mapping, exclusion set - or trim bounds even when the bucket name and aggregate row count stay the - same. The stats generator calls this same function on its selected manifest. - """ - - required = ("episode_index", "dataset_from_index", "length") - missing = [column for column in required if column not in eps_df.columns] - if missing: - raise DataContractError(f"InternDataA1 effective population is missing manifest columns {missing}") - - trim_spec = trim_spec or {} - records = [] - for _, row in eps_df.iterrows(): - episode_index = int(row["episode_index"]) - global_start = int(row["dataset_from_index"]) - length = int(row["length"]) - if episode_index < 0 or global_start < 0 or length <= 0: - raise DataContractError( - "InternDataA1 effective population contains an invalid manifest record: " - f"episode_index={episode_index}, dataset_from_index={global_start}, length={length}" - ) - entry = trim_spec.get(episode_index) - bounds = resolve_trim_bounds(entry, length, min_len) if entry is not None else None - head, tail = bounds if bounds is not None else (0, length) - records.append((episode_index, global_start, length, int(head), int(tail))) - - records.sort(key=lambda record: record[0]) - if len({record[0] for record in records}) != len(records): - raise DataContractError("InternDataA1 effective population has duplicate episode_index values") - - hasher = hashlib.sha256(b"openwam:interndata-a1-effective-population:v1\0") - num_rows = 0 - for record in records: - for value in record: - hasher.update(value.to_bytes(8, "little", signed=False)) - num_rows += record[4] - record[3] - return { - "schema_version": 1, - "sha256": hasher.hexdigest(), - "num_episodes": len(records), - "num_rows": int(num_rows), - } - - -def _validate_a1_root_stats_contributors( - root: Path, - stats_root: Path, - sub_dirs: List[Path], -) -> None: - """Bind every pooled per-embodiment stats file to the current root set. - - A leaf-level membership check cannot detect a contributor directory removed - after stats generation: all surviving leaves still find themselves in the - file even though its pooled values include the removed bucket. Root mode has - the authoritative discovered set, so compare it before bucket fan-out. - """ - - current: Dict[str, set[str]] = {} - for sub in sub_dirs: - try: - info = json.loads((sub / "meta" / "info.json").read_text()) - layout = detect_arm_layout(info.get("features", {}) or {}) - embodiment = embodiment_key(info.get("robot_type", "unknown"), layout) - dataset_id = str(sub.relative_to(root)) - except Exception as exc: - raise DataContractError(f"InternData-A1 cannot classify pooled stats contributor {sub}: {exc}") from exc - current.setdefault(embodiment, set()).add(dataset_id) - - for embodiment, current_buckets in sorted(current.items()): - stats_path = stats_root / "meta" / f"stats_{embodiment}.json" - try: - raw = json.loads(stats_path.read_text()) - population = raw["population"] - contributed = population["buckets"] - empty = population.get("empty_buckets", []) - except (KeyError, OSError, TypeError, ValueError) as exc: - raise DataContractError( - f"InternData-A1 pooled stats {stats_path} has no readable contributor population; regenerate stats." - ) from exc - if ( - not isinstance(contributed, list) - or not isinstance(empty, list) - or not all(isinstance(name, str) and name for name in [*contributed, *empty]) - or len(set(contributed)) != len(contributed) - or len(set(empty)) != len(empty) - or set(contributed).intersection(empty) - ): - raise DataContractError( - f"InternData-A1 pooled stats {stats_path} has a malformed contributor set; regenerate stats." - ) - recorded_buckets = set(contributed) | set(empty) - if recorded_buckets != current_buckets: - missing = sorted(current_buckets - recorded_buckets) - extra = sorted(recorded_buckets - current_buckets) - raise DataContractError( - f"InternData-A1 pooled stats {stats_path} contributor set no longer matches " - f"the current {embodiment} root population (missing={missing}, extra={extra}); " - "regenerate stats." - ) + if starts.size and (starts[0] != 0 or np.any(starts[1:] != ends[:-1])): + raise ValueError(f"{who}: Overlapping manifest ranges or missing rows in the cleaned dataset") class InternDataA1Dataset(LeRobotV3Reader): @@ -794,61 +311,15 @@ class InternDataA1Dataset(LeRobotV3Reader): STATS_DIM = _ACTION_DIM STATS_STRICT_MINMAX = False DEFAULT_NORMALIZE_MODE = "quantile" + STRICT_BUCKET_LOADING = True DEPLOY_ACTION_MODE = None - CONFIG_KEYS = LeRobotV3Reader.CONFIG_KEYS + ("trim_csv",) - - def __init__( - self, - dataset_dir, - *, - a1_stats_root: Optional[str] = None, - trim_csv: Optional[str] = None, - _trim_snapshot: Optional[_A1TrimSnapshot] = None, - **kwargs, - ): - """ - Args: - a1_stats_root: directory holding ``meta/stats_.json``. - Defaults to ``dataset_dir`` itself (single-bucket use); in root - mode :meth:`from_config` passes the dataset root so every bucket - shares one per-embodiment stats file. - trim_csv: optional path to a quality-audit trim list; see - :meth:`_filter_episodes`. ``None`` (default) disables trimming. - """ - - self._a1_stats_root = Path(a1_stats_root) if a1_stats_root else Path(dataset_dir) - self._trim_csv = trim_csv - trim_snapshot_was_provided = _trim_snapshot is not None - if _trim_snapshot is not None: - if trim_csv is None or not isinstance(_trim_snapshot, _A1TrimSnapshot): - raise ValueError("_trim_snapshot requires a matching non-null trim_csv") - if _trim_snapshot.path != str(trim_csv): - raise ValueError( - f"_trim_snapshot path {_trim_snapshot.path!r} does not match trim_csv {str(trim_csv)!r}" - ) - elif trim_csv is not None: - _trim_snapshot = _load_trim_snapshot(trim_csv) - self._trim_snapshot = _trim_snapshot - - self._a1_excluded_episode_indices = tuple(sorted(load_excluded_episodes(dataset_dir))) - super().__init__(dataset_dir, **kwargs) - current_exclusions = tuple(sorted(load_excluded_episodes(dataset_dir))) - if current_exclusions != self._a1_excluded_episode_indices: - raise DataContractError( - f"InternDataA1 bucket {self._dataset_id}: excluded_episodes.json changed during reader " - f"construction; expected {list(self._a1_excluded_episode_indices)}, " - f"got {list(current_exclusions)}." - ) - if self._trim_snapshot is not None and not trim_snapshot_was_provided: - _assert_trim_snapshot_current(self._trim_snapshot, context="reader construction") - - self._trim_snapshot = None + CONFIG_KEYS = LeRobotV3Reader.CONFIG_KEYS def _load_excluded_episode_indices(self) -> set[int]: - """Use the strict exclusion snapshot pinned before base construction.""" - return set(self._a1_excluded_episode_indices) + """The cleaned dataset already contains exactly the retained episodes.""" + return set() def _resolve_cameras(self, info: dict): """Detect the arm layout, then pick cameras and per-bucket columns. @@ -903,95 +374,21 @@ def _resolve_cameras(self, info: dict): return head, left_wrist, right_wrist def _add_data_offsets(self, eps) -> None: - """Resolve data-shard and row offsets from physical parquet lengths. - - Manifest file indices can be stale at shard boundaries. Global - ``dataset_from_index`` values are therefore located against cumulative - physical shard lengths. The method validates manifest ranges, total - row coverage, per-episode shard capacity, and the shard's advertised - episode envelope before publishing offsets. - """ - - from concurrent.futures import ThreadPoolExecutor - - import pyarrow.parquet as pq - - shards = iter_data_shards(self._dataset_dir) - if not shards: - raise FileNotFoundError(f"No data parquet files under {self._dataset_dir}/data") - - def _read_meta(entry): - chunk, file_idx, path = entry - pf = pq.ParquetFile(path) - return chunk, file_idx, pf.metadata.num_rows, _shard_episode_bounds(pf) - - with ThreadPoolExecutor(max_workers=min(len(shards), 4)) as pool: - data_files = list(pool.map(_read_meta, shards)) - - who = f"{self.DATASET_NAME}({self._dataset_id})" - + """Use canonical v3 file IDs and contiguous manifest row ranges.""" validate_manifest_ranges( eps["dataset_from_index"].to_numpy(), eps["dataset_to_index"].to_numpy(), eps["length"].to_numpy(), eps["episode_index"].to_numpy(), - who, + f"{self.DATASET_NAME}({self._dataset_id})", ) - - starts = np.concatenate([[0], np.cumsum([n for _, _, n, _ in data_files])]).astype(np.int64) - - manifest_end = int(eps["dataset_to_index"].to_numpy().max()) - if manifest_end != int(starts[-1]): - raise ValueError( - f"{who}: the data shards hold {int(starts[-1])} rows but the manifest ends at " - f"{manifest_end}. A shard is missing, truncated or out of order — resolving " - "offsets against this would map episodes onto another episode's rows." - ) - - global_starts = eps["dataset_from_index"].to_numpy().astype(np.int64) - file_pos = np.searchsorted(starts, global_starts, side="right") - 1 - if (file_pos < 0).any() or (file_pos >= len(data_files)).any(): - raise ValueError(f"{who}: dataset_from_index outside the data parquet row range") - - offsets = global_starts - starts[file_pos] - lengths = eps["dataset_to_index"].to_numpy().astype(np.int64) - global_starts - capacity = np.array([data_files[i][2] for i in file_pos], dtype=np.int64) - overflow = np.flatnonzero(offsets + lengths > capacity) - if overflow.size: - i = int(overflow[0]) - raise ValueError( - f"{who}: episode {int(eps['episode_index'].to_numpy()[i])} needs rows " - f"[{int(offsets[i])}, {int(offsets[i] + lengths[i])}) of a shard holding only " - f"{int(capacity[i])}. The shard set does not match the manifest — resolving " - "offsets against it would map episodes onto another episode's rows." - ) - - ep_vals = eps["episode_index"].to_numpy() - for i, pos in enumerate(file_pos): - bounds = data_files[pos][3] - if bounds is None: - continue - ep = int(ep_vals[i]) - if not bounds[0] <= ep <= bounds[1]: - raise ValueError( - f"{who}: the manifest puts episode {ep} in shard " - f"chunk-{data_files[pos][0]:03d}/file-{data_files[pos][1]:03d}, but that " - f"shard only holds episodes {bounds[0]}..{bounds[1]}. The shard set does " - "not match the manifest (a copied, reordered or substituted shard passes " - "the row-count checks above while holding another episode's data)." - ) - - eps["data/chunk_index"] = np.array([data_files[i][0] for i in file_pos], dtype=np.int64) - eps["data/file_index"] = np.array([data_files[i][1] for i in file_pos], dtype=np.int64) - eps["_data_row_offset"] = offsets + super()._add_data_offsets(eps) def _add_episode_offsets(self, eps) -> None: - """Rebuild camera-frame offsets from file-local timestamps. + """Read camera-frame offsets from the cleaned file-local timestamps. - ``round(from_timestamp * fps)`` is independent of which manifest rows - survived filtering, so dropped episodes cannot shift later video away - from its action/state rows. Missing timestamps and invalid fps values - fail explicitly rather than leaving one camera on a stale offset. + ``round(from_timestamp * fps)`` locates the same native frames as the + metadata. Missing timestamps and inconsistent spans fail before sampling. """ super()._add_episode_offsets(eps) @@ -1069,135 +466,6 @@ def _add_episode_offsets(self, eps) -> None: ) eps[self._video_offset_col(cam)] = off - def _trim_min_len(self) -> int: - """Minimum frames an episode must keep to be worth trimming. - - Named so `_load_stats` can quote the same number the stats generator - must have been given (`--min_keep`): the same trim CSV under a different - bound produces a different kept population. - """ - - return self._num_frames if self._split == "val" else self._train_min_window_len() - - def _match_bucket_key(self, keys, what: str) -> Optional[str]: - """Map this bucket onto a key in a mapping keyed by bucket path. - - Both the trim CSV's ``dataset`` column and the stats file's - ``exclusions`` map are keyed by the bucket path **relative to the - dataset root**, which is exactly what root mode already uses as - ``dataset_id`` (see ``_per_bucket`` in :meth:`from_config`) — so the - common case is a direct hit. - - Single-bucket mode is the awkward one: ``_dataset_id`` falls back to the - bare directory name. A1 bucket names repeat across embodiments - (``pick_beef_sandwich_on_conveyor`` exists under both lift2 and - split_aloha), so a bare name is accepted only when it resolves - uniquely — an ambiguous one is refused rather than guessed, since - matching the wrong bucket would apply another embodiment's numbers. - """ - - return resolve_bucket_key( - keys, self._dataset_id, self._dataset_dir, what=what, source=f"InternDataA1({self._dataset_id})" - ) - - def _trim_key(self) -> Optional[str]: - spec = self._get_trim_snapshot().spec - if not spec: - return None - return self._match_bucket_key(spec, "trim_csv") - - def _get_trim_snapshot(self) -> _A1TrimSnapshot: - snapshot = getattr(self, "_trim_snapshot", None) - if snapshot is None: - if not self._trim_csv: - raise ValueError("InternDataA1 trim snapshot requested while trim_csv is disabled") - snapshot = _load_trim_snapshot(self._trim_csv) - self._trim_snapshot = snapshot - return snapshot - - def _filter_episodes(self, eps_df): - """Trim motionless episode heads/tails from an optional trim artifact. - - Trimming adjusts data and video frame offsets together and shortens the - episode length; it does not splice mid-episode pauses. Bounds are - integer frame offsets, avoiding timestamp rounding skew. Entries whose - recorded length no longer matches the manifest are rejected as stale, - and the exact retained population is bound into normalization stats. - """ - - eps_df = super()._filter_episodes(eps_df) - spec = None - if self._trim_csv: - key = self._trim_key() - if key is not None: - spec = self._get_trim_snapshot().spec.get(key) - - min_len = self._trim_min_len() - - self._a1_effective_population = effective_a1_population_provenance( - eps_df, - spec, - min_len, - ) - if not spec: - return eps_df - - cam_cols = [c for c in eps_df.columns if c.startswith("_video_frame_offset/")] - - lengths = eps_df["length"].to_numpy().copy() - row_off = eps_df["_data_row_offset"].to_numpy().copy() - cam_off = {c: eps_df[c].to_numpy().copy() for c in cam_cols} - n_trim = n_skip = 0 - frames_before = int(lengths.sum()) - for pos, ep in enumerate(eps_df["episode_index"].to_numpy()): - entry = spec.get(int(ep)) - if entry is None: - continue - length = int(lengths[pos]) - bounds = resolve_trim_bounds(entry, length, min_len) - if bounds is None: - if not (entry[0] == 0 and entry[1] in (None, length)): - n_skip += 1 - continue - head, tail = bounds - lengths[pos] = tail - head - if head: - row_off[pos] += head - for c in cam_cols: - cam_off[c][pos] += head - n_trim += 1 - - if n_skip: - logger.warning( - "InternDataA1(%s): %d trim entries skipped — recorded total_frames disagrees " - "with the manifest length, or the trim would leave < %d frames. A total_frames " - "mismatch is what a STALE trim list looks like; regenerate it against this " - "corpus.", - self._dataset_id, - n_skip, - min_len, - ) - if not n_trim: - return eps_df - - eps_df = eps_df.copy() - eps_df["length"] = lengths - eps_df["_data_row_offset"] = row_off - for c in cam_cols: - eps_df[c] = cam_off[c] - after = int(lengths.sum()) - logger.info( - "InternDataA1(%s): trimmed %d/%d episodes, %d -> %d frames (-%.1f%%, %.2f h removed)", - self._dataset_id, - n_trim, - len(eps_df), - frames_before, - after, - 100.0 * (frames_before - after) / max(frames_before, 1), - (frames_before - after) / self._fps / 3600.0, - ) - return eps_df.reset_index(drop=True) - def _train_min_window_len(self) -> int: """Require >= 2 rows so every train window keeps >= 1 supervised step. @@ -1220,202 +488,38 @@ def _n_supervised_action_steps(self, actual_raw_len: int) -> int: return actual_raw_len return max(0, actual_raw_len - 1) - def _check_stats_population(self, raw: dict, stats_path) -> None: - """Refuse stats computed over a different population than this reader reads. - - Trimmed and untrimmed stats describe different distributions — trimming - removes the motionless head/tail, which pulls the positional means - toward the moving span — and NOTHING in the numbers says which is which. - The same is true of a val-derived file loaded by a train reader, of a - different ``--min_keep``, and of a bucket that dropped out of the scan - but still loads the shared per-embodiment file. Each was reproduced end - to end through the real generator and reader; each mis-scales every - action for the whole run with no error anywhere. - - The broad policy fields below catch split/trim/min-keep/coverage - mismatches. A trimming reader additionally requires schema-v2 - provenance: a hash of the exact trim CSV bytes and a per-bucket digest - of the selected manifest spans after exclusions and trim resolution. - The stats worker creates that certificate from the same immutable trim - snapshot and strict exclusion parser used to select its parquet rows. - - Missing block -> refuse. A stats file that cannot say what it covers is - exactly the unverifiable pairing this exists to prevent, and the fix is - one generator re-run. - """ - - pop = raw.get("population") - rerun = ( - "Re-run python -m openwam.dataloader.utils.stats_computation." - f"interndata_a1_stats_computation --dataset_dir {self._dataset_dir} " - f"--stats_root {self._a1_stats_root}" - ) - if not isinstance(pop, dict): - raise DataContractError( - f"InternData-A1 bucket {self._dataset_id}: {stats_path} has no 'population' " - f"block, so there is no way to tell whether it describes the rows this reader " - f"loads (split, trimming and keep-bound all change the distribution). {rerun}." - ) - - if pop.get("split") != "train": - raise DataContractError( - f"InternData-A1 bucket {self._dataset_id}: {stats_path} was computed over the " - f"{pop.get('split')!r} split. Normalization statistics must describe the " - f"training distribution — a val-derived file would scale training by numbers " - f"drawn from data the model never fits. {rerun} --split train." - ) - - want_trim = bool(self._trim_csv) - if bool(pop.get("trim_active")) != want_trim: - state = "with" if pop.get("trim_active") else "without" - mine = "is" if want_trim else "is not" - raise DataContractError( - f"InternData-A1 bucket {self._dataset_id}: {stats_path} was computed {state} a " - f"trim list but this reader {mine} trimming. Trimming removes the motionless " - f"head/tail, so the two describe different distributions. {rerun}" - + (f" --trim_csv {self._trim_csv}" if want_trim else "") - + "." - ) - - want_keep = self._train_min_window_len() - try: - recorded_keep = int(pop.get("min_keep", -1)) - except (TypeError, ValueError) as exc: - raise DataContractError( - f"InternData-A1 bucket {self._dataset_id}: {stats_path} has invalid " - f"population.min_keep={pop.get('min_keep')!r}. {rerun}." - ) from exc - if want_trim and recorded_keep != int(want_keep): - raise DataContractError( - f"InternData-A1 bucket {self._dataset_id}: {stats_path} used --min_keep=" - f"{pop.get('min_keep')} but this reader keeps episodes down to {want_keep} " - f"frames. The same trim CSV under a different bound keeps a different set of " - f"episodes. {rerun} --trim_csv {self._trim_csv} --min_keep {want_keep}." - ) - - buckets = pop.get("buckets") - empty_buckets = pop.get("empty_buckets", []) - if ( - not isinstance(buckets, list) - or not all(isinstance(key, str) and key for key in buckets) - or not isinstance(empty_buckets, list) - or not all(isinstance(key, str) and key for key in empty_buckets) - ): - raise DataContractError( - f"InternData-A1 bucket {self._dataset_id}: {stats_path} has no bucket list in " - f"its 'population' block, or carries malformed bucket identifiers. {rerun}." - ) - - known = list(buckets) + list(empty_buckets) - - try: - population_key = self._match_bucket_key(dict.fromkeys(known), "stats population") - except (AmbiguousBucketKey, TypeError, ValueError) as exc: - raise DataContractError( - f"InternData-A1 bucket {self._dataset_id}: cannot resolve this bucket in the " - f"stats population from {stats_path}: {exc}. {rerun}." - ) from exc - if population_key is None: - raise DataContractError( - f"InternData-A1 bucket {self._dataset_id}: {stats_path} covers {len(known)} " - f"buckets, none of them this one — the scan refused it (unreadable metadata, " - f"or shards that disagree with the manifest) while it still loads the shared " - f"per-embodiment file, so its actions would be scaled by other buckets' " - f"numbers. {rerun}." - ) - - if want_trim: - if pop.get("schema_version") != 2: - raise DataContractError( - f"InternData-A1 bucket {self._dataset_id}: {stats_path} predates fail-closed " - "trim/exclusion population provenance. Regenerate the normalization stats." - ) - expected_trim = self._get_trim_snapshot().provenance - if pop.get("trim_provenance") != expected_trim: - raise DataContractError( - f"InternData-A1 bucket {self._dataset_id}: {stats_path} trim provenance does " - f"not match {self._trim_csv}; expected {expected_trim}, got " - f"{pop.get('trim_provenance')}. Regenerate normalization stats." - ) - - bucket_provenance = pop.get("bucket_provenance") - if not isinstance(bucket_provenance, dict) or not all( - isinstance(key, str) and key for key in bucket_provenance - ): - raise DataContractError( - f"InternData-A1 bucket {self._dataset_id}: {stats_path} has no per-bucket " - "population provenance; regenerate normalization stats." - ) - try: - provenance_key = self._match_bucket_key(bucket_provenance, "stats bucket provenance") - except (AmbiguousBucketKey, TypeError, ValueError) as exc: - raise DataContractError( - f"InternData-A1 bucket {self._dataset_id}: cannot resolve this bucket in " - f"{stats_path} population provenance: {exc}. Regenerate normalization stats." - ) from exc - if provenance_key is None or not isinstance(bucket_provenance.get(provenance_key), dict): - raise DataContractError( - f"InternData-A1 bucket {self._dataset_id}: {stats_path} has no population " - "certificate for this bucket; regenerate normalization stats." - ) - recorded = bucket_provenance[provenance_key] - current_exclusions = list(self._a1_excluded_episode_indices) - if recorded.get("excluded_episode_indices") != current_exclusions: - raise DataContractError( - f"InternData-A1 bucket {self._dataset_id}: excluded_episodes.json changed after " - f"{stats_path} was generated; stats recorded " - f"{recorded.get('excluded_episode_indices')}, current={current_exclusions}." - ) - - if self._split == "train": - current_effective = getattr(self, "_a1_effective_population", None) - if recorded.get("effective_population") != current_effective: - raise DataContractError( - f"InternData-A1 bucket {self._dataset_id}: effective split/exclusion/trim " - f"population changed after {stats_path} was generated; regenerate stats." - ) - def _load_stats(self, info: dict): - """Load per-embodiment 20-D EEF stats (``meta/stats_.json``).""" + """Read the existing per-embodiment EEF statistics from the dataset tree.""" if not self._normalize_mode or self._normalize_mode in ("none", "null"): return None - stats_path = self._a1_stats_root / "meta" / f"stats_{self._embodiment}.json" - if not stats_path.exists(): - raise FileNotFoundError( - f"normalize_mode={self._normalize_mode!r} but stats file is missing: {stats_path}. " - "Run python -m openwam.dataloader.utils.stats_computation." - "interndata_a1_stats_computation --dataset_dir " - f"--stats_root {self._a1_stats_root}, or set normalize_mode=null." - ) + filename = f"stats_{self._embodiment}.json" + stats_path = find_dataset_stats(self._dataset_dir, filename, family_root=self._stats_root) with open(stats_path) as f: raw = json.load(f) - - self._check_stats_population(raw, stats_path) - - eef_raw = raw.get("eef", {}) - - for k in ("mean", "std", "min", "max", "q01", "q99"): - if k not in eef_raw: + if not isinstance(raw, dict) or not isinstance(raw.get("eef"), dict): + raise ValueError(f"InternData-A1: {stats_path} must contain an 'eef' statistics object") + if raw.get("embodiment", self._embodiment) != self._embodiment: + raise ValueError(f"InternData-A1: {stats_path} has statistics for a different embodiment") + eef_raw = raw["eef"] + for key in ("mean", "std", "min", "max", "q01", "q99"): + if key not in eef_raw: continue - width = len(eef_raw[k]) - if width != _ACTION_DIM: + values = np.asarray(eef_raw[key], dtype=np.float64) + if values.shape != (_ACTION_DIM,): raise ValueError( - f"InternData-A1 bucket {self._dataset_id}: 'eef' stats {k!r} width " - f"{width} in {stats_path} != expected {_ACTION_DIM}. " - "Re-run interndata_a1_stats_computation." + f"InternData-A1 bucket {self._dataset_id}: 'eef' stats {key!r} shape " + f"{values.shape} in {stats_path} != expected {_ACTION_DIM} dimensions" ) - stats = materialize_eef_stats( + if not np.isfinite(values).all(): + raise ValueError(f"InternData-A1: {stats_path}: eef.{key} contains non-finite values") + return materialize_eef_stats( eef_raw, self._normalize_mode, dim=_ACTION_DIM, strict_minmax=False, - source_hint=( - f"{stats_path}: eef.* — re-run python -m openwam.dataloader.utils." - "stats_computation.interndata_a1_stats_computation" - ), + source_hint=f"{stats_path}: eef.*", force_rot6d_identity=True, ) - return stats def _post_init(self, info: dict) -> None: """Sanity-check the quaternion magnitude on a small sample of rows. @@ -1495,38 +599,19 @@ def _multibucket_wrapper(cls): @classmethod def from_config(cls, config, split: str = "train") -> Any: - """Build one bucket, or recursively discover every bucket under a root. - - Overrides the base because A1 buckets sit at variable depth (2 or 3 - levels below the embodiment dir), which the base's single-level - ``iterdir`` scan cannot see. Everything else — the shared window kwargs, - the ``total_hours`` water-fill, per-bucket failure tolerance — is - delegated to the same :func:`build_multibucket` the other readers use. - """ - + """Build a cleaned bucket or all recursively discovered family buckets.""" from openwam.dataloader.utils import get_cfg as _get dataset_dir = _get(config, "dataset_dir") if dataset_dir is None: raise ValueError(f"{cls.__name__}: missing dataset_dir") root = Path(dataset_dir) - common: Dict[str, Any] = {"split": split} for key in cls.CONFIG_KEYS: - v = _get(config, key, _CONFIG_SENTINEL) - if v is _CONFIG_SENTINEL: + value = _get(config, key, _CONFIG_SENTINEL) + if value is _CONFIG_SENTINEL or (value is None and key != "normalize_mode"): continue - - if v is None and key != "normalize_mode": - continue - common[key] = v - - stats_root = _get(config, "stats_root") or str(root) - common["a1_stats_root"] = stats_root - trim_snapshot = None - if common.get("trim_csv"): - trim_snapshot = _load_trim_snapshot(common["trim_csv"]) - common["_trim_snapshot"] = trim_snapshot + common[key] = value if (root / "meta" / "info.json").is_file(): kwargs = dict(common) @@ -1537,32 +622,17 @@ def from_config(cls, config, split: str = "train") -> Any: if total_hours is not None: kwargs["max_hours"] = float(total_hours) kwargs["subsample_seed"] = int(_get(config, "seed", 42)) - dataset = cls(dataset_dir=str(root), **kwargs) - if trim_snapshot is not None: - _assert_trim_snapshot_current(trim_snapshot, context="from_config reader construction") - return dataset + return cls(dataset_dir=str(root), **kwargs) if not root.is_dir(): raise FileNotFoundError(f"{cls.__name__}: {root} does not exist") sub_dirs = discover_a1_buckets(root) if not sub_dirs: - raise FileNotFoundError( - f"{cls.__name__}: no buckets with meta/info.json found anywhere under {root}. " - "Did the tar.gz archives get extracted?" - ) + raise FileNotFoundError(f"{cls.__name__}: no buckets with meta/info.json found under {root}") logger.info("%s.from_config: root mode, %d buckets under %s", cls.__name__, len(sub_dirs), root) - normalize_mode = common.get("normalize_mode", cls.DEFAULT_NORMALIZE_MODE) - if normalize_mode and normalize_mode not in ("none", "null"): - _validate_a1_root_stats_contributors(root, Path(stats_root), sub_dirs) - - def _per_bucket(sub: Path) -> Dict[str, Any]: - try: - return {"dataset_id": str(sub.relative_to(root))} - except ValueError: - return {"dataset_id": sub.name} - - dataset = build_multibucket( + common["stats_root"] = str(root) + return build_multibucket( cls, sub_dirs, common, @@ -1570,11 +640,9 @@ def _per_bucket(sub: Path) -> Dict[str, Any]: total_hours=_get(config, "total_hours"), wrapper_cls=cls._multibucket_wrapper(), source_name=cls.__name__, - per_bucket_kwargs=_per_bucket, + strict=cls.STRICT_BUCKET_LOADING, + per_bucket_kwargs=lambda sub: {"dataset_id": sub.relative_to(root).as_posix()}, ) - if trim_snapshot is not None: - _assert_trim_snapshot_current(trim_snapshot, context="from_config reader construction") - return dataset class MultiInternDataA1Dataset(MultiLeRobotV3Reader): diff --git a/openwam/dataloader/mixture.py b/openwam/dataloader/mixture.py index 1e706969f..110abc9fb 100644 --- a/openwam/dataloader/mixture.py +++ b/openwam/dataloader/mixture.py @@ -98,9 +98,14 @@ def __init__( weights = [float(len(d)) for d in self._datasets] if len(weights) != len(self._datasets): raise ValueError(f"weights length ({len(weights)}) != datasets length ({len(self._datasets)})") + weights = [float(w) for w in weights] + if not all(np.isfinite(w) for w in weights): + raise ValueError(f"MixtureDataset: weights must be finite: {weights}") if any(w < 0 for w in weights): raise ValueError(f"MixtureDataset: negative weights are not allowed: {weights}") total_w = sum(weights) + if not np.isfinite(total_w): + raise ValueError("MixtureDataset: the sum of weights must be finite") if total_w <= 0: raise ValueError("MixtureDataset: all sub-sources are empty or zero-weight; nothing to sample.") self._weights = [w / total_w for w in weights] @@ -128,7 +133,7 @@ def set_epoch(self, epoch: int) -> None: =================================================================== ⚠️ CALLER REQUIREMENT — READ BEFORE USE ⚠️ - This method mutates ``self._index_map`` in place. Under + This method replaces ``self._index_map``. Under ``torch.utils.data.DataLoader`` with ``num_workers > 0``, each worker holds a forked copy of the dataset object; calling ``set_epoch`` on the parent only updates the parent's copy. To make @@ -171,27 +176,41 @@ def _build_index_map(self): """ total_real = sum(len(d) for d in self._datasets) - parts: List[np.ndarray] = [] + entries = [] + self._virtual_counts = [0] * len(self._datasets) for di, (ds, w) in enumerate(zip(self._datasets, self._weights)): n_real = len(ds) if n_real == 0 or w <= 0.0: logger.warning( - "MixtureDataset: skipping source '%s' (len=%d, weight=%.4f) — not sampled.", + "MixtureDataset: skipping source '%s' (len=%d, weight=%.4f); not sampled.", self._names[di], n_real, w, ) continue + if n_real - 1 > np.iinfo(np.int32).max: + raise ValueError(f"MixtureDataset: source '{self._names[di]}' exceeds the int32 sample-index range") virtual_n = max(1, round(total_real * w)) - sample_idx = (np.arange(virtual_n, dtype=np.int64) % n_real).astype(np.int32) - di_arr = np.full(virtual_n, di, dtype=np.int32) - parts.append(np.stack([di_arr, sample_idx], axis=1)) - if not parts: + self._virtual_counts[di] = virtual_n + entries.append((di, n_real, virtual_n)) + if not entries: raise RuntimeError("MixtureDataset: all sub-sources are empty or zero-weight; nothing to sample.") - combined = np.concatenate(parts, axis=0) - rng = np.random.RandomState(self._seed) - perm = rng.permutation(len(combined)) - self._index_map = combined[perm] + + # Four-source pretraining has hundreds of millions of windows. Fill one + # map with bounded temporary storage, then shuffle its rows in place; + # this keeps the legacy seeded order without a second map/permutation. + combined = np.empty((sum(self._virtual_counts), 2), dtype=np.int32) + offset = 0 + block_size = 1_000_000 + for di, n_real, virtual_n in entries: + combined[offset : offset + virtual_n, 0] = di + for start in range(0, virtual_n, block_size): + end = min(start + block_size, virtual_n) + combined[offset + start : offset + end, 1] = np.arange(start, end, dtype=np.int64) % n_real + offset += virtual_n + # One eight-byte row per shuffle item uses NumPy's fast 1-D path. + np.random.RandomState(self._seed).shuffle(combined.view(np.int64).reshape(-1)) + self._index_map = combined def __getitem__(self, idx: int) -> dict: di, si = self._index_map[idx] @@ -249,9 +268,7 @@ def get_dataset(self, name: str) -> BaseDataset: def dataset_sample_counts(self) -> Dict[str, int]: """Per-source virtual-sample counts, keyed by source name.""" - di_col = self._index_map[:, 0] - counts = np.bincount(di_col, minlength=len(self._datasets)) - return {self._names[i]: int(c) for i, c in enumerate(counts)} + return dict(zip(self._names, self._virtual_counts)) @classmethod def from_config(cls, config, split: str = "train") -> "MixtureDataset": @@ -259,7 +276,7 @@ def from_config(cls, config, split: str = "train") -> "MixtureDataset": Sub-sources live under the ``datasets:`` key. Two layouts: - **dict** (recommended, used by ``configs/dataloader/pretrain_data/mixture.yaml``). + **dict** (recommended, used by ``configs/dataloader/mixture.yaml``). Sub-source config is composed in via Hydra ``defaults`` (``@datasets.``); the per-source block under ``datasets:`` only carries ``enabled`` / ``weight`` overrides. The dict key is the @@ -346,6 +363,11 @@ def _normalize_entries(datasets_cfg): return entries weight_strategy = _get(config, "weight_strategy", "proportional") + valid_strategies = ("manual", "uniform", "inverse_size", "proportional") + if weight_strategy not in valid_strategies: + raise ValueError( + f"MixtureDataset: unknown weight_strategy={weight_strategy!r}. Valid options: {valid_strategies}." + ) split_manifest = _get(config, "split_manifest") mixture_seed = int(_get(config, "seed", 42)) @@ -380,12 +402,6 @@ def _normalize_entries(datasets_cfg): with ThreadPoolExecutor(max_workers=n_workers) as pool: sub_datasets = list(pool.map(lambda c: build_dataset(c, split=split), enabled_cfgs)) - valid_strategies = ("manual", "uniform", "inverse_size", "proportional") - if weight_strategy not in valid_strategies: - raise ValueError( - f"MixtureDataset: unknown weight_strategy={weight_strategy!r}. Valid options: {valid_strategies}." - ) - if weight_strategy == "uniform": weights = [1.0] * len(sub_datasets) logger.info("MixtureDataset: weight_strategy=uniform, all weights set to 1.0") diff --git a/openwam/dataloader/muka_franka.py b/openwam/dataloader/muka_franka.py deleted file mode 100644 index d3ba3a1c7..000000000 --- a/openwam/dataloader/muka_franka.py +++ /dev/null @@ -1,346 +0,0 @@ -"""MUKA Franka single-arm LeRobot v3 dataloader. - -The downloaded dataset stores row-aligned 7-D values as:: - - observation.state = achieved [xyz3, Euler-XYZ3, gripper_closedness1] - action = next achieved [xyz3, Euler-XYZ3, gripper_closedness1] - -Euler angles are extrinsic XYZ radians (``R = Rz @ Ry @ Rx``). Raw gripper -closedness is ``0 = open, 1 = closed``. This reader converts both streams to -the OpenWAM single-arm EEF10 contract:: - - [xyz3, rot6d6, gripper_open_scale1], -1 = closed, +1 = open - -The conversion is ``open_scale = 1 - 2 * closedness``. With -``unify_action: true``, EEF10 is normalized first and then scattered into the -shared 80-D left-arm slots 0..9. - -The source action at row ``t`` equals the state at row ``t+1``. The final -episode row repeats the terminal state, so it is excluded from action -supervision and from action normalization statistics. -""" - -from __future__ import annotations - -import json -import logging -import os -import time -from pathlib import Path -from typing import Any, ClassVar, Optional, Sequence, Tuple - -import numpy as np - -from openwam.dataloader.bases import LeRobotV3Reader -from openwam.dataloader.utils.eef import euler_xyz_to_rot6d -from openwam.dataloader.utils.normalization import ( - ROT6D_DIMS_ARM10, - apply_normalization, - load_stats_file, - load_stats_metadata, -) - -logger = logging.getLogger(__name__) - -_ACTION_MODE = "eef" -RAW_EEF7_DIM = 7 -EEF10_DIM = 10 -RAW_GRIPPER_CONVENTION = "zero_open_one_closed" -GRIPPER_CONVENTION = "minus1_closed_plus1_open" -GRIPPER_TRANSFORM = "1_minus_2_times_closedness" -ROTATION_CONVENTION = "extrinsic_xyz_radians_r_equals_rz_ry_rx" -NORMALIZATION_STATS_FILENAME = "normalization_stats.npy" -STATS_POPULATION = "all_episodes" - -_EXPECTED_FEATURE_NAMES = ( - "x", - "y", - "z", - "roll", - "pitch", - "yaw", - "gripper_closedness", -) -_EXPECTED_SOURCE_SCHEMA = { - "position_frame": "robot_base", - "rotation_storage": "Euler XYZ", - "rotation_convention": "extrinsic; R = Rz @ Ry @ Rx", - "rotation_unit": "radians", - "action_type": "absolute achieved EEF state", - "action_alignment": "action[t] = observation.state[t+1]; terminal row repeated and masked by reader", - "gripper": "continuous closedness; 0=open, 1=closed", -} - - -def _as_priority(value: Optional[Sequence[str]], default: Tuple[str, ...]) -> Tuple[str, ...]: - if value is None: - return default - if isinstance(value, str): - return (value,) - return tuple(str(item) for item in value) - - -def _pick_feature(features: dict, priorities: Sequence[str]) -> Optional[str]: - return next((key for key in priorities if key in features), None) - - -def closedness_to_open_scale(closedness: np.ndarray) -> np.ndarray: - """Map source gripper closedness to OpenWAM's ``[-1 closed, +1 open]``.""" - values = np.asarray(closedness, dtype=np.float32) - if not np.isfinite(values).all(): - raise ValueError("MUKA Franka gripper closedness contains NaN or infinity") - if np.any(values < -1e-6) or np.any(values > 1.0 + 1e-6): - observed_min = float(values.min()) - observed_max = float(values.max()) - raise ValueError( - f"MUKA Franka gripper closedness must be in [0, 1], got min={observed_min:.8g}, max={observed_max:.8g}" - ) - return (1.0 - 2.0 * values).astype(np.float32) - - -def euler7_to_eef10(values: np.ndarray) -> np.ndarray: - """Convert ``[xyz, extrinsic XYZ Euler, closedness]`` to canonical EEF10.""" - raw = np.asarray(values, dtype=np.float32) - if raw.ndim != 2 or raw.shape[1] != RAW_EEF7_DIM: - raise ValueError(f"MUKA Franka EEF input must be (T, {RAW_EEF7_DIM}), got {raw.shape}") - if not np.isfinite(raw).all(): - raise ValueError("MUKA Franka EEF input contains NaN or infinity") - return np.concatenate( - ( - raw[:, :3], - euler_xyz_to_rot6d(raw[:, 3:6]), - closedness_to_open_scale(raw[:, 6:7]), - ), - axis=-1, - ).astype(np.float32) - - -class MukaFrankaDataset(LeRobotV3Reader): - """Single-bucket reader for ``ewykric/muka_franka_lerobot_v3``.""" - - DATASET_NAME = "MUKAFranka" - ACTION_DIM = EEF10_DIM - NEEDED_COLS = ("action", "observation.state", "task_index") - PROMPT_FILE_REQUIRED = True - DEPLOY_ACTION_MODE = _ACTION_MODE - - HEAD_CAMERA_PRIORITY: ClassVar[Tuple[str, ...]] = ("observation.images.left",) - WRIST_CAMERA_PRIORITY: ClassVar[Tuple[str, ...]] = ("observation.images.left_wrist",) - CONFIG_KEYS: ClassVar[Tuple[str, ...]] = LeRobotV3Reader.CONFIG_KEYS + ( - "action_mode", - "gripper_convention", - "head_camera_priority", - "wrist_camera_priority", - "normalization_stats_path", - ) - - def __init__( - self, - dataset_dir: str, - *, - action_mode: str = _ACTION_MODE, - gripper_convention: str = GRIPPER_CONVENTION, - head_camera_priority: Optional[Sequence[str]] = None, - wrist_camera_priority: Optional[Sequence[str]] = None, - normalization_stats_path: Optional[str] = None, - unify_action: bool = False, - unify_action_map: Optional[Any] = None, - **kwargs: Any, - ): - mode = str(action_mode).strip().lower() - if mode != _ACTION_MODE: - raise ValueError(f"MUKA Franka supports only action_mode='eef', got {action_mode!r}") - convention = str(gripper_convention).strip() - if convention != GRIPPER_CONVENTION: - raise ValueError( - f"MUKA Franka requires gripper_convention={GRIPPER_CONVENTION!r}, got {gripper_convention!r}" - ) - unify_on = bool(unify_action) - if unify_on and unify_action_map is None: - raise ValueError( - "MUKA Franka unify_action=true requires an explicit unify_action_map; " - 'set ["0-9"] for the canonical single-arm left-slot mapping' - ) - self.action_mode = mode - self._head_priority = _as_priority(head_camera_priority, self.HEAD_CAMERA_PRIORITY) - self._wrist_priority = _as_priority(wrist_camera_priority, self.WRIST_CAMERA_PRIORITY) - self._source_stats_path = str(normalization_stats_path) if normalization_stats_path else None - self._resolved_stats_path: Optional[str] = None - super().__init__( - dataset_dir=dataset_dir, - unify_action=unify_on, - unify_action_map=unify_action_map, - **kwargs, - ) - - def _resolve_cameras(self, info: dict): - features = info.get("features", {}) or {} - if self._target_camera is not None: - return self._target_camera, None, None - return ( - _pick_feature(features, self._head_priority), - _pick_feature(features, self._wrist_priority), - None, - ) - - def _n_supervised_action_steps(self, actual_raw_len: int) -> int: - # action[t] targets state[t+1]; a window ending at the episode boundary - # includes one repeated terminal target that must stay out of the loss. - return max(0, actual_raw_len - 1) - - def _post_init(self, info: dict) -> None: - features = info.get("features", {}) or {} - expected_size = (384, 320) if self._multiview else (256, 320) - if (self._height, self._width) != expected_size: - mode = "multiview" if self._multiview else "single-view" - raise ValueError( - f"MUKA Franka {mode} requires height={expected_size[0]}, width={expected_size[1]}, " - f"got height={self._height}, width={self._width}" - ) - if info.get("robot_type") != "muka_franka_single_arm": - raise ValueError( - "MUKA Franka info.json must declare robot_type='muka_franka_single_arm', " - f"got {info.get('robot_type')!r}" - ) - for column in ("action", "observation.state"): - if column not in features: - raise KeyError(f"MUKA Franka requires {column!r}") - feature = features[column] - shape = tuple(feature.get("shape", ())) - if shape != (RAW_EEF7_DIM,): - raise ValueError(f"MUKA Franka {column} feature must have shape [{RAW_EEF7_DIM}], got {shape}") - names = tuple(feature.get("names") or ()) - if names != _EXPECTED_FEATURE_NAMES: - raise ValueError(f"MUKA Franka {column} names must be {_EXPECTED_FEATURE_NAMES}, got {names}") - self._validate_source_schema() - - def _validate_source_schema(self) -> None: - path = self._dataset_dir / "meta" / "muka_schema.json" - if not path.is_file(): - raise FileNotFoundError( - f"MUKA Franka source contract is missing: {path}; cannot safely infer rotation/gripper semantics" - ) - with path.open(encoding="utf-8") as handle: - schema = json.load(handle) - mismatches = { - key: (schema.get(key), expected) - for key, expected in _EXPECTED_SOURCE_SCHEMA.items() - if schema.get(key) != expected - } - if mismatches: - raise ValueError(f"MUKA Franka source contract mismatch in {path}: {mismatches}") - - @staticmethod - def _stats_contract_matches(path: Path) -> bool: - if not path.is_file(): - return False - try: - raw = np.load(path, allow_pickle=True).item() - block = raw.get(_ACTION_MODE, raw) if isinstance(raw, dict) else {} - return ( - block.get("gripper_convention") == GRIPPER_CONVENTION - and block.get("raw_gripper_convention") == RAW_GRIPPER_CONVENTION - and block.get("gripper_transform") == GRIPPER_TRANSFORM - and block.get("rotation_convention") == ROTATION_CONVENTION - and block.get("stats_population") == STATS_POPULATION - ) - except (OSError, ValueError, EOFError, AttributeError): - return False - - def _load_stats(self, info: dict): - if not self._normalize_mode or self._normalize_mode in ("none", "null"): - return None - if self._source_stats_path: - stats_path = Path(self._source_stats_path) - if not self._stats_contract_matches(stats_path): - raise ValueError(f"MUKA Franka stats {stats_path} do not match the required rot6d/gripper contract") - else: - stats_path = self._dataset_dir / "meta" / NORMALIZATION_STATS_FILENAME - if not self._stats_contract_matches(stats_path): - self._build_default_stats(stats_path) - self._resolved_stats_path = str(stats_path) - global_stats = load_stats_file( - stats_path, - action_mode=self.action_mode, - normalize_mode=str(self._normalize_mode), - dim=self._raw_action_dim, - ) - self._check_stats_contract(stats_path) - self.normalization_stats_path = str(stats_path) - return global_stats - - def _check_stats_contract(self, stats_path: Path) -> None: - metadata = load_stats_metadata(stats_path, action_mode=self.action_mode) - expected = { - "gripper_convention": GRIPPER_CONVENTION, - "raw_gripper_convention": RAW_GRIPPER_CONVENTION, - "gripper_transform": GRIPPER_TRANSFORM, - "rotation_convention": ROTATION_CONVENTION, - "stats_population": STATS_POPULATION, - } - mismatches = {key: (metadata.get(key), value) for key, value in expected.items() if metadata.get(key) != value} - if mismatches: - raise ValueError(f"MUKA Franka stats contract mismatch in {stats_path}: {mismatches}") - - def _build_default_stats(self, path: Path) -> None: - from openwam.dataloader.utils.stats_computation.muka_franka_stats_computation import ( - build_and_save_muka_franka_stats, - ) - - try: - import torch.distributed as dist - - dist_ready = dist.is_available() and dist.is_initialized() - except Exception: - dist_ready = False - rank = dist.get_rank() if dist_ready else int(os.environ.get("RANK", os.environ.get("LOCAL_RANK", 0))) - if rank == 0: - logger.info( - "MUKAFranka(%s): building OpenWAM normalization stats at %s", - self._dataset_id, - path, - ) - build_and_save_muka_franka_stats(self, path) - return - - deadline = time.monotonic() + float(os.environ.get("OPENWAM_STATS_WAIT_TIMEOUT_S", 12 * 60 * 60)) - poll_interval = float(os.environ.get("OPENWAM_STATS_POLL_INTERVAL_S", 10)) - while not self._stats_contract_matches(path): - if time.monotonic() >= deadline: - raise TimeoutError(f"timed out waiting for rank 0 to build MUKA Franka stats: {path}") - time.sleep(poll_interval) - - @staticmethod - def _read_euler7_column(win, column: str) -> np.ndarray: - values = np.stack(win[column].values).astype(np.float32) - return euler7_to_eef10(values) - - def _raw_action_eef10(self, win) -> np.ndarray: - return self._read_euler7_column(win, "action") - - def _raw_state_eef10(self, win) -> np.ndarray: - return self._read_euler7_column(win, "observation.state") - - def _action_20d(self, win) -> np.ndarray: - return apply_normalization(self._raw_action_eef10(win), self._normalization_stats, self._normalize_mode) - - def _proprio_20d(self, win) -> np.ndarray: - raw = self._raw_state_eef10(win)[0:1] - return apply_normalization(raw, self._normalization_stats, self._normalize_mode) - - -ROT6D_DIMS_EEF10 = ROT6D_DIMS_ARM10 - -__all__ = [ - "EEF10_DIM", - "GRIPPER_CONVENTION", - "GRIPPER_TRANSFORM", - "MukaFrankaDataset", - "RAW_EEF7_DIM", - "RAW_GRIPPER_CONVENTION", - "ROT6D_DIMS_EEF10", - "ROTATION_CONVENTION", - "STATS_POPULATION", - "closedness_to_open_scale", - "euler7_to_eef10", -] diff --git a/openwam/dataloader/oxe_droid.py b/openwam/dataloader/oxe_droid.py index c58c7a81c..5bd02eb64 100644 --- a/openwam/dataloader/oxe_droid.py +++ b/openwam/dataloader/oxe_droid.py @@ -1,64 +1,26 @@ -"""LeRobot v3 reader for the DROID portion of OXE. +"""Reader for the physically cleaned DROID LeRobot v3 export. -The reader uses the rigid arm-side wrist/gripper-mount frame: achieved state is -read from ``observation_gripper_pose6d`` and commanded action from -``action_wrist_pose``. The physically distinct task-TCP and clipped-delta -streams are intentionally not mixed into the same action slots. Euler XYZ -poses are converted to a 10-D single-arm EEF vector and placed in the left arm -of the shared schema. Raw gripper closedness is inverted to the project-wide -``0=closed, 1=open`` convention. - -Prompt text resolves from ``tasks.parquet`` and a deterministic per-row fallback -chain. Episodes without any resolvable prompt must be represented by the -canonical exclusion artifact. The reader validates prompt inputs, manifest -ranges, exclusions, and normalization provenance so stale artifacts fail -before training. +Prompts are resolved through the export's ``meta/tasks.parquet``. The rigid +arm-side wrist/gripper-mount poses retain the same base frame and Euler XYZ +conversion as the source loader; raw gripper closedness is inverted to the +shared ``0=closed, 1=open`` convention. Statistics are read from the dataset's +``meta/eef_stats.json`` and are never computed during loading. """ from __future__ import annotations -import hashlib import json -import re -from collections.abc import Mapping -from pathlib import Path -from typing import ClassVar, Optional, Tuple +from typing import Optional import numpy as np import pandas as pd -import pyarrow as pa -import pyarrow.compute as pc from openwam.dataloader.bases import LeRobotV3Reader from openwam.dataloader.utils.eef import EEF_POSE_FRAME_CONTRACT, LEFT_ARM_DIM_MASK, single_arm_20d -from openwam.dataloader.utils.lerobotv3 import ( - LeRobotV3DataPopulation, - LeRobotV3DataShard, - apply_info_splits, - digest_lerobot_v3_data_population, - read_lerobot_v3_population_shard, - resolve_lerobot_v3_data_population, -) +from openwam.dataloader.utils.lerobotv3 import DataContractError from openwam.dataloader.utils.normalization import materialize_eef_stats from openwam.dataloader.utils.oxe_schema import droid_euler7_to_arm10, droid_pose6_closedness_to_arm10 -# Present-but-empty instructions fall through to the next prompt source. -_PLACEHOLDER_RE = re.compile( - r"^(?:no[\s_-]*action\.?|not[\s_-]*action|no[\s_-]*instruction|n/?a|null|none|nothing|test" - r"|[.\-_/]+|pree|pm|op)$" -) - - -# Versioned contracts keep prompt exclusions and EEF stats tied to source data. -DROID_PROMPT_EXCLUSION_SCHEMA_VERSION = 3 -DROID_PROMPT_EXCLUSION_KEY = "droid_prompt_exclusions" -DROID_PROMPT_INDEPENDENT_EXCLUSIONS_KEY = "independently_owned_episode_indices" -DROID_PROMPT_INPUTS_DIGEST_KEY = "prompt_inputs_digest" -DROID_PROMPT_INPUTS_DIGEST_FORMAT_VERSION = 2 -DROID_DATA_POPULATION_DIGEST_KEY = "data_population_digest" -DROID_STATS_POPULATION_KEY = "stats_population" -DROID_STATS_POPULATION_SCHEMA_VERSION = 1 -DROID_STATS_POPULATION_POLICY = "info_split_then_prompt_exclusion" DROID_EEF_STATS_CONTRACT_KEY = "droid_eef_stats_contract" DROID_EEF_STATS_CONTRACT_VERSION = 2 DROID_EEF_STATS_CONTRACT = { @@ -71,454 +33,48 @@ "output_gripper_semantics": "openness:0=closed,1=open", "gripper_transform": "1-raw", } -DROID_PROMPT_FALLBACK_COLS = ( - "other_information.language_instruction_2", - "other_information.language_instruction_3", - "annotation.substask", - "annotation.instruction_add", -) -DROID_PROMPT_SOURCE_COLUMNS = ("episode_index", "task_index", *DROID_PROMPT_FALLBACK_COLS) - - -def _validate_droid_eef_stats_contract(value) -> None: - """Reject stats from a different pose frame or gripper conversion.""" - if not isinstance(value, dict) or set(value) != set(DROID_EEF_STATS_CONTRACT): - raise ValueError(f"expected exactly {DROID_EEF_STATS_CONTRACT!r}") - if type(value["schema_version"]) is not int or value["schema_version"] != DROID_EEF_STATS_CONTRACT_VERSION: - raise ValueError(f"schema_version must be integer {DROID_EEF_STATS_CONTRACT_VERSION}") - for key, expected in DROID_EEF_STATS_CONTRACT.items(): - if key == "schema_version": - continue - if not isinstance(value[key], str) or value[key] != expected: - raise ValueError(f"{key} must be {expected!r}") - - -def _clean_text(value) -> str: - """Strip a raw prompt candidate, or return ``""`` when it carries no instruction.""" - if value is None or not isinstance(value, str): - return "" - text = value.strip() - if not text or _PLACEHOLDER_RE.match(text.lower()): - return "" - return text - - -def _parse_episode_indices(value, field: str) -> set[int]: - if not isinstance(value, list) or any(type(i) is not int or i < 0 for i in value): - raise ValueError(f"{field} must be a list of non-negative integers") - return set(value) - - -def _parse_nonnegative_int(value, field: str) -> int: - if type(value) is not int or value < 0: - raise ValueError(f"{field} must be a non-negative integer") - return value - - -def _digest_framed_bytes(hasher, value: bytes | bytearray | memoryview) -> None: - hasher.update(len(value).to_bytes(8, "little")) - hasher.update(value) - - -def _digest_framed_text(hasher, value: str) -> None: - _digest_framed_bytes(hasher, value.encode("utf-8")) - - -def _sha256_file(path: Path) -> str: - hasher = hashlib.sha256() - with path.open("rb") as stream: - for chunk in iter(lambda: stream.read(1024 * 1024), b""): - hasher.update(chunk) - return hasher.hexdigest() - - -def resolve_droid_stats_population( - population: LeRobotV3DataPopulation, - info: dict, - excluded_episode_indices, - *, - split: str = "train", -) -> tuple[pd.DataFrame, dict]: - """Resolve and certify the exact DROID population used for normalization. - - The raw data-population digest deliberately describes the pre-split corpus. - This second certificate binds the stats to the selected info.json split and - prompt-exclusion set, including each episode's physical manifest range. - """ - - selected = apply_info_splits( - population.episodes, - split, - info.get("splits", {}) or {}, - source_name="OXE-DROID normalization stats", - ) - excluded = set(int(value) for value in excluded_episode_indices) - if excluded: - selected = selected[~selected["episode_index"].isin(excluded)].reset_index(drop=True) - - required = ( - "episode_index", - "dataset_from_index", - "length", - "data/chunk_index", - "data/file_index", - "_data_row_offset", - ) - missing = [column for column in required if column not in selected.columns] - if missing: - raise ValueError(f"DROID stats population is missing manifest columns {missing}") - - hasher = hashlib.sha256(b"openwam:droid-stats-population:v1\0") - ordered = selected.sort_values("episode_index", kind="stable") - for row in ordered.loc[:, required].itertuples(index=False, name=None): - for value in row: - integer = int(value) - if integer < 0: - raise ValueError("DROID stats population manifest values must be non-negative") - hasher.update(integer.to_bytes(8, "little", signed=False)) - provenance = { - "schema_version": DROID_STATS_POPULATION_SCHEMA_VERSION, - "split": split, - "policy": DROID_STATS_POPULATION_POLICY, - "effective_population_digest": hasher.hexdigest(), - "num_episodes": int(len(ordered)), - "num_rows": int(ordered["length"].sum()), - } - return selected, provenance - - -def read_droid_prompt_population_shard( - dataset_dir: str | Path, - shard: LeRobotV3DataShard, -) -> pa.Table: - """Read and validate the prompt rows addressed by one manifest shard.""" - try: - return read_lerobot_v3_population_shard( - Path(dataset_dir), - shard, - DROID_PROMPT_SOURCE_COLUMNS, - ) - except (OSError, TypeError, ValueError) as exc: - raise ValueError(f"prompt source population is invalid ({exc})") from exc - - -def digest_droid_prompt_shard(table: pa.Table) -> str: - """Hash the logical prompt-bearing values of one data shard. - - Numeric IDs are normalized to little-endian int64. Fallback text is - normalized to one combined large-string array with an explicit validity - vector and nulls filled to empty before hashing offsets/data. This keeps the - digest independent of parquet row groups, compression, and dictionary - encoding while preserving row order and null-vs-empty distinctions. - """ - - missing = [column for column in DROID_PROMPT_SOURCE_COLUMNS if column not in table.column_names] - if missing: - raise ValueError(f"prompt source table is missing columns {missing}") - table = table.select(list(DROID_PROMPT_SOURCE_COLUMNS)) - hasher = hashlib.sha256(b"openwam:droid-prompt-shard:v1\0") - hasher.update(table.num_rows.to_bytes(8, "little")) - for column_name in DROID_PROMPT_SOURCE_COLUMNS: - _digest_framed_text(hasher, column_name) - column = table.column(column_name) - if column_name in ("episode_index", "task_index"): - array = column.combine_chunks() - if array.null_count: - raise ValueError(f"{column_name} contains null values") - try: - array = pc.cast(array, pa.int64(), safe=True) - except (pa.ArrowInvalid, pa.ArrowNotImplementedError) as exc: - raise ValueError(f"{column_name} must contain integer values") from exc - values = np.asarray(array.to_numpy(zero_copy_only=False), dtype=" dict: - """Hash tasks and the exact prompt population addressed by the reader.""" - root = Path(dataset_dir) - resolved_population = population or resolve_lerobot_v3_data_population(root) - tasks_digest = tasks_sha256 if tasks_sha256 is not None else _sha256_file(root / "meta" / "tasks.parquet") - relative_paths = [shard.relative_path for shard in resolved_population.shards] - if shard_digests is None: - resolved_shard_digests = { - shard.relative_path: digest_droid_prompt_shard(read_droid_prompt_population_shard(root, shard)) - for shard in resolved_population.shards - } - else: - resolved_shard_digests = dict(shard_digests) - if set(resolved_shard_digests) != set(relative_paths): - raise ValueError("prompt shard digest paths do not match the manifest-addressed shard set") - - hasher = hashlib.sha256(b"openwam:droid-prompt-inputs:v2\0") - _digest_framed_text(hasher, "meta/tasks.parquet") - try: - tasks_digest_bytes = bytes.fromhex(tasks_digest) - except ValueError as exc: - raise ValueError("tasks_sha256 must be a 64-character hexadecimal SHA-256 digest") from exc - if len(tasks_digest_bytes) != hashlib.sha256().digest_size: - raise ValueError("tasks_sha256 must be a 64-character hexadecimal SHA-256 digest") - _digest_framed_bytes(hasher, tasks_digest_bytes) - population_digest = digest_lerobot_v3_data_population(resolved_population) - _digest_framed_text(hasher, DROID_DATA_POPULATION_DIGEST_KEY) - _digest_framed_bytes(hasher, bytes.fromhex(population_digest)) - for relative_path in relative_paths: - _digest_framed_text(hasher, relative_path) - digest = resolved_shard_digests[relative_path] - try: - digest_bytes = bytes.fromhex(digest) - except ValueError as exc: - raise ValueError(f"invalid SHA-256 digest for {relative_path}") from exc - if len(digest_bytes) != hashlib.sha256().digest_size: - raise ValueError(f"invalid SHA-256 digest for {relative_path}") - _digest_framed_bytes(hasher, digest_bytes) - return { - "algorithm": "sha256", - "format_version": DROID_PROMPT_INPUTS_DIGEST_FORMAT_VERSION, - "value": hasher.hexdigest(), - } - - -def _parse_prompt_inputs_digest(value) -> dict: - if not isinstance(value, dict): - raise ValueError(f"{DROID_PROMPT_INPUTS_DIGEST_KEY} must be an object") - if value.get("algorithm") != "sha256": - raise ValueError(f"{DROID_PROMPT_INPUTS_DIGEST_KEY}.algorithm must be 'sha256'") - format_version = value.get("format_version") - if type(format_version) is not int or format_version != DROID_PROMPT_INPUTS_DIGEST_FORMAT_VERSION: - raise ValueError( - f"{DROID_PROMPT_INPUTS_DIGEST_KEY}.format_version must be {DROID_PROMPT_INPUTS_DIGEST_FORMAT_VERSION}" - ) - digest = value.get("value") - if not isinstance(digest, str) or re.fullmatch(r"[0-9a-f]{64}", digest) is None: - raise ValueError(f"{DROID_PROMPT_INPUTS_DIGEST_KEY}.value must be a lowercase SHA-256 digest") - return value - - -def load_droid_prompt_exclusions( - dataset_dir: str | Path, - *, - population: LeRobotV3DataPopulation | None = None, -) -> tuple[dict, set[int]]: - """Load and validate the prompt-generator provenance required by this reader.""" - root = Path(dataset_dir) - path = root / "meta" / "excluded_episodes.json" - if not path.exists(): - raise FileNotFoundError( - f"{path} is missing. Provide the prompt-exclusion manifest before constructing OXE-DROID." - ) - try: - payload = json.loads(path.read_text(encoding="utf-8")) - canonical = _parse_episode_indices(payload["episode_indices"], "episode_indices") - prompt = payload[DROID_PROMPT_EXCLUSION_KEY] - if not isinstance(prompt, dict): - raise ValueError(f"{DROID_PROMPT_EXCLUSION_KEY} must be an object") - if type(prompt["schema_version"]) is not int or ( - prompt["schema_version"] != DROID_PROMPT_EXCLUSION_SCHEMA_VERSION - ): - raise ValueError( - f"{DROID_PROMPT_EXCLUSION_KEY}.schema_version={prompt['schema_version']!r}, " - f"expected {DROID_PROMPT_EXCLUSION_SCHEMA_VERSION}" - ) - if prompt["fallback_chain"] != list(DROID_PROMPT_FALLBACK_COLS): - raise ValueError( - f"{DROID_PROMPT_EXCLUSION_KEY}.fallback_chain={prompt['fallback_chain']!r}, " - f"expected {list(DROID_PROMPT_FALLBACK_COLS)!r}" - ) - prompt_owned = _parse_episode_indices( - prompt["episode_indices"], - f"{DROID_PROMPT_EXCLUSION_KEY}.episode_indices", - ) - independently_owned = _parse_episode_indices( - prompt[DROID_PROMPT_INDEPENDENT_EXCLUSIONS_KEY], - f"{DROID_PROMPT_EXCLUSION_KEY}.{DROID_PROMPT_INDEPENDENT_EXCLUSIONS_KEY}", - ) - if canonical != prompt_owned | independently_owned: - raise ValueError("episode_indices must equal the union of prompt-owned and independently-owned exclusions") - - latest_scan = prompt["latest_scan"] - if not isinstance(latest_scan, dict): - raise ValueError(f"{DROID_PROMPT_EXCLUSION_KEY}.latest_scan must be an object") - latest_scan_owned = _parse_episode_indices( - latest_scan["episode_indices"], - f"{DROID_PROMPT_EXCLUSION_KEY}.latest_scan.episode_indices", - ) - if latest_scan_owned != prompt_owned: - raise ValueError("latest_scan.episode_indices do not match the prompt-owned exclusions") - scan_stats = latest_scan["stats"] - if not isinstance(scan_stats, dict): - raise ValueError(f"{DROID_PROMPT_EXCLUSION_KEY}.latest_scan.stats must be an object") - if scan_stats["fallback_chain"] != list(DROID_PROMPT_FALLBACK_COLS): - raise ValueError("latest_scan.stats.fallback_chain does not match the reader fallback chain") - rows_scanned = _parse_nonnegative_int(scan_stats["rows_scanned"], "latest_scan.stats.rows_scanned") - unresolved_rows = _parse_nonnegative_int( - scan_stats["unresolved_rows"], - "latest_scan.stats.unresolved_rows", - ) - if rows_scanned == 0 or unresolved_rows > rows_scanned: - raise ValueError("latest_scan.stats row counts are inconsistent") - if _parse_nonnegative_int( - scan_stats["episodes_all_unresolved"], - "latest_scan.stats.episodes_all_unresolved", - ) != len(prompt_owned): - raise ValueError("latest_scan.stats.episodes_all_unresolved does not match prompt-owned exclusions") - for field in ("episodes_partially_unresolved", "task_index_missing_from_tasks_parquet"): - if _parse_nonnegative_int(scan_stats[field], f"latest_scan.stats.{field}") != 0: - raise ValueError(f"latest_scan.stats.{field} must be zero for a completed scan") - resolved_population = population or resolve_lerobot_v3_data_population(root) - if rows_scanned != resolved_population.total_rows: - raise ValueError( - f"latest_scan.stats.rows_scanned={rows_scanned} but the data manifest addresses " - f"{resolved_population.total_rows} rows" - ) - recorded_inputs_digest = _parse_prompt_inputs_digest(latest_scan[DROID_PROMPT_INPUTS_DIGEST_KEY]) - current_inputs_digest = compute_droid_prompt_inputs_digest( - root, - population=resolved_population, - ) - if current_inputs_digest != recorded_inputs_digest: - raise ValueError(f"{DROID_PROMPT_INPUTS_DIGEST_KEY} does not match tasks.parquet and prompt source columns") - except (KeyError, OSError, TypeError, ValueError) as exc: - raise ValueError(f"{path} is stale or malformed ({exc}). Regenerate the prompt-exclusion manifest.") from exc - return payload, canonical class OxeDroidDataset(LeRobotV3Reader): - DATASET_NAME = "OXE-DROID" + """Load DROID's exported rows directly, without further frame selection.""" + + STRICT_BUCKET_LOADING = True + DATASET_NAME = "DROID" HEAD_CAMERA = "observation.images.primary" LEFT_WRIST_CAMERA = "observation.images.wrist" RIGHT_WRIST_CAMERA = None - - PROMPT_FALLBACK_COLS: ClassVar[Tuple[str, ...]] = DROID_PROMPT_FALLBACK_COLS NEEDED_COLS = ( "state", "other_information.observation_gripper_pose6d", "other_information.action_wrist_pose", "task_index", - ) + PROMPT_FALLBACK_COLS + ) ACTION_DIM_MASK = LEFT_ARM_DIM_MASK - - PROMPT_SOURCE = "task_index" + PROMPT_FILE_REQUIRED = True DEFAULT_NORMALIZE_MODE = "quantile" STATS_FILENAME = "eef_stats.json" STATS_DIM = 10 STATS_STRICT_MINMAX = True - def _build_episode_index(self, info: dict) -> pd.DataFrame: - population = resolve_lerobot_v3_data_population(self._dataset_dir, info=info) - _, self._droid_excluded_episode_indices = load_droid_prompt_exclusions( - self._dataset_dir, - population=population, - ) - self._droid_data_population_digest = digest_lerobot_v3_data_population(population) - _, self._droid_stats_population = resolve_droid_stats_population( - population, - info, - self._droid_excluded_episode_indices, - split="train", - ) - eps = population.episodes.copy() - self._add_episode_offsets(eps) - info_splits = info.get("splits", {}) or {} - return apply_info_splits( - eps, - self._split, - info_splits, - source_name=f"{self.DATASET_NAME}({self._dataset_id})", - ) - def _load_excluded_episode_indices(self) -> set[int]: - """Reuse the canonical set validated with the prompt certificate.""" - return set(self._droid_excluded_episode_indices) + """The export has already physically removed excluded episodes.""" + return set() def _load_stats(self, info: dict) -> Optional[dict]: if not self._normalize_mode or self._normalize_mode in ("none", "null"): - return super()._load_stats(info) + return None stats_path = self._dataset_dir / "meta" / self.STATS_FILENAME - if not stats_path.exists(): - return super()._load_stats(info) - try: - raw = json.loads(stats_path.read_text(encoding="utf-8")) - stats_excluded = _parse_episode_indices( - raw["excluded_episode_indices"], - "excluded_episode_indices", + if not stats_path.is_file(): + raise FileNotFoundError(f"DROID requires its own precomputed statistics at {stats_path}.") + with stats_path.open(encoding="utf-8") as stream: + raw = json.load(stream) + if not isinstance(raw, dict): + raise DataContractError(f"{stats_path} must contain a statistics mapping.") + if raw.get(DROID_EEF_STATS_CONTRACT_KEY) != DROID_EEF_STATS_CONTRACT: + raise DataContractError( + f"{stats_path} does not match the rigid terminal-arm pose frame and canonical gripper convention. " + "Use the stats shipped with the cleaned DROID export." ) - stats_population_digest = raw[DROID_DATA_POPULATION_DIGEST_KEY] - stats_population = raw[DROID_STATS_POPULATION_KEY] - stats_contract = raw[DROID_EEF_STATS_CONTRACT_KEY] - if ( - not isinstance(stats_population_digest, str) - or re.fullmatch(r"[0-9a-f]{64}", stats_population_digest) is None - ): - raise ValueError(f"{DROID_DATA_POPULATION_DIGEST_KEY} must be a lowercase SHA-256 digest") - except (KeyError, TypeError, ValueError) as exc: - raise ValueError( - f"{stats_path} has missing or invalid DROID normalization provenance ({exc}). " - "Re-run oxe_stats_computation for DROID." - ) from exc - if stats_excluded != self._droid_excluded_episode_indices: - raise ValueError( - f"{stats_path} excluded_episode_indices do not match meta/excluded_episodes.json. " - "Re-run oxe_stats_computation for DROID after updating exclusions." - ) - if stats_population_digest != self._droid_data_population_digest: - raise ValueError( - f"{stats_path} {DROID_DATA_POPULATION_DIGEST_KEY} does not match the current data manifest. " - "Re-run oxe_stats_computation for DROID." - ) - if stats_population != self._droid_stats_population: - raise ValueError( - f"{stats_path} {DROID_STATS_POPULATION_KEY} does not match the current train " - "split/exclusion population. Re-run oxe_stats_computation for DROID." - ) - try: - _validate_droid_eef_stats_contract(stats_contract) - except (KeyError, TypeError, ValueError) as exc: - raise ValueError( - f"{stats_path} {DROID_EEF_STATS_CONTRACT_KEY} is stale or invalid; expected " - "the rigid terminal-arm pose frame and 0=closed, 1=open gripper transform. " - f"Re-run oxe_stats_computation for DROID ({exc})." - ) from exc stats = materialize_eef_stats( raw, self._normalize_mode, @@ -529,68 +85,24 @@ def _load_stats(self, info: dict) -> Optional[dict]: ) return stats - def _resolve_prompt(self, row, win: pd.DataFrame) -> str: - """tasks.parquet text, else the first non-placeholder fallback column. - - Raises when every source is blank. This is a **backstop, not a runtime - degradation path**: those episodes must already be gone from the index - via ``meta/excluded_episodes.json`` (see the module docstring). Relying - on ``_safe_get`` to retry past them does NOT work — its ``idx + 1`` walk - stays inside the same episode, and the condition is all-or-nothing per - episode, so any episode longer than ``_GETITEM_MAX_RETRIES`` frames - exhausts the retries and kills the DataLoader worker. - """ - - task_idx = int(win["task_index"].iloc[0]) - if task_idx not in self._task_idx_to_text: - raise KeyError( - f"{self.DATASET_NAME} prompt lookup failed: task_index={task_idx} not present " - "in this bucket's tasks.parquet." - ) - text = _clean_text(self._task_idx_to_text[task_idx]) - if text: - return text - for col in self.PROMPT_FALLBACK_COLS: - text = _clean_text(win[col].iloc[0]) - if text: - return text - raise ValueError( - f"{self.DATASET_NAME} task_index={task_idx} has a blank prompt and every fallback " - f"column {list(self.PROMPT_FALLBACK_COLS)} is blank too (episode_index=" - f"{int(row['episode_index'])}). Generate meta/excluded_episodes.json " - "before loading the dataset." - ) - def _action_20d(self, win: pd.DataFrame) -> np.ndarray: action = np.stack(win["other_information.action_wrist_pose"].values).astype(np.float32) + if action.ndim != 2 or action.shape[1] != 7: + raise DataContractError(f"DROID commanded wrist pose must have width 7, got {action.shape}.") return single_arm_20d(droid_euler7_to_arm10(action), self._normalization_stats, self._normalize_mode) def _proprio_20d(self, win: pd.DataFrame) -> Optional[np.ndarray]: pose = np.stack(win["other_information.observation_gripper_pose6d"].values[:1]).astype(np.float32) state = np.stack(win["state"].values[:1]).astype(np.float32) + if state.ndim != 2 or state.shape[1] < 7: + raise DataContractError(f"DROID state must provide gripper closedness at index 6, got {state.shape}.") arm10 = droid_pose6_closedness_to_arm10(pose, state[:, 6:7]) return single_arm_20d(arm10, self._normalization_stats, self._normalize_mode) __all__ = [ - "DROID_DATA_POPULATION_DIGEST_KEY", "DROID_EEF_STATS_CONTRACT", "DROID_EEF_STATS_CONTRACT_KEY", "DROID_EEF_STATS_CONTRACT_VERSION", - "DROID_PROMPT_EXCLUSION_KEY", - "DROID_PROMPT_EXCLUSION_SCHEMA_VERSION", - "DROID_PROMPT_FALLBACK_COLS", - "DROID_PROMPT_INDEPENDENT_EXCLUSIONS_KEY", - "DROID_PROMPT_INPUTS_DIGEST_FORMAT_VERSION", - "DROID_PROMPT_INPUTS_DIGEST_KEY", - "DROID_PROMPT_SOURCE_COLUMNS", - "DROID_STATS_POPULATION_KEY", - "DROID_STATS_POPULATION_POLICY", - "DROID_STATS_POPULATION_SCHEMA_VERSION", "OxeDroidDataset", - "compute_droid_prompt_inputs_digest", - "digest_droid_prompt_shard", - "load_droid_prompt_exclusions", - "read_droid_prompt_population_shard", - "resolve_droid_stats_population", ] diff --git a/openwam/dataloader/registry.py b/openwam/dataloader/registry.py index 084220cf4..3db5beb35 100644 --- a/openwam/dataloader/registry.py +++ b/openwam/dataloader/registry.py @@ -25,7 +25,7 @@ def register_dataset(name: str): """Decorator to register a dataset class by name. Args: - name: Config-level type name (e.g., "droid", "bridge_v2", "oxe"). + name: Config-level type name (e.g., "oxe_droid", "robocoin"). """ def wrapper(cls): @@ -79,7 +79,6 @@ def _register_builtins(): from openwam.dataloader.interndata_a1 import InternDataA1Dataset from openwam.dataloader.libero import LiberoDataset from openwam.dataloader.mixture import MixtureDataset - from openwam.dataloader.muka_franka import MukaFrankaDataset from openwam.dataloader.oxe_droid import OxeDroidDataset from openwam.dataloader.robocasa365 import MultiTaskRoboCasa365Dataset from openwam.dataloader.robocasa_gr1 import MultiRoboCasaGR1Dataset @@ -96,7 +95,6 @@ def _register_builtins(): register_dataset("robocoin")(MultiRobotCOINDataset) register_dataset("ebench")(EBenchDataset) register_dataset("libero")(LiberoDataset) - register_dataset("muka_franka")(MukaFrankaDataset) register_dataset("oxe_droid")(OxeDroidDataset) register_dataset("robocasa365")(MultiTaskRoboCasa365Dataset) # Registered on the SINGLE-bucket class: its from_config returns either one diff --git a/openwam/dataloader/robocoin.py b/openwam/dataloader/robocoin.py index ebff92e8b..eae50e707 100644 --- a/openwam/dataloader/robocoin.py +++ b/openwam/dataloader/robocoin.py @@ -1,149 +1,60 @@ -"""RoboCOIN LeRobot v3 reader with real action/state supervision. - -Grippered buckets consume the published ``eef_sim_pose_*`` Euler-XYZ poses and -``gripper_open_scale_*`` values, producing the canonical 20-D bimanual EEF -layout. The pose follows a rigid terminal-arm attachment in the robot base -frame and the gripper convention is ``0=closed, 1=open``. - -Dexterous-hand buckets have no scalar gripper columns. Without unification -they emit pose with masked grip slots; with unification they append left/right -finger joints and scatter pose/fingers into the corresponding 80-D slots while -leaving scalar-grip and unused hand slots masked. - -Rows are temporally aligned: action row ``t`` is supervised at row ``t``. -Normalization is pooled per robot type over action and state, with rot6d -dimensions pinned to identity. Optional episode trimming and exclusions are -validated against the exact manifest population before stats are accepted. +"""Read physically cleaned RoboCOIN LeRobot v3 task buckets. + +Action and state use the published Euler-XYZ EEF poses and gripper openness +(``0=closed, 1=open``). Dexterous-hand buckets retain their finger joints under +unified actions. Normalization reads the existing per-robot statistics from +``meta/stats_.json`` under the supplied family root; loading never +changes the dataset or applies an additional quality filter. """ from __future__ import annotations -import csv -import hashlib -import io import json import logging -import re -from concurrent.futures import ThreadPoolExecutor -from dataclasses import dataclass, field +from functools import lru_cache from pathlib import Path -from types import MappingProxyType from typing import List import numpy as np -import pyarrow.parquet as pq from openwam.dataloader.bases import LeRobotV3Reader, MultiLeRobotV3Reader from openwam.dataloader.utils.eef import EEF_DIM as _ACTION_DIM from openwam.dataloader.utils.eef import eef14_to_eef20 -from openwam.dataloader.utils.lerobotv3 import ( - DataContractError, - ExcludedEpisodesSnapshot, - apply_info_splits, - assert_excluded_episodes_snapshot_current, - build_multibucket, - load_episodes_parquet, - load_excluded_episodes_snapshot, - parse_info_json, +from openwam.dataloader.utils.normalization import ( + STAT_KEYS, + apply_normalization, + find_dataset_stats, + materialize_eef_stats, ) -from openwam.dataloader.utils.normalization import apply_normalization, materialize_eef_stats logger = logging.getLogger(__name__) -_STATE_DIM = _ACTION_DIM - -# One canonical denylist is shared by reader and stats discovery. -ROBOCOIN_WHOLE_BUCKET_EXCLUSIONS = MappingProxyType( - { - "Airbot_MMK2_storage_peach_pear": MappingProxyType( - { - "robot_type": "airbot_mmk2", - "reason": "fails the public action-unit contract and must not enter normalization statistics", - } - ), - } -) -_WHOLE_BUCKET_EXCLUSIONS_SCHEMA_VERSION = 1 -_WHOLE_BUCKET_EXCLUSIONS_POLICY = "drop_whole_bucket_before_reader_and_stats_discovery" - - -def is_robocoin_bucket_excluded(dataset_dir) -> bool: - """Whether ``dataset_dir`` is excluded as one indivisible task bucket.""" - return Path(dataset_dir).name in ROBOCOIN_WHOLE_BUCKET_EXCLUSIONS - - -def robocoin_bucket_exclusions_provenance(robot_type: str) -> dict: - """Canonical whole-bucket exclusion contract for one robot-type stats file.""" - datasets = { - dataset_id: dict(entry) - for dataset_id, entry in ROBOCOIN_WHOLE_BUCKET_EXCLUSIONS.items() - if entry["robot_type"] == str(robot_type) - } - return { - "schema_version": _WHOLE_BUCKET_EXCLUSIONS_SCHEMA_VERSION, - "policy": _WHOLE_BUCKET_EXCLUSIONS_POLICY, - "datasets": dict(sorted(datasets.items())), - } - - -def _validate_whole_bucket_exclusions_provenance(actual, *, robot_type, stats_path) -> None: - """Fail closed when stats for an affected robot type predate its denylist.""" - expected = robocoin_bucket_exclusions_provenance(robot_type) - if expected["datasets"]: - if actual != expected: - raise DataContractError( - f"RoboCOIN stats {stats_path} whole_bucket_exclusions_provenance does not " - f"match the repository exclusion contract for robot_type={robot_type!r}; " - f"expected {expected}, got {actual}. Regenerate stats." - ) - elif actual is not None and actual != expected: - raise DataContractError( - f"RoboCOIN stats {stats_path} has incompatible " - f"whole_bucket_exclusions_provenance for robot_type={robot_type!r}; " - f"expected {expected}, got {actual}. Regenerate stats." - ) +@lru_cache(maxsize=32) +def _read_robot_stats(path: str) -> dict: + """Share immutable on-disk stats across buckets of the same robot type.""" + with open(path, encoding="utf-8") as stream: + raw = json.load(stream) + if not isinstance(raw, dict): + raise ValueError(f"RoboCOIN stats {path} must contain a JSON object") + return raw -def _discover_data_parquets(dataset_dir: Path, data_path_template: str): - """Return canonical numeric RoboCOIN data shards in reader order.""" - if not isinstance(data_path_template, str) or not data_path_template: - raise DataContractError( - f"RoboCOIN({Path(dataset_dir).name}): info.json data_path must be a non-empty string, " - f"got {data_path_template!r}" - ) - data_files = {} - for path in sorted((Path(dataset_dir) / "data").glob("chunk-*/file-*.parquet")): - chunk_m = re.search(r"chunk-(\d+)$", path.parent.name) - file_m = re.search(r"file-(\d+)$", path.stem) - if chunk_m is not None and file_m is not None: - coordinates = (int(chunk_m.group(1)), int(file_m.group(1))) - previous = data_files.get(coordinates) - if previous is not None: - raise DataContractError( - f"RoboCOIN({Path(dataset_dir).name}): duplicate numeric data shard " - f"coordinates {coordinates}: {previous} and {path}" - ) - data_files[coordinates] = path - - result = [] - for (chunk_index, file_index), path in sorted(data_files.items()): - try: - expected = Path(dataset_dir) / data_path_template.format( - chunk_index=chunk_index, - file_index=file_index, - ) - except (AttributeError, IndexError, KeyError, TypeError, ValueError) as e: - raise DataContractError( - f"RoboCOIN({Path(dataset_dir).name}): invalid info.json data_path template {data_path_template!r}: {e}" - ) from e - if path != expected: - raise DataContractError( - f"RoboCOIN({Path(dataset_dir).name}): non-canonical numeric data shard {path}; " - f"info.json data_path resolves coordinates {(chunk_index, file_index)} to {expected}" +def _materialize_robot_stats(raw, *, dim: int, mode: str, source: str, rot6d: bool = False) -> dict: + """Validate stored vectors before applying the normalization contract.""" + if not isinstance(raw, dict): + raise ValueError(f"RoboCOIN stats {source} must contain a mapping") + for key in STAT_KEYS: + if key not in raw: + continue + values = np.asarray(raw[key], dtype=np.float32) + if values.shape != (dim,) or not np.isfinite(values).all(): + raise ValueError( + f"RoboCOIN stats {source}.{key} must be a finite vector with shape ({dim},), got {values.shape}" ) - result.append((path, chunk_index, file_index)) - return result + return materialize_eef_stats( + raw, mode, dim=dim, strict_minmax=False, source_hint=source, force_rot6d_identity=rot6d + ) _NEEDED_COLS = ( @@ -207,7 +118,7 @@ def dex_finger_layout(features: dict): ``None``. Single source of truth for the dex-bucket gate, shared by the reader - (:meth:`RoboCOINDataset.__init__`) and the stats script so they cannot drift. + (:meth:`RoboCOINDataset.__init__`) and all callers so their hand layouts agree. """ has_grip = all(c in features for c in _GRIP_COLS) @@ -245,9 +156,6 @@ def _build_dex_unify_map(k_left: int, k_right: int): ) -_eef14_to_eef20 = eef14_to_eef20 - - HEAD_CAMERA_PRIORITY = [ "observation.images.cam_high_rgb", "observation.images.cam_head_rgb", @@ -280,524 +188,6 @@ def _build_dex_unify_map(k_left: int, k_right: int): ] -_TRIM_SNAPSHOT_CACHE: dict = {} -_TRIM_SCHEMA_VERSION = 1 -_TRIM_MIN_LEN = 1 -_TRIM_ZERO_SPAN_POLICY = "drop" -_EXCLUDED_EPISODES_SCHEMA_VERSION = 1 -_EXCLUDED_EPISODES_POLICY = "drop_matching_episode_index_before_trim" -_STATS_POPULATION_SCHEMA_VERSION = 2 -_STATS_POPULATION_POLICY = "info_split_then_exclusion_then_trim" -_TRIM_CSV_COLUMNS = ( - "dataset", - "episode_index", - "total_frames", - "trim_head_to", - "trim_tail_from", -) - - -def _excluded_episodes_provenance( - snapshots: dict[str, ExcludedEpisodesSnapshot], -) -> dict: - """Serialize the exact effective blacklist for every stats contributor.""" - return { - "schema_version": _EXCLUDED_EPISODES_SCHEMA_VERSION, - "policy": _EXCLUDED_EPISODES_POLICY, - "datasets": { - dataset_id: {"episode_indices": list(snapshot.episode_indices)} - for dataset_id, snapshot in sorted(snapshots.items()) - }, - } - - -def _validate_excluded_episodes_provenance( - actual, - *, - dataset_dir: Path, - dataset_id: str, - current_snapshot: ExcludedEpisodesSnapshot, - num_datasets, - stats_path: Path, -) -> dict[str, ExcludedEpisodesSnapshot]: - """Validate every bucket population pooled into one robot-type stats file.""" - expected_keys = {"schema_version", "policy", "datasets"} - if ( - not isinstance(actual, dict) - or set(actual) != expected_keys - or type(actual.get("schema_version")) is not int - or actual.get("schema_version") != _EXCLUDED_EPISODES_SCHEMA_VERSION - or actual.get("policy") != _EXCLUDED_EPISODES_POLICY - or not isinstance(actual.get("datasets"), dict) - ): - raise DataContractError( - f"RoboCOIN bucket {dataset_id}: stats {stats_path} excluded_episodes_provenance " - "has a missing or incompatible schema/policy. Regenerate stats with the configured trim CSV." - ) - - datasets = actual["datasets"] - if type(num_datasets) is not int or num_datasets != len(datasets): - raise DataContractError( - f"RoboCOIN bucket {dataset_id}: stats {stats_path} declares num_datasets=" - f"{num_datasets!r} but excluded_episodes_provenance has {len(datasets)} dataset entries. " - "Regenerate stats with the configured trim CSV." - ) - if dataset_id not in datasets: - raise DataContractError( - f"RoboCOIN bucket {dataset_id}: stats {stats_path} excluded_episodes_provenance " - "does not include this bucket. Regenerate stats with the configured trim CSV." - ) - - root = Path(dataset_dir).parent - snapshots = {} - for contributor, entry in datasets.items(): - if not isinstance(contributor, str) or not contributor: - raise DataContractError( - f"RoboCOIN bucket {dataset_id}: stats {stats_path} has invalid contributor " - f"name {contributor!r} in excluded_episodes_provenance" - ) - contributor_path = Path(contributor) - if contributor in (".", "..") or contributor_path.is_absolute() or len(contributor_path.parts) != 1: - raise DataContractError( - f"RoboCOIN bucket {dataset_id}: stats {stats_path} has invalid contributor " - f"name {contributor!r} in excluded_episodes_provenance" - ) - bucket_dir = root / contributor - if not bucket_dir.is_dir(): - raise DataContractError( - f"RoboCOIN bucket {dataset_id}: stats {stats_path} references missing contributor " - f"directory {bucket_dir} in excluded_episodes_provenance" - ) - snapshot = current_snapshot if contributor == dataset_id else load_excluded_episodes_snapshot(bucket_dir) - snapshots[contributor] = snapshot - expected_entry = {"episode_indices": list(snapshot.episode_indices)} - if ( - not isinstance(entry, dict) - or set(entry) != {"episode_indices"} - or not isinstance(entry.get("episode_indices"), list) - or any(type(value) is not int for value in entry["episode_indices"]) - or entry != expected_entry - ): - raise DataContractError( - f"RoboCOIN bucket {dataset_id}: stats {stats_path} excluded_episodes_provenance " - f"for contributor {contributor!r} does not match the current exclusion population; " - f"expected {expected_entry}, got {entry}. Regenerate stats with the configured trim CSV." - ) - return snapshots - - -def _validate_stats_population_provenance( - actual, - *, - dataset_dir: Path, - current_dataset_id: str, - current_info: dict, - trim_snapshot: _TrimSnapshot, - exclusion_snapshots: dict[str, ExcludedEpisodesSnapshot], - num_datasets, - stats_path: Path, -) -> None: - """Validate every contributor to one pooled robot-type stats file. - - Checking only the leaf currently being constructed is insufficient: a - sibling bucket contributes to the same pooled numbers, so changing that - sibling's train split would stale normalization for this leaf too. Rebuild - the exact physical train spans for every contributor using the immutable - trim/exclusion snapshots already validated above. - """ - - expected_header = { - "schema_version": _STATS_POPULATION_SCHEMA_VERSION, - "split": "train", - "policy": _STATS_POPULATION_POLICY, - } - if ( - not isinstance(actual, dict) - or {key: actual.get(key) for key in expected_header} != expected_header - or set(actual) != {*expected_header, "datasets"} - or not isinstance(actual.get("datasets"), dict) - ): - raise DataContractError( - f"RoboCOIN bucket {current_dataset_id}: stats {stats_path} population must be " - "train-derived with the reader's split/exclusion/trim order; " - f"expected schema/header {expected_header}, got {actual}. Regenerate stats " - "with --split train and the configured trim CSV." - ) - datasets = actual["datasets"] - if type(num_datasets) is not int or len(datasets) != num_datasets or set(datasets) != set(exclusion_snapshots): - raise DataContractError( - f"RoboCOIN bucket {current_dataset_id}: stats {stats_path} population datasets " - "do not exactly match the pooled exclusion contributors; regenerate stats." - ) - - root = Path(dataset_dir).parent - - def _expected(contributor: str, snapshot: ExcludedEpisodesSnapshot) -> dict: - bucket_dir = root / contributor - info = current_info if contributor == current_dataset_id else parse_info_json(bucket_dir) - _spans, _manifest_end, expected_entry = _stats_population_spans( - bucket_dir, - dataset_id=contributor, - split="train", - dataset_spec=trim_snapshot.spec.get(contributor, {}), - excluded_episode_indices=snapshot.episode_indices, - info=info, - ) - return expected_entry - - current_expected = _expected(current_dataset_id, exclusion_snapshots[current_dataset_id]) - current_actual = datasets.get(current_dataset_id) - if current_actual != current_expected: - raise DataContractError( - f"RoboCOIN bucket {current_dataset_id}: stats {stats_path} effective train " - "population no longer matches its manifest/split/exclusion/trim spans; " - f"expected {current_expected}, got {current_actual}. Regenerate stats." - ) - - population_digest = hashlib.sha256( - json.dumps(actual, sort_keys=True, separators=(",", ":")).encode("utf-8") - ).hexdigest() - cache_key = (str(stats_path.resolve()), population_digest) - if cache_key in trim_snapshot.population_validation_cache: - return - - for contributor, snapshot in exclusion_snapshots.items(): - if contributor == current_dataset_id: - continue - expected_entry = _expected(contributor, snapshot) - actual_entry = datasets.get(contributor) - if actual_entry != expected_entry: - raise DataContractError( - f"RoboCOIN bucket {current_dataset_id}: stats {stats_path} effective train " - f"population for contributor {contributor!r} no longer matches its " - "manifest/split/exclusion/trim spans; " - f"expected {expected_entry}, got {actual_entry}. Regenerate stats." - ) - trim_snapshot.population_validation_cache.add(cache_key) - - -@dataclass(frozen=True) -class _TrimSnapshot: - """Parsed trim rules and provenance derived from one immutable byte snapshot.""" - - path: str - spec: dict - sha256: str - - population_validation_cache: set[tuple[str, str]] = field( - default_factory=set, - compare=False, - repr=False, - ) - - @property - def provenance(self) -> dict: - return { - "schema_version": _TRIM_SCHEMA_VERSION, - "sha256": self.sha256, - "min_len": _TRIM_MIN_LEN, - "zero_span_policy": _TRIM_ZERO_SPAN_POLICY, - } - - -class _TrimSnapshotConfig: - """Overlay one pinned snapshot onto dict, DictConfig, or namespace configs.""" - - def __init__(self, base, snapshot: _TrimSnapshot): - self._base = base - self._trim_snapshot = snapshot - - def __getattr__(self, key): - return getattr(self._base, key) - - def get(self, key, default=None): - if key == "_trim_snapshot": - return self._trim_snapshot - if hasattr(self._base, "get"): - return self._base.get(key, default) - return getattr(self._base, key, default) - - -def _trim_csv_int(row, name: str, *, path: str, line_number: int, required: bool = False): - """Parse one optional integer field, with actionable CSV error context.""" - raw = row.get(name) - value = raw.strip() if isinstance(raw, str) else "" - if not value: - if required: - raise ValueError(f"RoboCOIN trim_csv {path}, line {line_number}: '{name}' is required") - return None - try: - return int(value) - except ValueError as e: - raise ValueError( - f"RoboCOIN trim_csv {path}, line {line_number}: '{name}' must be an integer, got {value!r}" - ) from e - - -def _load_trim_snapshot(path) -> _TrimSnapshot: - """Read, parse, and hash one immutable trim CSV byte snapshot. - - Cached per unchanged path: MultiRobotCOINDataset builds one reader per - bucket, and each would otherwise re-read the same large CSV. - - Expected columns — ``dataset``, ``episode_index``, ``total_frames``, - ``trim_head_to``, ``trim_tail_from``. Empty trim fields mean "no trim on that - end". ``total_frames`` is required and carried through as a staleness check - (see :meth:`RoboCOINDataset._filter_episodes`). Invalid files fail fast; only - a completely parsed snapshot is cached. The parsed spec and SHA-256 always - come from the same bytes, so a path replacement cannot pair old spans with - new provenance inside one reader construction or stats scan. - """ - - key = str(path) - try: - stat = Path(key).stat() - except OSError as e: - raise OSError( - e.errno, - f"RoboCOIN trim_csv {key} could not be read: {e.strerror or e}", - key, - ) from e - signature = (stat.st_dev, stat.st_ino, stat.st_mtime_ns, stat.st_ctime_ns, stat.st_size) - cached = _TRIM_SNAPSHOT_CACHE.get(key) - if cached is not None and cached[0] == signature: - return cached[1] - try: - raw = Path(key).read_bytes() - except OSError as e: - raise OSError( - e.errno, - f"RoboCOIN trim_csv {key} could not be read: {e.strerror or e}", - key, - ) from e - try: - text = raw.decode("utf-8") - except UnicodeError as e: - raise ValueError(f"RoboCOIN trim_csv {key}, line unknown: invalid text encoding: {e}") from e - - spec: dict = {} - seen_entries = set() - n = 0 - try: - with io.StringIO(text, newline="") as fh: - reader = csv.DictReader(fh, strict=True) - fieldnames = reader.fieldnames - if fieldnames is None: - raise ValueError(f"RoboCOIN trim_csv {key}, line 1: missing CSV header") - missing = [name for name in _TRIM_CSV_COLUMNS if name not in fieldnames] - if missing: - raise ValueError(f"RoboCOIN trim_csv {key}, line 1: missing required column(s): {', '.join(missing)}") - duplicate = sorted({name for name in fieldnames if fieldnames.count(name) > 1}) - if duplicate: - raise ValueError(f"RoboCOIN trim_csv {key}, line 1: duplicate column(s): {', '.join(duplicate)}") - - for row in reader: - line_number = reader.line_num - if None in row: - raise ValueError( - f"RoboCOIN trim_csv {key}, line {line_number}: row has more values than the CSV header" - ) - dataset_raw = row.get("dataset") - dataset = dataset_raw.strip() if isinstance(dataset_raw, str) else "" - if not dataset: - raise ValueError(f"RoboCOIN trim_csv {key}, line {line_number}: 'dataset' is required") - episode_index = _trim_csv_int(row, "episode_index", path=key, line_number=line_number, required=True) - total = _trim_csv_int(row, "total_frames", path=key, line_number=line_number, required=True) - head = _trim_csv_int(row, "trim_head_to", path=key, line_number=line_number) - tail = _trim_csv_int(row, "trim_tail_from", path=key, line_number=line_number) - if episode_index < 0: - raise ValueError( - f"RoboCOIN trim_csv {key}, line {line_number}: " - f"'episode_index' must be >= 0, got {episode_index}" - ) - if total <= 0: - raise ValueError( - f"RoboCOIN trim_csv {key}, line {line_number}: 'total_frames' must be > 0, got {total}" - ) - for name, value in (("trim_head_to", head), ("trim_tail_from", tail)): - if value is not None and not 0 <= value <= total: - raise ValueError( - f"RoboCOIN trim_csv {key}, line {line_number}: " - f"'{name}' must be in [0, total_frames={total}], got {value}" - ) - head = head or 0 - effective_tail = total if tail is None else tail - if head > effective_tail: - raise ValueError( - f"RoboCOIN trim_csv {key}, line {line_number}: trim_head_to={head} " - f"must not exceed trim_tail_from={effective_tail}" - ) - entry_key = (dataset, episode_index) - if entry_key in seen_entries: - raise ValueError( - f"RoboCOIN trim_csv {key}, line {line_number}: duplicate entry for " - f"dataset={dataset!r}, episode_index={episode_index}" - ) - seen_entries.add(entry_key) - if not head and tail is None: - continue - dataset_spec = spec.setdefault(dataset, {}) - dataset_spec[episode_index] = (head, tail, total) - n += 1 - except csv.Error as e: - line_number = getattr(locals().get("reader"), "line_num", "unknown") - raise ValueError(f"RoboCOIN trim_csv {key}, line {line_number}: malformed CSV: {e}") from e - - logger.info("RoboCOIN: loaded %d trim entries across %d datasets from %s", n, len(spec), key) - immutable_spec = MappingProxyType({dataset: MappingProxyType(entries) for dataset, entries in spec.items()}) - snapshot = _TrimSnapshot(path=key, spec=immutable_spec, sha256=hashlib.sha256(raw).hexdigest()) - _TRIM_SNAPSHOT_CACHE[key] = (signature, snapshot) - return snapshot - - -def _load_trim_spec(path) -> dict: - """Return the parsed spec from one cached immutable trim snapshot.""" - return {dataset: dict(entries) for dataset, entries in _load_trim_snapshot(path).spec.items()} - - -def _assert_trim_snapshot_current(snapshot: _TrimSnapshot, *, context: str) -> None: - """Fail if the configured path changed after ``snapshot`` was pinned.""" - try: - actual_sha256 = hashlib.sha256(Path(snapshot.path).read_bytes()).hexdigest() - except OSError as e: - raise OSError( - e.errno, - f"RoboCOIN trim_csv {snapshot.path} could not be re-read after {context}: {e.strerror or e}", - snapshot.path, - ) from e - if actual_sha256 != snapshot.sha256: - raise DataContractError( - f"RoboCOIN trim_csv {snapshot.path} changed while {context}; " - f"started with sha256={snapshot.sha256}, now sha256={actual_sha256}. " - "Refusing mixed trim spans/provenance; retry with an immutable CSV." - ) - - -def _validate_trim_manifest(dataset_id: str, manifest, spec: dict) -> dict: - """Validate one bucket's trim spec against its full pre-split manifest. - - Returns ``{episode_index: (head, tail)}`` relative spans. Unknown IDs and - every ``total_frames`` mismatch are data-contract failures, including rows - outside the selected split or later excluded by the reader. - """ - - if not spec: - return {} - try: - episode_ids = [int(ep) for ep in manifest["episode_index"].to_numpy()] - lengths = [int(length) for length in manifest["length"].to_numpy()] - except (KeyError, TypeError, ValueError) as e: - raise DataContractError( - f"RoboCOIN({dataset_id}): full manifest must carry integer episode_index and length columns" - ) from e - manifest_lengths = dict(zip(episode_ids, lengths)) - if len(manifest_lengths) != len(episode_ids): - raise DataContractError(f"RoboCOIN({dataset_id}): full manifest has duplicate episode_index values") - - unknown = sorted(set(spec).difference(manifest_lengths)) - if unknown: - unknown_ids = ",".join(map(str, unknown[:10])) + ("..." if len(unknown) > 10 else "") - raise DataContractError( - f"RoboCOIN({dataset_id}): trim list references {len(unknown)} unknown episode_index " - f"value(s) absent from the full pre-split manifest (episode_index={unknown_ids}). " - "Refusing dataset construction; re-run " - "Regenerate the trim manifest against the current corpus before retrying." - ) - - stale = sorted(ep for ep, (_, _, total) in spec.items() if int(total) != manifest_lengths[ep]) - if stale: - stale_ids = ",".join(map(str, stale[:10])) + ("..." if len(stale) > 10 else "") - raise DataContractError( - f"RoboCOIN({dataset_id}): trim list is STALE for {len(stale)} matching episode(s) " - f"(episode_index={stale_ids}): recorded total_frames disagrees with the full manifest " - "length. Refusing dataset construction because episode indices shift after physical " - "deletion and same-length collisions cannot be detected individually. Re-run " - "Regenerate the trim manifest against the current corpus before retrying." - ) - - return { - ep: (int(head), manifest_lengths[ep] if tail is None else int(tail)) for ep, (head, tail, _) in spec.items() - } - - -def _stats_population_spans( - dataset_dir, - *, - dataset_id: str, - split: str, - dataset_spec: dict, - excluded_episode_indices, - info: dict | None = None, -): - """Return physical kept spans plus a digest of the exact stats population. - - The order deliberately mirrors the reader: validate trims against the full - pre-split manifest, apply ``info.json`` split membership, drop exclusions, - then apply the head/tail span. The digest binds each surviving episode to - its physical global ``[start, end)`` rows, so changing a split range cannot - be hidden by merely leaving a top-level ``split: train`` label in an old - stats file. - """ - - root = Path(dataset_dir) - manifest = load_episodes_parquet(root) - trim_spans = _validate_trim_manifest(dataset_id, manifest, dataset_spec) - if info is None: - info = parse_info_json(root) - selected = apply_info_splits( - manifest, - split, - info.get("splits", {}) or {}, - source_name=f"RoboCOIN stats population({dataset_id})", - ) - selected_ids = set(int(ep) for ep in selected["episode_index"].to_numpy()) - try: - episode_ids = [int(ep) for ep in manifest["episode_index"].to_numpy()] - starts = [int(start) for start in manifest["dataset_from_index"].to_numpy()] - lengths = [int(length) for length in manifest["length"].to_numpy()] - except (KeyError, TypeError, ValueError) as exc: - raise DataContractError( - f"RoboCOIN({dataset_id}): full manifest must carry integer " - "episode_index, dataset_from_index, and length columns" - ) from exc - - excluded = set(int(ep) for ep in excluded_episode_indices) - raw_spans = [] - kept = [] - for episode_id, start, length in zip(episode_ids, starts, lengths): - if start < 0 or length < 0: - raise DataContractError( - f"RoboCOIN({dataset_id}): invalid manifest span for episode_index={episode_id}: " - f"dataset_from_index={start}, length={length}" - ) - raw_spans.append((start, start + length, episode_id)) - if episode_id not in selected_ids or episode_id in excluded: - continue - head, tail = trim_spans.get(episode_id, (0, length)) - if tail - head >= _TRIM_MIN_LEN: - kept.append((start + head, start + tail, episode_id)) - - raw_spans.sort() - for previous, current in zip(raw_spans, raw_spans[1:]): - if current[0] < previous[1]: - raise DataContractError( - f"RoboCOIN({dataset_id}): overlapping global manifest spans for " - f"episode_index={previous[2]} and {current[2]}" - ) - kept.sort() - hasher = hashlib.sha256(b"openwam:robocoin-stats-population:v1\0") - for start, end, episode_id in kept: - for value in (episode_id, start, end): - hasher.update(int(value).to_bytes(8, "little", signed=False)) - entry = { - "effective_population_digest": hasher.hexdigest(), - "num_episodes": len(kept), - "num_rows": sum(end - start for start, end, _episode_id in kept), - } - manifest_end = max((end for _start, end, _episode_id in raw_spans), default=0) - return [(start, end) for start, end, _episode_id in kept], manifest_end, entry - - def _resolve_robocoin_cameras(features: dict) -> tuple: """Resolve head / left-wrist / right-wrist camera keys from info.json features. @@ -833,13 +223,12 @@ class RoboCOINDataset(LeRobotV3Reader): overridden below. """ + STRICT_BUCKET_LOADING = True DATASET_NAME = "RoboCOIN" NEEDED_COLS = _NEEDED_COLS PROMPT_FILE_REQUIRED = False - CONFIG_KEYS = LeRobotV3Reader.CONFIG_KEYS + ("trim_csv", "_trim_snapshot") - WRIST_DECODE_TOLERATED = (Exception,) def __init__( @@ -848,8 +237,6 @@ def __init__( *, unify_action: bool = False, unify_action_map=None, - trim_csv=None, - _trim_snapshot=None, **kwargs, ): """Detect dexterous-hand + unify buckets and wire finger encoding. @@ -861,30 +248,8 @@ def __init__( ``unify_action_map`` that scatters fingers into the dedicated 80-D hand slots [10:34)/[44:68) (overriding any config map). Requires unify so the variable-width raw collapses to the uniform 80-D head. - - ``trim_csv``: optional path to a quality-audit trim list; see - :meth:`_filter_episodes`. ``None`` (default) disables trimming. """ - if is_robocoin_bucket_excluded(dataset_dir): - exclusion = ROBOCOIN_WHOLE_BUCKET_EXCLUSIONS[Path(dataset_dir).name] - raise DataContractError( - f"RoboCOIN bucket {Path(dataset_dir).name!r} is excluded as a whole: {exclusion['reason']}" - ) - self._trim_csv = trim_csv - trim_snapshot_was_provided = _trim_snapshot is not None - if _trim_snapshot is not None: - if trim_csv is None or not isinstance(_trim_snapshot, _TrimSnapshot): - raise ValueError("_trim_snapshot requires a matching non-null trim_csv") - if _trim_snapshot.path != str(trim_csv): - raise ValueError( - f"_trim_snapshot path {_trim_snapshot.path!r} does not match trim_csv {str(trim_csv)!r}" - ) - if _trim_snapshot is None and trim_csv is not None: - _trim_snapshot = _load_trim_snapshot(trim_csv) - - _trim_snapshot.population_validation_cache.clear() - self._trim_snapshot = _trim_snapshot self._dex_unify = False self._k_left = 0 self._k_right = 0 @@ -920,38 +285,10 @@ def __init__( MAX_HAND_DOF, ) super().__init__(dataset_dir, unify_action=unify_action, unify_action_map=unify_action_map, **kwargs) - if self._trim_snapshot is not None and not trim_snapshot_was_provided: - _assert_trim_snapshot_current(self._trim_snapshot, context="reader construction") - if self._trim_csv is not None: - exclusion_snapshots = getattr( - self, - "_stats_exclusion_snapshots", - (self._excluded_episodes_snapshot,), - ) - for exclusion_snapshot in exclusion_snapshots: - assert_excluded_episodes_snapshot_current( - exclusion_snapshot, - context="reader construction", - ) - if hasattr(self, "_stats_exclusion_snapshots"): - del self._stats_exclusion_snapshots - - self._trim_snapshot = None - - def _get_trim_snapshot(self): - """Pin one snapshot for every stage of this reader construction.""" - if self._trim_csv is None: - return None - snapshot = getattr(self, "_trim_snapshot", None) - if snapshot is None: - snapshot = _load_trim_snapshot(self._trim_csv) - self._trim_snapshot = snapshot - return snapshot def _load_excluded_episode_indices(self) -> set[int]: - """Pin the exact exclusion population consumed by filtering and stats.""" - self._excluded_episodes_snapshot = load_excluded_episodes_snapshot(self._dataset_dir) - return set(self._excluded_episodes_snapshot.episode_indices) + """The processed manifest is the complete training population.""" + return set() def _resolve_cameras(self, info: dict): """Resolve cameras by info.features priority, record robot_type, and @@ -980,250 +317,44 @@ def _resolve_cameras(self, info: dict): self.ACTION_DIM_MASK = GRIP_EXCLUDED_DIM_MASK return head, left_wrist, right_wrist - def _filter_episodes(self, eps_df): - """Apply optional head/tail trim rules without breaking alignment. - - Data-row and every camera-frame offset move by the same integer head - count, and length becomes the retained ``[head, tail)`` span. The - method never removes mid-episode pauses. A trim entry is accepted only - when its recorded manifest identity and frame count still match; the - same retained-population contract is required by stats generation. - """ - - eps_df = super()._filter_episodes(eps_df) - if self._trim_csv is None: - return eps_df - spec = self._get_trim_snapshot().spec.get(self._dataset_id) - if not spec: - if hasattr(self, "_trim_spans"): - del self._trim_spans - return eps_df - - trim_spans = getattr(self, "_trim_spans", None) - if trim_spans is None: - trim_spans = _validate_trim_manifest(self._dataset_id, eps_df, spec) - else: - del self._trim_spans - - episode_indices = eps_df["episode_index"].to_numpy() - lengths = eps_df["length"].to_numpy().copy() - - cam_cols = [c for c in eps_df.columns if c.startswith("_video_frame_offset/")] - - row_off = eps_df["_data_row_offset"].to_numpy().copy() - cam_off = {c: eps_df[c].to_numpy().copy() for c in cam_cols} - - min_len = _TRIM_MIN_LEN - - keep = np.ones(len(eps_df), dtype=bool) - n_trim = n_degenerate = 0 - frames_before = int(lengths.sum()) - for pos, ep in enumerate(episode_indices): - span = trim_spans.get(int(ep)) - if span is None: - continue - head, tail = span - length = int(lengths[pos]) - if head == 0 and tail == length: - continue - if tail - head < min_len: - keep[pos] = False - n_degenerate += 1 - continue - lengths[pos] = tail - head - if head: - row_off[pos] += head - for c in cam_cols: - cam_off[c][pos] += head - n_trim += 1 - - if n_degenerate: - logger.warning( - "RoboCOIN(%s): dropped %d episode(s) whose trim leaves < %d frames.", - self._dataset_id, - n_degenerate, - min_len, - ) - if not n_trim and not n_degenerate: - return eps_df - - eps_df = eps_df.copy() - eps_df["length"] = lengths - eps_df["_data_row_offset"] = row_off - for c in cam_cols: - eps_df[c] = cam_off[c] - eps_df = eps_df[keep].reset_index(drop=True) - after = int(lengths[keep].sum()) - logger.info( - "RoboCOIN(%s): trimmed %d and dropped %d/%d episodes, %d -> %d frames (-%.1f%%, %.2f h removed)", - self._dataset_id, - n_trim, - n_degenerate, - len(keep), - frames_before, - after, - 100.0 * (frames_before - after) / max(frames_before, 1), - (frames_before - after) / self._fps / 3600.0, - ) - return eps_df - - def _add_data_offsets(self, eps) -> None: - if self._trim_csv is not None: - spec = self._get_trim_snapshot().spec.get(self._dataset_id, {}) - self._trim_spans = _validate_trim_manifest(self._dataset_id, eps, spec) - - self._add_data_offsets_from_files(eps) - - def _add_data_offsets_from_files(self, eps): - """Resolve data parquet file + row offset from actual parquet row counts. - - Some RoboCOIN episode metadata has stale ``data/file_index`` values at - file boundaries. The parquet files themselves are globally contiguous - and ``dataset_from_index`` is correct, so use real file row counts as - the source of truth for data lookup. - - Reads parquet metadata in parallel to avoid a serial startup bottleneck - when a bucket has many parquet shards. - """ - - paths = _discover_data_parquets(self._dataset_dir, self._data_path_template) - - def _read_meta(entry): - path, chunk_index, file_index = entry - return (chunk_index, file_index, pq.ParquetFile(path).metadata.num_rows) - - if not paths: - raise FileNotFoundError(f"No data parquet files under {self._dataset_dir}/data") - - n_workers = min(len(paths), 4) - with ThreadPoolExecutor(max_workers=n_workers) as pool: - data_files = list(pool.map(_read_meta, paths)) - - starts = np.concatenate([[0], np.cumsum([n for _, _, n in data_files])]).astype(np.int64) - global_starts = eps["dataset_from_index"].to_numpy().astype(np.int64) - file_pos = np.searchsorted(starts, global_starts, side="right") - 1 - if (file_pos < 0).any() or (file_pos >= len(data_files)).any(): - raise ValueError(f"{self._dataset_id}: dataset_from_index outside data parquet row range") - - chunks = np.array([data_files[i][0] for i in file_pos], dtype=np.int64) - files = np.array([data_files[i][1] for i in file_pos], dtype=np.int64) - eps["data/chunk_index"] = chunks - eps["data/file_index"] = files - eps["_data_row_offset"] = global_starts - starts[file_pos] - def _load_stats(self, info: dict): - """Load per-robot-type stats (``../meta/stats_.json`` 'eef').""" + """Read the supplied family's stored EEF and optional hand statistics.""" if not self._normalize_mode or self._normalize_mode in ("none", "null"): return None - stats_path = self._dataset_dir.parent / "meta" / f"stats_{self._robot_type}.json" - if not stats_path.exists(): - raise FileNotFoundError( - f"normalize_mode={self._normalize_mode!r} but stats file is missing: {stats_path}. " - f"Run python -m openwam.dataloader.utils.stats_computation.robocoin_stats_computation to generate it, or set normalize_mode=null." - ) - with open(stats_path) as f: - raw = json.load(f) - _validate_whole_bucket_exclusions_provenance( - raw.get("whole_bucket_exclusions_provenance"), - robot_type=self._robot_type, - stats_path=stats_path, - ) - has_trim_provenance = "trim_provenance" in raw - has_excluded_provenance = "excluded_episodes_provenance" in raw - if self._trim_csv is None: - if has_trim_provenance or has_excluded_provenance: - raise DataContractError( - f"RoboCOIN bucket {self._dataset_id}: stats {stats_path} were computed for a " - "trimmed/filtered population but trim_csv is disabled. Regenerate legacy " - "untrimmed stats or configure the matching trim_csv." - ) - else: - expected_provenance = self._get_trim_snapshot().provenance - actual_provenance = raw.get("trim_provenance") - if not has_trim_provenance or actual_provenance != expected_provenance: - raise DataContractError( - f"RoboCOIN bucket {self._dataset_id}: stats {stats_path} trim_provenance " - f"does not exactly match trim_csv {self._trim_csv}; expected " - f"{expected_provenance}, got {actual_provenance}. Regenerate stats with the " - "configured trim CSV." - ) - exclusion_snapshots = _validate_excluded_episodes_provenance( - raw.get("excluded_episodes_provenance"), - dataset_dir=self._dataset_dir, - dataset_id=self._dataset_id, - current_snapshot=self._excluded_episodes_snapshot, - num_datasets=raw.get("eef", {}).get("num_datasets"), - stats_path=stats_path, - ) - _validate_stats_population_provenance( - raw.get("population"), - dataset_dir=self._dataset_dir, - current_dataset_id=self._dataset_id, - current_info=info, - trim_snapshot=self._get_trim_snapshot(), - exclusion_snapshots=exclusion_snapshots, - num_datasets=raw.get("eef", {}).get("num_datasets"), - stats_path=stats_path, + path = find_dataset_stats(self._dataset_dir, f"stats_{self._robot_type}.json", family_root=self._stats_root) + raw = _read_robot_stats(str(path)) + eef = raw.get("eef") + stored_robot_type = raw.get("robot_type", eef.get("robot_type") if isinstance(eef, dict) else None) + if stored_robot_type is not None and stored_robot_type != self._robot_type: + raise ValueError( + f"RoboCOIN stats {path} robot_type={stored_robot_type!r} does not match " + f"bucket {self._dataset_id} robot_type={self._robot_type!r}" ) - self._stats_exclusion_snapshots = tuple(exclusion_snapshots.values()) - eef_stats = materialize_eef_stats( - raw.get("eef", {}), - self._normalize_mode, - dim=_ACTION_DIM, - strict_minmax=False, - source_hint=f"{stats_path}: eef.* — re-run python -m openwam.dataloader.utils.stats_computation.robocoin_stats_computation", - force_rot6d_identity=True, + eef_stats = _materialize_robot_stats( + eef, dim=_ACTION_DIM, mode=self._normalize_mode, source=f"{path}:eef", rot6d=True ) - - for k in ("mean", "std", "min", "max", "q01", "q99"): - if eef_stats[k].shape[0] != _ACTION_DIM: - raise ValueError( - f"RoboCOIN bucket {self._dataset_id}: 'eef' stats '{k}' width " - f"{eef_stats[k].shape[0]} in {stats_path} != expected {_ACTION_DIM}. " - f"Re-run robocoin_stats_computation." - ) if not self._dex_unify: return eef_stats - hand_raw = raw.get("hand") - if not hand_raw: - raise FileNotFoundError( - f"unify_action + dexterous-hand bucket {self._dataset_id} needs a 'hand' stats " - f"block in {stats_path}; re-run robocoin_stats_computation (it now emits hand stats)." - ) - kL, kR = self._k_left, self._k_right - - dof_l, dof_r = hand_raw.get("dof_left"), hand_raw.get("dof_right") - if dof_l is None or dof_r is None: - raise ValueError( - f"unify_action + dexterous-hand bucket {self._dataset_id}: 'hand' stats block in " - f"{stats_path} is missing dof_left/dof_right (got left={dof_l}, right={dof_r}); " - f"re-run robocoin_stats_computation." - ) - if dof_l != kL or dof_r != kR: + hand = raw.get("hand") + if not isinstance(hand, dict): + raise ValueError(f"RoboCOIN dexterous-hand bucket {self._dataset_id} requires hand stats in {path}") + k_left, k_right = self._k_left, self._k_right + dof = (hand.get("dof_left"), hand.get("dof_right")) + if dof != (k_left, k_right): raise ValueError( - f"unify_action + dexterous-hand bucket {self._dataset_id}: 'hand' stats DOF " - f"(left={dof_l}, right={dof_r}) in {stats_path} != this bucket's finger DOF " - f"(left={kL}, right={kR}). The per-robot-type 'hand' block is locked to the first " - f"dataset's DOF; re-run robocoin_stats_computation (it now hard-fails on mixed DOF), " - f"or split mismatched datasets into distinct robot_types." + f"RoboCOIN hand stats {path} DOF left/right={dof} do not match " + f"bucket {self._dataset_id} finger DOF {(k_left, k_right)}" ) - hand_stats = materialize_eef_stats( - hand_raw, self._normalize_mode, dim=kL + kR, strict_minmax=False, source_hint=f"{stats_path}: hand.*" + hand_stats = _materialize_robot_stats( + hand, dim=k_left + k_right, mode=self._normalize_mode, source=f"{path}:hand" ) - - for k in ("mean", "std", "min", "max", "q01", "q99"): - if hand_stats[k].shape[0] != kL + kR: - raise ValueError( - f"unify_action + dexterous-hand bucket {self._dataset_id}: 'hand' stats '{k}' " - f"width {hand_stats[k].shape[0]} in {stats_path} != expected kL+kR={kL + kR}. " - f"Re-run robocoin_stats_computation." - ) - combined = {} - for k in ("mean", "std", "min", "max", "q01", "q99"): - e, h = eef_stats[k], hand_stats[k] - combined[k] = np.concatenate([e[0:9], h[0:kL], e[10:19], h[kL : kL + kR]]).astype(np.float32) - return combined + return { + key: np.concatenate( + [eef_stats[key][0:9], hand_stats[key][0:k_left], eef_stats[key][10:19], hand_stats[key][k_left:]] + ).astype(np.float32) + for key in STAT_KEYS + } def _normalize_array(self, arr: np.ndarray) -> np.ndarray: """Apply per-bucket normalization to a (..., 20) array. @@ -1285,84 +416,6 @@ def _proprio_20d(self, win) -> np.ndarray: def robot_type(self): return self._robot_type - @classmethod - def from_config(cls, config, split: str = "train"): - """Validate global QC inputs and exclude whole buckets before fan-out.""" - from openwam.dataloader.utils import get_cfg - - dataset_dir = get_cfg(config, "dataset_dir") - trim_csv = get_cfg(config, "trim_csv") - snapshot = None - if trim_csv is not None and dataset_dir is not None: - root = Path(dataset_dir) - if root.is_dir(): - snapshot = _load_trim_snapshot(trim_csv) - snapshot.population_validation_cache.clear() - spec = snapshot.spec - if not (root / "meta" / "info.json").is_file(): - bucket_names = { - path.name - for path in root.iterdir() - if ( - path.is_dir() - and (path / "meta" / "info.json").is_file() - and not is_robocoin_bucket_excluded(path) - ) - } - if bucket_names and bucket_names.isdisjoint(spec): - raise ValueError( - f"RoboCOIN trim_csv {trim_csv}: none of its {len(spec)} dataset key(s) " - f"match any bucket directory under {root}; trimming would silently no-op." - ) - if snapshot is not None: - config = _TrimSnapshotConfig(config, snapshot) - - root = Path(dataset_dir) if dataset_dir is not None else None - - if root is None or (root / "meta" / "info.json").is_file(): - dataset = super().from_config(config, split=split) - else: - if not root.is_dir(): - raise FileNotFoundError(f"{cls.__name__}: {root} does not exist") - discovered = sorted( - path for path in root.iterdir() if path.is_dir() and (path / "meta" / "info.json").is_file() - ) - excluded = [path for path in discovered if is_robocoin_bucket_excluded(path)] - sub_dirs = [path for path in discovered if not is_robocoin_bucket_excluded(path)] - if excluded: - logger.info( - "%s.from_config: excluding %d whole bucket(s): %s", - cls.__name__, - len(excluded), - ", ".join(path.name for path in excluded), - ) - if not sub_dirs: - raise FileNotFoundError(f"{cls.__name__}: no non-excluded sub-buckets with meta/info.json under {root}") - - missing = object() - common = {"split": split} - for key in cls.CONFIG_KEYS: - value = get_cfg(config, key, missing) - if value is missing: - continue - if value is None and key != "normalize_mode": - continue - common[key] = value - base_seed = int(get_cfg(config, "seed", 42)) - total_hours = get_cfg(config, "total_hours") - dataset = build_multibucket( - cls, - sub_dirs, - common, - base_seed=base_seed, - total_hours=total_hours, - wrapper_cls=MultiRobotCOINDataset, - source_name=cls.__name__, - ) - if snapshot is not None: - _assert_trim_snapshot_current(snapshot, context="from_config reader construction") - return dataset - @classmethod def _multibucket_wrapper(cls): return MultiRobotCOINDataset @@ -1391,10 +444,4 @@ def from_config(cls, config, split: str = "train"): return RoboCOINDataset.from_config(config, split) -__all__ = [ - "RoboCOINDataset", - "MultiRobotCOINDataset", - "ROBOCOIN_WHOLE_BUCKET_EXCLUSIONS", - "is_robocoin_bucket_excluded", - "robocoin_bucket_exclusions_provenance", -] +__all__ = ["RoboCOINDataset", "MultiRobotCOINDataset"] diff --git a/openwam/dataloader/utils/lerobotv3.py b/openwam/dataloader/utils/lerobotv3.py index d41758c53..914a25089 100644 --- a/openwam/dataloader/utils/lerobotv3.py +++ b/openwam/dataloader/utils/lerobotv3.py @@ -1,81 +1,17 @@ -"""Shared helpers for LeRobot v3 single-bucket readers. - -Both ``RoboCOINDataset`` and ``EgoDexDataset`` follow the same setup -recipe at init: read ``meta/info.json``, concatenate the per-shard -``meta/episodes/*.parquet``, sort by ``episode_index``, compute file-local -row + frame offsets via lexsort + groupby cumsum, then maybe trim by an -info-declared train/val split. This module hosts the identical pieces so -each reader's ``__init__`` only has to wire them together + do its -data-specific work (action conversion, multiview canvas, normalization). - -These helpers are intentionally stateless module-level functions; they -take whatever they need as arguments and return plain values. The readers -remain ``BaseDataset`` subclasses (no shared base class wrapped -around them) so the inheritance graph stays flat and easy to follow. - -Functions ---------- -- parse_info_json(dataset_dir) - Read ``/meta/info.json`` and return the dict along with - a couple of normalised defaults (``data_path`` and ``video_path`` - templates filled in when missing). - -- load_episodes_parquet(dataset_dir) - Concatenate every ``meta/episodes/*.parquet`` shard into a single - sorted pandas DataFrame, dropping any ``stats/*`` columns. - -- compute_file_local_offsets(eps, chunk_col, file_col) - Return an int64 numpy array of per-episode file-local cumulative - offsets, suitable as a column appended to ``eps`` (e.g. - ``_data_row_offset`` or ``_video_frame_offset``). - -- resolve_lerobot_v3_data_population(dataset_dir, ...) - Resolve the complete pre-split manifest into the exact data paths and - file-local episode ranges addressed by the reader. - -- digest_lerobot_v3_data_population(population) - Hash the effective data-path template and resolved episode-range mapping. - -- read_lerobot_v3_population_shard(dataset_dir, shard, columns) - Read only a manifest-addressed shard population and verify every physical - ``episode_index`` row against its declared range. - -- apply_info_splits(eps_df, split, info_splits, *, source_name=...) - Honor ``info.json["splits"][split]`` when present; otherwise default - train to the full eps_df and val to empty. Raises with a clear - message on malformed split specs. - -- water_fill_hours(bucket_hours, total_budget) - Allocate ``total_budget`` hours across N buckets via water-filling. - Each bucket gets at most its own ``bucket_hours[i]``; surplus from - under-budgetable buckets is redistributed among the rest. - -- effective_episode_frames(eps_df, length_col="length") - Validate and return each row's post-trim sampleable span. - -- subsample_episodes_by_hours(eps_df, target_hours, fps, seed, ...) - Random subset of ``eps_df`` rows totalling >= ``target_hours`` of effective - post-trim footage. ``_valid_start/_valid_end`` take precedence when present; - otherwise the (possibly trim-adjusted) ``length`` column is used. Seeded - permutation + greedy prefix; preserves original row order so downstream - offset columns stay valid. - -Exceptions ----------- -- DataContractError - Raised by a reader that has PROVEN its data is broken (as opposed to an - IO / environment fault). ``build_multibucket`` tolerates every other - per-bucket exception by dropping the bucket with a warning; this one it - re-raises, so broken data cannot silently shrink the training set. +"""Shared metadata, split and footage-budget helpers for LeRobot v3 readers. + +Readers consume canonical episode manifests and compute file-local offsets +before applying train/validation splits or whole-episode footage budgets. +Final pretraining exports need no additional frame-selection artifacts. +``build_multibucket`` preserves strict pretraining failures while allowing +benchmark readers to retain their existing per-bucket error policy. """ from __future__ import annotations -import hashlib import json import logging from concurrent.futures import ThreadPoolExecutor -from dataclasses import dataclass, field from pathlib import Path from typing import Any, Callable, Dict, List, Optional @@ -111,79 +47,6 @@ class DataContractError(RuntimeError): """ -@dataclass(frozen=True) -class ExcludedEpisodesSnapshot: - """Immutable effective population from one excluded-episodes file read.""" - - path: str - episode_indices: tuple[int, ...] - - -def load_excluded_episodes_snapshot(dataset_dir: Path) -> ExcludedEpisodesSnapshot: - """Load the canonical per-bucket episode blacklist as sorted unique IDs.""" - path = Path(dataset_dir) / "meta" / "excluded_episodes.json" - if not path.exists(): - return ExcludedEpisodesSnapshot(path=str(path), episode_indices=()) - try: - payload = json.loads(path.read_bytes()) - except (UnicodeError, json.JSONDecodeError) as e: - raise DataContractError(f"Malformed LeRobot episode exclusions {path}: {e}") from e - if not isinstance(payload, dict) or "episode_indices" not in payload: - raise DataContractError( - f"Malformed LeRobot episode exclusions {path}: expected an object with 'episode_indices'" - ) - values = payload["episode_indices"] - if not isinstance(values, list) or any(type(value) is not int or value < 0 for value in values): - raise DataContractError( - f"Malformed LeRobot episode exclusions {path}: 'episode_indices' must be a list of non-negative integers" - ) - return ExcludedEpisodesSnapshot(path=str(path), episode_indices=tuple(sorted(set(values)))) - - -def assert_excluded_episodes_snapshot_current( - snapshot: ExcludedEpisodesSnapshot, - *, - context: str, -) -> None: - """Fail if a bucket's effective exclusion population changed after loading.""" - actual = load_excluded_episodes_snapshot(Path(snapshot.path).parent.parent) - if actual.episode_indices != snapshot.episode_indices: - raise DataContractError( - f"RoboCOIN episode exclusions {snapshot.path} changed while {context}; " - f"expected {list(snapshot.episode_indices)}, got {list(actual.episode_indices)}" - ) - - -@dataclass(frozen=True) -class LeRobotV3DataRange: - """One episode range addressed inside a LeRobot v3 data shard.""" - - episode_index: int - dataset_from_index: int - row_offset: int - length: int - - -@dataclass(frozen=True) -class LeRobotV3DataShard: - """One logical data shard and the exact ranges addressed by its manifest.""" - - chunk_index: int - file_index: int - relative_path: str - ranges: tuple[LeRobotV3DataRange, ...] - - -@dataclass(frozen=True) -class LeRobotV3DataPopulation: - """The complete pre-split/pre-exclusion data population a reader addresses.""" - - data_path_template: str - total_rows: int - shards: tuple[LeRobotV3DataShard, ...] - episodes: pd.DataFrame = field(compare=False, repr=False) - - def parse_info_json(dataset_dir: Path) -> dict: """Read ``/meta/info.json`` and fill in default templates. @@ -242,192 +105,6 @@ def compute_file_local_offsets(eps: pd.DataFrame, chunk_col: str, file_col: str) return out -def _manifest_int_column(eps: pd.DataFrame, column: str, *, minimum: int) -> np.ndarray: - if column not in eps.columns: - raise ValueError(f"episodes manifest is missing {column!r}") - values = [] - for value in eps[column].tolist(): - if isinstance(value, (bool, np.bool_)) or not isinstance(value, (int, np.integer)): - raise ValueError(f"episodes manifest {column!r} must contain integers") - integer = int(value) - if integer < minimum: - raise ValueError(f"episodes manifest {column!r} values must be >= {minimum}") - values.append(integer) - try: - return np.asarray(values, dtype=np.int64) - except OverflowError as exc: - raise ValueError(f"episodes manifest {column!r} values must fit int64") from exc - - -def resolve_lerobot_v3_data_population( - dataset_dir: Path, - *, - info: dict | None = None, - episodes: pd.DataFrame | None = None, -) -> LeRobotV3DataPopulation: - """Resolve the exact pre-split data ranges addressed by a v3 reader.""" - root = Path(dataset_dir) - resolved_info = parse_info_json(root) if info is None else info - data_path_template = resolved_info.get("data_path") - if not isinstance(data_path_template, str) or not data_path_template: - raise ValueError("info.json data_path must be a non-empty string") - - eps = load_episodes_parquet(root) if episodes is None else episodes.copy() - if eps.empty: - raise ValueError("episodes manifest must contain at least one episode") - int_columns = { - "episode_index": _manifest_int_column(eps, "episode_index", minimum=0), - "length": _manifest_int_column(eps, "length", minimum=1), - "dataset_from_index": _manifest_int_column(eps, "dataset_from_index", minimum=0), - "data/chunk_index": _manifest_int_column(eps, "data/chunk_index", minimum=0), - "data/file_index": _manifest_int_column(eps, "data/file_index", minimum=0), - } - for column, values in int_columns.items(): - eps[column] = values - if eps["episode_index"].duplicated().any(): - raise ValueError("episodes manifest episode_index values must be unique") - if eps.duplicated(["data/chunk_index", "data/file_index", "dataset_from_index"]).any(): - raise ValueError("episodes manifest has duplicate per-shard dataset_from_index values") - - eps["_data_row_offset"] = compute_file_local_offsets( - eps, - "data/chunk_index", - "data/file_index", - ) - total_rows = int(eps["length"].sum()) - declared_total = resolved_info.get("total_frames") - if declared_total is not None: - if type(declared_total) is not int or declared_total < 0: - raise ValueError("info.json total_frames must be a non-negative integer") - if declared_total != total_rows: - raise ValueError( - f"episodes manifest addresses {total_rows} rows but info.json declares total_frames={declared_total}" - ) - - grouped: dict[tuple[int, int, str], list[LeRobotV3DataRange]] = {} - path_owners: dict[str, tuple[int, int]] = {} - for _, row in eps.iterrows(): - chunk_index = int(row["data/chunk_index"]) - file_index = int(row["data/file_index"]) - try: - rendered_path = data_path_template.format( - chunk_index=chunk_index, - file_index=file_index, - ) - except (IndexError, KeyError, ValueError) as exc: - raise ValueError(f"info.json data_path cannot resolve chunk/file indices ({exc})") from exc - relative = Path(rendered_path) - if relative.is_absolute() or ".." in relative.parts: - raise ValueError(f"info.json data_path resolved outside the dataset root: {rendered_path!r}") - relative_path = relative.as_posix() - owner = (chunk_index, file_index) - if relative_path in path_owners and path_owners[relative_path] != owner: - raise ValueError(f"info.json data_path maps multiple chunk/file pairs to {relative_path!r}") - path_owners[relative_path] = owner - grouped.setdefault((chunk_index, file_index, relative_path), []).append( - LeRobotV3DataRange( - episode_index=int(row["episode_index"]), - dataset_from_index=int(row["dataset_from_index"]), - row_offset=int(row["_data_row_offset"]), - length=int(row["length"]), - ) - ) - - shards = [] - for (chunk_index, file_index, relative_path), ranges in grouped.items(): - ordered_ranges = tuple(sorted(ranges, key=lambda item: item.row_offset)) - expected_offset = 0 - for episode_range in ordered_ranges: - if episode_range.row_offset != expected_offset: - raise ValueError(f"episodes manifest ranges for {relative_path} overlap or contain a gap") - expected_offset += episode_range.length - shards.append( - LeRobotV3DataShard( - chunk_index=chunk_index, - file_index=file_index, - relative_path=relative_path, - ranges=ordered_ranges, - ) - ) - shards.sort(key=lambda shard: shard.relative_path) - return LeRobotV3DataPopulation( - data_path_template=data_path_template, - total_rows=total_rows, - shards=tuple(shards), - episodes=eps, - ) - - -def _digest_length_prefixed(hasher, value: bytes) -> None: - hasher.update(len(value).to_bytes(8, "little")) - hasher.update(value) - - -def digest_lerobot_v3_data_population(population: LeRobotV3DataPopulation) -> str: - """Hash the effective data template and every addressed episode range.""" - hasher = hashlib.sha256(b"openwam:lerobot-v3-data-population:v1\0") - _digest_length_prefixed(hasher, population.data_path_template.encode("utf-8")) - hasher.update(population.total_rows.to_bytes(8, "little")) - hasher.update(len(population.shards).to_bytes(8, "little")) - for shard in population.shards: - _digest_length_prefixed(hasher, shard.relative_path.encode("utf-8")) - hasher.update(shard.chunk_index.to_bytes(8, "little")) - hasher.update(shard.file_index.to_bytes(8, "little")) - hasher.update(len(shard.ranges).to_bytes(8, "little")) - for episode_range in shard.ranges: - for value in ( - episode_range.episode_index, - episode_range.dataset_from_index, - episode_range.row_offset, - episode_range.length, - ): - hasher.update(value.to_bytes(8, "little")) - return hasher.hexdigest() - - -def read_lerobot_v3_population_shard( - dataset_dir: Path, - shard: LeRobotV3DataShard, - columns: List[str] | tuple[str, ...], -) -> pa.Table: - """Read and validate only the rows addressed by one logical shard.""" - requested = list(dict.fromkeys(columns)) - read_columns = list(dict.fromkeys(["episode_index", *requested])) - path = Path(dataset_dir) / shard.relative_path - try: - table = pq.read_table(path, memory_map=True, columns=read_columns) - except pa.ArrowInvalid as exc: - if "Dot path" not in str(exc): - raise - table = pq.read_table(path, memory_map=True).select(read_columns) - - slices = [] - for episode_range in shard.ranges: - stop = episode_range.row_offset + episode_range.length - if episode_range.row_offset < 0 or stop > table.num_rows: - raise ValueError( - f"manifest range for episode_index={episode_range.episode_index} exceeds {shard.relative_path} " - f"({episode_range.row_offset}:{stop} of {table.num_rows} rows)" - ) - episode_table = table.slice(episode_range.row_offset, episode_range.length) - if episode_table.num_rows != episode_range.length: - raise ValueError( - f"manifest range for episode_index={episode_range.episode_index} is truncated in {shard.relative_path}" - ) - episode_column = episode_table.column("episode_index").combine_chunks() - if episode_column.null_count: - raise ValueError(f"episode_index contains nulls in {shard.relative_path}") - physical_ids = episode_column.to_numpy(zero_copy_only=False) - if not np.all(physical_ids == episode_range.episode_index): - raise ValueError( - f"manifest episode_index={episode_range.episode_index} does not match every physical row in " - f"{shard.relative_path}[{episode_range.row_offset}:{stop}]" - ) - slices.append(episode_table) - population_table = pa.concat_tables(slices) if len(slices) > 1 else slices[0] - return population_table.select(requested) - - def apply_info_splits( eps_df: pd.DataFrame, split: str, @@ -521,37 +198,11 @@ def water_fill_hours(bucket_hours: List[float], total_budget: float) -> List[flo def effective_episode_frames(eps_df: pd.DataFrame, length_col: str = "length") -> np.ndarray: - """Return validated post-trim frame counts for each episode row. - - Readers that physically shift offsets during trimming also replace - ``length`` and therefore need no special handling. Segment-based readers - retain nominal ``length`` and provide the sampleable half-open range through - the paired ``_valid_start`` / ``_valid_end`` columns. - """ - nominal = eps_df[length_col].to_numpy(dtype=np.int64) - if np.any(nominal < 0): + """Return validated physical frame counts for each episode row.""" + lengths = eps_df[length_col].to_numpy(dtype=np.int64) + if np.any(lengths < 0): raise ValueError(f"episode {length_col!r} values must be non-negative") - if length_col != "length": - return nominal - - has_start = "_valid_start" in eps_df.columns - has_end = "_valid_end" in eps_df.columns - if has_start != has_end: - missing = "_valid_end" if has_start else "_valid_start" - raise ValueError(f"episode table has only one valid-range column; missing {missing!r}") - if not has_start: - return nominal - - valid_start = eps_df["_valid_start"].to_numpy(dtype=np.int64) - valid_end = eps_df["_valid_end"].to_numpy(dtype=np.int64) - invalid = (valid_start < 0) | (valid_end < valid_start) | (valid_end > nominal) - if np.any(invalid): - bad_pos = np.flatnonzero(invalid)[:5] - preview = ", ".join( - f"row {int(i)}: [{int(valid_start[i])}, {int(valid_end[i])}) / length {int(nominal[i])}" for i in bad_pos - ) - raise ValueError(f"episode valid ranges must satisfy 0 <= start <= end <= length ({preview})") - return valid_end - valid_start + return lengths def subsample_episodes_by_hours( @@ -565,11 +216,7 @@ def subsample_episodes_by_hours( """Random subset totalling >= ``target_hours`` of effective footage. Selection: seeded permutation over rows, then greedy prefix until - cumulative effective length >= target_hours * fps * 3600. When the - standard ``length`` column is requested, optional ``_valid_start`` and - ``_valid_end`` columns define the count; this is how segment-trimmed - datasets budget only frames that can actually be sampled. Readers that - implement trimming by mutating ``length`` need no special case. Returns + cumulative episode length >= target_hours * fps * 3600. Returns rows in their *original* eps_df order — important because downstream offset columns (``_data_row_offset`` / ``_video_frame_offset``) are computed per-row and must stay aligned. @@ -582,9 +229,7 @@ def subsample_episodes_by_hours( target_hours: Target total duration in hours. Must be > 0. fps: Frame rate (from info.json). seed: Random seed for the permutation. - length_col: Name of the nominal episode-length column. For the default - ``"length"``, ``_valid_start/_valid_end`` override its sampled span - when both are present. + length_col: Name of the physical episode-length column. episode_frames: Optional explicit non-negative frame count per row. Reader-level callers use this to zero episodes that cannot produce a window. When omitted, :func:`effective_episode_frames` is used. @@ -632,48 +277,6 @@ def subsample_episodes_by_hours( return eps_df.iloc[sorted(selected)].reset_index(drop=True) -def load_tasks_annotated(dataset_dir: Path, *, source_name: str = "dataset") -> Dict[int, str]: - """Load ``meta/tasks_annotated.parquet`` → ``{episode_index: text}``. - - OXE prompt source: per-episode LLM-rewritten descriptions indexed by - ``episode_index``. Raises FileNotFoundError when absent and ValueError on a - wrong index name. Warns (does not raise) if any entry is empty. - """ - ann_path = Path(dataset_dir) / "meta" / "tasks_annotated.parquet" - if not ann_path.exists(): - raise FileNotFoundError( - f"{ann_path} missing — OXE readers require tasks_annotated.parquet for prompts. " - "Re-run the upstream OXE conversion if this file is absent." - ) - ann_df = pd.read_parquet(ann_path) - if ann_df.index.name != "episode_index": - raise ValueError(f"{ann_path}: expected index 'episode_index', got {ann_df.index.name!r}") - txt_col = "task" if "task" in ann_df.columns else ann_df.columns[0] - mapping: Dict[int, str] = {int(ep_idx): str(text) for ep_idx, text in ann_df[txt_col].items()} - n_empty = sum(1 for t in mapping.values() if not t.strip()) - if n_empty: - logger.warning( - "%s: %d/%d annotated prompts are empty (unexpected — coverage was 100%% at conversion).", - source_name, - n_empty, - len(mapping), - ) - return mapping - - -def resolve_prompt_by_episode(episode_idx_to_text: Dict[int, str], episode_index: int, ds_name: str) -> str: - """OXE prompt resolver: look up per-episode text by ``episode_index``. - - Raises KeyError on a missing episode_index — a real data inconsistency - that should not be silently swallowed. - """ - if int(episode_index) not in episode_idx_to_text: - raise KeyError( - f"{ds_name} prompt lookup failed: episode_index={episode_index} not present in tasks_annotated.parquet." - ) - return episode_idx_to_text[int(episode_index)] - - def build_multibucket( reader_cls: type, sub_dirs: List[Path], @@ -684,18 +287,16 @@ def build_multibucket( wrapper_cls: type, source_name: str = "dataset", per_bucket_kwargs: Optional[Callable[[Path], Dict[str, Any]]] = None, + strict: bool = False, ) -> Any: """Construct N per-bucket readers from ``sub_dirs`` and wrap them. - Shared multi-bucket root-mode builder for RoboCOIN / EgoDex (and any future - LeRobot v3 reader family). Honors an optional ``total_hours`` water-fill - budget over the readers' effective post-split/exclusion/trim population. - Leaves are initialized uncapped once, their exact effective capacities are - water-filled, and the resulting whole-episode subsets are applied in place. - This avoids estimating from the raw manifest, which overstates capacity for - trimmed datasets. Construction is **robust**: a bucket whose construction - raises is skipped (logged), empty buckets are filtered, and only an - all-failed result raises. + Honors an optional ``total_hours`` water-fill budget over the readers' + available post-split population. Leaves are initialized uncapped once, + their sampleable capacities are water-filled, and the resulting + whole-episode subsets are applied in place. Empty buckets are filtered. + With ``strict=True``, construction failures abort the launch; otherwise + ordinary failures are logged and skipped. DataContractError always aborts. Args: reader_cls: the per-bucket reader class to instantiate. @@ -707,6 +308,8 @@ def build_multibucket( source_name: label for log lines. per_bucket_kwargs: optional callable ``(sub_dir) -> dict`` for extra per-bucket kwargs (unused by current readers; kept for extension). + strict: propagate construction failures instead of silently dropping a + bucket. Final pretraining data uses this to preserve full coverage. """ if total_hours is not None: total_hours = float(total_hours) @@ -729,6 +332,8 @@ def _build_one(args): # the reader's own diagnostic abort the launch. raise except Exception as e: + if strict: + raise logger.warning("%s: skipping %s: %s", source_name, sub.name, e) return None @@ -748,7 +353,7 @@ def _build_one(args): # loss. Keep each initialized leaf unchanged. logger.warning( "%s total_hours=%.3f exceeds available effective footage %.3fh; " - "loading the full post-split/exclusion/trim population.", + "loading the full post-split population.", source_name, total_hours, total_avail, @@ -770,11 +375,7 @@ def _build_one(args): buckets = [reader for _, _, reader in loaded if len(reader) > 0] if not buckets: raise RuntimeError(f"All {source_name} buckets became empty after effective-hour subsampling") - # _build_one swallows per-bucket construction errors (missing/corrupt data, - # and — for RoboCOIN — stats-integrity validation raised in _load_stats) into - # a warning + None, so a misconfigured bucket is dropped rather than aborting - # the run. Surface the dropped set explicitly so that silent data loss (e.g. a - # whole robot_type lost to a stale stats file) is visible, not buried. + # Report empty buckets and failures tolerated by non-strict benchmark readers. dropped = [sub.name for sub, r in zip(sub_dirs, results) if r is None or len(r) == 0] if dropped: logger.warning( @@ -797,7 +398,5 @@ def _build_one(args): "water_fill_hours", "effective_episode_frames", "subsample_episodes_by_hours", - "load_tasks_annotated", - "resolve_prompt_by_episode", "build_multibucket", ] diff --git a/openwam/dataloader/utils/normalization.py b/openwam/dataloader/utils/normalization.py index 0363d7a26..c64316308 100644 --- a/openwam/dataloader/utils/normalization.py +++ b/openwam/dataloader/utils/normalization.py @@ -16,6 +16,7 @@ import functools import json import logging +import os from pathlib import Path from typing import Mapping, Optional, Sequence @@ -29,6 +30,44 @@ # readers hand to ``_write_deploy_normalizer_stats``). STAT_KEYS = ("mean", "std", "min", "max", "q01", "q99") + +def find_dataset_stats( + dataset_dir: str | Path, + filename: str, + *, + family_root: str | Path | None = None, +) -> Path: + """Find supplied statistics without escaping the dataset family. + + Root-mode readers pass their configured family root. A standalone bucket + can infer the family from its nearest ancestor with a ``meta`` directory; + that directory is the boundary even when the requested file is missing. + The caller's lexical tree is retained so symlinked buckets can use the + statistics belonging to their configured family, rather than another + physical source tree. No statistics are created or recomputed. + """ + if not filename or Path(filename).name != filename: + raise ValueError("The statistics filename must be a single file name.") + root = Path(os.path.abspath(Path(dataset_dir).expanduser())) + if family_root is None: + boundary = next((p for p in root.parents if (p / "meta").is_dir()), root) + else: + boundary = Path(os.path.abspath(Path(family_root).expanduser())) + if root != boundary and boundary not in root.parents: + raise ValueError(f"Dataset bucket {root} is outside its statistics family {boundary}.") + for parent in (root, *root.parents): + candidate = parent / "meta" / filename + if candidate.is_file(): + return candidate + if parent == boundary: + break + raise FileNotFoundError( + f"Precomputed normalization statistics {filename!r} were not found " + f"between {root}/meta and dataset family {boundary}/meta. " + "Keep the supplied statistics with the dataset." + ) + + # rot6d dims within the two EEF stat layouts: # 10-D single-arm [pos(0:3), rot6d(3:9), grip(9)] (OXE) # 20-D bimanual [L_pos, L_rot6d(3:9), L_grip, R_pos, R_rot6d(13:19), R_grip] (RoboCOIN) @@ -135,7 +174,7 @@ def apply_normalization( if "q01" not in stats or "q99" not in stats: raise KeyError( "quantile normalization requires stats['q01'] and stats['q99']; " - "rerun the OXE compute_stats script to populate quantile fields." + "provide the dataset's precomputed statistics with both quantile fields." ) q01 = stats["q01"] q99 = stats["q99"] @@ -201,6 +240,18 @@ def materialize_eef_stats( "q99": np.array(raw.get("q99", [1.0] * dim), dtype=np.float32), } + # Check supplied values before identity pinning can hide malformed fields. + for key, values in out.items(): + if values.shape != (dim,): + raise ValueError(f"Normalization stats {key!r} must have shape ({dim},), got {values.shape}. {source_hint}") + if not np.isfinite(values).all(): + raise ValueError(f"Normalization stats {key!r} contain non-finite values. {source_hint}") + if np.any(out["std"] < 0): + raise ValueError(f"Normalization standard deviation must be non-negative. {source_hint}") + for lower, upper in (("min", "max"), ("q01", "q99")): + if np.any(out[lower] > out[upper]): + raise ValueError(f"Normalization stats require {lower} <= {upper}. {source_hint}") + rot6d_dims = _ROT6D_DIMS_BY_WIDTH.get(dim) if force_rot6d_identity and rot6d_dims is not None: pin_rot6d_identity(out, rot6d_dims) @@ -217,7 +268,7 @@ def materialize_eef_stats( logger.warning( "rot6d dims %s of this stats file are not identity under mode=%r — the file " "likely predates rot6d identity pinning and normalization WILL distort the " - "rotation representation. Rerun the matching *_stats_computation script. %s", + "rotation representation. Provide statistics matching the dataset's action representation. %s", idx, mode, source_hint, @@ -291,6 +342,7 @@ def _load_raw_stats(path: str, action_mode: Optional[str]) -> Mapping: __all__ = [ "apply_normalization", + "find_dataset_stats", "load_stats_file", "load_stats_metadata", "materialize_eef_stats", diff --git a/openwam/dataloader/utils/oxe_schema.py b/openwam/dataloader/utils/oxe_schema.py index daf301c4b..ebf3e354c 100644 --- a/openwam/dataloader/utils/oxe_schema.py +++ b/openwam/dataloader/utils/oxe_schema.py @@ -1,63 +1,30 @@ -"""OXE → 10-D EEF schema converters. - -Each OXE dataset has its own raw state/action representation: - - * BC-Z / Bridge: ``state = [x,y,z,roll,pitch,yaw,pad,gripper] (8)``, - ``action = [x,y,z,roll,pitch,yaw,gripper] (7)`` (Euler XYZ). - * Fractal: ``state = [x,y,z,rx,ry,rz,rw,gripper] (8)`` (quat xyzw), - ``action = [x,y,z,roll,pitch,yaw,gripper] (7)`` (Euler XYZ). - * DROID: arm-side wrist / gripper-mount pose plus a raw gripper - **closedness** signal (``0=open``, ``1=closed``). The re-converted bucket - stores the achieved pose and gripper scalar in separate columns, while its - commanded wrist stream stores both in one 7-D column. DROID-specific - converters assemble the two layouts and invert only the final scalar to the - canonical open scale (``0=closed``, ``1=open``). - -These helpers normalize all four data sources to a single 10-D EEF -representation ``[pos(3) + rot6d(6) + grip(1)]`` so the reader and -stats-computation paths can share code. - -The 10-D output is then passed through :func:`assemble_single_arm_left` -to slot into the canonical bimanual 20-D EEF schema. -""" +"""Canonical Euler/wrist-pose converters for DROID and VLABench.""" from __future__ import annotations import numpy as np -from openwam.dataloader.utils.eef import euler_xyz_to_rot6d, quat_xyzw_to_rot6d +from openwam.dataloader.utils.eef import euler_xyz_to_rot6d ARM10_DIM = 10 -# --------------------------------------------------------------------------- -# BC-Z / Bridge: state has a 'pad' slot at index 6 -# --------------------------------------------------------------------------- - - -def bcz_state_to_arm10(state: np.ndarray) -> np.ndarray: - """``(..., 8)`` BC-Z / Bridge state → ``(..., 10)`` EEF. - - Layout: ``state[:, [0:3, 3:6, 7]]`` (drops pad at index 6). - """ - pos = state[..., 0:3] - euler = state[..., 3:6] - grip = state[..., 7:8] - rot6d = euler_xyz_to_rot6d(euler) - return np.concatenate([pos, rot6d, grip], axis=-1).astype(np.float32) - - def euler7_action_to_arm10(action: np.ndarray) -> np.ndarray: """``(..., 7)`` ``[x,y,z,roll,pitch,yaw,gripper]`` → ``(..., 10)`` EEF. - Used by BC-Z / Bridge / Fractal action streams. Current DROID has the same + Used by VLABench pose streams. Current DROID has the same pose layout but opposite raw gripper direction, so it must use :func:`droid_euler7_to_arm10` instead. """ + action = np.asarray(action) + if action.ndim == 0 or action.shape[-1] != 7: + raise ValueError(f"Euler pose/gripper input must have width 7, got shape {action.shape}") + if not np.isfinite(action).all(): + raise ValueError("Euler pose/gripper input must contain only finite values") pos = action[..., 0:3] euler = action[..., 3:6] grip = action[..., 6:7] - rot6d = euler_xyz_to_rot6d(euler) + rot6d = euler_xyz_to_rot6d(euler.reshape(-1, 3)).reshape(action.shape[:-1] + (6,)) return np.concatenate([pos, rot6d, grip], axis=-1).astype(np.float32) @@ -90,9 +57,9 @@ def droid_pose6_closedness_to_arm10(pose: np.ndarray, closedness: np.ndarray) -> """ pose = np.asarray(pose) closedness = np.asarray(closedness) - if pose.shape[-1] != 6: + if pose.ndim == 0 or pose.shape[-1] != 6: raise ValueError(f"DROID gripper-mount pose must have width 6, got shape {pose.shape}") - if closedness.shape[-1] != 1: + if closedness.ndim == 0 or closedness.shape[-1] != 1: raise ValueError(f"DROID closedness must have width 1, got shape {closedness.shape}") if pose.shape[:-1] != closedness.shape[:-1]: raise ValueError( @@ -101,57 +68,4 @@ def droid_pose6_closedness_to_arm10(pose: np.ndarray, closedness: np.ndarray) -> return droid_euler7_to_arm10(np.concatenate([pose, closedness], axis=-1)) -# --------------------------------------------------------------------------- -# Fractal: state has quaternion in xyzw layout at indices 3:7 -# --------------------------------------------------------------------------- - - -def fractal_state_to_arm10(state: np.ndarray) -> np.ndarray: - """``(..., 8)`` Fractal state → ``(..., 10)`` EEF. - - Layout: ``state[:, [0:3, 3:7, 7]]`` — quat is xyzw, no pad. - """ - pos = state[..., 0:3] - quat = state[..., 3:7] - grip = state[..., 7:8] - rot6d = quat_xyzw_to_rot6d(quat) - return np.concatenate([pos, rot6d, grip], axis=-1).astype(np.float32) - - -# --------------------------------------------------------------------------- -# Legacy DROID conversion: state assembled from two separate parquet columns -# --------------------------------------------------------------------------- - - -def droid_state_to_arm10(cartesian: np.ndarray, gripper: np.ndarray) -> np.ndarray: - """Legacy two-column DROID state assembly. - - The current re-converted bucket assembles - ``other_information.observation_gripper_pose6d`` with ``state[6]`` through - :func:`droid_pose6_closedness_to_arm10`; this helper remains for old - conversion tools. - - Args: - cartesian: ``(..., 6)`` from ``observation.state.cartesian_position`` - (Euler XYZ representation: ``[x,y,z,roll,pitch,yaw]``). - gripper: ``(..., 1)`` from ``observation.state.gripper_position``; - DROID closedness in ``[0, 1]`` (``0=open``, ``1=closed``). - - Returns: ``(..., 10)`` EEF tensor. - """ - pos = cartesian[..., 0:3] - euler = cartesian[..., 3:6] - rot6d = euler_xyz_to_rot6d(euler) - openness = 1.0 - gripper - return np.concatenate([pos, rot6d, openness], axis=-1).astype(np.float32) - - -__all__ = [ - "ARM10_DIM", - "bcz_state_to_arm10", - "droid_euler7_to_arm10", - "droid_pose6_closedness_to_arm10", - "euler7_action_to_arm10", - "fractal_state_to_arm10", - "droid_state_to_arm10", -] +__all__ = ["ARM10_DIM", "euler7_action_to_arm10", "droid_euler7_to_arm10", "droid_pose6_closedness_to_arm10"] diff --git a/openwam/dataloader/utils/stats_computation/accumulator.py b/openwam/dataloader/utils/stats_computation/accumulator.py new file mode 100644 index 000000000..3bf412249 --- /dev/null +++ b/openwam/dataloader/utils/stats_computation/accumulator.py @@ -0,0 +1,126 @@ +"""Streaming statistics accumulator shared by benchmark statistics tools.""" + +import operator + +import numpy as np + +RESERVOIR_CAP = 1_000_000 + + +class Accumulator: + """Online mean/std/min/max accumulator + reservoir for q01/q99.""" + + def __init__(self, dim: int = 20, reservoir_cap: int = RESERVOIR_CAP, seed: int = 0): + try: + dim = operator.index(dim) + reservoir_cap = operator.index(reservoir_cap) + except TypeError as error: + raise ValueError("dim and reservoir_cap must be positive integers") from error + if dim < 1 or reservoir_cap < 1: + raise ValueError("dim and reservoir_cap must be positive integers") + self.dim = dim + self.count = 0 + self.mean = np.zeros(dim, dtype=np.float64) + self.m2 = np.zeros(dim, dtype=np.float64) + self.min_val = np.full(dim, np.inf, dtype=np.float64) + self.max_val = np.full(dim, -np.inf, dtype=np.float64) + + self.cap = int(reservoir_cap) + self.rng = np.random.RandomState(seed) + self._res = np.empty((self.cap, dim), dtype=np.float32) + self._res_n = 0 + self._res_seen = 0 + + def _validate_batch(self, batch: np.ndarray) -> np.ndarray: + batch = np.asarray(batch, dtype=np.float64) + if batch.ndim != 2 or batch.shape[1] != self.dim: + raise ValueError(f"expected a (N, {self.dim}) statistics batch, got {batch.shape}") + if not np.isfinite(batch).all() or np.any(np.abs(batch) > np.finfo(np.float32).max): + raise ValueError("statistics batches must contain finite float32-representable values") + return batch + + def update(self, batch: np.ndarray): + """Update with (N, dim) array using Welford's online algorithm.""" + batch = self._validate_batch(batch) + for i in range(len(batch)): + x = batch[i].astype(np.float64) + self.count += 1 + delta = x - self.mean + self.mean += delta / self.count + delta2 = x - self.mean + self.m2 += delta * delta2 + self.min_val = np.minimum(self.min_val, x) + self.max_val = np.maximum(self.max_val, x) + self._reservoir_add(np.asarray(batch, dtype=np.float32)) + + def update_batch(self, batch: np.ndarray): + """Batch update (more efficient for large arrays).""" + batch = self._validate_batch(batch) + n = len(batch) + if n == 0: + return + self._reservoir_add(np.asarray(batch, dtype=np.float32)) + batch = batch.astype(np.float64) + batch_mean = batch.mean(axis=0) + batch_var = batch.var(axis=0) + batch_min = batch.min(axis=0) + batch_max = batch.max(axis=0) + + if self.count == 0: + self.mean = batch_mean + self.m2 = batch_var * n + self.min_val = batch_min + self.max_val = batch_max + self.count = n + else: + total = self.count + n + delta = batch_mean - self.mean + new_mean = self.mean + delta * n / total + self.m2 = self.m2 + batch_var * n + delta**2 * self.count * n / total + self.mean = new_mean + self.count = total + self.min_val = np.minimum(self.min_val, batch_min) + self.max_val = np.maximum(self.max_val, batch_max) + + def _reservoir_add(self, batch: np.ndarray): + """Feed (N, dim) rows into the reservoir (vectorized Algorithm R).""" + n = len(batch) + if n == 0: + return + + if self._res_n < self.cap: + take = min(self.cap - self._res_n, n) + self._res[self._res_n : self._res_n + take] = batch[:take] + self._res_n += take + self._res_seen += take + batch = batch[take:] + if len(batch) == 0: + return + + m = len(batch) + t = self._res_seen + np.arange(m) + p = self.rng.randint(0, t + 1) + keep = p < self.cap + self._res[p[keep]] = batch[keep] + self._res_seen += m + + def finalize(self): + if self.count == 0: + raise ValueError("cannot finalize statistics without any rows") + std = np.sqrt(self.m2 / max(self.count, 1)) + std = np.where(std < 1e-8, 1.0, std) + if self._res_n > 0: + res = self._res[: self._res_n] + q01 = np.quantile(res, 0.01, axis=0) + q99 = np.quantile(res, 0.99, axis=0) + else: + q01 = self.min_val + q99 = self.max_val + return { + "mean": self.mean.astype(np.float32).tolist(), + "std": std.astype(np.float32).tolist(), + "min": self.min_val.astype(np.float32).tolist(), + "max": self.max_val.astype(np.float32).tolist(), + "q01": q01.astype(np.float32).tolist(), + "q99": q99.astype(np.float32).tolist(), + } diff --git a/openwam/dataloader/utils/stats_computation/agibotworld_stats_computation.py b/openwam/dataloader/utils/stats_computation/agibotworld_stats_computation.py deleted file mode 100644 index 7b358c851..000000000 --- a/openwam/dataloader/utils/stats_computation/agibotworld_stats_computation.py +++ /dev/null @@ -1,375 +0,0 @@ -#!/usr/bin/env python3 -"""Generate flavor-aware normalization stats for AgiBotWorld-Beta. - -One stats set is pooled across task buckets of the shared embodiment. Each -column is accumulated only where it carries signal: pose everywhere, gripper -for grippered buckets, fingers for dexterous-hand buckets, and base velocity -independently for moving action/proprio streams. Grippers are converted to -``0=closed, 1=open`` before accumulation. - -State spans include every retained row; action spans omit the unsupervised -episode-terminal target. Means, variances and extrema merge exactly, while -quantiles use a bounded uniform reservoir. The output records the manifest, -exclusion, and effective segment population consumed by the reader. -""" - -import argparse -import json -import os -from concurrent.futures import ProcessPoolExecutor -from pathlib import Path - -import numpy as np -import pyarrow as pa - -from openwam.dataloader.agibotworld import ( - _DEX_BUCKET_IDS, - _GRIPPER_CONTRACT, - _STATS_SCHEMA_VERSION, - _action_gripper_to_open_convention, - _apply_segment_annotations, - _bucket_base_motion_flags, - _effective_segment_population_digest, - _state_gripper_to_open_convention, - _validate_trim_ratio, -) -from openwam.dataloader.utils.lerobotv3 import ( - apply_info_splits, - assert_excluded_episodes_snapshot_current, - digest_lerobot_v3_data_population, - load_excluded_episodes_snapshot, - parse_info_json, - read_lerobot_v3_population_shard, - resolve_lerobot_v3_data_population, -) -from openwam.dataloader.utils.stats_computation.robocoin_stats_computation import Accumulator - -# Native column widths; flavor/motion selection happens per bucket. -_COL_WIDTH = { - "action.ee_base": 18, - "observation.state.ee_base": 18, - "action.gripper": 2, - "observation.state.gripper": 2, - "action.dex": 12, - "observation.state.dex": 12, - "action.robot_velocity": 3, - "observation.state.robot_velocity": 3, -} - -OUT_FILENAME = "stats_g2a.json" -WORKER_CAP = 200_000 -GLOBAL_CAP = 1_000_000 - - -def _columns_for(name: str, action_moving: bool, state_moving: bool): - """Which _COL_WIDTH keys this bucket contributes to (flavor + motion aware).""" - is_dex = name in _DEX_BUCKET_IDS - cols = ["action.ee_base", "observation.state.ee_base"] - if is_dex: - cols += ["action.dex", "observation.state.dex"] - else: - cols += ["action.gripper", "observation.state.gripper"] - if action_moving: - cols.append("action.robot_velocity") - if state_moving: - cols.append("observation.state.robot_velocity") - return cols - - -def _partial_bucket( - bucket_dir: str, - split: str = "train", - use_segment_annotations: bool = True, - segment_max_trim_ratio: float | None = 0.7, -): - """Worker: stream one reader-identical bucket population into partial stats.""" - name = os.path.basename(bucket_dir.rstrip("/")) - action_moving, state_moving = _bucket_base_motion_flags(bucket_dir) - want = {c: _COL_WIDTH[c] for c in _columns_for(name, action_moving, state_moving)} - accs = {c: Accumulator(dim=d, reservoir_cap=WORKER_CAP) for c, d in want.items()} - - root = Path(bucket_dir).resolve() - info = parse_info_json(root) - data_population = resolve_lerobot_v3_data_population(root, info=info) - eps = apply_info_splits( - data_population.episodes, - split, - info.get("splits", {}) or {}, - source_name=f"AgiBotWorld stats({name})", - ) - exclusion_snapshot = load_excluded_episodes_snapshot(root) - if exclusion_snapshot.episode_indices: - eps = eps[~eps["episode_index"].isin(exclusion_snapshot.episode_indices)].reset_index(drop=True) - eps, _summary = _apply_segment_annotations( - eps, - use_segment_annotations=use_segment_annotations, - segment_max_trim_ratio=segment_max_trim_ratio, - ) - kept = { - int(row["episode_index"]): ( - int(row.get("_valid_start", 0)), - int(row.get("_valid_end", row["length"])), - ) - for _, row in eps.iterrows() - } - sampleable = {ep: span for ep, span in kept.items() if span[1] - span[0] >= 2} - expected_state_rows = sum(end - start for start, end in sampleable.values()) - expected_action_rows = sum(end - start - 1 for start, end in sampleable.values()) - observed_state_rows = 0 - observed_action_rows = 0 - - action_cols = tuple(c for c in want if c.startswith("action.")) - state_cols = tuple(c for c in want if c.startswith("observation.state.")) - - def _update_columns(slices, cols): - if not slices: - return - kept_table = pa.concat_tables(slices) if len(slices) > 1 else slices[0] - df = kept_table.to_pandas() - for c in cols: - values = np.stack(df[c].values).astype(np.float32) - if c == "action.gripper": - values = _action_gripper_to_open_convention(values) - elif c == "observation.state.gripper": - values = _state_gripper_to_open_convention(values) - accs[c].update_batch(values) - - for shard in data_population.shards: - table = read_lerobot_v3_population_shard(root, shard, tuple(want)) - action_slices = [] - state_slices = [] - cursor = 0 - for episode_range in shard.ranges: - span = sampleable.get(episode_range.episode_index) - if span is not None: - start, end = span - if not 0 <= start < end <= episode_range.length: - raise ValueError( - f"AgiBotWorld stats({name}): invalid kept span {start}:{end} for " - f"episode_index={episode_range.episode_index} length={episode_range.length}" - ) - state_slices.append(table.slice(cursor + start, end - start)) - action_slices.append(table.slice(cursor + start, end - start - 1)) - observed_state_rows += end - start - observed_action_rows += end - start - 1 - cursor += episode_range.length - - _update_columns(action_slices, action_cols) - _update_columns(state_slices, state_cols) - - if observed_action_rows != expected_action_rows or observed_state_rows != expected_state_rows: - raise ValueError( - f"AgiBotWorld stats({name}): manifest scan observed action/state rows " - f"{observed_action_rows}/{observed_state_rows}, expected " - f"{expected_action_rows}/{expected_state_rows}" - ) - assert_excluded_episodes_snapshot_current(exclusion_snapshot, context=f"AgiBotWorld stats scan for {name}") - - out = {} - for c, acc in accs.items(): - if acc.count == 0: - continue - out[c] = { - "count": int(acc.count), - "mean": acc.mean.astype(np.float64), - "m2": acc.m2.astype(np.float64), - "min": acc.min_val.astype(np.float64), - "max": acc.max_val.astype(np.float64), - "res": acc._res[: acc._res_n].astype(np.float32), - } - population = { - "data_population_digest": digest_lerobot_v3_data_population(data_population), - "effective_population_digest": _effective_segment_population_digest(eps), - "num_episodes": int(len(eps)), - "num_sampleable_episodes": int(len(sampleable)), - "num_rows": int(expected_state_rows), - "num_action_rows": int(expected_action_rows), - "num_state_rows": int(expected_state_rows), - "excluded_episode_indices": list(exclusion_snapshot.episode_indices), - } - return name, out, population - - -def _merge_reservoir(r1, n1, r2, n2, rng): - """Uniform sample (≤ GLOBAL_CAP) of the union of two streams from their uniform - reservoirs r1 (over n1 rows) and r2 (over n2 rows), weighted by row count.""" - - if len(r1) == 0: - return r2[:GLOBAL_CAP].copy() - if len(r2) == 0: - return r1[:GLOBAL_CAP].copy() - size = min(GLOBAL_CAP, n1 + n2) - from_g = rng.random(size) < (n1 / (n1 + n2)) - idx1 = rng.integers(0, len(r1), size) - idx2 = rng.integers(0, len(r2), size) - out = np.where(from_g[:, None], r1[idx1], r2[idx2]) - return out.astype(np.float32) - - -def _merge_into(g: dict, partial: dict, rng): - """Merge one bucket's partial into the running global state (per column).""" - for c, p in partial.items(): - if c not in g: - g[c] = {k: p[k] for k in ("count", "mean", "m2", "min", "max", "res")} - continue - s = g[c] - n1, n2 = s["count"], p["count"] - n = n1 + n2 - delta = p["mean"] - s["mean"] - mean = s["mean"] + delta * n2 / n - m2 = s["m2"] + p["m2"] + delta * delta * n1 * n2 / n - g[c] = { - "count": n, - "mean": mean, - "m2": m2, - "min": np.minimum(s["min"], p["min"]), - "max": np.maximum(s["max"], p["max"]), - "res": _merge_reservoir(s["res"], n1, p["res"], n2, rng), - } - - -def _finalize(g: dict, contrib: dict): - """Global state → per-column stats dict ready to dump.""" - out = {} - for c, s in g.items(): - std = np.sqrt(s["m2"] / max(s["count"], 1)) - std = np.where(std < 1e-8, 1.0, std) - res = s["res"] - q01 = np.quantile(res, 0.01, axis=0) - q99 = np.quantile(res, 0.99, axis=0) - out[c] = { - "mean": s["mean"].astype(np.float32).tolist(), - "std": std.astype(np.float32).tolist(), - "min": s["min"].astype(np.float32).tolist(), - "max": s["max"].astype(np.float32).tolist(), - "q01": q01.astype(np.float32).tolist(), - "q99": q99.astype(np.float32).tolist(), - "num_timesteps": int(s["count"]), - "num_buckets": int(contrib.get(c, 0)), - } - return out - - -def discover_buckets(root: str) -> list: - return sorted( - os.path.join(root, d) for d in os.listdir(root) if os.path.isfile(os.path.join(root, d, "meta", "info.json")) - ) - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--dataset_dir", required=True) - parser.add_argument("--workers", type=int, default=64) - parser.add_argument("--seed", type=int, default=0) - parser.add_argument("--split", default="train") - parser.add_argument( - "--segment-max-trim-ratio", - type=float, - default=0.7, - help="Drop episodes whose segment annotation trims at least this fraction (reader default: 0.7)", - ) - parser.add_argument( - "--no-segment-annotations", - action="store_true", - help="Pool full segments instead of applying segment_flag/segment_delta", - ) - args = parser.parse_args() - segment_max_trim_ratio = _validate_trim_ratio(args.segment_max_trim_ratio) - use_segment_annotations = not args.no_segment_annotations - if not use_segment_annotations: - segment_max_trim_ratio = None - - buckets = discover_buckets(args.dataset_dir) - print(f"Unified g2a stats over {len(buckets)} buckets with {args.workers} workers...", flush=True) - rng = np.random.default_rng(args.seed) - g: dict = {} - contrib: dict = {} - bucket_population: dict = {} - contributed, skipped, failed = [], [], [] - done = 0 - with ProcessPoolExecutor(max_workers=args.workers) as pool: - futures = { - b: pool.submit( - _partial_bucket, - b, - args.split, - use_segment_annotations, - segment_max_trim_ratio, - ) - for b in buckets - } - for b in buckets: - name = os.path.basename(b.rstrip("/")) - done += 1 - try: - _, partial, population = futures[b].result() - except Exception as e: - failed.append(name) - print(f"[{done}/{len(buckets)}] {name}: FAILED — {type(e).__name__}: {e}", flush=True) - continue - if not partial: - skipped.append(name) - print(f"[{done}/{len(buckets)}] {name}: SKIPPED — no data rows found", flush=True) - continue - contributed.append(name) - bucket_population[name] = population - _merge_into(g, partial, rng) - for c in partial: - contrib[c] = contrib.get(c, 0) + 1 - if done % 20 == 0 or done == len(buckets): - print(f"[{done}/{len(buckets)}] merged {name}: cols={sorted(partial)}", flush=True) - - if not contributed: - raise RuntimeError(f"No bucket contributed any data under {args.dataset_dir}; nothing to write.") - if skipped or failed: - raise RuntimeError( - f"Refusing to write partial AgiBotWorld stats: contributed={len(contributed)} " - f"skipped={sorted(skipped)} failed={sorted(failed)}" - ) - - result = { - "robot_type": "g2a", - "gripper_contract": _GRIPPER_CONTRACT, - "num_buckets": len(contributed), - "num_skipped": len(skipped), - "num_failed": len(failed), - "population": { - "schema_version": _STATS_SCHEMA_VERSION, - "split": args.split, - "use_segment_annotations": use_segment_annotations, - "segment_max_trim_ratio": segment_max_trim_ratio, - "buckets": bucket_population, - }, - } - result.update(_finalize(g, contrib)) - - out_dir = os.path.join(args.dataset_dir, "meta") - os.makedirs(out_dir, exist_ok=True) - out_path = os.path.join(out_dir, OUT_FILENAME) - tmp_path = f"{out_path}.{os.getpid()}.tmp" - with open(tmp_path, "w") as f: - json.dump(result, f, indent=2) - os.replace(tmp_path, out_path) - - print("\n" + "=" * 60, flush=True) - print(f"Wrote {out_path}", flush=True) - print( - f" contributed {len(contributed)} / skipped {len(skipped)} / failed {len(failed)} of {len(buckets)} buckets", - flush=True, - ) - if skipped: - print(f" SKIPPED (no data): {sorted(skipped)}", flush=True) - if failed: - print(f" FAILED: {sorted(failed)}", flush=True) - - for c in sorted(k for k, value in result.items() if isinstance(value, dict) and "num_timesteps" in value): - s = result[c] - print( - f" {c:34s} n={s['num_timesteps']:>12,} buckets={s['num_buckets']:>3} " - f"q01[0]={s['q01'][0]:+.3f} q99[0]={s['q99'][0]:+.3f}", - flush=True, - ) - - -if __name__ == "__main__": - main() diff --git a/openwam/dataloader/utils/stats_computation/ebench_stats_computation.py b/openwam/dataloader/utils/stats_computation/ebench_stats_computation.py index 716a724c1..5f8ae95b8 100644 --- a/openwam/dataloader/utils/stats_computation/ebench_stats_computation.py +++ b/openwam/dataloader/utils/stats_computation/ebench_stats_computation.py @@ -36,7 +36,7 @@ ) from openwam.dataloader.utils.eef import assert_unit_quaternion from openwam.dataloader.utils.normalization import ROT6D_DIMS_EEF20, pin_rot6d_identity -from openwam.dataloader.utils.stats_computation.robocoin_stats_computation import Accumulator +from openwam.dataloader.utils.stats_computation.accumulator import Accumulator def _validated_raw23(frame: pd.DataFrame, action_keys: Sequence[str], ctx: str) -> np.ndarray: diff --git a/openwam/dataloader/utils/stats_computation/interndata_a1_stats_computation.py b/openwam/dataloader/utils/stats_computation/interndata_a1_stats_computation.py deleted file mode 100644 index f3f97c430..000000000 --- a/openwam/dataloader/utils/stats_computation/interndata_a1_stats_computation.py +++ /dev/null @@ -1,640 +0,0 @@ -"""Generate per-embodiment 20-D EEF stats for InternData-A1. - -The generator uses the reader's schema detection, quaternion conversion, -gripper scaling, exclusions, split selection, and trim bounds so the emitted -vectors match training exactly. Both action and state streams contribute. -Stats are pooled per embodiment, not per task bucket, and rot6d plus any -single-arm padding dimensions are pinned to identity. - -Output files live under ``/meta/stats_.json`` and carry -a certificate for the exact source and retained population. -""" - -from __future__ import annotations - -import argparse -import json -import logging -from concurrent.futures import ProcessPoolExecutor -from pathlib import Path -from typing import Dict, List, Optional, Sequence, Tuple - -import numpy as np -import pyarrow.parquet as pq - -from openwam.dataloader.interndata_a1 import ( - _BIMANUAL_SIDES, - _SINGLE_ARM_SIDES, - _assert_trim_snapshot_current, - _load_trim_snapshot, - detect_arm_layout, - discover_a1_buckets, - effective_a1_population_provenance, - embodiment_key, - iter_data_shards, - load_excluded_episodes, - resolve_bucket_key, - resolve_gripper_scale, - resolve_trim_bounds, - validate_manifest_ranges, -) -from openwam.dataloader.utils.eef import ARM10_DIM, quat_wxyz_to_rot6d -from openwam.dataloader.utils.lerobotv3 import ( - apply_info_splits, - load_episodes_parquet, - parse_info_json, -) -from openwam.dataloader.utils.normalization import ROT6D_DIMS_EEF20, pin_rot6d_identity -from openwam.dataloader.utils.stats_computation.robocoin_stats_computation import Accumulator - -logger = logging.getLogger(__name__) - -EEF20_DIM = 20 - -# Single-arm readers reserve the right-arm half as masked zero padding. -RIGHT_ARM_DIMS_EEF20: Tuple[int, ...] = tuple(range(ARM10_DIM, EEF20_DIM)) - - -# Reuse reader-owned column tables so stats and runtime cannot drift. -_SIDES: Dict[str, Dict[str, Sequence]] = { - "bimanual": _BIMANUAL_SIDES, - "single_arm": _SINGLE_ARM_SIDES, -} - - -def _arm10(table, spec, grip_scale: float) -> np.ndarray: - """Build one arm's ``(N, 10)`` ``[xyz, rot6d, grip]`` block from a parquet table.""" - pose_col, grip_col = spec - pose = np.asarray(table[pose_col].to_pylist(), dtype=np.float32) - grip = np.asarray(table[grip_col].to_pylist(), dtype=np.float32).reshape(len(pose), 1) / grip_scale - rot6d = quat_wxyz_to_rot6d(pose[:, 3:7]) - return np.concatenate([pose[:, 0:3], rot6d, grip], axis=-1) - - -def _eef20(table, sides, kind: str, grip_scales) -> np.ndarray: - """Assemble the ``(N, 20)`` EEF for one stream; right half zero for single-arm. - - Must stay bit-identical to ``InternDataA1Dataset._eef20`` — including the - per-bucket gripper rescale — or the stats describe a different distribution - than the reader actually emits. - """ - - left_spec, right_spec = sides[kind] - n = table.num_rows - out = np.zeros((n, EEF20_DIM), dtype=np.float32) - out[:, :ARM10_DIM] = _arm10(table, left_spec, grip_scales[0]) - if right_spec is not None: - out[:, ARM10_DIM:] = _arm10(table, right_spec, grip_scales[1]) - return out - - -def _bucket_info(bucket: Path) -> Tuple[str, str, str]: - """Return ``(embodiment, robot_type, arm_layout)`` for a bucket.""" - with open(bucket / "meta" / "info.json") as f: - info = json.load(f) - layout = detect_arm_layout(info.get("features", {}) or {}) - robot_type = info.get("robot_type", "unknown") - return embodiment_key(robot_type, layout), robot_type, layout - - -def classify_buckets(buckets: Sequence[Path]) -> Dict[str, Dict]: - """Group every discovered bucket by embodiment, failing on unreadable metadata. - - A shared per-embodiment stats file is safe only when every discovered bucket - is either scanned successfully or is empty by construction under the chosen - split. Silently omitting a bucket whose ``info.json`` cannot be classified - would create the same partial-normalizer failure as dropping a bucket later - in the parquet scan, only before the worker pool has a chance to report it. - """ - - groups: Dict[str, Dict] = {} - for b in buckets: - try: - emb, robot_type, layout = _bucket_info(b) - except Exception as e: - raise RuntimeError( - f"Refusing to compute partial InternData-A1 stats: discovered bucket {b} " - f"has unreadable meta/info.json ({type(e).__name__}: {e})" - ) from e - g = groups.setdefault(emb, {"robot_type": robot_type, "arm_layout": layout, "dirs": []}) - g["dirs"].append(b) - return groups - - -def _kept_episodes(bucket: Path) -> Optional[set]: - """Episode indices the reader will actually emit for this bucket. - - ``meta/episodes`` minus ``meta/excluded_episodes.json``. Both halves matter - for a cleaned view: ``data/`` (and ``meta/episodes`` itself) are symlinks to - the source, so the parquet shards physically contain every deleted episode. - Deletions live only in the exclusion list — which is precisely why the - cleaned view cannot express them by dropping manifest rows: the reader - derives file-local offsets as a cumsum over surviving rows, so a shortened - manifest would slide every later episode onto the wrong frames. - - Returns ``None`` only when the manifest is unreadable — the caller treats - that as "cannot filter". An **empty set is meaningful and different**: it - means the manifest was read and nothing survives, so every row must be - excluded. Conflating the two would silently pool a fully-deleted bucket's - rows into the statistics. - - Raises when the exclusion list exists but cannot be parsed: continuing would - quietly include episodes the reader drops, and the caller already skips - (and logs) a bucket that raises. - """ - - eps = _manifest_episodes(bucket) - if eps is None: - return None - - return eps - load_excluded_episodes(bucket) - - -def _manifest_episodes(bucket: Path) -> Optional[set]: - """Every episode index listed in ``meta/episodes``, or ``None`` if unreadable.""" - files = sorted((bucket / "meta" / "episodes").rglob("*.parquet")) - if not files: - return None - eps: set = set() - try: - for f in files: - eps.update(int(x) for x in pq.read_table(f, columns=["episode_index"]).to_pydict()["episode_index"]) - except (OSError, KeyError, ValueError): - return None - return eps - - -def _assert_reader_can_load(bucket: Path, shards: Sequence) -> None: - """Refuse a bucket the reader would refuse, for the same reasons. - - ``len(rows) > 0`` proves rows were read; it does not prove a reader can - consume them. A shard truncated by one row, a manifest whose ranges overlap, - a missing middle shard — each of those still yields plenty of rows here - while :meth:`InternDataA1Dataset._add_data_offsets` raises. Statistics - pooled from a bucket no reader can open describe a population that will - never be trained on, and nothing downstream distinguishes them. - - Deliberately the reader's checks, called through the reader's functions, - rather than a parallel implementation of the same idea — a second - implementation is how the two sides diverged in the first place. - """ - - files = sorted((bucket / "meta" / "episodes").rglob("*.parquet")) - if not files: - return - cols = ["episode_index", "dataset_from_index", "dataset_to_index", "length"] - frm, to, length, epi = [], [], [], [] - for f in files: - names = set(pq.ParquetFile(f).schema_arrow.names) - if not set(cols) <= names: - return - d = pq.read_table(f, columns=cols).to_pydict() - epi += d["episode_index"] - frm += d["dataset_from_index"] - to += d["dataset_to_index"] - length += d["length"] - - validate_manifest_ranges(frm, to, length, epi, str(bucket)) - - total = sum(pq.ParquetFile(p).metadata.num_rows for _, _, p in shards) - manifest_end = int(max(to)) - if manifest_end != total: - raise ValueError( - f"{bucket}: the data shards hold {total} rows but the manifest ends at " - f"{manifest_end}. The reader refuses this bucket, so statistics pooled from " - "it would describe rows that are never trained on." - ) - - -def _split_episodes(bucket: Path, split: str) -> Optional[set]: - """Episode indices belonging to ``split``, or ``None`` when info.json has no splits. - - The reader applies `info.json`'s split ranges before anything else, so a - corpus whose splits actually partition the episodes would otherwise have its - val rows pooled into the train normalizer. A1's shipped splits happen to - cover every episode, which is exactly why this gap could sit unnoticed — - the equivalence is a property of the current data, not of the code. - - Resolved through the reader's own `apply_info_splits` rather than by parsing - the range here, so the two cannot disagree about what a split string means. - """ - - try: - with open(bucket / "meta" / "info.json") as f: - splits = json.load(f).get("splits") - except (OSError, ValueError) as e: - raise ValueError(f"{bucket}: cannot read meta/info.json for split resolution ({e})") from e - - eps = _manifest_episodes(bucket) - if eps is None: - raise ValueError(f"{bucket}: meta/episodes is unreadable, so the population is unknown") - if not splits: - return None if split == "train" else set() - import pandas as pd - - from openwam.dataloader.utils.lerobotv3 import apply_info_splits - - df = pd.DataFrame({"episode_index": sorted(eps)}) - try: - return set(apply_info_splits(df, split, splits, source_name=str(bucket))["episode_index"]) - except Exception as e: - raise ValueError(f"{bucket}: unusable splits {splits!r} ({e})") from e - - -def _row_mask(table, kept: Optional[set], trim: Optional[Dict[int, Tuple]], min_len: int) -> Optional[np.ndarray]: - """Which rows belong in the statistics: kept episodes, minus trimmed head/tail. - - ``kept is None`` means "manifest unreadable, cannot filter"; an **empty set - excludes everything** (see :func:`_kept_episodes`). Testing ``not kept`` - would collapse those two into "no filter" and pool a fully-deleted bucket - back in. - - Returns ``None`` only when there is genuinely nothing to filter, so the - plain path stays allocation-free. - """ - - if kept is None and not trim: - return None - ep = np.asarray(table.column("episode_index").to_pylist(), dtype=np.int64) - mask = np.ones(ep.shape[0], dtype=bool) - uniq = np.unique(ep) - - if kept is not None and not set(uniq.tolist()) <= kept: - mask &= ( - np.isin(ep, np.fromiter(kept, dtype=np.int64, count=len(kept))) - if kept - else np.zeros(ep.shape[0], dtype=bool) - ) - - if trim: - order = np.argsort(ep, kind="stable") - bounds = np.flatnonzero(np.diff(ep[order])) + 1 - for g in np.split(order, bounds): - entry = trim.get(int(ep[g[0]])) - if entry is None: - continue - n = g.shape[0] - - b = resolve_trim_bounds(entry, n, min_len) - if b is None: - continue - head, tail = b - if head: - mask[g[:head]] = False - if tail < n: - mask[g[tail:]] = False - return mask - - -def _scan_bucket(args) -> Tuple[str, np.ndarray, bool, dict]: - """Read one bucket's parquet shards and return its stacked (N, 20) rows. - - Runs in a worker process; returns the raw rows rather than a per-bucket - Accumulator so the parent can merge them into one reservoir in a fixed - (submission) order — see ``compute_stats_for_embodiment``. - - Per-bucket Accumulators are not ruled out by determinism — a fixed-order, - seeded, weighted reservoir merge would be deterministic too. They are ruled - out because no such merge exists on ``Accumulator`` and writing a - statistically correct one is not worth it here. Worth knowing if the parent's - memory ever needs work: it currently holds every bucket's array until the - pool drains. - """ - - bucket_str, layout, embodiment = args[:3] - dataset_id = args[3] if len(args) > 3 else None - trim_csv = args[4] if len(args) > 4 else None - min_len = args[5] if len(args) > 5 else 2 - split = args[6] if len(args) > 6 else "train" - bucket = Path(bucket_str) - sides = _SIDES[layout] - cols = set() - for kind in ("action", "state"): - for spec in sides[kind]: - if spec is not None: - cols.update(spec) - - grip_scales = tuple( - resolve_gripper_scale(bucket, embodiment, spec[1]) if spec is not None else 1.0 for spec in sides["state"] - ) - - trim = None - trim_snapshot = None - if trim_csv: - trim_snapshot = _load_trim_snapshot(trim_csv) - key = resolve_bucket_key( - trim_snapshot.spec, dataset_id or bucket.name, bucket, what="trim_csv", source=f"stats({dataset_id})" - ) - trim = trim_snapshot.spec.get(key) if key is not None else None - - manifest = load_episodes_parquet(bucket) - info = parse_info_json(bucket) - selected = apply_info_splits( - manifest, - split, - info.get("splits", {}) or {}, - source_name=f"InternDataA1 stats({dataset_id or bucket.name})", - ) - excluded_snapshot = tuple(sorted(load_excluded_episodes(bucket))) - if excluded_snapshot: - selected = selected[~selected["episode_index"].isin(excluded_snapshot)].reset_index(drop=True) - kept = set(int(ep) for ep in selected["episode_index"].to_numpy()) - bucket_provenance = { - "excluded_episode_indices": list(excluded_snapshot), - "effective_population": effective_a1_population_provenance(selected, trim, min_len), - } - need_ep = bool(trim) or kept is not None - - population_empty = kept is not None and len(kept) == 0 - - chunks: List[np.ndarray] = [] - - shards = iter_data_shards(bucket) - if not shards: - raise ValueError(f"{bucket}: no data shards under data/chunk-*/file-*.parquet") - - _assert_reader_can_load(bucket, shards) - - for _, _, pth in shards: - names = set(pq.ParquetFile(pth).schema_arrow.names) - use_ep = need_ep and "episode_index" in names - if need_ep and not use_ep: - raise ValueError( - f"{pth}: trim/exclusion filtering requested but the shard has no " - "`episode_index` column, so the rows cannot be selected." - ) - table = pq.read_table(pth, columns=sorted(cols | {"episode_index"}) if use_ep else sorted(cols)) - if table.num_rows == 0: - continue - mask = _row_mask(table, kept, trim, min_len) if use_ep else None - for kind in ("action", "state"): - rows = _eef20(table, sides, kind, grip_scales) - chunks.append(rows if mask is None else rows[mask]) - out = np.zeros((0, EEF20_DIM), dtype=np.float32) if not chunks else np.concatenate(chunks, axis=0) - - if len(out) == 0 and not population_empty: - raise ValueError( - f"{bucket}: the {split!r} split expects " - f"{'every episode' if kept is None else f'{len(kept)} episode(s)'} but no shard row " - "carries them. The shards do not hold the episodes the manifest assigns to this " - "split — pooling this as an empty population would let its reader load another " - "bucket's statistics and read another episode's rows." - ) - current_exclusions = tuple(sorted(load_excluded_episodes(bucket))) - if current_exclusions != excluded_snapshot: - raise ValueError( - f"{bucket}: excluded_episodes.json changed during stats scan; " - f"expected {list(excluded_snapshot)}, got {list(current_exclusions)}" - ) - if trim_snapshot is not None: - _assert_trim_snapshot_current(trim_snapshot, context=f"stats scan for {dataset_id or bucket.name}") - return bucket_str, out, bool(population_empty), bucket_provenance - - -def compute_stats_for_embodiment( - embodiment: str, - group: Dict, - *, - rot6d_identity: bool = True, - workers: int = 16, - root: Optional[Path] = None, - trim_csv: Optional[str] = None, - min_len: int = 2, - split: str = "train", -) -> dict: - """Pool every bucket of one embodiment into a single 20-D stats dict. - - ``trim_csv`` keys on the bucket path RELATIVE to ``root`` — the same id the - reader uses (bucket basenames repeat across embodiments). Without ``root`` - the relative id cannot be formed, so trimming is skipped rather than matched - on an ambiguous basename. - """ - - dirs: List[Path] = group["dirs"] - layout = group["arm_layout"] - acc = Accumulator(dim=EEF20_DIM) - n_rows = 0 - n_ok = 0 - - def _rel_id(d: Path) -> str: - """Bucket key, in the exact form the reader resolves against. - - `--dataset_dir` pointing straight at a bucket is a supported mode, and - there `relative_to(root)` is `'.'` — a key no trim CSV contains and one - the reader (whose id is then the bare directory name) cannot - suffix-match. Left as-is it silently skips the trim on the stats side - while the reader applies it, and writes `exclusions: {".": ...}` that the - reader rejects outright. - """ - - if root is None: - return d.name - try: - rel = str(d.relative_to(root)) - except ValueError: - return d.name - return d.name if rel == "." else rel - - if trim_csv and root is None: - logger.warning("trim_csv given without a root; bucket ids are ambiguous, not trimming.") - trim_csv = None - trim_snapshot = _load_trim_snapshot(trim_csv) if trim_csv else None - tasks = [(str(d), layout, embodiment, _rel_id(d), trim_csv, min_len, split) for d in dirs] - - scanned_buckets: List[str] = [] - empty_buckets: List[str] = [] - bucket_provenance: Dict[str, dict] = {} - scan_failures: List[Tuple[str, str]] = [] - with ProcessPoolExecutor(max_workers=min(workers, max(1, len(tasks)))) as pool: - futures = {pool.submit(_scan_bucket, t): t[0] for t in tasks} - - for i, fut in enumerate(futures, 1): - name = futures[fut] - try: - _, rows, population_empty, provenance = fut.result() - except Exception as e: - rel_id = _rel_id(Path(name)) - scan_failures.append((rel_id, f"{type(e).__name__}: {e}")) - logger.error(" [%s] failed %s (%s)", embodiment, name, e) - continue - if not len(rows): - if not population_empty: - rel_id = _rel_id(Path(name)) - message = "produced 0 rows without an empty-by-construction population" - scan_failures.append((rel_id, message)) - logger.error(" [%s] %s %s", embodiment, name, message) - continue - logger.info(" [%s] %s has no %s rows (val-only?)", embodiment, name, split) - rel_id = _rel_id(Path(name)) - empty_buckets.append(rel_id) - bucket_provenance[rel_id] = provenance - continue - rel_id = _rel_id(Path(name)) - scanned_buckets.append(rel_id) - bucket_provenance[rel_id] = provenance - acc.update_batch(rows) - n_rows += len(rows) - n_ok += 1 - if i % 20 == 0 or i == len(tasks): - logger.info(" [%s] %d/%d buckets, %d rows", embodiment, i, len(tasks), n_rows) - - if scan_failures: - details = "; ".join(f"{bucket}: {error}" for bucket, error in scan_failures) - raise RuntimeError( - f"Refusing to write partial InternData-A1 stats for embodiment {embodiment!r}: " - f"{len(scan_failures)} non-empty bucket scan(s) failed: {details}" - ) - if n_rows == 0: - raise RuntimeError(f"embodiment {embodiment!r}: every bucket yielded 0 rows") - - population = { - "schema_version": 2, - "split": split, - "trim_active": bool(trim_csv), - "min_keep": int(min_len), - "trim_provenance": trim_snapshot.provenance if trim_snapshot is not None else None, - "bucket_provenance": {key: bucket_provenance[key] for key in sorted(bucket_provenance)}, - "buckets": sorted(scanned_buckets), - "empty_buckets": sorted(empty_buckets), - } - if trim_snapshot is not None: - _assert_trim_snapshot_current(trim_snapshot, context=f"stats scan for embodiment {embodiment}") - - stats = acc.finalize() - if rot6d_identity: - pin_rot6d_identity(stats, ROT6D_DIMS_EEF20) - if layout == "single_arm": - pin_rot6d_identity(stats, RIGHT_ARM_DIMS_EEF20) - return { - "eef": stats, - "robot_type": group["robot_type"], - "arm_layout": layout, - "embodiment": embodiment, - "num_buckets": n_ok, - "num_rows": int(n_rows), - "population": population, - "scanned_buckets": population["buckets"], - "split": population["split"], - "rot6d_identity": bool(rot6d_identity), - "layout_doc": "[L_xyz(0:3), L_rot6d(3:9), L_grip(9), R_xyz(10:13), R_rot6d(13:19), R_grip(19)]", - } - - -def main(): - parser = argparse.ArgumentParser(description=__doc__.split("\n")[0]) - parser.add_argument("--dataset_dir", required=True, help="extracted InternData-A1 v3.0 root") - parser.add_argument( - "--stats_root", - default=None, - help="where to write meta/stats_.json (default: --dataset_dir). Set this to a " - "writable directory when the dataset mount is read-only, and pass the SAME path as " - "dataloader.stats_root in interndata_a1.yaml — the reader resolves the file as " - "{stats_root}/meta/stats_{embodiment}.json.", - ) - parser.add_argument("--embodiment", default=None, help="compute for a single embodiment only") - parser.add_argument( - "--trim_csv", - default=None, - help="quality-audit trim list (same file the dataloader takes). Head/tail frames it " - "names are excluded from the statistics, so the normalizer describes what the reader " - "actually feeds the model instead of including motionless frames it skips. Episodes " - "listed in a bucket's meta/excluded_episodes.json are dropped regardless of this " - "flag: a cleaned view symlinks data/ (and meta/episodes) at the source, so deleted " - "episodes are still physically present in the parquet and listed in the manifest.", - ) - parser.add_argument( - "--min_keep", - type=int, - default=2, - help="leave an episode untrimmed when the trim would leave fewer than this " - "many frames. MUST match the reader's minimum window length for the split " - "being trained (InternDataA1Dataset._train_min_window_len() == 2 for train; " - "num_frames for val) — otherwise the stats describe episodes the reader " - "never emits that way.", - ) - parser.add_argument( - "--split", - default="train", - help="which info.json split to pool. Normalization stats must come from the TRAINING " - "distribution, so this defaults to train and should rarely be changed — it exists so " - "the scan applies the same split ranges the reader does, rather than pooling every " - "episode on disk.", - ) - parser.add_argument("--workers", type=int, default=16, help="parallel bucket readers") - parser.add_argument( - "--no-rot6d-identity", - action="store_true", - help="do NOT pin rot6d (and single-arm padding) dims to identity — they then " - "normalize per-dim like pos/gripper, which distorts the rotation " - "representation; see pin_rot6d_identity.", - ) - args = parser.parse_args() - logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") - - root = Path(args.dataset_dir) - buckets = discover_a1_buckets(root) - if not buckets: - raise SystemExit(f"no buckets with meta/info.json under {root}") - logger.info("discovered %d buckets under %s", len(buckets), root) - - groups = classify_buckets(buckets) - logger.info("embodiments: %s", {k: len(v["dirs"]) for k, v in sorted(groups.items())}) - if args.embodiment: - if args.embodiment not in groups: - raise SystemExit(f"embodiment {args.embodiment!r} not found; have {sorted(groups)}") - groups = {args.embodiment: groups[args.embodiment]} - - stats_root = Path(args.stats_root) if args.stats_root else root - out_dir = stats_root / "meta" - if stats_root != root: - logger.info("writing stats to %s (dataset root %s left untouched)", out_dir, root) - results = {} - for emb in sorted(groups): - logger.info("=== %s (%d buckets) ===", emb, len(groups[emb]["dirs"])) - result = compute_stats_for_embodiment( - emb, - groups[emb], - rot6d_identity=not args.no_rot6d_identity, - workers=args.workers, - root=root, - trim_csv=args.trim_csv, - min_len=args.min_keep, - split=args.split, - ) - - result["trim_csv"] = args.trim_csv - result["trim_min_keep"] = args.min_keep if args.trim_csv else None - results[emb] = result - - out_dir.mkdir(parents=True, exist_ok=True) - for emb in sorted(results): - result = results[emb] - out = out_dir / f"stats_{emb}.json" - - tmp = out.with_suffix(".json.tmp") - with open(tmp, "w") as f: - json.dump(result, f, indent=1) - tmp.replace(out) - eef = result["eef"] - logger.info( - "wrote %s (%d rows) | L_grip q01=%.4f q99=%.4f | R_grip q01=%.4f q99=%.4f", - out, - result["num_rows"], - eef["q01"][9], - eef["q99"][9], - eef["q01"][19], - eef["q99"][19], - ) - logger.info("done: %s", sorted(f"stats_{e}.json" for e in groups)) - - -if __name__ == "__main__": - main() - - -__all__ = [ - "EEF20_DIM", - "RIGHT_ARM_DIMS_EEF20", - "classify_buckets", - "compute_stats_for_embodiment", -] diff --git a/openwam/dataloader/utils/stats_computation/muka_franka_stats_computation.py b/openwam/dataloader/utils/stats_computation/muka_franka_stats_computation.py deleted file mode 100644 index d8e68ffe7..000000000 --- a/openwam/dataloader/utils/stats_computation/muka_franka_stats_computation.py +++ /dev/null @@ -1,145 +0,0 @@ -"""Compute MUKA Franka EEF10 action/state normalization statistics. - -State rows and real next-state action targets are pooled after the exact reader -conversion (Euler XYZ -> rot6d and closedness -> ``[-1 closed, +1 open]``). -The repeated terminal action row of each episode is excluded. rot6d dimensions -3..8 are pinned to identity for every normalization mode. - -Example:: - - python -m openwam.dataloader.utils.stats_computation.muka_franka_stats_computation \ - --config configs/dataloader/pretrain_data/muka_franka.yaml \ - --output /path/to/muka_franka_lerobot_v3/meta/normalization_stats.npy -""" - -from __future__ import annotations - -import argparse -import os -import socket -import uuid -from pathlib import Path - -import numpy as np -from omegaconf import OmegaConf - -from openwam.dataloader.muka_franka import ( - GRIPPER_CONVENTION, - GRIPPER_TRANSFORM, - RAW_GRIPPER_CONVENTION, - ROT6D_DIMS_EEF10, - ROTATION_CONVENTION, - STATS_POPULATION, - MukaFrankaDataset, -) -from openwam.dataloader.utils.normalization import pin_rot6d_identity -from openwam.dataloader.utils.stats_computation.robocoin_stats_computation import Accumulator - - -def _iter_bucket_arrays(bucket: MukaFrankaDataset): - """Yield per-episode ``(real_action_eef10, state_eef10)`` arrays.""" - for pos, (_, row) in enumerate(bucket._eps_df.iterrows()): # noqa: SLF001 - table = bucket._load_data_table( # noqa: SLF001 - int(row["data/chunk_index"]), int(row["data/file_index"]) - ) - offset = int(bucket._ep_data_row_offset[pos]) # noqa: SLF001 - win = table.slice(offset, int(row["length"])).to_pandas() - action = bucket._raw_action_eef10(win)[:-1] # noqa: SLF001 - terminal row is repeated. - state = bucket._raw_state_eef10(win) # noqa: SLF001 - yield action, state - - -def _compute_global_stats(dataset: MukaFrankaDataset, reservoir_cap: int): - if not isinstance(dataset, MukaFrankaDataset): - raise TypeError(f"expected MukaFrankaDataset, got {type(dataset).__name__}") - raw_dim = dataset._raw_action_dim # noqa: SLF001 - accumulator = Accumulator(dim=raw_dim, reservoir_cap=reservoir_cap) - action_rows = 0 - state_rows = 0 - for action, state in _iter_bucket_arrays(dataset): - action = np.asarray(action, np.float32).reshape(-1, raw_dim) - state = np.asarray(state, np.float32).reshape(-1, raw_dim) - accumulator.update_batch(action) - accumulator.update_batch(state) - action_rows += action.shape[0] - state_rows += state.shape[0] - if action_rows == 0 or state_rows == 0: - raise ValueError("cannot compute MUKA Franka stats from an empty dataset") - - stats = accumulator.finalize() - stats.update( - { - "num_timesteps": accumulator.count, - "pool": "real_action_state", - "action_rows": action_rows, - "state_rows": state_rows, - "gripper_convention": GRIPPER_CONVENTION, - "raw_gripper_convention": RAW_GRIPPER_CONVENTION, - "gripper_transform": GRIPPER_TRANSFORM, - "rotation_convention": ROTATION_CONVENTION, - "rotation_transform": "euler_xyz_to_rot6d", - "action_alignment": "action[t]=state[t+1]; repeated terminal action excluded", - "stats_population": STATS_POPULATION, - "split": dataset._split, # noqa: SLF001 - "num_episodes": len(dataset._eps_df), # noqa: SLF001 - } - ) - pin_rot6d_identity(stats, ROT6D_DIMS_EEF10) - return dataset.action_mode, raw_dim, stats, action_rows, state_rows - - -def build_and_save_muka_franka_stats( - dataset: MukaFrankaDataset, - output: str | Path, - reservoir_cap: int = 1_000_000, -): - """Compute stats and atomically merge the EEF block into ``output``.""" - output = Path(output) - action_mode, raw_dim, global_stats, action_rows, state_rows = _compute_global_stats(dataset, reservoir_cap) - payload = {} - if output.exists(): - try: - previous = np.load(output, allow_pickle=True).item() - if isinstance(previous, dict): - payload.update(previous) - except (ValueError, EOFError): - pass - payload[action_mode] = global_stats - output.parent.mkdir(parents=True, exist_ok=True) - tmp_path = output.with_name(f".{output.name}.{socket.gethostname()}.{os.getpid()}.{uuid.uuid4().hex[:8]}.tmp") - try: - with tmp_path.open("wb") as handle: - np.save(handle, payload) - os.replace(tmp_path, output) - finally: - if tmp_path.exists(): - tmp_path.unlink() - return action_mode, raw_dim, action_rows, state_rows - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--config", default="configs/dataloader/pretrain_data/muka_franka.yaml") - parser.add_argument("--output", required=True, help="Deploy-compatible .npy output") - parser.add_argument("--reservoir-cap", type=int, default=1_000_000) - args = parser.parse_args() - - output = Path(args.output) - if output.suffix != ".npy": - raise ValueError("--output must end in .npy") - cfg = OmegaConf.load(args.config) - OmegaConf.update(cfg, "normalize_mode", None, merge=False) - dataset = MukaFrankaDataset.from_config(cfg, split=str(OmegaConf.select(cfg, "split", default="train"))) - action_mode, raw_dim, action_rows, state_rows = build_and_save_muka_franka_stats( - dataset, output, args.reservoir_cap - ) - print( - f"wrote {output} mode={action_mode} pool=real_action_state dim={raw_dim} " - f"action_rows={action_rows} state_rows={state_rows} " - f"total_rows={action_rows + state_rows}" - ) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/openwam/dataloader/utils/stats_computation/oxe_stats_computation.py b/openwam/dataloader/utils/stats_computation/oxe_stats_computation.py deleted file mode 100644 index dac20b678..000000000 --- a/openwam/dataloader/utils/stats_computation/oxe_stats_computation.py +++ /dev/null @@ -1,323 +0,0 @@ -"""Compute per-dataset 10-D EEF normalization stats for OXE readers. - -Each schema converts state and action to ``xyz(3) + rot6d(6) + grip(1)`` and -pools the two streams into ``meta/eef_stats.json``. For DROID, the scan uses -the arm-side wrist/gripper-mount streams, applies the canonical prompt -exclusions, inverts gripper closedness, and records data-population provenance. -Rot6d dimensions are pinned to identity before output. -""" - -from __future__ import annotations - -import argparse -import json -import logging -from pathlib import Path -from typing import Dict, List, Tuple - -import numpy as np -import pyarrow as pa -import pyarrow.parquet as pq - -from openwam.dataloader.oxe_droid import ( - DROID_DATA_POPULATION_DIGEST_KEY, - DROID_EEF_STATS_CONTRACT, - DROID_EEF_STATS_CONTRACT_KEY, - DROID_STATS_POPULATION_KEY, - load_droid_prompt_exclusions, - resolve_droid_stats_population, -) -from openwam.dataloader.utils.eef import assert_unit_quaternion -from openwam.dataloader.utils.lerobotv3 import ( - digest_lerobot_v3_data_population, - parse_info_json, - read_lerobot_v3_population_shard, - resolve_lerobot_v3_data_population, -) -from openwam.dataloader.utils.normalization import ROT6D_DIMS_ARM10, pin_rot6d_identity -from openwam.dataloader.utils.oxe_schema import ( - bcz_state_to_arm10, - droid_euler7_to_arm10, - droid_pose6_closedness_to_arm10, - droid_state_to_arm10, - euler7_action_to_arm10, - fractal_state_to_arm10, -) - -logging.basicConfig(format="%(asctime)s [%(levelname)s] %(message)s", level=logging.INFO) -logger = logging.getLogger("oxe_stats_computation") - - -SCHEMA: Dict[str, Dict] = { - "BC-Z": { - "state_cols": ["observation.state"], - "action_cols": ["action"], - "state_fn": "bcz_state", - "action_fn": "euler7_action", - }, - "Bridge": { - "state_cols": ["observation.state"], - "action_cols": ["action"], - "state_fn": "bcz_state", - "action_fn": "euler7_action", - }, - "Fractal": { - "state_cols": ["observation.state"], - "action_cols": ["action"], - "state_fn": "fractal_state", - "action_fn": "euler7_action", - }, - "DROID": { - "state_cols": ["other_information.observation_gripper_pose6d", "state"], - "action_cols": ["other_information.action_wrist_pose"], - "state_fn": "droid_gripper_pose6", - "action_fn": "droid_euler7", - }, -} - - -def _convert_state(rows: Dict[str, np.ndarray], state_fn: str) -> np.ndarray: - if state_fn == "bcz_state": - return bcz_state_to_arm10(rows["observation.state"]) - if state_fn == "fractal_state": - quat = rows["observation.state"][:, 3:7] - assert_unit_quaternion(quat, tol=0.05, sample_n=min(64, len(quat))) - return fractal_state_to_arm10(rows["observation.state"]) - if state_fn == "droid_state": - return droid_state_to_arm10( - rows["observation.state.cartesian_position"], - rows["observation.state.gripper_position"], - ) - if state_fn == "euler7_state": - return euler7_action_to_arm10(rows[list(rows.keys())[0]]) - if state_fn == "droid_euler7": - return droid_euler7_to_arm10(rows[list(rows.keys())[0]]) - if state_fn == "droid_gripper_pose6": - return droid_pose6_closedness_to_arm10( - rows["other_information.observation_gripper_pose6d"], - rows["state"][:, 6:7], - ) - raise ValueError(f"unknown state_fn={state_fn}") - - -def _convert_action(rows: Dict[str, np.ndarray], action_fn: str) -> np.ndarray: - if action_fn == "euler7_action": - return euler7_action_to_arm10(rows[list(rows.keys())[0]]) - if action_fn == "droid_euler7": - return droid_euler7_to_arm10(rows[list(rows.keys())[0]]) - raise ValueError(f"unknown action_fn={action_fn}") - - -def _table_rows(table: pa.Table, cols: List[str]) -> Dict[str, np.ndarray]: - """Convert selected Arrow columns to stacked numpy rows.""" - out: Dict[str, np.ndarray] = {} - for c in cols: - col_data = table.column(c).to_pylist() - - if col_data and not isinstance(col_data[0], (list, np.ndarray)): - out[c] = np.asarray(col_data, dtype=np.float32).reshape(-1, 1) - else: - out[c] = np.asarray(col_data, dtype=np.float32) - return out - - -def _load_shard(path: Path, cols: List[str]) -> Dict[str, np.ndarray]: - """Load one parquet shard, returning a dict {col: ndarray-of-stacked-rows}.""" - return _table_rows(pq.read_table(path, memory_map=True, columns=cols), cols) - - -def compute_dataset_stats( - dataset_dir: Path, - dataset_name: str, - rot6d_identity: bool = True, - *, - split: str = "train", -) -> Tuple[dict, int, int]: - """Walk a dataset's data parquets, convert to 10-D EEF, aggregate stats. - - Returns: - (stats_dict, n_state_samples, n_action_samples) - """ - - spec = SCHEMA[dataset_name] - excluded_episode_indices: set[int] = set() - droid_population = None - droid_stats_population = None - droid_stats_episode_indices = None - if dataset_name == "DROID": - info = parse_info_json(dataset_dir) - droid_population = resolve_lerobot_v3_data_population(dataset_dir, info=info) - _, excluded_episode_indices = load_droid_prompt_exclusions( - dataset_dir, - population=droid_population, - ) - selected, droid_stats_population = resolve_droid_stats_population( - droid_population, - info, - excluded_episode_indices, - split=split, - ) - droid_stats_episode_indices = selected["episode_index"].to_numpy(dtype=np.int64, copy=False) - shard_inputs = list(droid_population.shards) - logger.info("%s: scanning %d manifest-addressed data shards", dataset_name, len(shard_inputs)) - else: - shard_inputs = sorted((dataset_dir / "data").rglob("*.parquet")) - if not shard_inputs: - raise FileNotFoundError(f"No parquet shards under {dataset_dir}/data") - logger.info("%s: scanning %d parquet shards under %s/data", dataset_name, len(shard_inputs), dataset_dir) - - state_arrs: List[np.ndarray] = [] - action_arrs: List[np.ndarray] = [] - for i, shard_input in enumerate(shard_inputs, start=1): - if droid_population is not None: - columns = list(dict.fromkeys(["episode_index", *spec["state_cols"], *spec["action_cols"]])) - table = read_lerobot_v3_population_shard(dataset_dir, shard_input, columns) - episode_indices = table.column("episode_index").combine_chunks().to_numpy(zero_copy_only=False) - keep = np.isin(episode_indices, droid_stats_episode_indices) - if not keep.any(): - continue - table = table.filter(pa.array(keep)) - state_rows = _table_rows(table, spec["state_cols"]) - action_rows = _table_rows(table, spec["action_cols"]) - else: - p = shard_input - state_rows = _load_shard(p, spec["state_cols"]) - action_rows = _load_shard(p, spec["action_cols"]) - if not state_rows or not action_rows: - continue - state10 = _convert_state(state_rows, spec["state_fn"]) - action10 = _convert_action(action_rows, spec["action_fn"]) - state_arrs.append(state10) - action_arrs.append(action10) - if i % 50 == 0 or i == len(shard_inputs): - logger.info(" %s: processed %d/%d shards", dataset_name, i, len(shard_inputs)) - - if not state_arrs: - raise ValueError(f"{dataset_name}: every parquet row is excluded; cannot compute stats") - state_all = np.concatenate(state_arrs, axis=0) - action_all = np.concatenate(action_arrs, axis=0) - n_state = int(len(state_all)) - n_action = int(len(action_all)) - - merged = np.concatenate([state_all, action_all], axis=0) - logger.info( - "%s: merged %d state + %d action rows = %d total samples for stats", - dataset_name, - n_state, - n_action, - len(merged), - ) - - stats = { - "n_samples": int(len(merged)), - "n_state_samples": n_state, - "n_action_samples": n_action, - "min": merged.min(axis=0).astype(np.float64).tolist(), - "max": merged.max(axis=0).astype(np.float64).tolist(), - "mean": merged.mean(axis=0, dtype=np.float64).tolist(), - "std": merged.std(axis=0, dtype=np.float64).tolist(), - "q01": np.quantile(merged, 0.01, axis=0).astype(np.float64).tolist(), - "q99": np.quantile(merged, 0.99, axis=0).astype(np.float64).tolist(), - } - if dataset_name == "DROID": - if n_state != droid_stats_population["num_rows"] or n_action != droid_stats_population["num_rows"]: - raise ValueError( - f"DROID stats scan selected {n_state} state/{n_action} action rows but the " - f"effective train population declares {droid_stats_population['num_rows']} rows" - ) - stats["excluded_episode_indices"] = sorted(excluded_episode_indices) - stats[DROID_DATA_POPULATION_DIGEST_KEY] = digest_lerobot_v3_data_population(droid_population) - stats[DROID_STATS_POPULATION_KEY] = droid_stats_population - stats[DROID_EEF_STATS_CONTRACT_KEY] = dict(DROID_EEF_STATS_CONTRACT) - if rot6d_identity: - pin_rot6d_identity(stats, ROT6D_DIMS_ARM10) - return stats, n_state, n_action - - -def _print_stats_table(stats: dict, name: str) -> None: - """Per-dim summary for human eyeballing.""" - dim_names = ["x", "y", "z", "r6_0", "r6_1", "r6_2", "r6_3", "r6_4", "r6_5", "grip"] - print(f"\n {name} (n_samples={stats['n_samples']:,})") - print(f" {'dim':<6} {'min':>10} {'max':>10} {'q01':>10} {'q99':>10} {'mean':>10} {'std':>10}") - for i, dn in enumerate(dim_names): - print( - f" {dn:<6} " - f"{stats['min'][i]:>10.3f} " - f"{stats['max'][i]:>10.3f} " - f"{stats['q01'][i]:>10.3f} " - f"{stats['q99'][i]:>10.3f} " - f"{stats['mean'][i]:>10.3f} " - f"{stats['std'][i]:>10.3f}" - ) - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument( - "--root", - type=str, - default="/path/to/OXE", - help="OXE dataset root (each dataset is in {root}/-Dataset/)", - ) - parser.add_argument( - "--dataset", - type=str, - choices=list(SCHEMA.keys()), - help="Single dataset to process (mutually exclusive with --all)", - ) - parser.add_argument("--all", action="store_true", help="Process all 4 OXE datasets") - parser.add_argument( - "--dataset-dir", - type=str, - default=None, - help="Bucket path to scan, overriding {root}/{dataset}-Dataset. Requires --dataset " - "(the schema to read it with) and is incompatible with --all. Use when a bucket lives " - "outside the OXE root, e.g. --dataset DROID --dataset-dir /path/to/pretrain_dataset/Droid", - ) - parser.add_argument( - "--dry-run", - action="store_true", - help="Compute and print stats but do not write meta/eef_stats.json", - ) - parser.add_argument( - "--split", - default="train", - help="info.json split used for DROID normalization statistics (default: train)", - ) - parser.add_argument( - "--no-rot6d-identity", - action="store_true", - help="Disable pinning rot6d stats to identity (rot6d would then be per-dim normalized " - "like pos/gripper — generally undesirable; see pin_rot6d_identity).", - ) - args = parser.parse_args() - - if not args.dataset and not args.all: - parser.error("must specify either --dataset NAME or --all") - if args.dataset_dir and args.all: - parser.error("--dataset-dir applies to a single bucket; use --dataset NAME, not --all") - targets = list(SCHEMA.keys()) if args.all else [args.dataset] - root = Path(args.root) - for name in targets: - ds_dir = Path(args.dataset_dir) if args.dataset_dir else root / f"{name}-Dataset" - if not ds_dir.is_dir(): - logger.warning("%s: directory %s missing, skipping", name, ds_dir) - continue - stats, n_state, n_action = compute_dataset_stats( - ds_dir, - name, - rot6d_identity=not args.no_rot6d_identity, - split=args.split, - ) - _print_stats_table(stats, name) - if not args.dry_run: - out_path = ds_dir / "meta" / "eef_stats.json" - with open(out_path, "w") as f: - json.dump(stats, f, indent=2) - logger.info("%s: wrote %s", name, out_path) - else: - logger.info("%s: --dry-run, no file written", name) - - -if __name__ == "__main__": - main() diff --git a/openwam/dataloader/utils/stats_computation/robocasa_gr1_stats_computation.py b/openwam/dataloader/utils/stats_computation/robocasa_gr1_stats_computation.py index a24bfaa7a..3978cfb3e 100644 --- a/openwam/dataloader/utils/stats_computation/robocasa_gr1_stats_computation.py +++ b/openwam/dataloader/utils/stats_computation/robocasa_gr1_stats_computation.py @@ -36,7 +36,7 @@ ) from openwam.dataloader.utils.gr1_kinematics import HAND_DIMS_EEF33, ROT6D_DIMS_EEF33 from openwam.dataloader.utils.normalization import pin_rot6d_identity -from openwam.dataloader.utils.stats_computation.robocoin_stats_computation import Accumulator +from openwam.dataloader.utils.stats_computation.accumulator import Accumulator def _iter_buckets(dataset) -> Iterable[RoboCasaGR1Dataset]: diff --git a/openwam/dataloader/utils/stats_computation/robocoin_stats_computation.py b/openwam/dataloader/utils/stats_computation/robocoin_stats_computation.py deleted file mode 100644 index 248a0739f..000000000 --- a/openwam/dataloader/utils/stats_computation/robocoin_stats_computation.py +++ /dev/null @@ -1,556 +0,0 @@ -#!/usr/bin/env python3 -"""Compute per-robot-type EEF and hand stats for RoboCOIN. - -Published Euler-XYZ pose and gripper columns are converted to the reader's -20-D EEF layout, then action and state rows are pooled together. Dexterous -buckets additionally produce a left-then-right hand-joint stats block while -their absent scalar-grip slots remain masked. Means, variances and extrema -are streamed exactly; quantiles use a bounded uniform reservoir. - -Rot6d dimensions are pinned to identity, and each output records contributor, -exclusion, trim, and retained-population provenance required by the reader. -""" - -import argparse -import json -import os -from pathlib import Path - -import numpy as np -import pandas as pd -import pyarrow.parquet as pq - -from openwam.dataloader.robocoin import ( - MAX_HAND_DOF, - _assert_trim_snapshot_current, - _discover_data_parquets, - _eef14_to_eef20, - _excluded_episodes_provenance, - _finger_indices, - _load_trim_snapshot, - _stats_population_spans, - dex_finger_layout, - is_robocoin_bucket_excluded, - robocoin_bucket_exclusions_provenance, -) -from openwam.dataloader.utils.lerobotv3 import ( - DataContractError, - assert_excluded_episodes_snapshot_current, - load_excluded_episodes_snapshot, - parse_info_json, -) -from openwam.dataloader.utils.normalization import ROT6D_DIMS_EEF20, pin_rot6d_identity - -RESERVOIR_CAP = 1_000_000 - - -class Accumulator: - """Online mean/std/min/max accumulator + reservoir for q01/q99.""" - - def __init__(self, dim: int = 20, reservoir_cap: int = RESERVOIR_CAP, seed: int = 0): - self.dim = dim - self.count = 0 - self.mean = np.zeros(dim, dtype=np.float64) - self.m2 = np.zeros(dim, dtype=np.float64) - self.min_val = np.full(dim, np.inf, dtype=np.float64) - self.max_val = np.full(dim, -np.inf, dtype=np.float64) - - self.cap = int(reservoir_cap) - self.rng = np.random.RandomState(seed) - self._res = np.empty((self.cap, dim), dtype=np.float32) - self._res_n = 0 - self._res_seen = 0 - - def update(self, batch: np.ndarray): - """Update with (N, dim) array using Welford's online algorithm.""" - for i in range(len(batch)): - x = batch[i].astype(np.float64) - self.count += 1 - delta = x - self.mean - self.mean += delta / self.count - delta2 = x - self.mean - self.m2 += delta * delta2 - self.min_val = np.minimum(self.min_val, x) - self.max_val = np.maximum(self.max_val, x) - self._reservoir_add(np.asarray(batch, dtype=np.float32)) - - def update_batch(self, batch: np.ndarray): - """Batch update (more efficient for large arrays).""" - n = len(batch) - if n == 0: - return - self._reservoir_add(np.asarray(batch, dtype=np.float32)) - batch = batch.astype(np.float64) - batch_mean = batch.mean(axis=0) - batch_var = batch.var(axis=0) - batch_min = batch.min(axis=0) - batch_max = batch.max(axis=0) - - if self.count == 0: - self.mean = batch_mean - self.m2 = batch_var * n - self.min_val = batch_min - self.max_val = batch_max - self.count = n - else: - total = self.count + n - delta = batch_mean - self.mean - new_mean = self.mean + delta * n / total - self.m2 = self.m2 + batch_var * n + delta**2 * self.count * n / total - self.mean = new_mean - self.count = total - self.min_val = np.minimum(self.min_val, batch_min) - self.max_val = np.maximum(self.max_val, batch_max) - - def _reservoir_add(self, batch: np.ndarray): - """Feed (N, dim) rows into the reservoir (vectorized Algorithm R).""" - n = len(batch) - if n == 0: - return - - if self._res_n < self.cap: - take = min(self.cap - self._res_n, n) - self._res[self._res_n : self._res_n + take] = batch[:take] - self._res_n += take - self._res_seen += take - batch = batch[take:] - if len(batch) == 0: - return - - m = len(batch) - t = self._res_seen + np.arange(m) - p = self.rng.randint(0, t + 1) - keep = p < self.cap - self._res[p[keep]] = batch[keep] - self._res_seen += m - - def finalize(self): - std = np.sqrt(self.m2 / max(self.count, 1)) - std = np.where(std < 1e-8, 1.0, std) - if self._res_n > 0: - res = self._res[: self._res_n] - q01 = np.quantile(res, 0.01, axis=0) - q99 = np.quantile(res, 0.99, axis=0) - else: - q01 = self.min_val - q99 = self.max_val - return { - "mean": self.mean.astype(np.float32).tolist(), - "std": std.astype(np.float32).tolist(), - "min": self.min_val.astype(np.float32).tolist(), - "max": self.max_val.astype(np.float32).tolist(), - "q01": q01.astype(np.float32).tolist(), - "q99": q99.astype(np.float32).tolist(), - } - - -def discover_datasets_by_robot_type(root: str) -> dict: - """Group datasets by robot_type.""" - groups = {} - for name in sorted(os.listdir(root)): - if is_robocoin_bucket_excluded(name): - print(f" Excluding whole RoboCOIN bucket from stats discovery: {name}") - continue - info_path = os.path.join(root, name, "meta", "info.json") - if not os.path.isfile(info_path): - continue - with open(info_path) as f: - info = json.load(f) - - rtype = str(info.get("robot_type", "unknown")) - groups.setdefault(rtype, []).append(os.path.join(root, name)) - return groups - - -_NEEDED_COLS = [ - "eef_sim_pose_action", - "gripper_open_scale_action", - "eef_sim_pose_state", - "gripper_open_scale_state", -] -_EEF_COLS = ["eef_sim_pose_action", "eef_sim_pose_state"] -_GRIP_COLS = ("gripper_open_scale_action", "gripper_open_scale_state") - -_HAND_RAW_COLS = ["action", "observation.state"] - - -def _classify_dataset(ds_dir: str, *, fail_closed: bool = False): - """Classify one dataset and return its dex-hand finger layout if any. - - Reads ``meta/info.json`` ONCE and returns ``(kind, layout)``: - * ``kind`` is one of: - ``"grip"`` — has ``gripper_open_scale_*`` (real gripper); - ``"dex"`` — pose, no gripper, a valid finger layout; - ``"nogrip"`` — pose, no gripper, but the finger gate failed - (single-handed / >MAX_HAND_DOF / action-state disagree); - ``"other"`` — no pose column at all (its files won't pool into `eef`). - * ``layout`` is ``(idx_action_LR, idx_state_LR, kL, kR)`` for ``"dex"`` (the - left-then-right ``*_hand_joint_*`` indices within the raw ``action`` / - ``observation.state`` arrays), else ``None``. - - ``"dex"`` and ``"nogrip"`` both zero-fill the gripper slots (9/19) when pooled, - so both must be kept out of a robot_type that also has ``"grip"`` datasets (the - caller's mixed-bucket guard treats them alike). Uses the shared - :func:`dex_finger_layout` gate so the stats producer and the reader agree. - """ - - try: - with open(os.path.join(ds_dir, "meta", "info.json")) as f: - feats = json.load(f).get("features", {}) - except (OSError, ValueError): - if fail_closed: - raise - return "other", None - if all(c in feats for c in _GRIP_COLS): - return "grip", None - layout = dex_finger_layout(feats) - if layout is not None: - aL, aR, sL, sR = layout - return "dex", (aL + aR, sL + sR, len(aL), len(aR)) - if "eef_sim_pose_action" in feats: - aL, aR = _finger_indices(feats.get("action", {})) - sL, sR = _finger_indices(feats.get("observation.state", {})) - print( - f" Warning: {ds_dir} looks dexterous-hand (pose, no gripper) but finger layout " - f"failed the gate (action L/R={len(aL)}/{len(aR)}, state L/R={len(sL)}/{len(sR)}, " - f"max={MAX_HAND_DOF}); no finger stats emitted." - ) - return "nogrip", None - return "other", None - - -def _slice_file_to_global_spans(df, file_start: int, file_end: int, kept_spans): - """Select intersections with kept global spans using physical file offsets.""" - pieces = [] - for span_start, span_end in kept_spans: - if span_end <= file_start: - continue - if span_start >= file_end: - break - local_start = max(span_start, file_start) - file_start - local_end = min(span_end, file_end) - file_start - if local_start < local_end: - pieces.append(df.iloc[local_start:local_end]) - if not pieces: - return df.iloc[0:0] - return pd.concat(pieces, ignore_index=True) - - -def compute_stats_for_robot_type( - rtype: str, - dataset_dirs: list, - rot6d_identity: bool = True, - trim_csv=None, - *, - split: str = "train", - _trim_snapshot=None, -) -> dict: - """Compute unified action+state EEF stats across all datasets of one robot type. - - Each parquet row contributes TWO 20-D points to the pool: one from action - columns and one from state columns. The accumulator sees their union, so - min/max end up as the per-dim envelope across both streams and mean/std - reflect the joint distribution. - - For dexterous-hand types (no gripper), a second accumulator pools the raw - ``*_hand_joint_*`` finger values (action + state) into a ``hand`` stats block - of width ``kL + kR`` (left then right), used to normalize the finger dims the - reader scatters into the 80-D hand slots under unify_action. - - With ``trim_csv`` set, every stream is restricted to the same split, - exclusion and manifest-based trim spans as the reader and all - file/manifest errors fail closed. - - Returns a dict ready to dump: ``{"eef": <20-D stats>, "hand": ?}``. - """ - - dataset_dirs = list(dataset_dirs) - excluded_dirs = [ds_dir for ds_dir in dataset_dirs if is_robocoin_bucket_excluded(ds_dir)] - dataset_dirs = [ds_dir for ds_dir in dataset_dirs if not is_robocoin_bucket_excluded(ds_dir)] - if excluded_dirs: - print( - " Excluding whole RoboCOIN bucket(s) from stats population: " - + ", ".join(sorted(Path(path).name for path in excluded_dirs)) - ) - if not dataset_dirs: - raise DataContractError(f"RoboCOIN stats for robot_type {rtype!r}: no non-excluded dataset directories") - - trim_enabled = trim_csv is not None - if _trim_snapshot is not None and (not trim_enabled or _trim_snapshot.path != str(trim_csv)): - raise ValueError("_trim_snapshot requires a matching non-null trim_csv") - trim_snapshot = _trim_snapshot - if trim_enabled and trim_snapshot is None: - trim_snapshot = _load_trim_snapshot(trim_csv) - trim_spec = trim_snapshot.spec if trim_enabled else {} - exclusion_snapshots = {} - population_datasets = {} - acc = Accumulator(dim=20) - total_files = 0 - grip_files = 0 - - hand_acc = None - hand_dims = None - hand_files = 0 - - grip_example = None - nogrip_example = None - - for ds_dir in dataset_dirs: - dataset_id = Path(ds_dir).name - data_dir = Path(ds_dir) / "data" - if not data_dir.is_dir(): - if trim_enabled: - raise FileNotFoundError(f"RoboCOIN stats: missing data directory {data_dir}") - continue - exclusion_snapshot = None - if trim_enabled: - if dataset_id in exclusion_snapshots: - raise DataContractError( - f"RoboCOIN stats for robot_type {rtype!r}: duplicate dataset directory " - f"name {dataset_id!r} cannot be represented in exclusion provenance" - ) - exclusion_snapshot = load_excluded_episodes_snapshot(Path(ds_dir)) - exclusion_snapshots[dataset_id] = exclusion_snapshot - kind, layout = _classify_dataset(ds_dir, fail_closed=trim_enabled) - if kind == "grip": - grip_example = grip_example or ds_dir - elif kind in ("dex", "nogrip"): - nogrip_example = nogrip_example or ds_dir - if grip_example and nogrip_example: - raise ValueError( - f"robot_type {rtype!r} mixes a grippered dataset ({grip_example}) with a " - f"no-gripper (dexterous-hand) dataset ({nogrip_example}). They share one 20-D " - f"'eef' stats block, so the no-gripper zero-filled gripper slots (9/19) would " - f"corrupt the grippered reader's gripper normalization. Split these into distinct " - f"robot_types." - ) - if layout is not None: - idx_act_lr, idx_state_lr, kL, kR = layout - if hand_acc is None: - hand_dims = (kL, kR) - hand_acc = Accumulator(dim=kL + kR) - elif hand_dims != (kL, kR): - raise ValueError( - f"robot_type {rtype!r} has heterogeneous dexterous-hand finger DOF: " - f"{ds_dir} has {(kL, kR)} but an earlier dataset had {hand_dims}. " - f"A per-robot-type 'hand' stats block cannot describe both; the reader " - f"slices it by each bucket's own kL/kR. Split these into distinct robot_types." - ) - if trim_enabled: - info = parse_info_json(Path(ds_dir)) - file_paths = [ - path - for path, _, _ in _discover_data_parquets( - Path(ds_dir), - info["data_path"], - ) - ] - else: - file_paths = [] - for chunk in sorted(os.listdir(data_dir)): - chunk_path = data_dir / chunk - if not chunk_path.is_dir(): - continue - file_paths.extend( - chunk_path / fname for fname in sorted(os.listdir(chunk_path)) if fname.endswith(".parquet") - ) - if not file_paths: - if trim_enabled: - raise FileNotFoundError(f"RoboCOIN stats: no data parquet files under {data_dir}") - continue - - kept_spans = None - manifest_end = 0 - if trim_enabled: - kept_spans, manifest_end, population_entry = _stats_population_spans( - ds_dir, - dataset_id=dataset_id, - dataset_spec=trim_spec.get(dataset_id, {}), - excluded_episode_indices=exclusion_snapshot.episode_indices, - split=split, - info=info, - ) - population_datasets[dataset_id] = population_entry - - file_entries = [] - physical_end = 0 - for fpath in file_paths: - try: - parquet_file = pq.ParquetFile(fpath) - present = set(parquet_file.schema_arrow.names) - n_rows = int(parquet_file.metadata.num_rows) - file_entries.append((fpath, physical_end, physical_end + n_rows, present)) - physical_end += n_rows - except Exception as e: - if trim_enabled: - raise - print(f" Warning: skipping {fpath}: {e}") - if trim_enabled and manifest_end > physical_end: - raise DataContractError( - f"RoboCOIN({Path(ds_dir).name}): manifest ends at global row {manifest_end}, " - f"past the physical parquet row count {physical_end}" - ) - - for fpath, file_start, file_end, present in file_entries: - try: - has_grip = all(c in present for c in _GRIP_COLS) - cols = list(_NEEDED_COLS) if has_grip else list(_EEF_COLS) - if layout is not None: - cols = cols + _HAND_RAW_COLS - df = pd.read_parquet(fpath, columns=cols) - if trim_enabled: - expected_rows = file_end - file_start - if len(df) != expected_rows: - raise DataContractError( - f"RoboCOIN({Path(ds_dir).name}): {fpath} metadata declares " - f"{expected_rows} rows but reading returned {len(df)}" - ) - df = _slice_file_to_global_spans(df, file_start, file_end, kept_spans) - - total_files += 1 - grip_files += int(has_grip) - if layout is not None: - hand_files += 1 - if len(df) == 0: - continue - - eef_a = np.stack(df["eef_sim_pose_action"].values).astype(np.float32) - eef_s = np.stack(df["eef_sim_pose_state"].values).astype(np.float32) - if has_grip: - grip_a = np.stack(df["gripper_open_scale_action"].values).astype(np.float32) - grip_s = np.stack(df["gripper_open_scale_state"].values).astype(np.float32) - else: - grip_a = np.zeros((len(eef_a), 2), dtype=np.float32) - grip_s = np.zeros((len(eef_s), 2), dtype=np.float32) - - action_20d = _eef14_to_eef20(eef_a, grip_a) - state_20d = _eef14_to_eef20(eef_s, grip_s) - acc.update_batch(np.concatenate([action_20d, state_20d], axis=0)) - - if layout is not None: - act_arr = np.stack(df["action"].values).astype(np.float32) - state_arr = np.stack(df["observation.state"].values).astype(np.float32) - fingers = np.concatenate([act_arr[:, idx_act_lr], state_arr[:, idx_state_lr]], axis=0) - hand_acc.update_batch(fingers) - except Exception as e: - if trim_enabled: - raise - print(f" Warning: skipping {fpath}: {e}") - - if trim_enabled and acc.count == 0: - raise DataContractError(f"RoboCOIN stats for robot_type {rtype!r}: trimming left no EEF rows") - - stats = acc.finalize() - if rot6d_identity: - pin_rot6d_identity(stats, ROT6D_DIMS_EEF20) - stats["num_timesteps"] = int(acc.count) - stats["num_datasets"] = len(dataset_dirs) - stats["num_files"] = total_files - stats["robot_type"] = rtype - stats["pool"] = "action+state" - - stats["grip_present"] = total_files > 0 and grip_files == total_files - - result = { - "eef": stats, - "whole_bucket_exclusions_provenance": robocoin_bucket_exclusions_provenance(rtype), - } - if hand_acc is not None and hand_acc.count > 0: - hand = hand_acc.finalize() - hand["dof_left"] = hand_dims[0] - hand["dof_right"] = hand_dims[1] - hand["num_timesteps"] = int(hand_acc.count) - hand["num_files"] = hand_files - hand["pool"] = "action+state" - hand["layout"] = "left_fingers + right_fingers" - result["hand"] = hand - if trim_enabled: - result["population"] = { - "schema_version": 2, - "split": split, - "policy": "info_split_then_exclusion_then_trim", - "datasets": dict(sorted(population_datasets.items())), - } - result["trim_provenance"] = trim_snapshot.provenance - result["excluded_episodes_provenance"] = _excluded_episodes_provenance(exclusion_snapshots) - _assert_trim_snapshot_current(trim_snapshot, context=f"stats scan for robot_type {rtype!r}") - for snapshot in exclusion_snapshots.values(): - assert_excluded_episodes_snapshot_current( - snapshot, - context=f"stats scan for robot_type {rtype!r}", - ) - return result - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--dataset_dir", required=True) - parser.add_argument("--robot_type", default=None, help="Compute stats for a single robot type only") - parser.add_argument( - "--split", - default="train", - help="info.json split used for normalization statistics (default: train)", - ) - parser.add_argument( - "--trim_csv", - default=None, - help="Apply the same leading/trailing episode trims as the RoboCOIN reader", - ) - parser.add_argument( - "--no-rot6d-identity", - action="store_true", - help="Disable pinning rot6d stats to identity (rot6d would then be per-dim normalized " - "like pos/gripper — generally undesirable; see pin_rot6d_identity).", - ) - args = parser.parse_args() - - groups = discover_datasets_by_robot_type(args.dataset_dir) - print(f"Found {len(groups)} robot types: {sorted(groups.keys())}") - - trim_snapshot = None - if args.trim_csv is not None: - trim_snapshot = _load_trim_snapshot(args.trim_csv) - trim_spec = trim_snapshot.spec - bucket_names = {Path(ds_dir).name for ds_list in groups.values() for ds_dir in ds_list} - if bucket_names and bucket_names.isdisjoint(trim_spec): - raise ValueError( - f"RoboCOIN trim_csv {args.trim_csv}: none of its {len(trim_spec)} dataset key(s) " - f"match any bucket directory under {args.dataset_dir}; stats would be untrimmed." - ) - - out_dir = os.path.join(args.dataset_dir, "meta") - os.makedirs(out_dir, exist_ok=True) - - for rtype in sorted(groups.keys()): - if args.robot_type and rtype != args.robot_type: - continue - ds_list = groups[rtype] - print(f"\n{'=' * 60}") - print(f"Computing stats for {rtype} ({len(ds_list)} datasets)...") - result = compute_stats_for_robot_type( - rtype, - ds_list, - rot6d_identity=not args.no_rot6d_identity, - trim_csv=args.trim_csv, - split=args.split, - _trim_snapshot=trim_snapshot, - ) - stats = result["eef"] - - out_path = os.path.join(out_dir, f"stats_{rtype}.json") - with open(out_path, "w") as f: - json.dump(result, f, indent=2) - - print(f" timesteps: {stats['num_timesteps']:,}") - print(f" mean[:5]: {stats['mean'][:5]}") - print(f" std[:5]: {stats['std'][:5]}") - print(f" min[:5]: {stats['min'][:5]}") - print(f" max[:5]: {stats['max'][:5]}") - if "hand" in result: - h = result["hand"] - print(f" hand: dof L/R={h['dof_left']}/{h['dof_right']}, timesteps={h['num_timesteps']:,}") - print(f" Saved to: {out_path}") - - -if __name__ == "__main__": - main() diff --git a/openwam/dataloader/utils/stats_computation/robodojo_stats_computation.py b/openwam/dataloader/utils/stats_computation/robodojo_stats_computation.py index 35b2fe293..58be53521 100644 --- a/openwam/dataloader/utils/stats_computation/robodojo_stats_computation.py +++ b/openwam/dataloader/utils/stats_computation/robodojo_stats_computation.py @@ -49,7 +49,7 @@ STAT_KEYS, pin_rot6d_identity, ) -from openwam.dataloader.utils.stats_computation.robocoin_stats_computation import ( +from openwam.dataloader.utils.stats_computation.accumulator import ( Accumulator, ) diff --git a/tests/dataloader/test_agibotworld.py b/tests/dataloader/test_agibotworld.py index 7f6188295..0873bba1d 100644 --- a/tests/dataloader/test_agibotworld.py +++ b/tests/dataloader/test_agibotworld.py @@ -1,15 +1,4 @@ -"""Tests for the AgiBotWorld dataloader (openwam/dataloader/agibotworld.py). - -Focus: the ``segment_flag`` / ``segment_delta`` boundary-cleanup path and the -``segment_max_trim_ratio`` drop rule, both of which reach into the shared -``LeRobotV3Reader`` window index via the ``_valid_start`` / ``_valid_end`` -columns. The end-to-end cases assert that the trimmed offset feeds BOTH the -parquet slice and the video decode — a prefix trim applied to only one of the -two would silently pair frame t with action t+delta. - -Synthetic buckets are written to tmp_path; video decode is monkeypatched to -return frame-index-encoded images so alignment is checkable without mp4s. -""" +"""Clean AgiBotWorld exports retain physical semantics without read-time filtering.""" from __future__ import annotations @@ -25,12 +14,9 @@ from openwam.dataloader.agibotworld import ( _GRIPPER_CONTRACT, - _STATS_SCHEMA_VERSION, AgiBotWorldDataset, - _validate_trim_ratio, ) from openwam.dataloader.utils.lerobotv3 import DataContractError -from openwam.dataloader.utils.stats_computation.agibotworld_stats_computation import _partial_bucket FPS = 15.0 HEAD = "observation.images.head" @@ -41,44 +27,20 @@ VIDEO_STRIDE = 4 -# --------------------------------------------------------------------------- -# _validate_trim_ratio -# --------------------------------------------------------------------------- - - -class TestValidateTrimRatio: - def test_none_passes_through(self): - assert _validate_trim_ratio(None) is None - - @pytest.mark.parametrize("v", [0.7, 1.0, "0.5", 0.001]) - def test_valid(self, v): - assert _validate_trim_ratio(v) == float(v) - - @pytest.mark.parametrize("v", [0.0, -0.1, 1.0001, 70, float("nan"), float("inf"), "abc", object()]) - def test_rejected(self, v): - # 70 (percent) and 0 (would drop everything) must raise, not clamp. - with pytest.raises(ValueError): - _validate_trim_ratio(v) - - -# --------------------------------------------------------------------------- -# Synthetic bucket -# --------------------------------------------------------------------------- - - def _make_bucket( bucket: Path, episodes: list[dict], *, - with_segment_cols: bool = True, action_velocity=None, state_velocity=None, action_gripper=None, state_gripper=None, + action_dex=None, + state_dex=None, ) -> Path: """Write a minimal grippered AgiBotWorld bucket. - ``episodes`` items: ``{"length": int, "flag": int, "delta": int}``. All + ``episodes`` items contain each exported episode length. All episodes live in one data shard and one video file per camera, so the per-episode offsets are plain cumulative sums — which is exactly the layout that makes an offset bug read into the NEXT episode instead of erroring. @@ -92,7 +54,7 @@ def _make_bucket( "robot_type": "g2a", "data_path": "data/chunk-{chunk_index:03d}/file-{file_index:03d}.parquet", "video_path": "videos/{video_key}/chunk-{chunk_index:03d}/file-{file_index:03d}.mp4", - "features": {"segment_flag": {"dtype": "int64"}, "segment_delta": {"dtype": "int64"}}, + "features": {}, } ) ) @@ -109,9 +71,6 @@ def _make_bucket( for cam in CAMS: row[f"videos/{cam}/chunk_index"] = 0 row[f"videos/{cam}/file_index"] = 0 - if with_segment_cols: - row["segment_flag"] = int(spec["flag"]) - row["segment_delta"] = int(spec["delta"]) rows.append(row) cum += int(spec["length"]) pq.write_table(pa.Table.from_pandas(pd.DataFrame(rows)), meta / "episodes" / "chunk-000.parquet") @@ -136,6 +95,8 @@ def _values(value, dim): state_vel = _values(state_velocity, 3) action_grip = _values(action_gripper, 2) state_grip = _values(state_gripper, 2) + action_fingers = _values(action_dex, 12) + state_fingers = _values(state_dex, 12) (meta / "stats.json").write_text( json.dumps( { @@ -169,6 +130,8 @@ def _values(value, dim): "observation.state.ee_base": list(ee), "action.gripper": list(action_grip), "observation.state.gripper": list(state_grip), + "action.dex": list(action_fingers), + "observation.state.dex": list(state_fingers), "action.robot_velocity": list(action_vel), "observation.state.robot_velocity": list(state_vel), } @@ -213,200 +176,6 @@ def _reader(bucket: Path, **kw) -> AgiBotWorldDataset: ) -# --------------------------------------------------------------------------- -# Flag semantics -# --------------------------------------------------------------------------- - - -class TestSegmentFlags: - def test_flag0_untouched(self, tmp_path): - ds = _reader(_make_bucket(tmp_path / "b", [{"length": 100, "flag": 0, "delta": 0}])) - assert ds._ep_valid_start.tolist() == [0] - assert ds._ep_valid_end.tolist() == [100] - assert len(ds) == 99 # final one-row start has no real next-state target - - def test_flag1_trims_prefix(self, tmp_path): - ds = _reader(_make_bucket(tmp_path / "b", [{"length": 100, "flag": 1, "delta": 30}])) - assert ds._ep_valid_start.tolist() == [30] - assert ds._ep_valid_end.tolist() == [100] - assert len(ds) == 69 - - def test_flag2_trims_suffix(self, tmp_path): - ds = _reader(_make_bucket(tmp_path / "b", [{"length": 100, "flag": 2, "delta": 30}])) - assert ds._ep_valid_start.tolist() == [0] - assert ds._ep_valid_end.tolist() == [70] - assert len(ds) == 69 - - def test_flag3_dropped(self, tmp_path): - ds = _reader( - _make_bucket( - tmp_path / "b", - [ - {"length": 100, "flag": 0, "delta": 0}, - {"length": 100, "flag": 3, "delta": 0}, - {"length": 50, "flag": 0, "delta": 0}, - ], - ) - ) - assert ds._eps_df["episode_index"].tolist() == [0, 2] - assert len(ds) == 148 - - def test_delta_at_or_over_length_drops_episode(self, tmp_path): - # delta >= length leaves nothing: must drop, never produce a negative span. - ds = _reader( - _make_bucket( - tmp_path / "b", - [ - {"length": 40, "flag": 1, "delta": 40}, - {"length": 40, "flag": 2, "delta": 90}, - {"length": 40, "flag": 0, "delta": 0}, - ], - ) - ) - assert ds._eps_df["episode_index"].tolist() == [2] - assert (ds._ep_valid_end >= ds._ep_valid_start).all() - - def test_disabled_keeps_full_segments(self, tmp_path): - eps = [{"length": 100, "flag": 1, "delta": 30}, {"length": 100, "flag": 3, "delta": 0}] - ds = _reader(_make_bucket(tmp_path / "b", eps), use_segment_annotations=False) - assert len(ds._eps_df) == 2 - assert ds._ep_valid_start.tolist() == [0, 0] - assert len(ds) == 198 - - def test_missing_columns_falls_back(self, tmp_path, caplog): - eps = [{"length": 100, "flag": 0, "delta": 0}] - ds = _reader(_make_bucket(tmp_path / "b", eps, with_segment_cols=False)) - assert len(ds) == 99 - assert "segment annotation columns" in caplog.text - - -# --------------------------------------------------------------------------- -# segment_max_trim_ratio -# --------------------------------------------------------------------------- - - -class TestMaxTrimRatio: - EPS = [ - {"length": 100, "flag": 1, "delta": 10}, # 10% trimmed -> keep - {"length": 100, "flag": 2, "delta": 70}, # 70% trimmed -> boundary - {"length": 100, "flag": 1, "delta": 90}, # 90% trimmed -> drop at 0.7 - {"length": 100, "flag": 0, "delta": 0}, # untrimmed -> always keep - ] - - def test_none_keeps_everything(self, tmp_path): - ds = _reader(_make_bucket(tmp_path / "b", self.EPS), segment_max_trim_ratio=None) - assert ds._eps_df["episode_index"].tolist() == [0, 1, 2, 3] - - def test_threshold_is_inclusive(self, tmp_path): - # "70% or more trimmed" — an episode trimmed to exactly the threshold goes. - ds = _reader(_make_bucket(tmp_path / "b", self.EPS), segment_max_trim_ratio=0.7) - assert ds._eps_df["episode_index"].tolist() == [0, 3] - assert len(ds) == 89 + 99 - - def test_higher_threshold_keeps_more(self, tmp_path): - # 0.8 spares the 70%-trimmed episode and still drops the 90% one. - ds = _reader(_make_bucket(tmp_path / "b", self.EPS), segment_max_trim_ratio=0.8) - assert ds._eps_df["episode_index"].tolist() == [0, 1, 3] - # above every trim ratio present, nothing is dropped. - ds = _reader(_make_bucket(tmp_path / "b2", self.EPS), segment_max_trim_ratio=0.95) - assert ds._eps_df["episode_index"].tolist() == [0, 1, 2, 3] - - def test_untrimmed_survives_ratio_one(self, tmp_path): - # ratio=1.0 may only remove episodes trimmed to nothing (already dropped), - # never a flag-0 episode. - ds = _reader(_make_bucket(tmp_path / "b", self.EPS), segment_max_trim_ratio=1.0) - assert 3 in ds._eps_df["episode_index"].tolist() - - def test_ignored_when_annotations_off(self, tmp_path, caplog): - ds = _reader( - _make_bucket(tmp_path / "b", self.EPS), - use_segment_annotations=False, - segment_max_trim_ratio=0.7, - ) - assert len(ds._eps_df) == 4 - assert "ignored because" in caplog.text - - def test_config_keys_expose_both_knobs(self): - assert "use_segment_annotations" in AgiBotWorldDataset.CONFIG_KEYS - assert "segment_max_trim_ratio" in AgiBotWorldDataset.CONFIG_KEYS - - -# --------------------------------------------------------------------------- -# End-to-end: the trimmed offset must drive parquet AND video identically -# --------------------------------------------------------------------------- - - -class TestTrimmedOffsetAlignment: - EPS = [ - {"length": 120, "flag": 1, "delta": 25}, # prefix trim - {"length": 120, "flag": 2, "delta": 25}, # suffix trim - {"length": 120, "flag": 0, "delta": 0}, # control - ] - - def _ds(self, tmp_path): - return _reader(_make_bucket(tmp_path / "b", self.EPS)) - - def test_first_window_starts_after_the_trimmed_prefix(self, tmp_path, patch_decode): - ds = self._ds(tmp_path) - s = ds[0] # first window of episode 0, whose valid range is [25, 120) - # action payload encodes the global parquet row index; episode 0 starts at 0 - assert s["action"][0, 0].item() == pytest.approx(25.0) - # video frames must come from the SAME offset, not from frame 0 - assert _decoded_index(s["video"][0]) == 25 - assert [_decoded_index(f) for f in s["video"]] == list(range(25, 25 + NUM_FRAMES, VIDEO_STRIDE)) - - def test_window_count_matches_trimmed_span(self, tmp_path): - ds = self._ds(tmp_path) - assert len(ds) == 94 + 94 + 119 - - def test_last_window_of_suffix_trimmed_episode_stops_at_valid_end(self, tmp_path, patch_decode): - ds = self._ds(tmp_path) - ep1_first = 94 # episode 0 contributed 94 supervised-window starts - last = ep1_first + 93 # last window of episode 1 - s = ds[last] - # episode 1 valid range is [0, 95); its last window starts at 93 and has - # two real rows (one real next-state action), so the video mask samples - # one frame and the payload - # must stay inside episode 1 (global rows 120..239). - assert int(s["video_mask"].sum()) == 1 - assert int(s["action_mask"].any(dim=1).sum()) == 1 - assert s["action"][0, 0].item() == pytest.approx(120.0 + 93.0) - assert _decoded_index(s["video"][0]) == 120 + 93 - - def test_no_window_reads_past_its_own_episode(self, tmp_path, patch_decode): - ds = self._ds(tmp_path) - bounds = [(0, 120), (120, 240), (240, 360)] - starts = ds._cum_n_starts - for i in range(len(ds)): - ep = int(np.searchsorted(starts, i, side="right") - 1) - lo, hi = bounds[ep] - s = ds[i] - rows = s["action"][s["action_mask"].any(dim=1)][:, 0].numpy() - assert rows.min() >= lo and rows.max() < hi, f"window {i} (episode {ep}) leaked outside {lo}..{hi}" - frames = [_decoded_index(f) for f, m in zip(s["video"], s["video_mask"].tolist()) if m] - assert min(frames) >= lo and max(frames) < hi - - def test_trimmed_frames_are_never_served(self, tmp_path, patch_decode): - ds = self._ds(tmp_path) - seen = set() - for i in range(len(ds)): - s = ds[i] - seen.update(int(r) for r in s["action"][s["action_mask"].any(dim=1)][:, 0].numpy()) - # episode 0 rows 0..24 (trimmed prefix) and episode 1 rows 215..239 - # (trimmed suffix, global 120+95 .. 120+119) must be absent. - assert not (seen & set(range(0, 25))) - assert not (seen & set(range(120 + 95, 240))) - # Every kept row with a real successor must appear. The last kept row of - # each episode (119/214/359) is a clamped action and stays unsupervised. - expected = set(range(25, 119)) | set(range(120, 214)) | set(range(240, 359)) - assert seen == expected - - -# --------------------------------------------------------------------------- -# Independent action/state movement masks + canonical gripper direction -# --------------------------------------------------------------------------- - - class TestPhysicalSemantics: @staticmethod def _unified_reader(bucket: Path) -> AgiBotWorldDataset: @@ -467,6 +236,26 @@ def test_state_only_base_motion_masks_action_independently(self, tmp_path, patch assert sample["proprio"][0, 68].item() == pytest.approx(0.25) assert sample["proprio"][0, 70].item() == pytest.approx(0.4) + def test_dex_fingers_keep_distinct_left_and_right_unified_slots(self, tmp_path, patch_decode): + bucket = _make_bucket( + tmp_path / "475", + [{"length": 40}], + action_dex=np.arange(12, dtype=np.float32), + state_dex=np.arange(12, dtype=np.float32) / 10, + ) + sample = self._unified_reader(bucket)[0] + np.testing.assert_allclose(sample["action"][0, 10:16], np.arange(6)) + np.testing.assert_allclose(sample["action"][0, 44:50], np.arange(6, 12)) + np.testing.assert_allclose(sample["proprio"][0, 44:50], np.arange(6, 12) / 10) + assert not sample["action_mask"][:, 9].any() + assert not sample["action_mask"][:, 43].any() + + def test_unified_path_requires_bucket_motion_stats(self, tmp_path): + bucket = _make_bucket(tmp_path / "b", [{"length": 40}]) + (bucket / "meta" / "stats.json").unlink() + with pytest.raises(DataContractError, match="movement statistics"): + self._unified_reader(bucket) + def test_non_unified_path_does_not_require_base_motion_stats(self, tmp_path, patch_decode): bucket = _make_bucket( tmp_path / "b", @@ -485,22 +274,6 @@ def test_non_unified_path_does_not_require_base_motion_stats(self, tmp_path, pat assert sample["action"].shape[-1] == 20 assert sample["proprio"].shape[-1] == 20 - def test_stats_use_independent_motion_population_and_invert_action_gripper(self, tmp_path): - bucket = _make_bucket( - tmp_path / "b", - [{"length": 10, "flag": 0, "delta": 0}], - action_velocity=[1.0, 0.0, -1.0], - state_velocity=[0.0, 0.0, 0.0], - action_gripper=[0.2, 0.8], - state_gripper=[0.035, 0.125], - ) - _, partial, _ = _partial_bucket(str(bucket), "train", True, 0.7) - - assert "action.robot_velocity" in partial - assert "observation.state.robot_velocity" not in partial - np.testing.assert_allclose(partial["action.gripper"]["mean"], [0.8, 0.2]) - np.testing.assert_allclose(partial["observation.state.gripper"]["mean"], [1.0, 0.0]) - class TestTemporalAlignment: def test_truncated_tail_drops_clamped_final_target(self, tmp_path, patch_decode): @@ -543,143 +316,94 @@ def test_one_row_episode_is_not_sampleable(self, tmp_path, patch_decode): assert all(ds[i]["action_mask"].any() for i in range(len(ds))) -# --------------------------------------------------------------------------- -# Normalization stats must use and attest to the same trimmed population -# --------------------------------------------------------------------------- +class TestCleanExport: + def test_every_exported_episode_is_kept_without_annotation_files(self, tmp_path, patch_decode): + bucket = _make_bucket(tmp_path / "family" / "bucket", [{"length": 20}, {"length": 40}]) + ds = _reader(bucket) + assert ds._eps_df["episode_index"].tolist() == [0, 1] + assert ds._eps_df["length"].tolist() == [20, 40] + assert ds[0]["action"][0, 0].item() == 0 + ep1_start = int(ds._cum_n_starts[1]) + assert ds[ep1_start]["action"][0, 0].item() == 20 + assert _decoded_index(ds[ep1_start]["video"][0]) == 20 + def test_selection_knobs_are_not_part_of_the_config_surface(self): + assert "use_segment_annotations" not in AgiBotWorldDataset.CONFIG_KEYS + assert "segment_max_trim_ratio" not in AgiBotWorldDataset.CONFIG_KEYS -class TestTrimAwareStats: @staticmethod - def _stats_block(dim: int) -> dict: - return { - "mean": [0.0] * dim, - "std": [1.0] * dim, - "min": [-1.0] * dim, - "max": [1.0] * dim, - "q01": [-1.0] * dim, - "q99": [1.0] * dim, - } + def _write_stats(bucket, **overrides): + def block(dim): + return { + "mean": [0.0] * dim, + "std": [1.0] * dim, + "min": [-2.0] * dim, + "max": [2.0] * dim, + "q01": [-2.0] * dim, + "q99": [2.0] * dim, + } - def test_partial_bucket_excludes_trimmed_and_blacklisted_rows(self, tmp_path): - bucket = _make_bucket( - tmp_path / "root" / "b", - [ - {"length": 10, "flag": 1, "delta": 2}, # keep global rows 2..9 - {"length": 10, "flag": 2, "delta": 3}, # keep global rows 10..16 - {"length": 10, "flag": 3, "delta": 0}, # static: drop - {"length": 10, "flag": 0, "delta": 0}, # generic exclusion: drop - {"length": 10, "flag": 1, "delta": 8}, # >=70% trimmed: drop - ], - ) - (bucket / "meta" / "excluded_episodes.json").write_text(json.dumps({"episode_indices": [3]})) - - name, partial, population = _partial_bucket(str(bucket), "train", True, 0.7) - - assert name == "b" - assert population["num_episodes"] == 2 - assert population["num_rows"] == 15 - assert population["num_action_rows"] == 13 - assert population["num_state_rows"] == 15 - assert population["excluded_episode_indices"] == [3] - assert partial["action.ee_base"]["count"] == 13 - assert partial["observation.state.ee_base"]["count"] == 15 - # ee_base[0] encodes the global physical row in the fixture. These - # bounds prove that neither prefix/suffix outside the kept spans leaked. - assert partial["action.ee_base"]["min"][0] == pytest.approx(2.0) - assert partial["action.ee_base"]["max"][0] == pytest.approx(15.0) - assert partial["observation.state.ee_base"]["max"][0] == pytest.approx(16.0) - - def test_reader_accepts_matching_population_and_rejects_trim_policy_drift(self, tmp_path): - root = tmp_path / "root" - bucket = _make_bucket( - root / "b", - [{"length": 40, "flag": 1, "delta": 4}, {"length": 40, "flag": 0, "delta": 0}], - ) - _, _partial, bucket_population = _partial_bucket(str(bucket), "train", True, 0.7) - stats = { - "robot_type": "g2a", - "gripper_contract": _GRIPPER_CONTRACT, - "population": { - "schema_version": _STATS_SCHEMA_VERSION, - "split": "train", - "use_segment_annotations": True, - "segment_max_trim_ratio": 0.7, - "buckets": {"b": bucket_population}, - }, - "action.ee_base": self._stats_block(18), - "observation.state.ee_base": self._stats_block(18), - "action.gripper": self._stats_block(2), - "observation.state.gripper": self._stats_block(2), - } - (root / "meta").mkdir(exist_ok=True) - stats_path = root / "meta" / "stats_g2a.json" - stats_path.write_text(json.dumps(stats)) - - ds = _reader(bucket, normalize_mode="quantile", segment_max_trim_ratio=0.7) - # Reader-side materialization pins rot6d fields to identity regardless - # of the affine stats stored on disk. - np.testing.assert_array_equal(ds._action_norm_stats["mean"][3:9], np.zeros(6)) - np.testing.assert_array_equal(ds._action_norm_stats["std"][3:9], np.ones(6)) - - stats["gripper_contract"] = {**_GRIPPER_CONTRACT, "open_endpoint_m": 0.036} - stats_path.write_text(json.dumps(stats)) + stats = {"gripper_contract": _GRIPPER_CONTRACT} + for prefix in ("action", "observation.state"): + stats[f"{prefix}.ee_base"] = block(18) + stats[f"{prefix}.gripper"] = block(2) + stats.update(overrides) + path = bucket.parent / "meta" / "stats_g2a.json" + path.parent.mkdir(exist_ok=True) + path.write_text(json.dumps(stats)) + return path + + def test_loads_precomputed_family_root_stats(self, tmp_path, patch_decode): + bucket = _make_bucket(tmp_path / "family" / "bucket", [{"length": 40}]) + self._write_stats(bucket) + ds = _reader(bucket, normalize_mode="quantile") + sample = ds[1] + assert sample["action"][0, 0].item() == pytest.approx(0.5) + assert ds._action_norm_stats["q01"][3] == -1 + assert ds._action_norm_stats["q99"][3] == 1 + + def test_missing_stats_fail_without_generating_them(self, tmp_path): + bucket = _make_bucket(tmp_path / "family" / "bucket", [{"length": 40}]) + with pytest.raises(FileNotFoundError, match="stats_g2a.json"): + _reader(bucket, normalize_mode="quantile") + assert not (bucket.parent / "meta" / "stats_g2a.json").exists() + + def test_invalid_gripper_stats_contract_is_rejected(self, tmp_path): + bucket = _make_bucket(tmp_path / "family" / "bucket", [{"length": 40}]) + self._write_stats(bucket, gripper_contract={}) with pytest.raises(DataContractError, match="gripper_contract"): - _reader(bucket, normalize_mode="quantile", segment_max_trim_ratio=0.7) - - stats["gripper_contract"] = _GRIPPER_CONTRACT - stats["population"]["segment_max_trim_ratio"] = 0.8 - stats_path.write_text(json.dumps(stats)) - with pytest.raises(DataContractError, match="generated with"): - _reader(bucket, normalize_mode="quantile", segment_max_trim_ratio=0.7) - - def test_reader_rejects_non_train_stats(self, tmp_path): - root = tmp_path / "root" - bucket = _make_bucket(root / "b", [{"length": 40, "flag": 0, "delta": 0}]) - _, _partial, bucket_population = _partial_bucket(str(bucket), "train", True, 0.7) - stats = { - "robot_type": "g2a", - "gripper_contract": _GRIPPER_CONTRACT, - "population": { - "schema_version": _STATS_SCHEMA_VERSION, - "split": "val", - "use_segment_annotations": True, - "segment_max_trim_ratio": 0.7, - "buckets": {"b": bucket_population}, - }, - "action.ee_base": self._stats_block(18), - "observation.state.ee_base": self._stats_block(18), - "action.gripper": self._stats_block(2), - "observation.state.gripper": self._stats_block(2), - } - (root / "meta").mkdir(exist_ok=True) - (root / "meta" / "stats_g2a.json").write_text(json.dumps(stats)) - - with pytest.raises(DataContractError, match="train-derived"): - _reader(bucket, normalize_mode="quantile", segment_max_trim_ratio=0.7) - - def test_reader_rejects_removed_pooled_contributor(self, tmp_path): - root = tmp_path / "root" - bucket = _make_bucket(root / "b", [{"length": 40, "flag": 0, "delta": 0}]) - _, _partial, bucket_population = _partial_bucket(str(bucket), "train", True, 0.7) - stats = { - "robot_type": "g2a", - "gripper_contract": _GRIPPER_CONTRACT, - "population": { - "schema_version": _STATS_SCHEMA_VERSION, - "split": "train", - "use_segment_annotations": True, - "segment_max_trim_ratio": 0.7, - # ``removed`` no longer exists on disk but still polluted the - # pooled numbers in this stale file. - "buckets": {"b": bucket_population, "removed": bucket_population}, - }, - "action.ee_base": self._stats_block(18), - "observation.state.ee_base": self._stats_block(18), - "action.gripper": self._stats_block(2), - "observation.state.gripper": self._stats_block(2), - } - (root / "meta").mkdir(exist_ok=True) - (root / "meta" / "stats_g2a.json").write_text(json.dumps(stats)) - - with pytest.raises(DataContractError, match="pooled contributor set"): - _reader(bucket, normalize_mode="quantile", segment_max_trim_ratio=0.7) + _reader(bucket, normalize_mode="quantile") + + def test_stats_width_and_finiteness_are_validated(self, tmp_path): + bucket = _make_bucket(tmp_path / "family" / "bucket", [{"length": 40}]) + path = self._write_stats(bucket) + stats = json.loads(path.read_text()) + stats["action.ee_base"]["mean"] = [0.0] + path.write_text(json.dumps(stats)) + with pytest.raises((DataContractError, ValueError), match="18|shape"): + _reader(bucket, normalize_mode="quantile") + + +def test_each_camera_uses_its_own_timestamp_after_a_video_gap(tmp_path, monkeypatch): + bucket = _make_bucket(tmp_path / "bucket", [{"length": 8}, {"length": 8}]) + manifest = bucket / "meta" / "episodes" / "chunk-000.parquet" + episodes = pd.read_parquet(manifest) + for number, camera in enumerate(CAMS): + starts = np.array([100, 150]) + number * 100 + episodes[f"videos/{camera}/from_timestamp"] = starts / FPS + episodes[f"videos/{camera}/to_timestamp"] = (starts + 8) / FPS + episodes.to_parquet(manifest, index=False) + requested = {} + + def decode(path, frame_indices, height, width): + requested[Path(path).parts[-3]] = frame_indices + return [Image.new("RGB", (width, height)) for _ in frame_indices] + + monkeypatch.setattr("openwam.dataloader.bases.lerobot_v3_reader._decode_video_frames", decode) + dataset = AgiBotWorldDataset( + dataset_dir=str(bucket), normalize_mode=None, multiview=True, num_frames=5, video_stride=1 + ) + sample = dataset._getitem_impl(int(dataset._cum_n_starts[1])) + assert sample["action"][0, 0].item() == 8 + for number, camera in enumerate(CAMS): + assert requested[camera] == list(range(150 + number * 100, 155 + number * 100)) diff --git a/tests/dataloader/test_interndata_a1.py b/tests/dataloader/test_interndata_a1.py index c8d084071..e14a83796 100644 --- a/tests/dataloader/test_interndata_a1.py +++ b/tests/dataloader/test_interndata_a1.py @@ -15,7 +15,6 @@ import json import logging import os -import shutil from pathlib import Path import numpy as np @@ -32,10 +31,6 @@ detect_arm_layout, discover_a1_buckets, embodiment_key, - iter_data_shards, - load_excluded_episodes, - parse_shard_path, - resolve_trim_bounds, ) from openwam.dataloader.utils.eef import ( ARM10_DIM, @@ -44,7 +39,6 @@ quat_wxyz_to_rot6d, quat_xyzw_to_rot6d, ) -from openwam.dataloader.utils.lerobotv3 import DataContractError from openwam.dataloader.utils.normalization import ROT6D_DIMS_EEF20 HEAD = "images.rgb.head" @@ -138,7 +132,7 @@ def _make_bucket( rng = np.random.RandomState(seed) cols: dict = { "task_index": np.zeros(total, dtype=np.int64), - # Real A1 shards carry this; the cleaned-view filters key on it. + # Physical episode IDs are aligned with the cleaned manifest. "episode_index": np.repeat(np.arange(n_eps, dtype=np.int64), ep_len), } for i, side in enumerate(sides): @@ -671,110 +665,87 @@ def test_never_opened_gripper_falls_back_to_the_stroke_and_stays_zero(self, tmp_ class TestStats: - def _write_stats( - self, - root: Path, - embodiment: str, - *, - pin_rot6d: bool = True, - buckets=("cat/split_aloha/task",), - split="train", - trim_active=False, - min_keep=2, - ): + @staticmethod + def _write_stats(root: Path, embodiment: str, *, width: int = EEF_DIM): (root / "meta").mkdir(parents=True, exist_ok=True) eef = { - "mean": [0.0] * EEF_DIM, - "std": [1.0] * EEF_DIM, - "min": [-2.0] * EEF_DIM, - "max": [2.0] * EEF_DIM, - "q01": [-2.0] * EEF_DIM, - "q99": [2.0] * EEF_DIM, + "mean": [0.0] * width, + "std": [1.0] * width, + "min": [-2.0] * width, + "max": [2.0] * width, + "q01": [-2.0] * width, + "q99": [2.0] * width, } - if pin_rot6d: - for dim in ROT6D_DIMS_EEF20: - eef["q01"][dim] = -1.0 - eef["q99"][dim] = 1.0 - (root / "meta" / f"stats_{embodiment}.json").write_text( - json.dumps( - { - "eef": eef, - "population": { - "split": split, - "trim_active": trim_active, - "min_keep": min_keep, - "buckets": list(buckets), - }, - } - ) - ) + path = root / "meta" / f"stats_{embodiment}.json" + path.write_text(json.dumps({"eef": eef})) + return path - def test_missing_stats_file_raises_actionable_error(self, tmp_path, patch_decode): - d = _make_bucket(tmp_path, "cat/split_aloha/task") - with pytest.raises(FileNotFoundError, match="interndata_a1_stats_computation"): - InternDataA1Dataset(str(d), a1_stats_root=str(tmp_path), normalize_mode="quantile") - - def test_missing_stats_error_prescribes_the_stats_root_it_looked_in(self, tmp_path, patch_decode): - """The suggested command must carry --stats_root, pointing at the SAME - directory the failed lookup used. Otherwise the one scenario stats_root - exists for — a read-only dataset mount — hands the user a command that - writes where this lookup does not read, reproducing the same error.""" - d = _make_bucket(tmp_path, "cat/split_aloha/task") - elsewhere = tmp_path / "scratch" - with pytest.raises(FileNotFoundError) as e: - InternDataA1Dataset(str(d), a1_stats_root=str(elsewhere), normalize_mode="quantile") - msg = str(e.value) - assert f"--stats_root {elsewhere}" in msg - # And it names the path it actually looked for, so the two agree. - assert str(elsewhere / "meta" / "stats_split_aloha.json") in msg + def test_missing_precomputed_stats_fails_without_writing(self, tmp_path, patch_decode): + d = _make_bucket(tmp_path, "family/cat/split_aloha/task") + before = {p.relative_to(tmp_path) for p in tmp_path.rglob("*")} + with pytest.raises(FileNotFoundError, match="stats_split_aloha.json"): + InternDataA1Dataset(str(d), normalize_mode="quantile") + assert {p.relative_to(tmp_path) for p in tmp_path.rglob("*")} == before def test_normalize_mode_null_needs_no_stats(self, tmp_path, patch_decode): d = _make_bucket(tmp_path, "cat/split_aloha/task") ds = InternDataA1Dataset(str(d), normalize_mode=None, num_frames=9, video_stride=4) assert ds._normalization_stats is None - def test_shared_stats_root_is_used_not_the_bucket_dir(self, tmp_path, patch_decode): - """Buckets sit at variable depth; stats live once at the dataset root.""" - d = _make_bucket(tmp_path, "cat/split_aloha/task") - self._write_stats(tmp_path, "split_aloha") - ds = InternDataA1Dataset( - str(d), a1_stats_root=str(tmp_path), normalize_mode="quantile", num_frames=9, video_stride=4 - ) + def test_leaf_uses_stats_in_renamed_family_ancestor(self, tmp_path, patch_decode): + family = tmp_path / "my_portable_family" + d = _make_bucket(family, "cat/split_aloha/task") + stats = self._write_stats(family, "split_aloha") + before = stats.read_bytes() + ds = InternDataA1Dataset(str(d), normalize_mode="quantile", num_frames=9, video_stride=4) assert ds._normalization_stats is not None + assert stats.read_bytes() == before + assert not list(tmp_path.rglob("normalization_stats.npy")) + + def test_nearest_dataset_stats_are_used(self, tmp_path, patch_decode): + family = tmp_path / "family" + d = _make_bucket(family, "cat/split_aloha/task") + outer = self._write_stats(tmp_path, "split_aloha") + inner = self._write_stats(family, "split_aloha") + payload = json.loads(outer.read_text()) + payload["eef"]["q99"][0] = 77.0 + outer.write_text(json.dumps(payload)) + ds = InternDataA1Dataset(str(d), normalize_mode="quantile") + assert ds._normalization_stats["q99"][0] == 2.0 + assert inner.is_file() def test_rot6d_dims_pass_through_normalization(self, tmp_path, patch_decode): - """Pinned rot6d stats must leave the rotation representation untouched, - while pos/gripper dims are rescaled.""" d = _make_bucket(tmp_path, "cat/split_aloha/task") self._write_stats(tmp_path, "split_aloha") raw = InternDataA1Dataset(str(d), normalize_mode=None, num_frames=9, video_stride=4) - norm = InternDataA1Dataset( - str(d), a1_stats_root=str(tmp_path), normalize_mode="quantile", num_frames=9, video_stride=4 - ) - a_raw = raw[0]["action"].numpy() - a_norm = norm[0]["action"].numpy() + norm = InternDataA1Dataset(str(d), normalize_mode="quantile", num_frames=9, video_stride=4) + a_raw, a_norm = raw[0]["action"].numpy(), norm[0]["action"].numpy() np.testing.assert_allclose(a_norm[:, list(ROT6D_DIMS_EEF20)], a_raw[:, list(ROT6D_DIMS_EEF20)], atol=1e-6) - # xyz dims used q01/q99 = +-2 -> genuinely rescaled assert not np.allclose(a_norm[:, 0:3], a_raw[:, 0:3], atol=1e-6) def test_wrong_width_stats_rejected(self, tmp_path, patch_decode): d = _make_bucket(tmp_path, "cat/split_aloha/task") - (tmp_path / "meta").mkdir(parents=True, exist_ok=True) - (tmp_path / "meta" / "stats_split_aloha.json").write_text( - json.dumps( - { - "eef": {k: [0.0] * 10 for k in ("mean", "std", "min", "max", "q01", "q99")}, - "population": { - "split": "train", - "trim_active": False, - "min_keep": 2, - "buckets": ["cat/split_aloha/task"], - }, - } - ) - ) + self._write_stats(tmp_path, "split_aloha", width=10) with pytest.raises(ValueError, match="!= expected 20"): - InternDataA1Dataset(str(d), a1_stats_root=str(tmp_path), normalize_mode="quantile") + InternDataA1Dataset(str(d), normalize_mode="quantile") + + def test_nonfinite_stats_rejected(self, tmp_path, patch_decode): + d = _make_bucket(tmp_path, "cat/split_aloha/task") + path = self._write_stats(tmp_path, "split_aloha") + payload = json.loads(path.read_text()) + payload["eef"]["q99"][0] = float("nan") + path.write_text(json.dumps(payload)) + with pytest.raises(ValueError, match="non-finite"): + InternDataA1Dataset(str(d), normalize_mode="quantile") + + def test_required_quantile_fields_are_not_computed(self, tmp_path, patch_decode): + d = _make_bucket(tmp_path, "cat/split_aloha/task") + path = self._write_stats(tmp_path, "split_aloha") + payload = json.loads(path.read_text()) + del payload["eef"]["q99"] + path.write_text(json.dumps(payload)) + with pytest.raises(KeyError, match="q99"): + InternDataA1Dataset(str(d), normalize_mode="quantile") class TestUnifyScatter: @@ -827,25 +798,22 @@ def test_root_mode_discovers_and_aggregates_mixed_embodiments(self, tmp_path, pa assert len(ds) == sum(len(b) for b in ds.buckets) assert ds[0]["action"].shape == (8, EEF_DIM) - def test_root_mode_rejects_removed_pooled_contributor(self, tmp_path, patch_decode): + def test_root_mode_uses_existing_family_stats_without_historical_provenance(self, tmp_path, patch_decode): _make_bucket(tmp_path, "cat/split_aloha/current") - (tmp_path / "meta").mkdir(exist_ok=True) - (tmp_path / "meta" / "stats_split_aloha.json").write_text( - json.dumps( - { - "population": { - "buckets": ["cat/split_aloha/current", "cat/split_aloha/removed"], - "empty_buckets": [], - } - } - ) + TestStats._write_stats(tmp_path, "split_aloha") + ds = InternDataA1Dataset.from_config( + {"dataset_dir": str(tmp_path), "normalize_mode": "quantile", "num_frames": 9}, + split="train", ) + assert len(ds.buckets) == 1 + assert ds.buckets[0]._normalization_stats is not None - with pytest.raises(DataContractError, match="contributor set no longer matches"): - InternDataA1Dataset.from_config( - {"dataset_dir": str(tmp_path), "normalize_mode": "quantile"}, - split="train", - ) + def test_root_missing_embodiment_stats_cannot_silently_drop_a_bucket(self, tmp_path, patch_decode): + _make_bucket(tmp_path, "cat/split_aloha/present") + _make_bucket(tmp_path, "cat/franka/missing", layout="single_arm", robot_type="Franka") + TestStats._write_stats(tmp_path, "split_aloha") + with pytest.raises(FileNotFoundError, match="stats_franka.json"): + InternDataA1Dataset.from_config({"dataset_dir": str(tmp_path), "normalize_mode": "quantile"}) def test_single_bucket_mode(self, tmp_path, patch_decode): d = _make_bucket(tmp_path, "cat/split_aloha/task") @@ -866,8 +834,8 @@ def test_bucket_ids_are_root_relative_paths(self, tmp_path, patch_decode): ids = sorted(b._dataset_id for b in ds.buckets) assert ids == ["cat/franka/taskA/google_scan-book", "cat/franka/taskB/google_scan-book"] - def test_empty_root_raises_pointing_at_extraction(self, tmp_path): - with pytest.raises(FileNotFoundError, match="extracted"): + def test_empty_root_raises_with_missing_metadata(self, tmp_path): + with pytest.raises(FileNotFoundError, match="meta/info.json"): InternDataA1Dataset.from_config({"dataset_dir": str(tmp_path)}, split="train") def test_registered_under_interndata_a1(self): @@ -877,188 +845,43 @@ def test_registered_under_interndata_a1(self): # --------------------------------------------------------------------------- -# Cleaned-view / trim regression coverage +# Cleaned-manifest offset and timestamp coverage # --------------------------------------------------------------------------- class TestCleanedViewOffsets: - """A cleaned view must not express deletions by dropping meta/episodes rows. - - ``_data_row_offset`` is a ``groupby(chunk, file).cumsum()`` over the rows - currently in ``eps_df``. A cleaned view symlinks ``data/`` at the untouched - source shards, so a shortened manifest makes every deleted episode's length - vanish from that sum and silently slides later episodes onto earlier frames. - ``meta/excluded_episodes.json`` is applied after the offsets are computed, - so it does not have this failure mode. - """ + """Clean manifests and file IDs are authoritative; no source repair runs.""" - def test_excluded_first_episode_leaves_the_second_at_its_physical_offset(self, tmp_path, patch_decode): + def test_cleaned_data_does_not_apply_an_old_external_exclusion_file(self, tmp_path, patch_decode): d = _make_bucket(tmp_path, "cat/split_aloha/task", n_eps=2, ep_len=4) (d / "meta" / "excluded_episodes.json").write_text(json.dumps({"episode_indices": [0]})) - ds = InternDataA1Dataset(str(d), normalize_mode=None, num_frames=2, video_stride=1) + assert ds._eps_df["episode_index"].tolist() == [0, 1] + assert ds._ep_data_row_offset.tolist() == [0, 4] + for off in ds._ep_video_frame_offsets.values(): + assert off.tolist() == [0, 4] - assert ds._eps_df["episode_index"].tolist() == [1] - # Episode 1 physically starts at row 4 of the shard; the exclusion must - # not renumber it to 0. - assert int(ds._ep_data_row_offset[0]) == 4 - assert ds._ep_video_frame_offsets, "no camera offsets resolved" - for cam, off in ds._ep_video_frame_offsets.items(): - assert int(off[0]) == 4, f"{cam} offset collapsed to {int(off[0])}" - - def test_dropping_a_manifest_row_no_longer_shifts_the_survivor(self, tmp_path, patch_decode): - """The failure mode this class was written for, now closed at the source. - - `_add_data_offsets` rebuilds each offset from `dataset_from_index` and the - shards' real row counts, so it no longer depends on a cumsum over the - surviving rows — a shortened manifest cannot displace anything. This once - asserted the WRONG offset (0 instead of 4) to pin the hazard; it now - asserts the right one, so the guarantee is pinned rather than the bug. - - excluded_episodes.json remains the correct way to express deletions — - it keeps meta/episodes intact and self-describing — but offsets are no - longer the reason why. - """ + def test_incomplete_manifest_is_rejected_without_repair(self, tmp_path, patch_decode): d = _make_bucket(tmp_path, "cat/split_aloha/task", n_eps=2, ep_len=4) - man = d / "meta" / "episodes" / "chunk-000" / "file-000.parquet" - t = pq.read_table(man) - pq.write_table(t.slice(1, 1), man) # keep only episode 1 - - ds = InternDataA1Dataset(str(d), normalize_mode=None, num_frames=2, video_stride=1) - - assert ds._eps_df["episode_index"].tolist() == [1] - assert int(ds._ep_data_row_offset[0]) == 4 # physical start, not 0 - # And EVERY camera, not just the parquet rows. Checking only the data - # offset once let a half-fix advertise alignment safety while the video - # offsets still collapsed to 0 — pairing episode 1's actions with - # episode 0's frames, which is worse than either error alone. - assert ds._ep_video_frame_offsets, "no camera offsets resolved" - for cam, off in ds._ep_video_frame_offsets.items(): - assert int(off[0]) == 4, f"{cam} offset collapsed to {int(off[0])}" - - -class TestTrimTooShortIsLeftWhole: - """A trim that would leave less than one window keeps the episode intact — - and the stats path must make the identical call (they share - :func:`resolve_trim_bounds`), or the normalizer describes episodes the - reader never emits that way.""" - - def test_reader_leaves_a_too_short_trim_whole(self, tmp_path, patch_decode): - d = _make_bucket(tmp_path, "cat/split_aloha/task", n_eps=1, ep_len=4) - trim = tmp_path / "trim.csv" - trim.write_text( - "dataset,episode_index,total_frames,trim_head_to,trim_tail_from\n" - "cat/split_aloha/task,0,4,3,\n" # would leave 1 frame < min_len 2 - ) - ds = InternDataA1Dataset( - str(d), - dataset_id="cat/split_aloha/task", - normalize_mode=None, - trim_csv=str(trim), - num_frames=2, - video_stride=1, - ) - assert int(ds._eps_df["length"].iloc[0]) == 4 - assert int(ds._ep_data_row_offset[0]) == 0 - - def test_resolve_trim_bounds_rejects_the_short_case_and_accepts_a_valid_one(self): - assert resolve_trim_bounds((3, None, 4), 4, 2) is None # leaves 1 < 2 - assert resolve_trim_bounds((1, None, 4), 4, 2) == (1, 4) # leaves 3 - assert resolve_trim_bounds((1, None, 99), 4, 2) is None # stale total_frames - assert resolve_trim_bounds((0, None, 4), 4, 2) is None # no-op - - -class TestStaleShardIndex: - """`data/file_index` goes stale at shard boundaries: the episode that starts a - new shard keeps the previous file's index. - - In affected multi-shard buckets, the base - `groupby(chunk,file).cumsum()` points boundary episodes into the previous - shard — a valid row range that reads back real numbers, so it pairs an - episode with another one's frames and raises nothing. - """ - - def _two_shard_bucket(self, tmp_path): - d = _make_bucket(tmp_path, "cat/split_aloha/task", n_eps=2, ep_len=4) - src = d / "data" / "chunk-000" / "file-000.parquet" - t = pq.read_table(src) - pq.write_table(t.slice(0, 4), src) - pq.write_table(t.slice(4, 4), d / "data" / "chunk-000" / "file-001.parquet") - man = d / "meta" / "episodes" / "chunk-000" / "file-000.parquet" - m = pq.read_table(man).to_pydict() - m["data/file_index"] = [0, 0] # ep1 really lives in file-001 - m["dataset_from_index"] = [0, 4] - m["dataset_to_index"] = [4, 8] - pq.write_table(pa.Table.from_pydict(m), man) - return d - - def test_the_boundary_episode_resolves_to_its_real_shard(self, tmp_path, patch_decode): - d = self._two_shard_bucket(tmp_path) - ds = InternDataA1Dataset(str(d), normalize_mode=None, num_frames=2, video_stride=1) - i = ds._eps_df.index[ds._eps_df["episode_index"] == 1][0] - pos = list(ds._eps_df.index).index(i) - assert int(ds._eps_df["data/file_index"].loc[i]) == 1 - assert int(ds._ep_data_row_offset[pos]) == 0 - - def test_the_unaffected_episode_is_untouched(self, tmp_path, patch_decode): - d = self._two_shard_bucket(tmp_path) - ds = InternDataA1Dataset(str(d), normalize_mode=None, num_frames=2, video_stride=1) - i = ds._eps_df.index[ds._eps_df["episode_index"] == 0][0] - pos = list(ds._eps_df.index).index(i) - assert int(ds._eps_df["data/file_index"].loc[i]) == 0 - assert int(ds._ep_data_row_offset[pos]) == 0 - - -class TestShardCompleteness: - """A missing shard must fail loudly rather than resolve onto another episode. - - The cumulative boundaries close over whatever files exist, so with a shard - gone every episode start still lands inside the total found on disk, and each - later episode maps onto a plausible row of the wrong file. - """ - - def test_a_missing_middle_shard_is_rejected(self, tmp_path, patch_decode): - d = _make_bucket(tmp_path, "cat/split_aloha/task", n_eps=3, ep_len=4) - src = d / "data" / "chunk-000" / "file-000.parquet" - t = pq.read_table(src) - pq.write_table(t.slice(0, 4), src) - pq.write_table(t.slice(8, 4), d / "data" / "chunk-000" / "file-002.parquet") - man = d / "meta" / "episodes" / "chunk-000" / "file-000.parquet" - m = pq.read_table(man).to_pydict() - m["dataset_from_index"] = [0, 4, 8] - m["dataset_to_index"] = [4, 8, 12] - pq.write_table(pa.Table.from_pydict(m), man) - with pytest.raises(ValueError, match="shard is missing"): + path = d / "meta" / "episodes" / "chunk-000" / "file-000.parquet" + pq.write_table(pq.read_table(path).slice(1, 1), path) + with pytest.raises(ValueError, match="missing rows"): InternDataA1Dataset(str(d), normalize_mode=None, num_frames=2, video_stride=1) - def test_an_episode_overrunning_its_shard_is_rejected(self, tmp_path, patch_decode): - """Equal totals are not enough — 3+5 physical rows satisfy a 4+4 manifest.""" - d = _make_bucket(tmp_path, "cat/split_aloha/task", n_eps=2, ep_len=4) - src = d / "data" / "chunk-000" / "file-000.parquet" - t = pq.read_table(src) - pq.write_table(t.slice(0, 3), src) - pq.write_table(t.slice(3, 5), d / "data" / "chunk-000" / "file-001.parquet") - man = d / "meta" / "episodes" / "chunk-000" / "file-000.parquet" - m = pq.read_table(man).to_pydict() - m["dataset_from_index"] = [0, 4] - m["dataset_to_index"] = [4, 8] - pq.write_table(pa.Table.from_pydict(m), man) - with pytest.raises(ValueError, match="shard holding only"): - InternDataA1Dataset(str(d), normalize_mode=None, num_frames=2, video_stride=1) - - def test_complete_shards_are_accepted(self, tmp_path, patch_decode): - d = _make_bucket(tmp_path, "cat/split_aloha/task", n_eps=2, ep_len=4) - src = d / "data" / "chunk-000" / "file-000.parquet" - t = pq.read_table(src) - pq.write_table(t.slice(0, 4), src) - pq.write_table(t.slice(4, 4), d / "data" / "chunk-000" / "file-001.parquet") - man = d / "meta" / "episodes" / "chunk-000" / "file-000.parquet" - m = pq.read_table(man).to_pydict() - m["dataset_from_index"] = [0, 4] - m["dataset_to_index"] = [4, 8] - pq.write_table(pa.Table.from_pydict(m), man) + def test_canonical_file_ids_and_file_local_offsets_are_preserved(self, tmp_path, patch_decode): + d = _make_bucket(tmp_path, "cat/split_aloha/task", n_eps=3, ep_len=4) + manifest = d / "meta" / "episodes" / "chunk-000" / "file-000.parquet" + rows = pq.read_table(manifest).to_pydict() + rows["data/file_index"] = [0, 0, 1] + pq.write_table(pa.Table.from_pydict(rows), manifest) + path = d / "data" / "chunk-000" / "file-000.parquet" + data = pq.read_table(path) + pq.write_table(data.slice(0, 8), path) + pq.write_table(data.slice(8, 4), path.with_name("file-001.parquet")) ds = InternDataA1Dataset(str(d), normalize_mode=None, num_frames=2, video_stride=1) - assert len(ds._eps_df) == 2 + assert ds._eps_df["data/file_index"].tolist() == [0, 0, 1] + assert ds._ep_data_row_offset.tolist() == [0, 4, 0] + assert len(ds) > 0 class TestVideoOffsetValidation: @@ -1080,7 +903,7 @@ def test_non_finite_timestamps_are_rejected(self, tmp_path, patch_decode): m = pq.read_table(man).to_pydict() m["videos/images.rgb.head/from_timestamp"] = [float("nan"), 0.1333] pq.write_table(pa.Table.from_pydict(m), man) - with pytest.raises(ValueError, match="non-finite or negative"): + with pytest.raises(ValueError, match="invalid from_timestamp"): InternDataA1Dataset(str(d), normalize_mode=None, num_frames=2, video_stride=1) def test_two_episodes_sharing_a_frame_offset_are_rejected(self, tmp_path, patch_decode): @@ -1091,6 +914,7 @@ def test_two_episodes_sharing_a_frame_offset_are_rejected(self, tmp_path, patch_ k = f"videos/{cam}/from_timestamp" if k in m: m[k] = [0.0, 0.0] # both episodes start at frame 0 of one shard + m[f"videos/{cam}/to_timestamp"] = [4 / 30, 4 / 30] pq.write_table(pa.Table.from_pydict(m), man) with pytest.raises(ValueError, match="overlaps episode"): InternDataA1Dataset(str(d), normalize_mode=None, num_frames=2, video_stride=1) @@ -1148,7 +972,7 @@ def test_a_video_span_disagreeing_with_length_is_rejected(self, tmp_path, patch_ pytest.skip("fixture carries no to_timestamp column") m[tk] = [m["videos/images.rgb.head/from_timestamp"][0] + 9 / 30, m[tk][1]] pq.write_table(pa.Table.from_pydict(m), man) - with pytest.raises(ValueError, match="describes two different episodes"): + with pytest.raises(ValueError, match="span does not match the declared episode length"): InternDataA1Dataset(str(d), normalize_mode=None, num_frames=2, video_stride=1) def test_a_non_finite_video_span_is_rejected_not_skipped(self, tmp_path, patch_decode): @@ -1161,7 +985,7 @@ def test_a_non_finite_video_span_is_rejected_not_skipped(self, tmp_path, patch_d pytest.skip("fixture carries no to_timestamp column") m[tk] = [float("nan"), m[tk][1]] pq.write_table(pa.Table.from_pydict(m), man) - with pytest.raises(ValueError, match="describes two different episodes"): + with pytest.raises(ValueError, match="span does not match the declared episode length"): InternDataA1Dataset(str(d), normalize_mode=None, num_frames=2, video_stride=1) @@ -1194,254 +1018,20 @@ def test_a_range_disagreeing_with_length_is_rejected(self, tmp_path, patch_decod with pytest.raises(ValueError, match="declares length=3"): InternDataA1Dataset(str(d), normalize_mode=None, num_frames=2, video_stride=1) - def test_a_gap_is_still_accepted(self, tmp_path, patch_decode): - """Non-overlap is required; a perfect tiling is NOT. - - A manifest gap means rows belong to no episode, which is how a dropped - row expresses "skip this one" — handled correctly by the offset rebuild - and asserted by TestCleanedViewOffsets. Demanding contiguity would - reject that working configuration to catch nothing. - """ - d = _make_bucket(tmp_path, "cat/split_aloha/task", n_eps=3, ep_len=4) - man = d / "meta" / "episodes" / "chunk-000" / "file-000.parquet" - m = pq.read_table(man).to_pydict() - keep = [i for i in range(3) if i != 1] # drop the middle episode - pq.write_table(pa.Table.from_pydict({k: [v[i] for i in keep] for k, v in m.items()}), man) - r = InternDataA1Dataset(str(d), normalize_mode=None, num_frames=2, video_stride=1) - assert len(r._eps_df) == 2 - - def test_a_substituted_equal_length_shard_is_rejected(self, tmp_path, patch_decode): - """An equal-length copy of another shard passes every count-based check. - - file-001 replaced by a copy of file-000: the totals still balance, every - episode still fits its resolved shard, and the reader constructs — while - manifest episodes [0,1] both resolve onto physical episode 0. The - shard's own episode_index range (from parquet row-group statistics, so - no column is read) is what distinguishes them. - """ - d = _make_bucket(tmp_path, "cat/split_aloha/task", n_eps=2, ep_len=4) - data_dir = d / "data" / "chunk-000" - src = pq.read_table(data_dir / "file-000.parquet").to_pydict() - half = {k: v[:4] for k, v in src.items()} - pq.write_table(pa.Table.from_pydict(half), data_dir / "file-000.parquet") - pq.write_table(pa.Table.from_pydict(half), data_dir / "file-001.parquet") - with pytest.raises(ValueError, match="only holds episodes 0..0"): - InternDataA1Dataset(str(d), normalize_mode=None, num_frames=2, video_stride=1) - - def test_a_stray_backup_parquet_is_not_a_shard(self, tmp_path, patch_decode): - """`file-000.backup.parquet` matches the glob but is not a shard. - - The reader's exact parse always dropped it; the stats generator's glob - did not, so a copy left in place doubled the statistics population while - the reader read half of it. Both sides now enumerate through - `iter_data_shards`. - """ - d = _make_bucket(tmp_path, "cat/split_aloha/task", n_eps=2, ep_len=4) - data_dir = d / "data" / "chunk-000" - shutil.copy(data_dir / "file-000.parquet", data_dir / "file-000.backup.parquet") - assert parse_shard_path(data_dir / "file-000.backup.parquet") is None - assert [p.name for _, _, p in iter_data_shards(d)] == ["file-000.parquet"] - r = InternDataA1Dataset(str(d), normalize_mode=None, num_frames=2, video_stride=1) - assert len(r) > 0 - - def test_shards_are_ordered_by_index_not_lexically(self, tmp_path): - """`file-1000` follows `file-999`; a lexical sort puts it first. - - Both names are canonical (`{:03d}` widens past three digits), so this is - reachable on a corpus with over a thousand shards. - """ - b = tmp_path / "b" - (b / "data" / "chunk-000").mkdir(parents=True) - for i in (0, 999, 1000): - pq.write_table(pa.table({"x": [i]}), b / "data" / "chunk-000" / f"file-{i:03d}.parquet") - assert [f for _, f, _ in iter_data_shards(b)] == [0, 999, 1000] - - def test_a_non_canonical_shard_name_is_not_a_shard(self, tmp_path, patch_decode): - """`file-0.parquet` is a name the data_path template cannot produce. - - Sampling rebuilds the path from `info.json[data_path]` with `{:03d}`, so - an unpadded file is enumerable but unreadable: the scan pooled its rows - and the reader built windows over them, and only the first real data - load failed — on a different, non-existent path. Rejecting it at - enumeration turns that into an immediate, accurate error. - """ - d = _make_bucket(tmp_path, "cat/split_aloha/task", n_eps=2, ep_len=4) - shard = d / "data" / "chunk-000" / "file-000.parquet" - shard.rename(shard.with_name("file-0.parquet")) - assert parse_shard_path(shard.with_name("file-0.parquet")) is None - assert iter_data_shards(d) == [] - with pytest.raises(FileNotFoundError, match="No data parquet files"): - InternDataA1Dataset(str(d), normalize_mode=None, num_frames=2, video_stride=1) - - def test_an_aliased_shard_name_is_not_a_shard(self, tmp_path): - """`file-00.parquet` would otherwise claim shard 0 alongside `file-000`.""" - b = tmp_path / "b" - (b / "data" / "chunk-000").mkdir(parents=True) - for name in ("file-000.parquet", "file-00.parquet"): - pq.write_table(pa.table({"x": [0]}), b / "data" / "chunk-000" / name) - assert [f for _, f, _ in iter_data_shards(b)] == [0] - - def test_a_canonical_shard_under_a_non_canonical_chunk_is_rejected(self, tmp_path): - b = tmp_path / "b" - (b / "data" / "chunk-0").mkdir(parents=True) - pq.write_table(pa.table({"x": [0]}), b / "data" / "chunk-0" / "file-000.parquet") - assert iter_data_shards(b) == [] +def test_root_mode_cannot_borrow_stats_from_outside_configured_family(tmp_path, patch_decode): + family = tmp_path / "family" + _make_bucket(family, "cat/split_aloha/task") + TestStats._write_stats(tmp_path, "split_aloha") + with pytest.raises(FileNotFoundError, match="stats_split_aloha.json"): + InternDataA1Dataset.from_config({"dataset_dir": str(family), "normalize_mode": "quantile"}) -class TestExclusionParser: - """One strict parser, because the two casts disagreed.""" - def test_quoted_indices_are_refused_rather_than_reinterpreted(self, tmp_path): - """`["0"]` excluded episode 0 in the generator and nothing in the reader. - - The base reader keeps the JSON values verbatim and matches them against - an integer column, so a quoted index excludes nothing there, while an - `int()` cast in the generator excluded it. Refusing is the one answer - that is identical on both sides by construction. - """ - b = tmp_path / "b" - (b / "meta").mkdir(parents=True) - (b / "meta" / "excluded_episodes.json").write_text(json.dumps({"episode_indices": ["0"]})) - with pytest.raises(ValueError, match="must be JSON integers"): - load_excluded_episodes(b) - - def test_booleans_are_refused(self, tmp_path): - """`True` passes isinstance(x, int) and would silently become episode 1.""" - b = tmp_path / "b" - (b / "meta").mkdir(parents=True) - (b / "meta" / "excluded_episodes.json").write_text(json.dumps({"episode_indices": [True]})) - with pytest.raises(ValueError, match="must be JSON integers"): - load_excluded_episodes(b) - - def test_the_actual_reader_refuses_quoted_indices_too(self, tmp_path, patch_decode): - """The helper being strict is worthless if the reader never calls it. - - The base keeps the JSON values verbatim and matches them against an - integer column, so `["0"]` excluded nothing there while the generator - excluded episode 0. Asserting on the helper alone did not catch that — - this instantiates the reader. - """ - d = _make_bucket(tmp_path, "cat/split_aloha/task", n_eps=2, ep_len=4) - (d / "meta" / "excluded_episodes.json").write_text(json.dumps({"episode_indices": ["0"]})) - with pytest.raises(ValueError, match="must be JSON integers"): - InternDataA1Dataset(str(d), normalize_mode=None, num_frames=2, video_stride=1) - - def test_the_actual_reader_still_applies_integer_exclusions(self, tmp_path, patch_decode): - d = _make_bucket(tmp_path, "cat/split_aloha/task", n_eps=2, ep_len=4) - (d / "meta" / "excluded_episodes.json").write_text(json.dumps({"episode_indices": [0]})) - r = InternDataA1Dataset(str(d), normalize_mode=None, num_frames=2, video_stride=1) - assert r._eps_df["episode_index"].tolist() == [1] - - def test_plain_integers_are_accepted(self, tmp_path): - b = tmp_path / "b" - (b / "meta").mkdir(parents=True) - (b / "meta" / "excluded_episodes.json").write_text(json.dumps({"episode_indices": [0, 3]})) - assert load_excluded_episodes(b) == {0, 3} - - def test_a_missing_file_means_nothing_excluded(self, tmp_path): - assert load_excluded_episodes(tmp_path) == set() - - -class TestStatsPopulationContract: - """The four ways the generator and the reader were seen to disagree.""" - - def _bucket_and_stats(self, tmp_path, **pop): - d = _make_bucket(tmp_path, "cat/split_aloha/task", n_eps=2, ep_len=8) - block = {"split": "train", "trim_active": False, "min_keep": 2, "buckets": ["cat/split_aloha/task"]} - block.update(pop) - (tmp_path / "meta").mkdir(parents=True, exist_ok=True) - eef = {k: [0.0] * EEF_DIM for k in ("mean", "min", "q01")} - eef.update({k: [1.0] * EEF_DIM for k in ("std", "max", "q99")}) - (tmp_path / "meta" / "stats_split_aloha.json").write_text(json.dumps({"eef": eef, "population": block})) - return d - - def _read(self, d, tmp_path, **kw): - return InternDataA1Dataset( - str(d), a1_stats_root=str(tmp_path), normalize_mode="quantile", num_frames=2, video_stride=1, **kw - ) - - def test_val_derived_stats_are_refused(self, tmp_path, patch_decode): - d = self._bucket_and_stats(tmp_path, split="val") - with pytest.raises(DataContractError, match="computed over the 'val' split"): - self._read(d, tmp_path) - - def test_train_derived_stats_are_accepted_by_a_val_reader(self, tmp_path, patch_decode): - """A val reader must LOAD train stats, not demand val-derived ones. - - There is one stats file per embodiment and it describes the training - distribution by definition. Comparing the recorded split against the - reader's own split refused every val run against a correctly generated - file — a regression with no failing test until this one. - """ - d = self._bucket_and_stats(tmp_path, split="train") - info = json.loads((d / "meta" / "info.json").read_text()) - info["splits"] = {"train": "0:1", "val": "1:2"} - (d / "meta" / "info.json").write_text(json.dumps(info)) - r = self._read(d, tmp_path, split="val") - assert r._normalization_stats is not None - - def test_untrimmed_stats_are_refused_by_a_trimming_reader(self, tmp_path, patch_decode): - d = self._bucket_and_stats(tmp_path, trim_active=False) - csv = tmp_path / "trim.csv" - csv.write_text("dataset,episode_index,total_frames,trim_head_to,trim_tail_from\ncat/split_aloha/task,0,8,2,6\n") - with pytest.raises(DataContractError, match="without a trim list but this reader is trimming"): - self._read(d, tmp_path, trim_csv=str(csv)) - - def test_trimmed_stats_are_refused_by_an_untrimmed_reader(self, tmp_path, patch_decode): - d = self._bucket_and_stats(tmp_path, trim_active=True) - with pytest.raises(DataContractError, match="with a trim list but this reader is not"): - self._read(d, tmp_path) - - def test_a_different_keep_bound_is_refused(self, tmp_path, patch_decode): - d = self._bucket_and_stats(tmp_path, trim_active=True, min_keep=33) - csv = tmp_path / "trim.csv" - csv.write_text("dataset,episode_index,total_frames,trim_head_to,trim_tail_from\ncat/split_aloha/task,0,8,2,6\n") - with pytest.raises(DataContractError, match="--min_keep=33"): - self._read(d, tmp_path, trim_csv=str(csv)) - - def test_a_bucket_absent_from_the_scan_is_refused(self, tmp_path, patch_decode): - """The bucket dropped out of the scan but still loads the shared file.""" - d = self._bucket_and_stats(tmp_path, buckets=["cat/split_aloha/other"]) - with pytest.raises(DataContractError, match="none of them this one"): - self._read(d, tmp_path) - - def test_a_stats_file_without_the_block_is_refused(self, tmp_path, patch_decode): - d = _make_bucket(tmp_path, "cat/split_aloha/task", n_eps=2, ep_len=8) - (tmp_path / "meta").mkdir(parents=True, exist_ok=True) - eef = {k: [0.0] * EEF_DIM for k in ("mean", "min", "q01")} - eef.update({k: [1.0] * EEF_DIM for k in ("std", "max", "q99")}) - (tmp_path / "meta" / "stats_split_aloha.json").write_text(json.dumps({"eef": eef})) - with pytest.raises(DataContractError, match="no 'population' block"): - self._read(d, tmp_path) - - def test_root_mode_does_not_silently_drop_a_bucket_missing_from_stats(self, tmp_path, patch_decode): - """DataContractError must escape build_multibucket's tolerant wrapper.""" - _make_bucket(tmp_path, "cat/split_aloha/good", n_eps=2, ep_len=8) - _make_bucket(tmp_path, "cat/split_aloha/bad", n_eps=2, ep_len=8) - (tmp_path / "meta").mkdir(parents=True, exist_ok=True) - eef = {k: [0.0] * EEF_DIM for k in ("mean", "min", "q01")} - eef.update({k: [1.0] * EEF_DIM for k in ("std", "max", "q99")}) - population = { - "split": "train", - "trim_active": False, - "min_keep": 2, - "buckets": ["cat/split_aloha/good"], - "empty_buckets": [], - } - (tmp_path / "meta" / "stats_split_aloha.json").write_text(json.dumps({"eef": eef, "population": population})) - - with pytest.raises(DataContractError, match="contributor set no longer matches"): - InternDataA1Dataset.from_config( - { - "dataset_dir": str(tmp_path), - "stats_root": str(tmp_path), - "normalize_mode": "quantile", - "num_frames": 2, - "video_stride": 1, - }, - split="train", - ) - - def test_a_matching_population_loads(self, tmp_path, patch_decode): - d = self._bucket_and_stats(tmp_path) - assert len(self._read(d, tmp_path)) > 0 +def test_stats_embodiment_marker_must_match_bucket(tmp_path, patch_decode): + bucket = _make_bucket(tmp_path, "cat/split_aloha/task") + path = TestStats._write_stats(tmp_path, "split_aloha") + payload = json.loads(path.read_text()) + payload["embodiment"] = "franka" + path.write_text(json.dumps(payload)) + with pytest.raises(ValueError, match="different embodiment"): + InternDataA1Dataset(str(bucket), normalize_mode="quantile") diff --git a/tests/dataloader/test_interndata_a1_stats_computation.py b/tests/dataloader/test_interndata_a1_stats_computation.py deleted file mode 100644 index 42ac2283d..000000000 --- a/tests/dataloader/test_interndata_a1_stats_computation.py +++ /dev/null @@ -1,946 +0,0 @@ -"""Tests for the InternData-A1 stats script. - -(openwam/dataloader/utils/stats_computation/interndata_a1_stats_computation.py) - -Mirrors the coverage its siblings already have (test_ebench_stats_computation.py, -test_robocoin_compute_stats.py) and pins the three behaviours the script's own -comments flag but nothing enforced: - -1. **Reader/stats parity.** ``_SIDES`` / ``_arm10`` / ``_eef20`` are duplicated - from the reader ("must stay bit-identical to ``InternDataA1Dataset._eef20``"), - so an edit to either copy would surface only as silently skewed normalization - — the stats would describe a distribution the reader never emits. -2. **Identity pinning at generation time.** ``materialize_eef_stats`` only - *warns* on an unpinned file, so dropping ``pin_rot6d_identity`` would distort - every rotation with a green suite. Same for the single-arm right-half pin - that keeps franka's padding at exactly 0 instead of the -1 boundary. -3. **Merge determinism.** q01/q99 come from a seeded reservoir whose surviving - rows depend on insertion order, so the parent must merge in a fixed order. - -The bucket builder is shared with the reader tests (same precedent as -test_ebench_stats_computation.py importing from test_ebench_dataset). -""" - -from __future__ import annotations - -import json -from concurrent.futures import Future -from pathlib import Path - -import numpy as np -import pyarrow.parquet as pq -import pytest - -import openwam.dataloader.utils.stats_computation.interndata_a1_stats_computation as a1s -from openwam.dataloader.interndata_a1 import ( - _BIMANUAL_SIDES, - _SINGLE_ARM_SIDES, - InternDataA1Dataset, - resolve_gripper_scale, -) -from openwam.dataloader.utils.lerobotv3 import DataContractError -from openwam.dataloader.utils.normalization import ROT6D_DIMS_EEF20 -from openwam.dataloader.utils.stats_computation.robocoin_stats_computation import Accumulator -from tests.dataloader.test_interndata_a1 import _make_bucket - -BIMANUAL = ("bimanual", "AgileX Split Aloha", "split_aloha") -SINGLE_ARM = ("single_arm", "Franka", "franka") - - -def _group(dirs, layout: str, robot_type: str) -> dict: - return {"robot_type": robot_type, "arm_layout": layout, "dirs": list(dirs)} - - -class _InlinePool: - """ProcessPoolExecutor stub that runs each task inline at submit time. - - Keeps ``_scan_bucket`` in-process (covered, and no process-spawn cost) and - leaves every future already-done by the time the merge loop runs — which is - exactly the state in which ``as_completed``'s yield order stops tracking - submission order. - """ - - def __init__(self, max_workers=None): - pass - - def __enter__(self): - return self - - def __exit__(self, *exc): - return False - - def submit(self, fn, arg): - fut: Future = Future() - try: - fut.set_result(fn(arg)) - except BaseException as e: # noqa: BLE001 - mirror the real pool's contract - fut.set_exception(e) - return fut - - -@pytest.fixture -def inline_pool(monkeypatch): - monkeypatch.setattr(a1s, "ProcessPoolExecutor", _InlinePool) - - -# --------------------------------------------------------------------------- -# 1. Reader / stats parity -# --------------------------------------------------------------------------- - - -class TestReaderParity: - """The stats script and the reader must assemble byte-identical vectors.""" - - @pytest.mark.parametrize("layout,robot_type,_emb", [BIMANUAL, SINGLE_ARM]) - @pytest.mark.parametrize("kind", ["action", "state"]) - def test_eef20_is_bit_identical_to_the_reader(self, tmp_path, layout, robot_type, _emb, kind): - d = _make_bucket(tmp_path, "cat/emb/task", layout=layout, robot_type=robot_type) - ds = InternDataA1Dataset(str(d), normalize_mode=None, num_frames=9, video_stride=4) - win = ds._load_data_table(0, 0).to_pandas() - reader_out = ds._eef20(win, kind, len(win)) - - sides = a1s._SIDES[layout] - table = pq.read_table(d / "data" / "chunk-000" / "file-000.parquet") - grip_scales = tuple( - resolve_gripper_scale(d, _emb, spec[1]) if spec is not None else 1.0 for spec in sides["state"] - ) - stats_out = a1s._eef20(table, sides, kind, grip_scales) - - assert stats_out.shape == reader_out.shape - np.testing.assert_array_equal(stats_out, reader_out) - - def test_gripper_rescale_is_shared_not_reimplemented(self, tmp_path): - """A franka Robotiq bucket: both sides must pick the same 1.0 stroke, or - the stats describe a distribution 12.5x off what the reader emits.""" - d = _make_bucket(tmp_path, "cat/franka/task", layout="single_arm", robot_type="Franka") - (d / "meta" / "stats.json").write_text( - json.dumps({"states.gripper.position": {"min": [0.0], "max": [1.0], "mean": [0.45]}}) - ) - ds = InternDataA1Dataset(str(d), normalize_mode=None, num_frames=9, video_stride=4) - win = ds._load_data_table(0, 0).to_pandas() - sides = a1s._SIDES["single_arm"] - grip_scales = tuple( - resolve_gripper_scale(d, "franka", spec[1]) if spec is not None else 1.0 for spec in sides["state"] - ) - assert grip_scales[0] == pytest.approx(ds._grip_scale[0]) == 1.0 - table = pq.read_table(d / "data" / "chunk-000" / "file-000.parquet") - np.testing.assert_array_equal( - a1s._eef20(table, sides, "action", grip_scales), ds._eef20(win, "action", len(win)) - ) - - def test_each_side_gets_its_own_scale_in_both_implementations(self, tmp_path, monkeypatch): - """Every bimanual embodiment currently declares one stroke, so the two - sides always resolve to the SAME divisor — which means a left/right - scale mix-up in either `_eef20` is invisible to the parity test above. - Force distinct per-side scales and assert each arm's ABSOLUTE value, so a - mix-up applied to both copies at once — which the parity assert cannot - see, since it only compares the two against each other — still fails.""" - d = _make_bucket(tmp_path, "cat/emb/task") - scales = (0.1, 0.4) - ds = InternDataA1Dataset(str(d), normalize_mode=None, num_frames=9, video_stride=4) - monkeypatch.setattr(ds, "_grip_scale", scales) - win = ds._load_data_table(0, 0).to_pandas() - table = pq.read_table(d / "data" / "chunk-000" / "file-000.parquet") - - reader_out = ds._eef20(win, "action", len(win)) - stats_out = a1s._eef20(table, a1s._SIDES["bimanual"], "action", scales) - np.testing.assert_array_equal(stats_out, reader_out) - - # Each arm's gripper must be its OWN raw column over its OWN divisor. - # (Asserting only that cols 9 and 19 differ is unfalsifiable here: - # _make_bucket draws each side's gripper independently, so they differ - # whatever the routing does.) - raw = { - side: np.stack(win[f"actions.{side}_gripper.position"].values).astype(np.float32).ravel() - for side in ("left", "right") - } - for col, side, scale in ((9, "left", scales[0]), (19, "right", scales[1])): - np.testing.assert_allclose(reader_out[:, col], raw[side] / scale, rtol=0, atol=1e-6) - # Sanity: the two divisors are distinct, so the assertions above are not - # accidentally identical. - assert scales[0] != scales[1] - - def test_sides_tables_are_the_readers_own(self): - """The column tables are no longer restated in the stats script — it - imports the reader's. `is` rather than `==`: a re-introduced local copy - would compare equal on the day it was written and drift later, which is - exactly the failure the script's old "must stay in sync" comment named.""" - assert a1s._SIDES["bimanual"] is _BIMANUAL_SIDES - assert a1s._SIDES["single_arm"] is _SINGLE_ARM_SIDES - - -# --------------------------------------------------------------------------- -# 2. Identity pinning at generation time -# --------------------------------------------------------------------------- - - -class TestIdentityPinning: - def test_rot6d_dims_are_pinned(self, tmp_path, inline_pool): - d = _make_bucket(tmp_path, "cat/emb/task") - out = a1s.compute_stats_for_embodiment("split_aloha", _group([d], "bimanual", "AgileX Split Aloha")) - eef = out["eef"] - assert out["rot6d_identity"] is True - for i in ROT6D_DIMS_EEF20: - assert (eef["min"][i], eef["max"][i]) == (-1.0, 1.0) - assert (eef["q01"][i], eef["q99"][i]) == (-1.0, 1.0) - assert (eef["mean"][i], eef["std"][i]) == (0.0, 1.0) - - def test_no_rot6d_identity_leaves_the_measured_values(self, tmp_path, inline_pool): - d = _make_bucket(tmp_path, "cat/emb/task") - out = a1s.compute_stats_for_embodiment( - "split_aloha", _group([d], "bimanual", "AgileX Split Aloha"), rot6d_identity=False - ) - eef = out["eef"] - assert out["rot6d_identity"] is False - # Random unit quaternions never produce exactly +-1 on all 12 rot6d dims. - assert not all(eef["q01"][i] == -1.0 and eef["q99"][i] == 1.0 for i in ROT6D_DIMS_EEF20) - - def test_single_arm_right_half_is_pinned_so_padding_stays_zero(self, tmp_path, inline_pool): - """franka fills [0:10) only. Without the right-half pin the zero padding - would normalize onto the -1 boundary instead of staying at 0.""" - d = _make_bucket(tmp_path, "cat/franka/task", layout="single_arm", robot_type="Franka") - out = a1s.compute_stats_for_embodiment("franka", _group([d], "single_arm", "Franka")) - eef = out["eef"] - assert out["arm_layout"] == "single_arm" - for i in range(10, 20): # xyz, rot6d AND gripper of the padded right arm - assert (eef["q01"][i], eef["q99"][i]) == (-1.0, 1.0) - assert (eef["mean"][i], eef["std"][i]) == (0.0, 1.0) - # quantile-normalizing a 0 with q01/q99 = -/+1 leaves it at 0. - assert (0.0 - eef["q01"][10]) / (eef["q99"][10] - eef["q01"][10]) * 2 - 1 == 0.0 - - def test_bimanual_right_half_is_data_not_padding(self, tmp_path, inline_pool): - """The right-half pin must fire for single_arm ONLY — pinning a bimanual - bucket's real right arm would disable its normalization.""" - d = _make_bucket(tmp_path, "cat/emb/task") - eef = a1s.compute_stats_for_embodiment("split_aloha", _group([d], "bimanual", "AgileX Split Aloha"))["eef"] - assert (eef["q01"][10], eef["q99"][10]) != (-1.0, 1.0) # right xyz - assert (eef["q01"][19], eef["q99"][19]) != (-1.0, 1.0) # right gripper - - def test_written_file_round_trips_into_the_reader(self, tmp_path, inline_pool, monkeypatch): - """End-to-end: generated stats must satisfy the reader's own width check - and load under the default quantile mode.""" - d = _make_bucket(tmp_path, "cat/split_aloha/task") - monkeypatch.setattr("sys.argv", ["prog", "--dataset_dir", str(tmp_path)]) - a1s.main() - ds = InternDataA1Dataset( - str(d), a1_stats_root=str(tmp_path), normalize_mode="quantile", num_frames=9, video_stride=4 - ) - assert ds._normalization_stats is not None - - -# --------------------------------------------------------------------------- -# 3. Merge determinism -# --------------------------------------------------------------------------- - - -class TestMergeDeterminism: - """q01/q99 come from a seeded Algorithm-R reservoir: which rows survive - depends on INSERTION order, so consuming as_completed made the file drift run - to run on identical data (mean/std/min/max are order-invariant and stayed - put, which is what let it go unnoticed).""" - - @staticmethod - def _small_acc_factory(cap: int): - class _SmallAccumulator(Accumulator): - def __init__(self, dim: int = 20, **kw): - super().__init__(dim=dim, reservoir_cap=cap, seed=0) - - return _SmallAccumulator - - def _buckets(self, tmp_path): - # Distinct row counts so the merge order is identifiable from batch sizes. - return [ - _make_bucket(tmp_path, f"cat/emb/task{i}", n_eps=1, ep_len=ep, seed=i) - for i, ep in enumerate((10, 12, 14, 16, 18, 20)) - ] - - def test_merges_in_submission_order_not_completion_order(self, tmp_path, inline_pool, monkeypatch): - dirs = self._buckets(tmp_path) - seen: list[int] = [] - base = self._small_acc_factory(64) - - class _Recording(base): # type: ignore[valid-type,misc] - def update_batch(self, batch): - seen.append(len(batch)) - return super().update_batch(batch) - - monkeypatch.setattr(a1s, "Accumulator", _Recording) - a1s.compute_stats_for_embodiment("split_aloha", _group(dirs, "bimanual", "AgileX Split Aloha")) - # _scan_bucket returns action rows ++ state rows -> 2 * ep_len per bucket, - # and dirs are submitted in the order given. - assert seen == [2 * ep for ep in (10, 12, 14, 16, 18, 20)] - - def test_quantiles_match_a_sequential_merge_of_the_same_buckets(self, tmp_path, inline_pool, monkeypatch): - """The reservoir is deliberately shrunk below the row count here — at the - 1M production cap a small fixture never evicts, so order could not bite.""" - dirs = self._buckets(tmp_path) - small = self._small_acc_factory(64) - monkeypatch.setattr(a1s, "Accumulator", small) - got = a1s.compute_stats_for_embodiment( - "split_aloha", _group(dirs, "bimanual", "AgileX Split Aloha"), rot6d_identity=False - )["eef"] - - ref = small(dim=a1s.EEF20_DIM) - for d in dirs: # same fixed order the script must use - _, rows, _, _ = a1s._scan_bucket((str(d), "bimanual", "split_aloha")) - ref.update_batch(rows) - expected = ref.finalize() - for key in ("mean", "std", "min", "max", "q01", "q99"): - np.testing.assert_allclose(got[key], expected[key], atol=0, rtol=0, err_msg=key) - - def test_regenerating_reproduces_the_file_byte_for_byte(self, tmp_path, inline_pool, monkeypatch): - _make_bucket(tmp_path, "cat/split_aloha/task") - monkeypatch.setattr("sys.argv", ["prog", "--dataset_dir", str(tmp_path)]) - out = tmp_path / "meta" / "stats_split_aloha.json" - a1s.main() - first = out.read_bytes() - a1s.main() - assert out.read_bytes() == first - - -# --------------------------------------------------------------------------- -# 4. Output location (read-only dataset mounts) -# --------------------------------------------------------------------------- - - -class TestStatsRoot: - def test_defaults_to_the_dataset_dir(self, tmp_path, inline_pool, monkeypatch): - _make_bucket(tmp_path, "cat/split_aloha/task") - monkeypatch.setattr("sys.argv", ["prog", "--dataset_dir", str(tmp_path)]) - a1s.main() - assert (tmp_path / "meta" / "stats_split_aloha.json").is_file() - - def test_stats_root_redirects_the_write_off_the_dataset_mount(self, tmp_path, inline_pool, monkeypatch): - """interndata_a1.yaml advertises stats_root for read-only mounts and the - reader wires a1_stats_root for it; without this flag the shipped tool - could not produce the files in the one scenario stats_root exists for.""" - root = tmp_path / "readonly_mount" - elsewhere = tmp_path / "scratch" - d = _make_bucket(root, "cat/split_aloha/task") - monkeypatch.setattr("sys.argv", ["prog", "--dataset_dir", str(root), "--stats_root", str(elsewhere)]) - a1s.main() - assert (elsewhere / "meta" / "stats_split_aloha.json").is_file() - assert not (root / "meta").exists() # dataset mount untouched - # And the reader resolves the same path from dataloader.stats_root. - ds = InternDataA1Dataset( - str(d), a1_stats_root=str(elsewhere), normalize_mode="quantile", num_frames=9, video_stride=4 - ) - assert ds._normalization_stats is not None - - -# --------------------------------------------------------------------------- -# 5. Bucket classification -# --------------------------------------------------------------------------- - - -class TestClassifyBuckets: - def test_groups_by_embodiment(self, tmp_path): - _make_bucket(tmp_path, "cat/split_aloha/a") - _make_bucket(tmp_path, "cat/split_aloha/b") - _make_bucket(tmp_path, "cat/franka/c", layout="single_arm", robot_type="Franka") - _make_bucket(tmp_path, "cat/lift2/d", robot_type="ARX Lift-2") - groups = a1s.classify_buckets(a1s.discover_a1_buckets(tmp_path)) - assert {k: len(v["dirs"]) for k, v in groups.items()} == {"split_aloha": 2, "franka": 1, "lift2": 1} - assert groups["franka"]["arm_layout"] == "single_arm" - assert groups["split_aloha"]["arm_layout"] == "bimanual" - - def test_one_unreadable_bucket_aborts_classification(self, tmp_path): - _make_bucket(tmp_path, "cat/split_aloha/good") - bad = tmp_path / "cat" / "split_aloha" / "bad" - (bad / "meta").mkdir(parents=True) - (bad / "meta" / "info.json").write_text("{ not json") - with pytest.raises(RuntimeError, match="Refusing to compute partial.*unreadable meta/info.json"): - a1s.classify_buckets(a1s.discover_a1_buckets(tmp_path)) - - def test_nonempty_manifest_with_zero_physical_rows_fails_closed(self, tmp_path, inline_pool): - """A corrupt non-empty bucket must not become a degenerate stats file.""" - d = _make_bucket(tmp_path, "cat/emb/task", n_eps=1, ep_len=1) - pq.write_table( - pq.read_table(d / "data" / "chunk-000" / "file-000.parquet").slice(0, 0), - d / "data" / "chunk-000" / "file-000.parquet", - ) - with pytest.raises(RuntimeError, match="Refusing to write partial.*non-empty bucket"): - a1s.compute_stats_for_embodiment("split_aloha", _group([d], "bimanual", "AgileX Split Aloha")) - - -def test_scan_bucket_reads_action_and_state_rows(tmp_path): - """Both streams are pooled — actions[t] == states[t+1], the same signal - offset by one row — so the row count is 2x the parquet length.""" - d = _make_bucket(tmp_path, "cat/emb/task", n_eps=2, ep_len=20) - name, rows, _, _ = a1s._scan_bucket((str(d), "bimanual", "split_aloha")) - assert name == str(d) - assert rows.shape == (2 * 40, a1s.EEF20_DIM) - assert rows.dtype == np.float32 - - -def test_scan_bucket_of_a_single_arm_leaves_the_right_half_zero(tmp_path): - d = _make_bucket(tmp_path, "cat/franka/task", layout="single_arm", robot_type="Franka") - _, rows, _, _ = a1s._scan_bucket((str(d), "single_arm", "franka")) - np.testing.assert_array_equal(rows[:, 10:], 0.0) - assert np.abs(rows[:, :10]).sum() > 0 - - -def test_module_exports_stay_importable(): - for name in a1s.__all__: - assert hasattr(a1s, name), name - assert Path(a1s.__file__).name == "interndata_a1_stats_computation.py" - - -# --------------------------------------------------------------------------- -# Cleaned-view filtering: the stats must describe exactly what the reader emits -# --------------------------------------------------------------------------- - - -def _trim_file(path: Path, rows: str) -> str: - path.write_text("dataset,episode_index,total_frames,trim_head_to,trim_tail_from\n" + rows) - return str(path) - - -class TestCleanedViewFiltering: - """A cleaned view symlinks ``data/``, so deleted episodes are still - physically in the parquet. The scanner has to honour - ``meta/excluded_episodes.json`` or the normalizer pools rows the reader - never emits.""" - - def test_excluded_episodes_are_dropped_from_the_scan(self, tmp_path): - d = _make_bucket(tmp_path, "cat/emb/task", n_eps=2, ep_len=20) - (d / "meta" / "excluded_episodes.json").write_text(json.dumps({"episode_indices": [0]})) - _, rows, _, _ = a1s._scan_bucket((str(d), "bimanual", "split_aloha", "cat/emb/task", None)) - assert rows.shape[0] == 2 * 20 # only episode 1, both streams - - def test_an_empty_kept_set_excludes_everything(self, tmp_path): - """``set()`` means "manifest read, nothing survives" — distinct from - ``None`` ("unreadable, cannot filter"). Collapsing the two would pool a - fully-deleted bucket's rows back into the statistics.""" - d = _make_bucket(tmp_path, "cat/emb/task", n_eps=2, ep_len=20) - (d / "meta" / "excluded_episodes.json").write_text(json.dumps({"episode_indices": [0, 1]})) - assert a1s._kept_episodes(d) == set() - _, rows, _, _ = a1s._scan_bucket((str(d), "bimanual", "split_aloha", "cat/emb/task", None)) - assert rows.shape[0] == 0 - - def test_kept_episodes_is_none_only_when_the_manifest_is_unreadable(self, tmp_path): - d = _make_bucket(tmp_path, "cat/emb/task", n_eps=2, ep_len=20) - assert a1s._kept_episodes(d) == {0, 1} - for f in (d / "meta" / "episodes").rglob("*.parquet"): - f.unlink() - assert a1s._kept_episodes(d) is None - - -class TestStatsTrimMatchesReader: - """Both paths call :func:`resolve_trim_bounds`, so a trim that the reader - declines to apply must not be applied here either.""" - - def test_a_too_short_trim_is_left_whole_like_the_reader_does(self, tmp_path): - d = _make_bucket(tmp_path, "cat/emb/task", n_eps=1, ep_len=4) - trim = _trim_file(tmp_path / "trim.csv", "cat/emb/task,0,4,3,\n") # leaves 1 < min_len 2 - _, rows, _, _ = a1s._scan_bucket((str(d), "bimanual", "split_aloha", "cat/emb/task", trim, 2)) - assert rows.shape[0] == 2 * 4, "stats trimmed an episode the reader keeps whole" - - def test_a_valid_trim_is_applied(self, tmp_path): - d = _make_bucket(tmp_path, "cat/emb/task", n_eps=1, ep_len=20) - trim = _trim_file(tmp_path / "trim.csv", "cat/emb/task,0,20,4,18\n") - _, rows, _, _ = a1s._scan_bucket((str(d), "bimanual", "split_aloha", "cat/emb/task", trim, 2)) - assert rows.shape[0] == 2 * 14 # 18 - 4 - - def test_a_stale_total_frames_disables_the_entry(self, tmp_path): - d = _make_bucket(tmp_path, "cat/emb/task", n_eps=1, ep_len=20) - trim = _trim_file(tmp_path / "trim.csv", "cat/emb/task,0,999,4,18\n") - _, rows, _, _ = a1s._scan_bucket((str(d), "bimanual", "split_aloha", "cat/emb/task", trim, 2)) - assert rows.shape[0] == 2 * 20 - - -class TestDirectBucketParity: - """`--dataset_dir` pointed straight at a bucket is a supported mode, and it - is the one where the generator and the reader can silently disagree. - - There `relative_to(root)` is `'.'` — a key no trim CSV holds and one the - reader (whose id is then the bare directory name) cannot suffix-match. Left - unnormalized the generator skips the trim while the reader applies it, and - writes `exclusions: {".": ...}` the reader then rejects. These go through the - real `main()` rather than the helpers, because that is the seam the unit - tests were blind to. - """ - - def _setup(self, tmp_path): - d = _make_bucket(tmp_path, "cat/emb/task", n_eps=2, ep_len=40) - (d / "meta" / "excluded_episodes.json").write_text(json.dumps({"episode_indices": [0]})) - trim = tmp_path / "trim.csv" - trim.write_text( - "dataset,episode_index,total_frames,trim_head_to,trim_tail_from\n" - "task,1,40,10,30\n" # keyed by the bare bucket name, as the reader resolves it - ) - return d, trim - - def _run(self, monkeypatch, d, trim, out): - import sys - - monkeypatch.setattr( - sys, - "argv", - [ - "prog", - "--dataset_dir", - str(d), - "--stats_root", - str(out), - "--trim_csv", - str(trim), - "--workers", - "1", - "--min_keep", - "2", - ], - ) - a1s.main() - return json.load(open(out / "meta" / "stats_split_aloha.json")) - - def test_generator_applies_the_trim_and_the_exclusion(self, tmp_path, monkeypatch): - d, trim = self._setup(tmp_path) - out = tmp_path / "stats" - res = self._run(monkeypatch, d, trim, out) - # ep0 excluded; ep1 trimmed to [10, 30) = 20 frames, over action+state. - assert res["num_rows"] == 40, "generator ignored the trim or the exclusion" - - def test_the_reader_accepts_what_the_generator_produced(self, tmp_path, monkeypatch): - d, trim = self._setup(tmp_path) - out = tmp_path / "stats" - self._run(monkeypatch, d, trim, out) - ds = InternDataA1Dataset( # no dataset_id: id is the bare bucket name - str(d), - a1_stats_root=str(out), - trim_csv=str(trim), - normalize_mode="quantile", - num_frames=9, - video_stride=4, - ) - assert ds._normalization_stats is not None - assert ds._eps_df["episode_index"].tolist() == [1] - assert int(ds._eps_df["length"].iloc[0]) == 20 - - def test_the_emitted_key_is_not_a_dot(self, tmp_path, monkeypatch): - d, trim = self._setup(tmp_path) - out = tmp_path / "stats" - res = self._run(monkeypatch, d, trim, out) - assert "." not in res["scanned_buckets"], "direct-bucket key leaked as '.'" - assert res["scanned_buckets"] == ["task"] - - def test_reader_rejects_stats_after_trim_csv_content_changes(self, tmp_path, monkeypatch): - d, trim = self._setup(tmp_path) - out = tmp_path / "stats" - self._run(monkeypatch, d, trim, out) - - # Keep the same path and bucket key, but change the effective span after - # the stats were generated. A path-only or trim-enabled boolean check - # would accept the stale normalizer. - trim.write_text("dataset,episode_index,total_frames,trim_head_to,trim_tail_from\ntask,1,40,5,30\n") - with pytest.raises(DataContractError, match="trim provenance"): - InternDataA1Dataset( - str(d), - a1_stats_root=str(out), - trim_csv=str(trim), - normalize_mode="quantile", - num_frames=9, - video_stride=4, - ) - - def test_reader_rejects_stats_after_exclusions_change(self, tmp_path, monkeypatch): - d, trim = self._setup(tmp_path) - out = tmp_path / "stats" - self._run(monkeypatch, d, trim, out) - - # The old stats exclude episode 0. Removing that exclusion changes the - # population without changing the shared stats path or trim artifact. - (d / "meta" / "excluded_episodes.json").write_text(json.dumps({"episode_indices": []})) - with pytest.raises(DataContractError, match="excluded_episodes.json changed"): - InternDataA1Dataset( - str(d), - a1_stats_root=str(out), - trim_csv=str(trim), - normalize_mode="quantile", - num_frames=9, - video_stride=4, - ) - - def test_capped_reader_still_checks_the_full_train_population(self, tmp_path, monkeypatch): - d = _make_bucket(tmp_path, "cat/emb/task", n_eps=3, ep_len=40) - trim = tmp_path / "trim.csv" - trim.write_text("dataset,episode_index,total_frames,trim_head_to,trim_tail_from\ntask,1,40,5,30\n") - out = tmp_path / "stats" - self._run(monkeypatch, d, trim, out) - - # Narrow the train split after stats generation. max_hours is applied - # only after A1 captures its full pre-subsample train certificate, so it - # must not exempt this reader from the population check. - info_path = d / "meta" / "info.json" - info = json.loads(info_path.read_text()) - info["splits"] = {"train": "0:2", "val": "2:3"} - info_path.write_text(json.dumps(info)) - with pytest.raises(DataContractError, match="effective.*population changed"): - InternDataA1Dataset( - str(d), - a1_stats_root=str(out), - trim_csv=str(trim), - normalize_mode="quantile", - num_frames=9, - video_stride=4, - max_hours=1.0, - ) - - -class TestFailClosedCoverage: - """A real bucket failure aborts the shared stats instead of shrinking it.""" - - def test_a_failed_nonempty_bucket_aborts_the_embodiment(self, tmp_path): - good = _make_bucket(tmp_path, "cat/emb/good", n_eps=2, ep_len=8) - bad = _make_bucket(tmp_path, "cat/emb/bad", n_eps=2, ep_len=8) - # Strip episode_index so the scan of `bad` raises (trim/exclusions cannot - # select rows without it). - (bad / "meta" / "excluded_episodes.json").write_text(json.dumps({"episode_indices": [0]})) - p = bad / "data" / "chunk-000" / "file-000.parquet" - t = pq.read_table(p) - pq.write_table(t.drop(["episode_index"]), p) - - with pytest.raises(RuntimeError, match="Refusing to write partial.*cat/emb/bad"): - a1s.compute_stats_for_embodiment( - "split_aloha", - _group([good, bad], "bimanual", "AgileX Split Aloha"), - workers=1, - root=tmp_path, - ) - - def test_cli_publishes_no_embodiment_when_a_later_one_fails(self, tmp_path, monkeypatch, inline_pool): - """The CLI is transactional across its requested embodiment set.""" - _make_bucket( - tmp_path, - "cat/franka/good", - n_eps=2, - ep_len=8, - layout="single_arm", - robot_type="Franka", - ) - bad = _make_bucket(tmp_path, "cat/split_aloha/bad", n_eps=2, ep_len=8) - p = bad / "data" / "chunk-000" / "file-000.parquet" - pq.write_table(pq.read_table(p).slice(0, 15), p) - - out = tmp_path / "published" - meta = out / "meta" - meta.mkdir(parents=True) - old = { - meta / "stats_franka.json": "old-franka\n", - meta / "stats_split_aloha.json": "old-split\n", - } - for path, content in old.items(): - path.write_text(content) - monkeypatch.setattr( - "sys.argv", - [ - "prog", - "--dataset_dir", - str(tmp_path), - "--stats_root", - str(out), - "--workers", - "1", - ], - ) - - with pytest.raises(RuntimeError, match="Refusing to write partial.*cat/split_aloha/bad"): - a1s.main() - assert {path: path.read_text() for path in old} == old - assert not list(meta.glob("*.json.tmp")) - - -class TestSplitFailuresDoNotFailOpen: - """ "Cannot determine the population" must never degrade into "use every row". - - The reader raises on the same metadata and refuses the bucket, so pooling it - here would put rows into the normalizer that training can never load. - """ - - def test_a_malformed_split_spec_raises(self, tmp_path): - d = _make_bucket(tmp_path, "cat/emb/task", n_eps=2, ep_len=8) - info = json.loads((d / "meta" / "info.json").read_text()) - info["splits"] = {"train": "not-a-range"} - (d / "meta" / "info.json").write_text(json.dumps(info)) - with pytest.raises(ValueError, match="unusable splits"): - a1s._split_episodes(d, "train") - - def test_an_unreadable_manifest_raises(self, tmp_path): - d = _make_bucket(tmp_path, "cat/emb/task", n_eps=2, ep_len=8) - for f in (d / "meta" / "episodes").rglob("*.parquet"): - f.unlink() - with pytest.raises(ValueError, match="population is unknown"): - a1s._split_episodes(d, "train") - - def test_absent_splits_are_empty_for_val_not_everything(self, tmp_path): - d = _make_bucket(tmp_path, "cat/emb/task", n_eps=2, ep_len=8) - info = json.loads((d / "meta" / "info.json").read_text()) - info.pop("splits", None) - (d / "meta" / "info.json").write_text(json.dumps(info)) - assert a1s._split_episodes(d, "train") is None # unrestricted - assert a1s._split_episodes(d, "val") == set() # not "everything" - - -class TestScannerRefusesWhatTheReaderRefuses: - """`len(rows) > 0` proves rows were read, not that a reader can consume them. - - Each case below yielded a healthy row count here while - ``InternDataA1Dataset`` raised on the same bucket — statistics describing a - population that never reaches training, with nothing downstream able to tell. - Both sides now enumerate shards and validate the manifest through the same - functions, so these are parity assertions rather than a second rule set. - """ - - def _scan(self, d, **kw): - return ( - a1s._scan_bucket((str(d), "bimanual", "split_aloha", kw.pop("dsid", "task")), **kw) - if kw - else a1s._scan_bucket((str(d), "bimanual", "split_aloha")) - ) - - def test_a_truncated_shard_is_refused_like_the_reader_refuses_it(self, tmp_path): - d = _make_bucket(tmp_path, "cat/emb/task", n_eps=2, ep_len=4) - p = d / "data" / "chunk-000" / "file-000.parquet" - t = pq.read_table(p) - pq.write_table(t.slice(0, t.num_rows - 1), p) # manifest says 8, shard holds 7 - - with pytest.raises(ValueError, match="manifest ends at 8"): - InternDataA1Dataset(str(d), normalize_mode=None, num_frames=2, video_stride=1) - with pytest.raises(ValueError, match="manifest ends at 8"): - a1s._scan_bucket((str(d), "bimanual", "split_aloha")) - - def test_a_stray_backup_parquet_does_not_double_the_population(self, tmp_path): - """`file-000.backup.parquet` matched the old glob; the reader never read it.""" - import shutil - - d = _make_bucket(tmp_path, "cat/emb/task", n_eps=2, ep_len=4) - src = d / "data" / "chunk-000" / "file-000.parquet" - _, rows_before, _, _ = a1s._scan_bucket((str(d), "bimanual", "split_aloha")) - shutil.copy(src, src.with_name("file-000.backup.parquet")) - _, rows_after, _, _ = a1s._scan_bucket((str(d), "bimanual", "split_aloha")) - assert len(rows_after) == len(rows_before), "a non-shard file entered the statistics" - - def test_overlapping_manifest_ranges_are_refused(self, tmp_path): - d = _make_bucket(tmp_path, "cat/emb/task", n_eps=3, ep_len=4) - man = d / "meta" / "episodes" / "chunk-000" / "file-000.parquet" - m = pq.read_table(man).to_pydict() - m["dataset_from_index"], m["dataset_to_index"] = [0, 2, 8], [4, 6, 12] - m["length"] = [4, 4, 4] - import pyarrow as pa - - pq.write_table(pa.Table.from_pydict(m), man) - with pytest.raises(ValueError, match="Overlapping manifest ranges"): - a1s._scan_bucket((str(d), "bimanual", "split_aloha")) - - def test_quoted_exclusion_indices_are_refused_on_both_sides(self, tmp_path): - """`["0"]` excluded episode 0 here and nothing in the reader.""" - d = _make_bucket(tmp_path, "cat/emb/task", n_eps=2, ep_len=4) - (d / "meta" / "excluded_episodes.json").write_text(json.dumps({"episode_indices": ["0"]})) - with pytest.raises(ValueError, match="must be JSON integers"): - a1s._scan_bucket((str(d), "bimanual", "split_aloha")) - - -class TestPopulationContractEndToEnd: - """Through the real `main()` and the real reader — no hand-written stats. - - Every provenance defect so far survived a green suite because the tests - wrote the stats file themselves and so could not observe what the generator - actually emits. These run the CLI and hand its output to a reader. - """ - - def _run(self, monkeypatch, d, out, *extra): - import sys - - monkeypatch.setattr( - sys, "argv", ["prog", "--dataset_dir", str(d), "--stats_root", str(out), "--workers", "1", *extra] - ) - a1s.main() - return json.load(open(out / "meta" / "stats_split_aloha.json")) - - def test_the_generator_records_what_it_scanned(self, tmp_path, monkeypatch): - d = _make_bucket(tmp_path, "cat/emb/task", n_eps=2, ep_len=8) - res = self._run(monkeypatch, d, tmp_path / "stats") - pop = res["population"] - assert {key: pop[key] for key in ("split", "trim_active", "min_keep", "buckets", "empty_buckets")} == { - "split": "train", - "trim_active": False, - "min_keep": 2, - "buckets": ["task"], - "empty_buckets": [], - } - assert pop["schema_version"] == 2 - assert pop["trim_provenance"] is None - assert pop["bucket_provenance"]["task"]["excluded_episode_indices"] == [] - assert pop["bucket_provenance"]["task"]["effective_population"]["num_episodes"] == 2 - assert pop["bucket_provenance"]["task"]["effective_population"]["num_rows"] == 16 - - def test_val_stats_are_refused_by_a_train_reader(self, tmp_path, monkeypatch): - d = _make_bucket(tmp_path, "cat/emb/task", n_eps=2, ep_len=8) - out = tmp_path / "stats" - # info.json ships splits {"train": "0:N"}, so give it a val range to scan. - info = json.loads((d / "meta" / "info.json").read_text()) - info["splits"] = {"train": "0:1", "val": "1:2"} - (d / "meta" / "info.json").write_text(json.dumps(info)) - res = self._run(monkeypatch, d, out, "--split", "val") - assert res["population"]["split"] == "val" - with pytest.raises(DataContractError, match="computed over the 'val' split"): - InternDataA1Dataset(str(d), a1_stats_root=str(out), normalize_mode="quantile", num_frames=2, video_stride=1) - - def test_untrimmed_stats_are_refused_by_a_trimming_reader(self, tmp_path, monkeypatch): - d = _make_bucket(tmp_path, "cat/emb/task", n_eps=2, ep_len=8) - out = tmp_path / "stats" - self._run(monkeypatch, d, out) # generated WITHOUT --trim_csv - trim = tmp_path / "trim.csv" - trim.write_text("dataset,episode_index,total_frames,trim_head_to,trim_tail_from\ntask,0,8,2,6\n") - with pytest.raises(DataContractError, match="without a trim list but this reader is"): - InternDataA1Dataset( - str(d), - a1_stats_root=str(out), - trim_csv=str(trim), - normalize_mode="quantile", - num_frames=2, - video_stride=1, - ) - - def test_cli_refuses_to_publish_when_one_bucket_scan_fails(self, tmp_path, monkeypatch): - """One good bucket must not turn a failed embodiment into partial stats.""" - _make_bucket(tmp_path, "cat/emb/good", n_eps=2, ep_len=8) - bad = _make_bucket(tmp_path, "cat/emb/bad", n_eps=2, ep_len=8) - p = bad / "data" / "chunk-000" / "file-000.parquet" - pq.write_table(pq.read_table(p).slice(0, 15), p) # truncated: scan raises - - out = tmp_path / "stats" - with pytest.raises(RuntimeError, match="Refusing to write partial.*cat/emb/bad"): - self._run(monkeypatch, tmp_path, out) - assert not (out / "meta" / "stats_split_aloha.json").exists() - - -class TestValOnlyBucketReusesTrainStats: - """A val-only bucket must borrow the train distribution, not be refused. - - Establishing "stats always come from train" while also demanding the bucket - appear among the TRAIN contributors is contradictory: a bucket whose - `info.json` declares `train=0:0, val=0:2` cannot be a train contributor and - is still a perfectly valid bucket to evaluate. Refusing it rejected the - configuration the train-derived rule exists to support. - """ - - def _two_buckets(self, tmp_path): - a = _make_bucket(tmp_path, "cat/emb/a", n_eps=2, ep_len=8) - b = _make_bucket(tmp_path, "cat/emb/b", n_eps=2, ep_len=8) - for bucket, splits in ((a, {"train": "0:2"}), (b, {"train": "0:0", "val": "0:2"})): - info = json.loads((bucket / "meta" / "info.json").read_text()) - info["splits"] = splits - (bucket / "meta" / "info.json").write_text(json.dumps(info)) - return a, b - - def _generate(self, monkeypatch, root, out): - import sys - - monkeypatch.setattr( - sys, "argv", ["prog", "--dataset_dir", str(root), "--stats_root", str(out), "--workers", "1"] - ) - a1s.main() - return json.load(open(out / "meta" / "stats_split_aloha.json")) - - def test_the_val_only_bucket_is_recorded_as_empty_not_dropped(self, tmp_path, monkeypatch): - self._two_buckets(tmp_path) - pop = self._generate(monkeypatch, tmp_path, tmp_path / "stats")["population"] - assert pop["buckets"] == ["cat/emb/a"] - assert pop["empty_buckets"] == ["cat/emb/b"] - - def test_the_val_only_bucket_has_windows_without_normalization(self, tmp_path): - """Pins the premise: B is a real bucket, not an empty one.""" - _, b = self._two_buckets(tmp_path) - r = InternDataA1Dataset( - str(b), dataset_id="cat/emb/b", normalize_mode=None, num_frames=2, video_stride=1, split="val" - ) - assert len(r) > 0 - - def test_the_val_only_bucket_loads_the_train_stats(self, tmp_path, monkeypatch): - _, b = self._two_buckets(tmp_path) - out = tmp_path / "stats" - self._generate(monkeypatch, tmp_path, out) - r = InternDataA1Dataset( - str(b), - dataset_id="cat/emb/b", - a1_stats_root=str(out), - normalize_mode="quantile", - num_frames=2, - video_stride=1, - split="val", - ) - assert r._normalization_stats is not None - assert len(r) > 0 - - def test_a_corrupt_val_only_bucket_still_aborts_generation(self, tmp_path, monkeypatch): - """Empty-by-split is allowed only when the bucket itself validates.""" - _, b = self._two_buckets(tmp_path) - p = b / "data" / "chunk-000" / "file-000.parquet" - pq.write_table(pq.read_table(p).slice(0, 15), p) # truncated -> scan raises - out = tmp_path / "stats" - with pytest.raises(RuntimeError, match="Refusing to write partial.*cat/emb/b"): - self._generate(monkeypatch, tmp_path, out) - assert not (out / "meta" / "stats_split_aloha.json").exists() - - -class TestZeroRowsIsNotProofOfAnEmptyPopulation: - """A substituted physical slice survives every count-and-envelope check. - - Ranges, grand total and the shard's `episode_index` min/max are all - preserved when one episode's rows are overwritten with another's, so the - only symptom is that the mask for the expected episode selects nothing. - Read as "empty population" that bucket gets certified as covered, and its - reader then loads another bucket's statistics while reading another - episode's rows. - """ - - def _substituted(self, tmp_path): - """3 episodes x 4 rows, split train=1:2, rows [4:8] replaced by ep2's.""" - import pyarrow as pa - - d = _make_bucket(tmp_path, "cat/emb/bad", n_eps=3, ep_len=4) - info = json.loads((d / "meta" / "info.json").read_text()) - info["splits"] = {"train": "1:2"} - (d / "meta" / "info.json").write_text(json.dumps(info)) - p = d / "data" / "chunk-000" / "file-000.parquet" - t = pq.read_table(p).to_pydict() - for k, v in t.items(): - t[k] = v[:4] + v[8:12] + v[8:12] # ep0, ep2, ep2 — total and envelope intact - pq.write_table(pa.Table.from_pydict(t), p) - return d - - def test_the_premise_holds_total_and_envelope_are_unchanged(self, tmp_path): - d = self._substituted(tmp_path) - t = pq.read_table(d / "data" / "chunk-000" / "file-000.parquet") - eps = t.column("episode_index").to_pylist() - assert t.num_rows == 12, "grand total must still match the manifest" - assert (min(eps), max(eps)) == (0, 2), "envelope must still contain episode 1" - assert 1 not in eps, "episode 1's rows are the ones that were replaced" - - def test_the_scanner_refuses_it_instead_of_calling_it_empty(self, tmp_path): - d = self._substituted(tmp_path) - with pytest.raises(ValueError, match="no shard row carries them"): - a1s._scan_bucket((str(d), "bimanual", "split_aloha", "cat/emb/bad")) - - def test_substituted_rows_abort_generation(self, tmp_path, monkeypatch): - import sys - - _make_bucket(tmp_path, "cat/emb/good", n_eps=2, ep_len=8) - self._substituted(tmp_path) - out = tmp_path / "stats" - monkeypatch.setattr( - sys, "argv", ["prog", "--dataset_dir", str(tmp_path), "--stats_root", str(out), "--workers", "1"] - ) - with pytest.raises(RuntimeError, match="Refusing to write partial.*cat/emb/bad"): - a1s.main() - assert not (out / "meta" / "stats_split_aloha.json").exists() - - def test_a_genuinely_empty_train_split_is_still_reported_empty(self, tmp_path): - """The distinction must not collapse the other way.""" - d = _make_bucket(tmp_path, "cat/emb/valonly", n_eps=2, ep_len=8) - info = json.loads((d / "meta" / "info.json").read_text()) - info["splits"] = {"train": "0:0", "val": "0:2"} - (d / "meta" / "info.json").write_text(json.dumps(info)) - _, rows, population_empty, _ = a1s._scan_bucket((str(d), "bimanual", "split_aloha", "cat/emb/valonly")) - assert len(rows) == 0 and population_empty is True diff --git a/tests/dataloader/test_lerobotv3_utils.py b/tests/dataloader/test_lerobotv3_utils.py index a7d74aa56..886876403 100644 --- a/tests/dataloader/test_lerobotv3_utils.py +++ b/tests/dataloader/test_lerobotv3_utils.py @@ -188,17 +188,6 @@ def test_target_hit_bounds(self): assert actual_hours >= target assert actual_hours < target + (100 / 30.0 / 3600.0) + 1e-9 - def test_effective_valid_range_drives_budget(self): - """Segment-trimmed episodes are charged only for sampleable frames.""" - df = _make_eps_df(20, ep_length=100) - df["_valid_start"] = 10 - df["_valid_end"] = 60 # 50 effective frames per episode, not 100 - target = 250.0 / 30.0 / 3600.0 - out = subsample_episodes_by_hours(df, target_hours=target, fps=30.0, seed=42) - assert len(out) == 5 - effective_frames = (out["_valid_end"] - out["_valid_start"]).sum() - assert effective_frames == 250 - def test_zero_window_episode_cannot_satisfy_tiny_positive_budget(self): df = _make_eps_df(2, ep_length=10) explicit = np.array([0, 10], dtype=np.int64) @@ -222,17 +211,11 @@ def test_all_zero_sampleable_frames_raise(self): episode_frames=np.zeros(2, dtype=np.int64), ) - def test_unpaired_or_out_of_bounds_valid_range_raises(self): - unpaired = _make_eps_df(2, ep_length=10) - unpaired["_valid_start"] = 1 - with pytest.raises(ValueError, match="missing '_valid_end'"): - subsample_episodes_by_hours(unpaired, target_hours=0.1, fps=30.0, seed=42) - - out_of_bounds = _make_eps_df(2, ep_length=10) - out_of_bounds["_valid_start"] = 0 - out_of_bounds["_valid_end"] = [11, 10] - with pytest.raises(ValueError, match="start <= end <= length"): - subsample_episodes_by_hours(out_of_bounds, target_hours=0.1, fps=30.0, seed=42) + def test_negative_episode_length_raises(self): + df = _make_eps_df(2, ep_length=10) + df.loc[0, "length"] = -1 + with pytest.raises(ValueError, match="values must be non-negative"): + subsample_episodes_by_hours(df, target_hours=0.1, fps=30.0, seed=42) def test_negative_target_raises(self): df = _make_eps_df(10) diff --git a/tests/dataloader/test_mixture.py b/tests/dataloader/test_mixture.py index 629d0dc49..a5e04221a 100644 --- a/tests/dataloader/test_mixture.py +++ b/tests/dataloader/test_mixture.py @@ -198,3 +198,34 @@ def test_set_epoch_idempotent_with_same_arg(self, fake_dataset_factory): epoch_3_map = m._index_map m.set_epoch(3) assert m._index_map is epoch_3_map + + +@pytest.mark.parametrize("seed", [0, 42]) +def test_mixture_preserves_legacy_seeded_order_without_extra_maps(fake_dataset_factory, seed): + datasets = [fake_dataset_factory(7), fake_dataset_factory(11)] + mixture = MixtureDataset(datasets, weights=[1.0, 4.0], seed=seed) + for epoch in (0, 3): + mixture.set_epoch(epoch) + parts = [] + for di, (dataset, weight) in enumerate(zip(datasets, mixture.weights)): + count = max(1, round(18 * weight)) + parts.append( + np.stack( + [ + np.full(count, di, dtype=np.int32), + (np.arange(count, dtype=np.int64) % len(dataset)).astype(np.int32), + ], + axis=1, + ) + ) + legacy = np.concatenate(parts) + legacy = legacy[np.random.RandomState(seed + epoch * 7919).permutation(len(legacy))] + np.testing.assert_array_equal(mixture._index_map, legacy) + assert mixture._index_map.dtype == np.int32 + assert mixture.dataset_sample_counts() == {"source_0": 4, "source_1": 14} + + +@pytest.mark.parametrize("weights", [[float("nan"), 1.0], [float("inf"), 1.0], [1e308, 1e308]]) +def test_mixture_rejects_nonfinite_weights(fake_dataset_factory, weights): + with pytest.raises(ValueError, match="finite"): + MixtureDataset([fake_dataset_factory(3), fake_dataset_factory(4)], weights=weights) diff --git a/tests/dataloader/test_mixture_from_config.py b/tests/dataloader/test_mixture_from_config.py index 683ea14c2..2ff52e5a2 100644 --- a/tests/dataloader/test_mixture_from_config.py +++ b/tests/dataloader/test_mixture_from_config.py @@ -264,3 +264,12 @@ def test_action_dim_override_field_present_uses_legacy(self, fake_registry): } m = MixtureDataset.from_config(cfg) assert m.action_dim == 20 + + +def test_invalid_strategy_fails_before_loading_sources(monkeypatch): + def unexpected_load(*args, **kwargs): + raise AssertionError("invalid mixture configuration must not scan datasets") + + monkeypatch.setattr(registry_mod, "build_dataset", unexpected_load) + with pytest.raises(ValueError, match="unknown weight_strategy"): + MixtureDataset.from_config({"weight_strategy": "typo", "datasets": {"a": {"type": "robocoin"}}}) diff --git a/tests/dataloader/test_muka_franka.py b/tests/dataloader/test_muka_franka.py deleted file mode 100644 index af2bbd1a9..000000000 --- a/tests/dataloader/test_muka_franka.py +++ /dev/null @@ -1,253 +0,0 @@ -from __future__ import annotations - -import json -from contextlib import contextmanager -from pathlib import Path -from unittest.mock import patch - -import numpy as np -import pandas as pd -import pyarrow as pa -import pyarrow.parquet as pq -from omegaconf import OmegaConf -from PIL import Image - -from openwam.dataloader.muka_franka import ( - EEF10_DIM, - GRIPPER_CONVENTION, - GRIPPER_TRANSFORM, - RAW_GRIPPER_CONVENTION, - ROT6D_DIMS_EEF10, - ROTATION_CONVENTION, - STATS_POPULATION, - MukaFrankaDataset, - closedness_to_open_scale, - euler7_to_eef10, -) -from openwam.dataloader.registry import build_dataset, list_registered_datasets -from openwam.dataloader.utils.stats_computation.muka_franka_stats_computation import ( - _compute_global_stats, -) -from openwam.dataloader.utils.unify_action import UNIFY_DIM, parse_unify_spec, unmap_from_unify - -EP_LENGTH = 6 -HEAD = "observation.images.left" -WRIST = "observation.images.left_wrist" -UNIFY_MAP = ["0-9"] - - -def _source_state() -> np.ndarray: - state = np.zeros((EP_LENGTH, 7), dtype=np.float32) - state[:, :3] = np.arange(EP_LENGTH, dtype=np.float32)[:, None] * np.array([0.01, -0.02, 0.03], dtype=np.float32) - state[:, 3:6] = np.array( - [ - [0.0, 0.0, 0.0], - [0.0, 0.0, np.pi / 2], - [0.1, -0.2, 0.3], - [-0.3, 0.2, -0.1], - [0.4, 0.1, -0.2], - [0.4, 0.1, -0.2], - ], - dtype=np.float32, - ) - state[:, 6] = np.linspace(0.0, 1.0, EP_LENGTH, dtype=np.float32) - return state - - -def _write_bucket(bucket: Path) -> tuple[np.ndarray, np.ndarray]: - (bucket / "meta" / "episodes" / "chunk-000").mkdir(parents=True) - (bucket / "data" / "chunk-000").mkdir(parents=True) - for camera in (HEAD, WRIST): - video_dir = bucket / "videos" / camera / "chunk-000" - video_dir.mkdir(parents=True) - (video_dir / "file-000.mp4").write_bytes(b"") - - names = ["x", "y", "z", "roll", "pitch", "yaw", "gripper_closedness"] - info = { - "fps": 15, - "robot_type": "muka_franka_single_arm", - "data_path": "data/chunk-{chunk_index:03d}/file-{file_index:03d}.parquet", - "video_path": "videos/{video_key}/chunk-{chunk_index:03d}/file-{file_index:03d}.mp4", - "features": { - "observation.state": {"dtype": "float32", "shape": [7], "names": names}, - "action": {"dtype": "float32", "shape": [7], "names": names}, - HEAD: {"dtype": "video", "shape": [480, 640, 3]}, - WRIST: {"dtype": "video", "shape": [480, 640, 3]}, - "task_index": {"dtype": "int64", "shape": [1]}, - }, - "splits": {"train": "0:1"}, - } - (bucket / "meta" / "info.json").write_text(json.dumps(info), encoding="utf-8") - source_schema = { - "position_frame": "robot_base", - "rotation_storage": "Euler XYZ", - "rotation_convention": "extrinsic; R = Rz @ Ry @ Rx", - "rotation_unit": "radians", - "action_type": "absolute achieved EEF state", - "action_alignment": "action[t] = observation.state[t+1]; terminal row repeated and masked by reader", - "gripper": "continuous closedness; 0=open, 1=closed", - } - (bucket / "meta" / "muka_schema.json").write_text(json.dumps(source_schema), encoding="utf-8") - - episode = { - "episode_index": 0, - "length": EP_LENGTH, - "dataset_from_index": 0, - "data/chunk_index": 0, - "data/file_index": 0, - f"videos/{HEAD}/chunk_index": 0, - f"videos/{HEAD}/file_index": 0, - f"videos/{WRIST}/chunk_index": 0, - f"videos/{WRIST}/file_index": 0, - } - pq.write_table( - pa.Table.from_pandas(pd.DataFrame([episode])), - bucket / "meta" / "episodes" / "chunk-000" / "file-000.parquet", - ) - pd.DataFrame({"task_index": [0]}, index=pd.Index(["place the object"], name="task")).to_parquet( - bucket / "meta" / "tasks.parquet" - ) - - state = _source_state() - action = np.concatenate((state[1:], state[-1:]), axis=0) - pq.write_table( - pa.Table.from_pandas( - pd.DataFrame( - { - "observation.state": list(state), - "action": list(action), - "task_index": np.zeros(EP_LENGTH, dtype=np.int64), - } - ) - ), - bucket / "data" / "chunk-000" / "file-000.parquet", - ) - return state, action - - -@contextmanager -def _mock_decoder(): - def _fake(path, frame_indices, height, width): - color = (255, 0, 0) if HEAD in str(path) else (0, 255, 0) - return [Image.new("RGB", (width, height), color) for _ in frame_indices] - - with patch("openwam.dataloader.bases.lerobot_v3_reader._decode_video_frames", side_effect=_fake): - yield - - -def _dataset(bucket: Path, **overrides) -> MukaFrankaDataset: - config = { - "dataset_dir": str(bucket), - "num_frames": 5, - "video_stride": 1, - "height": 384, - "width": 320, - "multiview": True, - "normalize_mode": None, - } - config.update(overrides) - return MukaFrankaDataset.from_config(OmegaConf.create(config), split="train") - - -def test_closedness_and_euler_conversion_match_pretrain_contract(): - np.testing.assert_array_equal( - closedness_to_open_scale(np.array([[0.0], [0.5], [1.0]], dtype=np.float32)), - np.array([[1.0], [0.0], [-1.0]], dtype=np.float32), - ) - converted = euler7_to_eef10(_source_state()) - assert converted.shape == (EP_LENGTH, EEF10_DIM) - np.testing.assert_allclose(converted[0, 3:9], [1, 0, 0, 0, 1, 0], atol=1e-6) - np.testing.assert_allclose(converted[1, 3:9], [0, 1, 0, -1, 0, 0], atol=1e-6) - np.testing.assert_allclose(np.linalg.norm(converted[:, 3:6], axis=-1), 1.0, atol=1e-6) - np.testing.assert_allclose(np.linalg.norm(converted[:, 6:9], axis=-1), 1.0, atol=1e-6) - np.testing.assert_allclose(np.sum(converted[:, 3:6] * converted[:, 6:9], axis=-1), 0.0, atol=1e-6) - - -def test_reader_converts_both_streams_and_masks_repeated_terminal_action(tmp_path: Path): - raw_state, raw_action = _write_bucket(tmp_path) - dataset = _dataset(tmp_path) - with _mock_decoder(): - first = dataset[0] - terminal = dataset[len(dataset) - 1] - - assert dataset.action_dim == EEF10_DIM - np.testing.assert_allclose(first["action"].numpy(), euler7_to_eef10(raw_action[:4])) - np.testing.assert_allclose(first["proprio"].numpy(), euler7_to_eef10(raw_state[:1])) - assert first["action_mask"].all() - assert first["proprio_mask"].all() - assert not terminal["action_mask"].any() - assert first["prompt"] == "place the object" - assert first["video"][0].size == (320, 384) - - -def test_unify_scatter_uses_only_left_arm_slots_zero_through_nine(tmp_path: Path): - _write_bucket(tmp_path) - raw_dataset = _dataset(tmp_path) - unified_dataset = _dataset(tmp_path, unify_action=True, unify_action_map=UNIFY_MAP) - with _mock_decoder(): - raw = raw_dataset[0] - unified = unified_dataset[0] - - dst = parse_unify_spec(UNIFY_MAP, UNIFY_DIM) - assert unified_dataset.action_dim == UNIFY_DIM - np.testing.assert_array_equal(dst, np.arange(10)) - np.testing.assert_allclose(unmap_from_unify(unified["action"].numpy(), dst), raw["action"].numpy()) - np.testing.assert_allclose(unmap_from_unify(unified["proprio"].numpy(), dst), raw["proprio"].numpy()) - assert unified["action_mask"][:, :10].all() - assert not unified["action_mask"][:, 10:].any() - - -def test_stats_pool_real_actions_and_states_and_pin_rot6d_identity(tmp_path: Path): - raw_state, raw_action = _write_bucket(tmp_path) - dataset = _dataset(tmp_path) - mode, dim, stats, action_rows, state_rows = _compute_global_stats(dataset, reservoir_cap=100) - expected = np.concatenate((euler7_to_eef10(raw_action[:-1]), euler7_to_eef10(raw_state)), axis=0) - ordinary = np.setdiff1d(np.arange(EEF10_DIM), ROT6D_DIMS_EEF10) - - assert mode == "eef" - assert dim == EEF10_DIM - assert action_rows == EP_LENGTH - 1 - assert state_rows == EP_LENGTH - np.testing.assert_allclose(np.asarray(stats["mean"])[ordinary], expected.mean(0)[ordinary], atol=1e-7) - for key, value in { - "mean": 0.0, - "std": 1.0, - "min": -1.0, - "max": 1.0, - "q01": -1.0, - "q99": 1.0, - }.items(): - np.testing.assert_array_equal(np.asarray(stats[key])[list(ROT6D_DIMS_EEF10)], value) - assert stats["gripper_convention"] == GRIPPER_CONVENTION - assert stats["raw_gripper_convention"] == RAW_GRIPPER_CONVENTION - assert stats["gripper_transform"] == GRIPPER_TRANSFORM - assert stats["rotation_convention"] == ROTATION_CONVENTION - assert stats["stats_population"] == STATS_POPULATION - assert stats["split"] == "train" - assert stats["num_episodes"] == 1 - - -def test_registry_and_training_yaml_select_muka_franka(tmp_path: Path): - _write_bucket(tmp_path) - dataset = build_dataset( - OmegaConf.create( - { - "type": "muka_franka", - "dataset_dir": str(tmp_path), - "num_frames": 5, - "height": 384, - "width": 320, - "multiview": True, - "normalize_mode": None, - } - ) - ) - assert isinstance(dataset, MukaFrankaDataset) - assert "muka_franka" in list_registered_datasets() - - config = OmegaConf.load("configs/dataloader/pretrain_data/muka_franka.yaml") - assert config.type == "muka_franka" - assert config.dataset_dir == "/path/to/muka_franka_lerobot_v3" - assert config.gripper_convention == GRIPPER_CONVENTION - assert list(config.unify_action_map) == UNIFY_MAP - assert config.unify_action is True diff --git a/tests/dataloader/test_normalization_utils.py b/tests/dataloader/test_normalization_utils.py index c074931dc..15e8c7527 100644 --- a/tests/dataloader/test_normalization_utils.py +++ b/tests/dataloader/test_normalization_utils.py @@ -8,6 +8,7 @@ from __future__ import annotations import numpy as np +import pytest from openwam.dataloader.utils.normalization import apply_normalization @@ -134,3 +135,104 @@ def test_quantile_within_bounds_no_clip(self): # No clipping for values within [q01, q99] expected = np.array([[0.5, 0.5], [-0.5, -0.5]], dtype=np.float32) np.testing.assert_allclose(out, expected, atol=1e-6) + + +class TestDatasetStatsBoundary: + def test_root_mode_does_not_borrow_statistics_outside_its_family(self, tmp_path): + from openwam.dataloader.utils.normalization import find_dataset_stats + + family = tmp_path / "family" + bucket = family / "category" / "robot" / "task" + (bucket / "meta").mkdir(parents=True) + (tmp_path / "meta").mkdir() + (tmp_path / "meta" / "stats_robot.json").write_text("{}") + with pytest.raises(FileNotFoundError, match="dataset family"): + find_dataset_stats(bucket, "stats_robot.json", family_root=family) + + def test_standalone_bucket_stops_at_nearest_family_metadata(self, tmp_path): + from openwam.dataloader.utils.normalization import find_dataset_stats + + family = tmp_path / "family" + bucket = family / "nested" / "task" + (bucket / "meta").mkdir(parents=True) + (family / "meta").mkdir() + (tmp_path / "meta").mkdir() + (tmp_path / "meta" / "stats_robot.json").write_text("{}") + with pytest.raises(FileNotFoundError, match="dataset family"): + find_dataset_stats(bucket, "stats_robot.json") + + def test_symlinked_bucket_uses_configured_family_statistics(self, tmp_path): + from openwam.dataloader.utils.normalization import find_dataset_stats + + physical = tmp_path / "physical" / "bucket" + (physical / "meta").mkdir(parents=True) + (physical.parent / "meta").mkdir() + (physical.parent / "meta" / "stats_robot.json").write_text("physical") + selected = tmp_path / "selected" + (selected / "meta").mkdir(parents=True) + expected = selected / "meta" / "stats_robot.json" + expected.write_text("selected") + alias = selected / "bucket" + try: + alias.symlink_to(physical, target_is_directory=True) + except OSError: + pytest.skip("directory symlinks are unavailable") + assert find_dataset_stats(alias, "stats_robot.json", family_root=selected) == expected + + def test_local_statistics_take_precedence_within_the_family(self, tmp_path): + from openwam.dataloader.utils.normalization import find_dataset_stats + + bucket = tmp_path / "bucket" + (bucket / "meta").mkdir(parents=True) + expected = bucket / "meta" / "stats_robot.json" + expected.write_text("local") + (tmp_path / "meta").mkdir() + (tmp_path / "meta" / "stats_robot.json").write_text("pooled") + assert find_dataset_stats(bucket, "stats_robot.json", family_root=tmp_path) == expected + + def test_bucket_must_belong_to_explicit_family(self, tmp_path): + from openwam.dataloader.utils.normalization import find_dataset_stats + + with pytest.raises(ValueError, match="outside its statistics family"): + find_dataset_stats(tmp_path / "elsewhere", "stats.json", family_root=tmp_path / "family") + + def test_filename_cannot_escape_metadata_directory(self, tmp_path): + from openwam.dataloader.utils.normalization import find_dataset_stats + + with pytest.raises(ValueError, match="single file name"): + find_dataset_stats(tmp_path, "../stats.json") + + +class TestMaterializedStatsValidation: + @pytest.mark.parametrize( + "field,index,value,match", + [ + ("std", 0, -1.0, "standard deviation"), + ("std", 3, -1.0, "standard deviation"), + ("min", 0, 2.0, "min <= max"), + ("q01", 0, 2.0, "q01 <= q99"), + ("q01", 3, 2.0, "q01 <= q99"), + ("mean", 0, float("nan"), "non-finite"), + ], + ) + def test_corrupt_stats_are_rejected_before_rotation_identity_pinning(self, field, index, value, match): + from openwam.dataloader.utils.normalization import materialize_eef_stats + + raw = { + "mean": [0.0] * 10, + "std": [1.0] * 10, + "min": [-1.0] * 10, + "max": [1.0] * 10, + "q01": [-1.0] * 10, + "q99": [1.0] * 10, + } + raw[field][index] = value + with pytest.raises(ValueError, match=match): + materialize_eef_stats(raw, "quantile", dim=10, strict_minmax=False, force_rot6d_identity=True) + + def test_optional_defaults_still_support_benchmark_min_max_stats(self): + from openwam.dataloader.utils.normalization import materialize_eef_stats + + stats = materialize_eef_stats({"min": [-2.0] * 10, "max": [2.0] * 10}, "min-max", dim=10, strict_minmax=False) + np.testing.assert_array_equal(stats["mean"], np.zeros(10)) + np.testing.assert_array_equal(stats["std"], np.ones(10)) diff --git a/tests/dataloader/test_oxe_droid.py b/tests/dataloader/test_oxe_droid.py new file mode 100644 index 000000000..954f76dec --- /dev/null +++ b/tests/dataloader/test_oxe_droid.py @@ -0,0 +1,170 @@ +"""DROID clean-export loading uses shipped stats without rescanning source rows.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +from PIL import Image + +from openwam.dataloader.oxe_droid import ( + DROID_EEF_STATS_CONTRACT, + DROID_EEF_STATS_CONTRACT_KEY, + OxeDroidDataset, +) +from openwam.dataloader.utils.lerobotv3 import DataContractError + + +@pytest.fixture +def droid_export(tmp_path, monkeypatch): + root = tmp_path / "Droid" + meta = root / "meta" + (meta / "episodes").mkdir(parents=True) + (root / "data" / "chunk-000").mkdir(parents=True) + (meta / "info.json").write_text( + json.dumps( + { + "fps": 15, + "total_episodes": 2, + "total_frames": 48, + "data_path": "data/chunk-{chunk_index:03d}/file-{file_index:03d}.parquet", + "video_path": "videos/{video_key}/chunk-{chunk_index:03d}/file-{file_index:03d}.mp4", + "features": {}, + } + ) + ) + rows = [] + for episode in range(2): + row = { + "episode_index": episode, + "length": 24, + "dataset_from_index": 24 * episode, + "data/chunk_index": 0, + "data/file_index": 0, + } + for camera in ("observation.images.primary", "observation.images.wrist"): + row[f"videos/{camera}/chunk_index"] = 0 + row[f"videos/{camera}/file_index"] = 0 + rows.append(row) + pd.DataFrame(rows).to_parquet(meta / "episodes" / "chunk-000.parquet") + pd.DataFrame({"task_index": [0, 1]}, index=pd.Index(["pick up the cup", "place the cup"], name="task")).to_parquet( + meta / "tasks.parquet" + ) + action = np.zeros((48, 7), dtype=np.float32) + action[:, 0] = np.arange(48) / 100 + action[:, 6] = 0.25 + pose = np.zeros((48, 6), dtype=np.float32) + pose[:, 0] = 0.75 + state = np.zeros((48, 7), dtype=np.float32) + state[:, :3] = 100 # TCP values must never replace the rigid mount pose. + state[:, 6] = 0.8 + pd.DataFrame( + { + "episode_index": np.repeat([0, 1], 24), + "frame_index": np.tile(np.arange(24), 2), + "task_index": np.repeat([0, 1], 24), + "state": list(state), + "other_information.observation_gripper_pose6d": list(pose), + "other_information.action_wrist_pose": list(action), + } + ).to_parquet(root / "data" / "chunk-000" / "file-000.parquet") + stats = { + "mean": [0.0] * 10, + "std": [1.0] * 10, + "min": [-2.0] * 10, + "max": [2.0] * 10, + "q01": [-2.0] * 10, + "q99": [2.0] * 10, + DROID_EEF_STATS_CONTRACT_KEY: DROID_EEF_STATS_CONTRACT, + } + (meta / "eef_stats.json").write_text(json.dumps(stats)) + monkeypatch.setattr( + "openwam.dataloader.bases.lerobot_v3_reader._decode_video_frames", + lambda path, frame_indices, height, width: [Image.new("RGB", (width, height)) for _ in frame_indices], + ) + return root + + +def _reader(root, **kwargs): + return OxeDroidDataset(str(root), num_frames=9, video_stride=4, height=32, width=32, **kwargs) + + +def test_clean_droid_loads_without_exclusion_certificate_or_fallback_columns(droid_export): + ds = _reader(droid_export, normalize_mode=None) + assert ds._eps_df["episode_index"].tolist() == [0, 1] + sample = ds[0] + assert sample["prompt"] == "pick up the cup" + assert sample["action"].shape == (8, 20) + assert sample["proprio"].shape == (1, 20) + assert sample["action_mask"][:, :10].all() + assert not sample["action_mask"][:, 10:].any() + assert sample["proprio"][0, 0].item() == pytest.approx(0.75) + assert sample["action"][0, 9].item() == pytest.approx(0.75) + assert sample["proprio"][0, 9].item() == pytest.approx(0.2) + np.testing.assert_allclose(sample["action"][0, 3:9], [1, 0, 0, 0, 1, 0]) + assert ds[int(ds._cum_n_starts[1])]["prompt"] == "place the cup" + + +def test_droid_shipped_stats_are_used_without_mutating_export(droid_export): + before = {str(p.relative_to(droid_export)): p.stat().st_mtime_ns for p in droid_export.rglob("*") if p.is_file()} + ds = _reader(droid_export) + sample = ds[0] + assert sample["action"][0, 9].item() == pytest.approx(0.375) + np.testing.assert_allclose(sample["action"][0, 3:9], [1, 0, 0, 0, 1, 0]) + after = {str(p.relative_to(droid_export)): p.stat().st_mtime_ns for p in droid_export.rglob("*") if p.is_file()} + assert before == after + + +def test_droid_startup_does_not_scan_numeric_or_prompt_source_shards(droid_export, monkeypatch): + read_parquet = pd.read_parquet + + def metadata_only(path, *args, **kwargs): + assert "data" not in Path(path).relative_to(droid_export).parts + return read_parquet(path, *args, **kwargs) + + monkeypatch.setattr(pd, "read_parquet", metadata_only) + assert len(_reader(droid_export)) == 48 + + +def test_legacy_external_exclusions_cannot_reapply_filtering(droid_export): + (droid_export / "meta" / "excluded_episodes.json").write_text("not a JSON artifact") + assert _reader(droid_export)._eps_df["episode_index"].tolist() == [0, 1] + + +def test_missing_droid_stats_fail_without_creating_output(droid_export): + path = droid_export / "meta" / "eef_stats.json" + path.unlink() + with pytest.raises(FileNotFoundError, match="eef_stats.json"): + _reader(droid_export) + assert not path.exists() + + +def test_droid_stats_reject_incompatible_pose_contract(droid_export): + path = droid_export / "meta" / "eef_stats.json" + stats = json.loads(path.read_text()) + stats[DROID_EEF_STATS_CONTRACT_KEY]["action_pose_source"] = "other_information.action_tcp_pose" + path.write_text(json.dumps(stats)) + with pytest.raises(DataContractError, match="terminal-arm"): + _reader(droid_export) + + +def test_droid_stats_reject_wrong_width(droid_export): + path = droid_export / "meta" / "eef_stats.json" + stats = json.loads(path.read_text()) + stats["mean"] = [0.0] * 11 + path.write_text(json.dumps(stats)) + with pytest.raises((DataContractError, ValueError), match="10|shape"): + _reader(droid_export) + + +def test_droid_cannot_borrow_another_datasets_parent_statistics(droid_export): + local = droid_export / "meta" / "eef_stats.json" + outer = droid_export.parent / "meta" / "eef_stats.json" + outer.parent.mkdir() + outer.write_bytes(local.read_bytes()) + local.unlink() + with pytest.raises(FileNotFoundError, match="DROID requires its own"): + _reader(droid_export) diff --git a/tests/dataloader/test_oxe_mixture_integration.py b/tests/dataloader/test_oxe_mixture_integration.py index 5e21e5b60..5541a947c 100644 --- a/tests/dataloader/test_oxe_mixture_integration.py +++ b/tests/dataloader/test_oxe_mixture_integration.py @@ -31,7 +31,7 @@ class TestMixtureYamlComposition: def test_mixture_includes_expected_entries(self): from hydra import compose, initialize_config_dir - config_dir = os.path.abspath("configs/dataloader/pretrain_data") + config_dir = os.path.abspath("configs/dataloader") with initialize_config_dir(config_dir=config_dir, version_base=None): cfg = compose(config_name="mixture") assert tuple(cfg.datasets.keys()) == MIXTURE_ENTRIES @@ -39,7 +39,7 @@ def test_mixture_includes_expected_entries(self): def test_blocks_inherit_dataset_dir(self): from hydra import compose, initialize_config_dir - config_dir = os.path.abspath("configs/dataloader/pretrain_data") + config_dir = os.path.abspath("configs/dataloader") with initialize_config_dir(config_dir=config_dir, version_base=None): cfg = compose(config_name="mixture") # Each entry inherits dataset_dir + type from its standalone yaml. @@ -51,7 +51,7 @@ def test_blocks_inherit_dataset_dir(self): def test_proportional_weight_strategy_default(self): from hydra import compose, initialize_config_dir - config_dir = os.path.abspath("configs/dataloader/pretrain_data") + config_dir = os.path.abspath("configs/dataloader") with initialize_config_dir(config_dir=config_dir, version_base=None): cfg = compose(config_name="mixture") assert cfg.weight_strategy == "proportional" diff --git a/tests/dataloader/test_oxe_schema.py b/tests/dataloader/test_oxe_schema.py index 24cdc554b..f8c64285d 100644 --- a/tests/dataloader/test_oxe_schema.py +++ b/tests/dataloader/test_oxe_schema.py @@ -1,54 +1,15 @@ -"""Tests for the OXE schema conversion helpers in ``utils.oxe_schema``. - -Pins the per-dataset state/action layouts against the canonical 10-D EEF -representation ``[pos(3) + rot6d(6) + grip(1)]`` consumed by the OXE -readers. -""" - -from __future__ import annotations +"""Canonical DROID and shared Euler pose conversion regressions.""" import numpy as np +import pytest from openwam.dataloader.utils.oxe_schema import ( - bcz_state_to_arm10, droid_euler7_to_arm10, droid_pose6_closedness_to_arm10, - droid_state_to_arm10, euler7_action_to_arm10, - fractal_state_to_arm10, ) -class TestBczStateToArm10: - def test_shape_and_dtype(self): - state = np.zeros((4, 8), dtype=np.float32) - out = bcz_state_to_arm10(state) - assert out.shape == (4, 10) - assert out.dtype == np.float32 - - def test_pad_index_6_dropped(self): - # state[:, 6] = pad is ignored. Use distinct values for grip vs pad. - state = np.zeros((1, 8), dtype=np.float32) - state[0, 6] = 99.0 # pad value — should NOT appear in output - state[0, 7] = 0.5 # gripper - out = bcz_state_to_arm10(state) - # grip slot is the last (index 9) - assert out[0, 9] == np.float32(0.5) - # No 99.0 anywhere in the output - assert not (out == 99.0).any() - - def test_position_passthrough(self): - state = np.array([[1.5, 2.5, 3.5, 0, 0, 0, 0, 0.7]], dtype=np.float32) - out = bcz_state_to_arm10(state) - np.testing.assert_allclose(out[0, :3], [1.5, 2.5, 3.5], atol=1e-6) - - def test_identity_euler_yields_identity_rot6d(self): - state = np.zeros((1, 8), dtype=np.float32) - out = bcz_state_to_arm10(state) - # euler = 0 → rot6d = [1, 0, 0, 0, 1, 0] - np.testing.assert_allclose(out[0, 3:9], [1, 0, 0, 0, 1, 0], atol=1e-6) - - class TestEuler7ActionToArm10: def test_shape(self): action = np.zeros((4, 7), dtype=np.float32) @@ -63,69 +24,6 @@ def test_pos_grip_passthrough(self): assert out[0, 9] == np.float32(0.9) -class TestFractalStateToArm10: - def test_shape(self): - state = np.zeros((4, 8), dtype=np.float32) - state[:, 6] = 1.0 # w = 1 → identity quat - out = fractal_state_to_arm10(state) - assert out.shape == (4, 10) - assert out.dtype == np.float32 - - def test_identity_quat_yields_identity_rot6d(self): - # state[3:7] = [0, 0, 0, 1] is identity in xyzw - state = np.array([[1.0, 2.0, 3.0, 0, 0, 0, 1, 0.5]], dtype=np.float32) - out = fractal_state_to_arm10(state) - # position - np.testing.assert_allclose(out[0, :3], [1, 2, 3], atol=1e-6) - # rot6d - np.testing.assert_allclose(out[0, 3:9], [1, 0, 0, 0, 1, 0], atol=1e-6) - # gripper - assert out[0, 9] == np.float32(0.5) - - def test_quat_vs_euler_consistency(self): - # Build BC-Z-style with euler [0, 0, pi/4]; build Fractal-style with - # equivalent quaternion. Both should give the same rot6d. - from scipy.spatial.transform import Rotation as _R - - T = 5 - euler = np.random.RandomState(0).randn(T, 3).astype(np.float32) - quat = _R.from_euler("xyz", euler).as_quat().astype(np.float32) - - bcz_state = np.zeros((T, 8), dtype=np.float32) - bcz_state[:, 3:6] = euler - fractal_state = np.zeros((T, 8), dtype=np.float32) - fractal_state[:, 3:7] = quat - - bcz_out = bcz_state_to_arm10(bcz_state) - fractal_out = fractal_state_to_arm10(fractal_state) - # Rot6d portion should match - np.testing.assert_allclose(bcz_out[:, 3:9], fractal_out[:, 3:9], atol=1e-5) - - -class TestDroidStateToArm10: - def test_shape(self): - cart = np.zeros((4, 6), dtype=np.float32) - grip = np.zeros((4, 1), dtype=np.float32) - out = droid_state_to_arm10(cart, grip) - assert out.shape == (4, 10) - assert out.dtype == np.float32 - - def test_layout(self): - cart = np.array([[0.1, 0.2, 0.3, 0, 0, 0]], dtype=np.float32) - grip = np.array([[0.8]], dtype=np.float32) - out = droid_state_to_arm10(cart, grip) - np.testing.assert_allclose(out[0, :3], [0.1, 0.2, 0.3], atol=1e-6) - np.testing.assert_allclose(out[0, 3:9], [1, 0, 0, 0, 1, 0], atol=1e-6) - np.testing.assert_allclose(out[0, 9], 0.2, atol=1e-7) - - def test_scalar_gripper_shape(self): - # Even if gripper is shape (T, 1), the output works. - cart = np.zeros((3, 6), dtype=np.float32) - grip = np.array([[0.1], [0.5], [0.9]], dtype=np.float32) - out = droid_state_to_arm10(cart, grip) - np.testing.assert_allclose(out[:, 9], [0.9, 0.5, 0.1], atol=1e-6) - - class TestDroidEuler7ToArm10: def test_closedness_is_inverted_to_canonical_openness(self): value = np.zeros((3, 7), dtype=np.float32) @@ -155,3 +53,29 @@ def test_assembles_arm_side_pose_and_inverts_only_closedness(self): np.testing.assert_allclose(out, reference, atol=1e-7) np.testing.assert_allclose(out[:, :3], pose[:, :3], atol=1e-7) np.testing.assert_allclose(out[:, 9], [0.75], atol=1e-7) + + +@pytest.mark.parametrize("leading", [(), (2,), (2, 3)]) +def test_euler_pose_conversion_preserves_arbitrary_leading_dimensions(leading): + value = np.arange(int(np.prod(leading) or 1) * 7, dtype=np.float32).reshape(leading + (7,)) / 10 + expected = euler7_action_to_arm10(value.reshape(-1, 7)).reshape(leading + (10,)) + np.testing.assert_array_equal(euler7_action_to_arm10(value), expected) + + +@pytest.mark.parametrize("value", [np.zeros(8), np.zeros((2, 6)), np.array(0.0)]) +def test_euler_pose_conversion_rejects_wrong_width_instead_of_dropping_fields(value): + with pytest.raises(ValueError, match="width 7"): + euler7_action_to_arm10(value) + + +@pytest.mark.parametrize("index", [0, 3, 6]) +def test_droid_conversion_rejects_nonfinite_position_rotation_or_gripper(index): + value = np.zeros((1, 7), dtype=np.float32) + value[0, index] = np.nan + with pytest.raises(ValueError, match="finite"): + droid_euler7_to_arm10(value) + + +def test_droid_scalar_pose_is_rejected_with_an_actionable_error(): + with pytest.raises(ValueError, match="width 6"): + droid_pose6_closedness_to_arm10(np.array(0.0), np.zeros(1)) diff --git a/tests/dataloader/test_oxe_stats_computation.py b/tests/dataloader/test_oxe_stats_computation.py deleted file mode 100644 index 607cab352..000000000 --- a/tests/dataloader/test_oxe_stats_computation.py +++ /dev/null @@ -1,304 +0,0 @@ -from __future__ import annotations - -import json - -import numpy as np -import pandas as pd -import pytest - -from openwam.dataloader.oxe_droid import ( - DROID_DATA_POPULATION_DIGEST_KEY, - DROID_EEF_STATS_CONTRACT, - DROID_EEF_STATS_CONTRACT_KEY, - DROID_PROMPT_EXCLUSION_SCHEMA_VERSION, - DROID_PROMPT_INPUTS_DIGEST_KEY, - OxeDroidDataset, - compute_droid_prompt_inputs_digest, -) -from openwam.dataloader.utils.lerobotv3 import ( - digest_lerobot_v3_data_population, - resolve_lerobot_v3_data_population, -) -from openwam.dataloader.utils.stats_computation.oxe_stats_computation import compute_dataset_stats - - -def _write_manifest(root, episode_lengths): - (root / "meta" / "episodes").mkdir() - total_frames = sum(episode_lengths) - (root / "meta" / "info.json").write_text( - json.dumps( - { - "fps": 10, - "total_frames": total_frames, - "data_path": "data/chunk-{chunk_index:03d}/file-{file_index:03d}.parquet", - } - ) - ) - starts = np.cumsum([0, *episode_lengths[:-1]]).tolist() - pd.DataFrame( - { - "episode_index": list(range(len(episode_lengths))), - "length": episode_lengths, - "dataset_from_index": starts, - "data/chunk_index": [0] * len(episode_lengths), - "data/file_index": [0] * len(episode_lengths), - } - ).to_parquet(root / "meta" / "episodes" / "chunk-000.parquet") - - -def _write_exclusions(root, canonical, prompt_owned=None, independently_owned=None): - prompt_owned = prompt_owned or [] - independently_owned = independently_owned if independently_owned is not None else canonical - fallback_chain = list(OxeDroidDataset.PROMPT_FALLBACK_COLS) - total_frames = resolve_lerobot_v3_data_population(root).total_rows - (root / "meta" / "excluded_episodes.json").write_text( - json.dumps( - { - "episode_indices": canonical, - "droid_prompt_exclusions": { - "schema_version": DROID_PROMPT_EXCLUSION_SCHEMA_VERSION, - "fallback_chain": fallback_chain, - "episode_indices": prompt_owned, - "independently_owned_episode_indices": independently_owned, - "latest_scan": { - "episode_indices": prompt_owned, - "stats": { - "rows_scanned": total_frames, - "unresolved_rows": 0, - "episodes_all_unresolved": len(prompt_owned), - "episodes_partially_unresolved": 0, - "task_index_missing_from_tasks_parquet": 0, - "fallback_chain": fallback_chain, - }, - DROID_PROMPT_INPUTS_DIGEST_KEY: compute_droid_prompt_inputs_digest(root), - }, - }, - } - ) - ) - - -def test_droid_stats_exclude_reader_blacklist_population(tmp_path): - root = tmp_path / "Droid" - (root / "meta").mkdir(parents=True) - (root / "data" / "chunk-000").mkdir(parents=True) - n_rows = 60 - kept_state = np.zeros((n_rows, 7), dtype=np.float32) - kept_state[:, :6] = 1_000.0 # Old task-TCP state pose: must not enter stats. - kept_pose = np.zeros((n_rows, 6), dtype=np.float32) - kept_pose[:, 0] = 1.0 - kept_action = np.zeros((n_rows, 7), dtype=np.float32) - kept_action[:, 0] = 1.0 - kept_tcp_action = np.zeros((n_rows, 7), dtype=np.float32) - kept_tcp_action[:, :6] = 2_000.0 # Old task-TCP action: must not enter stats. - # Deliberately malformed 2-D poses: filtering must happen in Arrow before - # numpy conversion, matching the reader's episode-level exclusion behavior. - excluded_pose = np.zeros((n_rows, 2), dtype=np.float32) - excluded_pose[:, 0] = 100.0 - excluded_state = np.zeros((n_rows, 7), dtype=np.float32) - pd.DataFrame( - { - "episode_index": [0] * n_rows + [1] * n_rows, - "task_index": [0] * (n_rows * 2), - "state": list(kept_state) + list(excluded_state), - "other_information.observation_gripper_pose6d": list(kept_pose) + list(excluded_pose), - "other_information.action_wrist_pose": list(kept_action) + list(excluded_pose), - "other_information.action_tcp_pose": list(kept_tcp_action) + list(excluded_pose), - **{column: [""] * (n_rows * 2) for column in OxeDroidDataset.PROMPT_FALLBACK_COLS}, - } - ).to_parquet(root / "data" / "chunk-000" / "file-000.parquet") - pd.DataFrame({"task_index": [0]}, index=pd.Index(["task"], name="task")).to_parquet(root / "meta" / "tasks.parquet") - _write_manifest(root, [n_rows, n_rows]) - _write_exclusions(root, canonical=[1]) - - stats, n_state, n_action = compute_dataset_stats(root, "DROID", rot6d_identity=False) - - assert n_state == n_rows - assert n_action == n_rows - assert stats["n_samples"] == n_rows * 2 - assert stats["max"][0] == 1.0 - assert stats["excluded_episode_indices"] == [1] - assert stats[DROID_EEF_STATS_CONTRACT_KEY] == DROID_EEF_STATS_CONTRACT - assert stats[DROID_DATA_POPULATION_DIGEST_KEY] == digest_lerobot_v3_data_population( - resolve_lerobot_v3_data_population(root) - ) - - -def test_droid_stats_aggregate_canonical_openness_not_raw_closedness(tmp_path): - root = tmp_path / "Droid" - (root / "meta").mkdir(parents=True) - (root / "data" / "chunk-000").mkdir(parents=True) - state_values = np.zeros((4, 7), dtype=np.float32) - state_values[:, :6] = 1_000.0 # Legacy task-TCP state pose sentinel. - state_values[:, 6] = [0.1, 0.2, 0.3, 0.4] - observation_pose_values = np.zeros((4, 6), dtype=np.float32) - action_values = np.zeros((4, 7), dtype=np.float32) - action_values[:, 6] = [0.15, 0.25, 0.35, 0.45] - tcp_action_values = np.full((4, 7), 2_000.0, dtype=np.float32) - pd.DataFrame( - { - "episode_index": [0] * 4, - "task_index": [0] * 4, - "state": list(state_values), - "other_information.observation_gripper_pose6d": list(observation_pose_values), - "other_information.action_wrist_pose": list(action_values), - "other_information.action_tcp_pose": list(tcp_action_values), - **{column: [""] * 4 for column in OxeDroidDataset.PROMPT_FALLBACK_COLS}, - } - ).to_parquet(root / "data" / "chunk-000" / "file-000.parquet") - pd.DataFrame({"task_index": [0]}, index=pd.Index(["task"], name="task")).to_parquet(root / "meta" / "tasks.parquet") - _write_manifest(root, [4]) - _write_exclusions(root, canonical=[]) - - stats, _, _ = compute_dataset_stats(root, "DROID", rot6d_identity=False) - - assert stats["min"][0] == 0.0 - assert stats["max"][0] == 0.0 - assert stats["min"][9] == pytest.approx(0.55) - assert stats["max"][9] == pytest.approx(0.9) - assert stats["mean"][9] == pytest.approx(0.725) - assert stats["q01"][9] == pytest.approx(0.5535) - assert stats["q99"][9] == pytest.approx(0.8965) - assert stats[DROID_EEF_STATS_CONTRACT_KEY] == DROID_EEF_STATS_CONTRACT - - -def test_droid_stats_fail_clearly_when_every_row_is_excluded(tmp_path): - root = tmp_path / "Droid" - (root / "meta").mkdir(parents=True) - (root / "data" / "chunk-000").mkdir(parents=True) - state_values = list(np.zeros((120, 7), dtype=np.float32)) - pose_values = list(np.zeros((120, 6), dtype=np.float32)) - action_values = list(np.zeros((120, 7), dtype=np.float32)) - pd.DataFrame( - { - "episode_index": [0] * 120, - "task_index": [0] * 120, - "state": state_values, - "other_information.observation_gripper_pose6d": pose_values, - "other_information.action_wrist_pose": action_values, - **{column: [""] * 120 for column in OxeDroidDataset.PROMPT_FALLBACK_COLS}, - } - ).to_parquet(root / "data" / "chunk-000" / "file-000.parquet") - pd.DataFrame({"task_index": [0]}, index=pd.Index(["task"], name="task")).to_parquet(root / "meta" / "tasks.parquet") - _write_manifest(root, [120]) - _write_exclusions(root, canonical=[0]) - - with pytest.raises(ValueError, match="every parquet row is excluded"): - compute_dataset_stats(root, "DROID", rot6d_identity=False) - - -def test_droid_stats_ignore_unreferenced_backup_parquet(tmp_path): - root = tmp_path / "Droid" - (root / "meta").mkdir(parents=True) - (root / "data" / "chunk-000").mkdir(parents=True) - state_values = list(np.zeros((2, 7), dtype=np.float32)) - pose_values = list(np.zeros((2, 6), dtype=np.float32)) - action_values = list(np.zeros((2, 7), dtype=np.float32)) - frame = pd.DataFrame( - { - "episode_index": [0, 0], - "task_index": [0, 0], - "state": state_values, - "other_information.observation_gripper_pose6d": pose_values, - "other_information.action_wrist_pose": action_values, - **{column: [""] * 2 for column in OxeDroidDataset.PROMPT_FALLBACK_COLS}, - } - ) - frame.to_parquet(root / "data" / "chunk-000" / "file-000.parquet") - backup = frame.copy() - backup["episode_index"] = 999 - backup.to_parquet(root / "data" / "chunk-000" / "backup.parquet") - pd.DataFrame({"task_index": [0]}, index=pd.Index(["task"], name="task")).to_parquet(root / "meta" / "tasks.parquet") - _write_manifest(root, [2]) - _write_exclusions(root, canonical=[]) - - stats, n_state, n_action = compute_dataset_stats(root, "DROID", rot6d_identity=False) - - assert n_state == 2 - assert n_action == 2 - assert stats["n_samples"] == 4 - - -def test_droid_stats_ignore_unaddressed_tail_rows_in_referenced_shard(tmp_path): - root = tmp_path / "Droid" - (root / "meta").mkdir(parents=True) - (root / "data" / "chunk-000").mkdir(parents=True) - valid_state = np.zeros(7, dtype=np.float32) - valid_pose = np.zeros(6, dtype=np.float32) - valid_action = np.zeros(7, dtype=np.float32) - malformed_tail = np.zeros(2, dtype=np.float32) - pd.DataFrame( - { - "episode_index": [0, 999], - "task_index": [0, 0], - "state": [valid_state, malformed_tail], - "other_information.observation_gripper_pose6d": [valid_pose, malformed_tail], - "other_information.action_wrist_pose": [valid_action, malformed_tail], - **{column: ["", ""] for column in OxeDroidDataset.PROMPT_FALLBACK_COLS}, - } - ).to_parquet(root / "data" / "chunk-000" / "file-000.parquet") - pd.DataFrame({"task_index": [0]}, index=pd.Index(["task"], name="task")).to_parquet(root / "meta" / "tasks.parquet") - _write_manifest(root, [1]) - _write_exclusions(root, canonical=[]) - - stats, n_state, n_action = compute_dataset_stats(root, "DROID", rot6d_identity=False) - - assert n_state == 1 - assert n_action == 1 - assert stats["n_samples"] == 2 - - -def test_droid_stats_do_not_substitute_backup_for_missing_manifest_file(tmp_path): - root = tmp_path / "Droid" - (root / "meta").mkdir(parents=True) - (root / "data" / "chunk-000").mkdir(parents=True) - state_values = list(np.zeros((2, 7), dtype=np.float32)) - pose_values = list(np.zeros((2, 6), dtype=np.float32)) - action_values = list(np.zeros((2, 7), dtype=np.float32)) - pd.DataFrame( - { - "episode_index": [0, 0], - "task_index": [0, 0], - "state": state_values, - "other_information.observation_gripper_pose6d": pose_values, - "other_information.action_wrist_pose": action_values, - **{column: [""] * 2 for column in OxeDroidDataset.PROMPT_FALLBACK_COLS}, - } - ).to_parquet(root / "data" / "chunk-000" / "file-000.parquet") - pd.DataFrame({"task_index": [0]}, index=pd.Index(["task"], name="task")).to_parquet(root / "meta" / "tasks.parquet") - _write_manifest(root, [2]) - _write_exclusions(root, canonical=[]) - data_path = root / "data" / "chunk-000" / "file-000.parquet" - data_path.rename(data_path.with_name("backup.parquet")) - - with pytest.raises((FileNotFoundError, ValueError), match="file-000.parquet"): - compute_dataset_stats(root, "DROID", rot6d_identity=False) - - -def test_droid_stats_reject_manifest_episode_mapping_mismatch(tmp_path): - root = tmp_path / "Droid" - (root / "meta").mkdir(parents=True) - (root / "data" / "chunk-000").mkdir(parents=True) - state_values = list(np.zeros((2, 7), dtype=np.float32)) - pose_values = list(np.zeros((2, 6), dtype=np.float32)) - action_values = list(np.zeros((2, 7), dtype=np.float32)) - pd.DataFrame( - { - "episode_index": [0, 1], - "task_index": [0, 0], - "state": state_values, - "other_information.observation_gripper_pose6d": pose_values, - "other_information.action_wrist_pose": action_values, - **{column: [""] * 2 for column in OxeDroidDataset.PROMPT_FALLBACK_COLS}, - } - ).to_parquet(root / "data" / "chunk-000" / "file-000.parquet") - pd.DataFrame({"task_index": [0]}, index=pd.Index(["task"], name="task")).to_parquet(root / "meta" / "tasks.parquet") - _write_manifest(root, [1, 1]) - _write_exclusions(root, canonical=[]) - manifest_path = root / "meta" / "episodes" / "chunk-000.parquet" - manifest = pd.read_parquet(manifest_path) - manifest["episode_index"] = [1, 0] - manifest.to_parquet(manifest_path) - - with pytest.raises(ValueError, match="prompt source population"): - compute_dataset_stats(root, "DROID", rot6d_identity=False) diff --git a/tests/dataloader/test_robocoin_compute_stats.py b/tests/dataloader/test_robocoin_compute_stats.py deleted file mode 100644 index ace223d85..000000000 --- a/tests/dataloader/test_robocoin_compute_stats.py +++ /dev/null @@ -1,773 +0,0 @@ -"""Tests for RoboCOIN stats accumulation and trim-population provenance. - -mean/std/min/max are streamed exactly; q01/q99 come from a bounded reservoir -sample. Synthetic bucket tests additionally pin trim-aware population selection -and the on-disk provenance contract consumed by ``RoboCOINDataset``. -""" - -from __future__ import annotations - -import hashlib -import json - -import numpy as np -import pandas as pd -import pytest - -from openwam.dataloader.robocoin import ( - ROBOCOIN_WHOLE_BUCKET_EXCLUSIONS, - RoboCOINDataset, - robocoin_bucket_exclusions_provenance, -) -from openwam.dataloader.utils.lerobotv3 import DataContractError -from openwam.dataloader.utils.normalization import ROT6D_DIMS_EEF20 -from openwam.dataloader.utils.stats_computation import robocoin_stats_computation as stats_module -from openwam.dataloader.utils.stats_computation.robocoin_stats_computation import ( - Accumulator, - compute_stats_for_robot_type, - discover_datasets_by_robot_type, -) -from tests.dataloader.test_robocoin_trim import ( - _make_bucket, - _set_info_splits, - _trim_row, - _write_trim_csv, -) - - -def _make_encoded_stats_bucket(tmp_path): - root = tmp_path / "root" - bucket = _make_bucket(root / "bucket", [10]) - data_path = bucket / "data" / "chunk-000" / "file-000.parquet" - df = pd.read_parquet(data_path) - source_values = np.array([100, 100, 2, 3, 4, 5, 6, 7, 100, 100], dtype=np.float32) - for column in ("eef_sim_pose_action", "eef_sim_pose_state"): - eef = np.stack(df[column].values).astype(np.float32) - eef[:, 0] = source_values - df[column] = list(eef) - df.to_parquet(data_path, index=False) - return root, bucket - - -def _make_trim_csv(tmp_path, *, episode_index=0, total_frames=10): - return _write_trim_csv( - tmp_path / "trim.csv", - [ - _trim_row( - dataset="bucket", - episode_index=episode_index, - total_frames=total_frames, - trim_head_to=2, - trim_tail_from=8, - ) - ], - ) - - -def _write_stats(root, payload): - meta = root / "meta" - meta.mkdir(exist_ok=True) - path = meta / "stats_test_robot.json" - path.write_text(json.dumps(payload)) - return path - - -def _write_exclusions(bucket, episode_indices): - path = bucket / "meta" / "excluded_episodes.json" - path.write_text(json.dumps({"episode_indices": episode_indices})) - return path - - -def _set_robot_type(bucket, robot_type): - info_path = bucket / "meta" / "info.json" - info = json.loads(info_path.read_text()) - info["robot_type"] = robot_type - info_path.write_text(json.dumps(info)) - - -@pytest.fixture -def trimmed_stats_case(tmp_path): - root, bucket = _make_encoded_stats_bucket(tmp_path) - trim_csv = _make_trim_csv(tmp_path) - result = compute_stats_for_robot_type( - "test_robot", - [str(bucket)], - rot6d_identity=False, - trim_csv=str(trim_csv), - ) - return root, bucket, trim_csv, result - - -class TestAccumulator: - def test_exact_stats_and_complete_schema(self): - rng = np.random.RandomState(0) - data = rng.uniform(-3.0, 3.0, size=(50000, 4)).astype(np.float32) - acc = Accumulator(dim=4, reservoir_cap=10000, seed=1) - for i in range(0, len(data), 1000): - acc.update_batch(data[i : i + 1000]) - out = acc.finalize() - # mean/std/min/max are streamed over every row → exact. - np.testing.assert_allclose(out["mean"], data.mean(0), atol=1e-3) - np.testing.assert_allclose(out["std"], data.std(0), atol=1e-3) - np.testing.assert_allclose(out["min"], data.min(0), atol=1e-5) - np.testing.assert_allclose(out["max"], data.max(0), atol=1e-5) - # schema carries all six fields the readers may need. - assert {"mean", "std", "min", "max", "q01", "q99"}.issubset(out) - assert len(out["q01"]) == 4 and len(out["q99"]) == 4 - - def test_reservoir_quantiles_approximate_truth(self): - rng = np.random.RandomState(0) - data = rng.uniform(-3.0, 3.0, size=(50000, 4)).astype(np.float32) - acc = Accumulator(dim=4, reservoir_cap=10000, seed=1) - for i in range(0, len(data), 1000): - acc.update_batch(data[i : i + 1000]) - out = acc.finalize() - # cap (10k) < N (50k) → reservoir is a uniform subsample; q01/q99 are - # close to the true quantiles within sampling error. - np.testing.assert_allclose(out["q01"], np.quantile(data, 0.01, axis=0), atol=0.15) - np.testing.assert_allclose(out["q99"], np.quantile(data, 0.99, axis=0), atol=0.15) - - def test_reservoir_holds_all_when_under_cap_gives_exact_quantiles(self): - rng = np.random.RandomState(2) - data = rng.uniform(0.0, 1.0, size=(500, 3)).astype(np.float32) - acc = Accumulator(dim=3, reservoir_cap=10000, seed=0) - acc.update_batch(data) - out = acc.finalize() - # N (500) < cap → reservoir holds every row → quantiles are exact. - np.testing.assert_allclose(out["q01"], np.quantile(data, 0.01, axis=0), atol=1e-5) - np.testing.assert_allclose(out["q99"], np.quantile(data, 0.99, axis=0), atol=1e-5) - - -class TestWholeBucketExclusion: - def test_discovery_and_direct_compute_share_the_same_exclusion(self, tmp_path): - root = tmp_path / "root" - kept = _make_bucket(root / "kept-airbot", [10]) - excluded_name = next(iter(ROBOCOIN_WHOLE_BUCKET_EXCLUSIONS)) - excluded = _make_bucket(root / excluded_name, [10]) - _set_robot_type(kept, "airbot_mmk2") - _set_robot_type(excluded, "airbot_mmk2") - - groups = discover_datasets_by_robot_type(str(root)) - assert groups == {"airbot_mmk2": [str(kept)]} - - result = compute_stats_for_robot_type( - "airbot_mmk2", - [str(kept), str(excluded)], - rot6d_identity=False, - ) - assert result["eef"]["num_datasets"] == 1 - assert result["eef"]["num_timesteps"] == 20 # action + state - assert result["whole_bucket_exclusions_provenance"] == (robocoin_bucket_exclusions_provenance("airbot_mmk2")) - - def test_affected_robot_type_rejects_stats_without_exclusion_provenance(self, tmp_path): - root = tmp_path / "root" - bucket = _make_bucket(root / "kept-airbot", [40]) - _set_robot_type(bucket, "airbot_mmk2") - payload = compute_stats_for_robot_type("airbot_mmk2", [str(bucket)], rot6d_identity=False) - payload.pop("whole_bucket_exclusions_provenance") - meta = root / "meta" - meta.mkdir(exist_ok=True) - (meta / "stats_airbot_mmk2.json").write_text(json.dumps(payload)) - - with pytest.raises(DataContractError, match="whole_bucket_exclusions_provenance"): - RoboCOINDataset( - dataset_dir=str(bucket), - normalize_mode="quantile", - num_frames=5, - video_stride=1, - ) - - -class TestTrimmedPopulationStats: - def test_stats_use_only_kept_rows_and_serialize_trim_provenance(self, trimmed_stats_case): - _, _, trim_csv, result = trimmed_stats_case - eef = result["eef"] - - # Six kept frames in [2, 8), pooled once from action and once from state. - assert eef["num_timesteps"] == 12 - assert eef["min"][0] == pytest.approx(2.0) - assert eef["max"][0] == pytest.approx(7.0) - - # The complete output must survive the CLI's JSON serialization without - # losing either the exact input digest or the trim-policy parameters. - round_tripped = json.loads(json.dumps(result)) - provenance = round_tripped["trim_provenance"] - assert provenance == result["trim_provenance"] - assert provenance["schema_version"] == 1 - assert provenance["sha256"] == hashlib.sha256(trim_csv.read_bytes()).hexdigest() - assert provenance["min_len"] == 1 - assert provenance["zero_span_policy"] == "drop" - assert round_tripped["excluded_episodes_provenance"] == { - "schema_version": 1, - "policy": "drop_matching_episode_index_before_trim", - "datasets": {"bucket": {"episode_indices": []}}, - } - population = round_tripped["population"] - assert {key: population[key] for key in ("schema_version", "split", "policy")} == { - "schema_version": 2, - "split": "train", - "policy": "info_split_then_exclusion_then_trim", - } - assert population["datasets"]["bucket"]["num_episodes"] == 1 - assert population["datasets"]["bucket"]["num_rows"] == 6 - assert len(population["datasets"]["bucket"]["effective_population_digest"]) == 64 - - def test_stats_apply_train_split_before_exclusion_and_trim(self, tmp_path): - root = tmp_path / "root" - bucket = _make_bucket(root / "bucket", [10, 10]) - _set_info_splits(bucket, {"train": "0:1", "val": "1:2"}) - data_path = bucket / "data" / "chunk-000" / "file-000.parquet" - df = pd.read_parquet(data_path) - encoded = np.array([100, 100, 2, 3, 4, 5, 6, 7, 100, 100] + [999] * 10, dtype=np.float32) - for column in ("eef_sim_pose_action", "eef_sim_pose_state"): - values = np.stack(df[column].values).astype(np.float32) - values[:, 0] = encoded - df[column] = list(values) - df.to_parquet(data_path, index=False) - - result = compute_stats_for_robot_type( - "test_robot", - [str(bucket)], - rot6d_identity=False, - trim_csv=str(_make_trim_csv(tmp_path)), - split="train", - ) - - # Six kept train rows, pooled from action and state. The val episode's - # sentinel value must not enter any statistic. - assert result["eef"]["num_timesteps"] == 12 - assert result["eef"]["min"][0] == pytest.approx(2.0) - assert result["eef"]["max"][0] == pytest.approx(7.0) - - def test_stats_exclude_blacklisted_episode_and_serialize_population(self, tmp_path): - root = tmp_path / "root" - bucket = _make_bucket(root / "bucket", [10, 10]) - data_path = bucket / "data" / "chunk-000" / "file-000.parquet" - df = pd.read_parquet(data_path) - source_values = np.array( - [100, 100, 2, 3, 4, 5, 6, 7, 100, 100] + [999] * 10, - dtype=np.float32, - ) - for column in ("eef_sim_pose_action", "eef_sim_pose_state"): - eef = np.stack(df[column].values).astype(np.float32) - eef[:, 0] = source_values - df[column] = list(eef) - df.to_parquet(data_path, index=False) - _write_exclusions(bucket, [1]) - - result = compute_stats_for_robot_type( - "test_robot", - [str(bucket)], - rot6d_identity=False, - trim_csv=str(_make_trim_csv(tmp_path)), - ) - - assert result["eef"]["num_timesteps"] == 12 - assert result["eef"]["min"][0] == pytest.approx(2.0) - assert result["eef"]["max"][0] == pytest.approx(7.0) - assert result["excluded_episodes_provenance"]["datasets"] == {"bucket": {"episode_indices": [1]}} - - @pytest.mark.parametrize( - ("episode_index", "total_frames"), - [(0, 9), (99, 10)], - ids=["stale-length", "unknown-episode"], - ) - def test_stale_or_unknown_trim_spec_fails_closed(self, tmp_path, episode_index, total_frames): - _, bucket = _make_encoded_stats_bucket(tmp_path) - trim_csv = _make_trim_csv( - tmp_path, - episode_index=episode_index, - total_frames=total_frames, - ) - - with pytest.raises(DataContractError): - compute_stats_for_robot_type( - "test_robot", - [str(bucket)], - rot6d_identity=False, - trim_csv=str(trim_csv), - ) - - def test_trim_span_crossing_parquet_boundary_uses_physical_global_rows(self, tmp_path): - _, bucket = _make_encoded_stats_bucket(tmp_path) - first_path = bucket / "data" / "chunk-000" / "file-000.parquet" - second_path = bucket / "data" / "chunk-000" / "file-001.parquet" - df = pd.read_parquet(first_path) - df.iloc[:5].to_parquet(first_path, index=False) - df.iloc[5:].to_parquet(second_path, index=False) - trim_csv = _make_trim_csv(tmp_path) - - result = compute_stats_for_robot_type( - "test_robot", - [str(bucket)], - rot6d_identity=False, - trim_csv=str(trim_csv), - ) - assert result["eef"]["num_timesteps"] == 12 - assert result["eef"]["min"][0] == pytest.approx(2.0) - assert result["eef"]["max"][0] == pytest.approx(7.0) - - def test_stats_ignore_sorting_earlier_non_numeric_backup_shard(self, tmp_path): - _, bucket = _make_encoded_stats_bucket(tmp_path) - real_path = bucket / "data" / "chunk-000" / "file-000.parquet" - backup_path = bucket / "data" / "chunk--backup" / "file-old.parquet" - backup_path.parent.mkdir() - backup = pd.read_parquet(real_path) - for column in ("eef_sim_pose_action", "eef_sim_pose_state"): - values = np.stack(backup[column].values).astype(np.float32) - values[:, 0] = 999 - backup[column] = list(values) - backup.to_parquet(backup_path, index=False) - - result = compute_stats_for_robot_type( - "test_robot", - [str(bucket)], - rot6d_identity=False, - trim_csv=str(_make_trim_csv(tmp_path)), - ) - - assert result["eef"]["num_files"] == 1 - assert result["eef"]["min"][0] == pytest.approx(2.0) - assert result["eef"]["max"][0] == pytest.approx(7.0) - - def test_stats_reject_duplicate_numeric_shard_coordinates(self, tmp_path): - _, bucket = _make_encoded_stats_bucket(tmp_path) - canonical = bucket / "data" / "chunk-000" / "file-000.parquet" - alias = canonical.with_name("file-0.parquet") - pd.read_parquet(canonical).to_parquet(alias, index=False) - - with pytest.raises(DataContractError, match="duplicate numeric data shard coordinates"): - compute_stats_for_robot_type( - "test_robot", - [str(bucket)], - rot6d_identity=False, - trim_csv=str(_make_trim_csv(tmp_path)), - ) - - def test_stats_fail_if_trim_csv_changes_during_scan(self, tmp_path, monkeypatch): - _, bucket = _make_encoded_stats_bucket(tmp_path) - trim_csv = _make_trim_csv(tmp_path) - replacement = _write_trim_csv( - tmp_path / "replacement.csv", - [_trim_row(trim_head_to=3, trim_tail_from=7)], - ) - original_read_parquet = stats_module.pd.read_parquet - replaced = False - - def replace_before_read(*args, **kwargs): - nonlocal replaced - if not replaced: - replacement.replace(trim_csv) - replaced = True - return original_read_parquet(*args, **kwargs) - - monkeypatch.setattr(stats_module.pd, "read_parquet", replace_before_read) - with pytest.raises(DataContractError, match="changed while.*stats scan"): - compute_stats_for_robot_type( - "test_robot", - [str(bucket)], - rot6d_identity=False, - trim_csv=str(trim_csv), - ) - - def test_stats_fail_if_exclusions_change_during_scan(self, tmp_path, monkeypatch): - _, bucket = _make_encoded_stats_bucket(tmp_path) - exclusions = _write_exclusions(bucket, []) - replacement = bucket / "meta" / "replacement-exclusions.json" - replacement.write_text(json.dumps({"episode_indices": [0]})) - original_read_parquet = stats_module.pd.read_parquet - replaced = False - - def replace_before_read(*args, **kwargs): - nonlocal replaced - if not replaced: - replacement.replace(exclusions) - replaced = True - return original_read_parquet(*args, **kwargs) - - monkeypatch.setattr(stats_module.pd, "read_parquet", replace_before_read) - with pytest.raises(DataContractError, match="changed while.*stats scan"): - compute_stats_for_robot_type( - "test_robot", - [str(bucket)], - rot6d_identity=False, - trim_csv=str(_make_trim_csv(tmp_path)), - ) - - -def test_untrimmed_stats_keep_legacy_nonstandard_parquet_discovery(tmp_path): - _, bucket = _make_encoded_stats_bucket(tmp_path) - original = bucket / "data" / "chunk-000" / "file-000.parquet" - legacy_dir = bucket / "data" / "legacy" - legacy_dir.mkdir() - original.rename(legacy_dir / "shard.parquet") - - result = compute_stats_for_robot_type( - "test_robot", - [str(bucket)], - rot6d_identity=False, - trim_csv=None, - ) - assert result["eef"]["num_timesteps"] == 20 - - -class TestReaderTrimStatsProvenance: - def test_matching_provenance_is_accepted(self, trimmed_stats_case): - root, bucket, trim_csv, result = trimmed_stats_case - _write_stats(root, result) - - reader = RoboCOINDataset( - dataset_dir=str(bucket), - normalize_mode="min-max", - trim_csv=str(trim_csv), - ) - assert reader._eps_df["length"].tolist() == [6] - assert reader._normalization_stats is not None - dims = list(ROT6D_DIMS_EEF20) - np.testing.assert_array_equal(reader._normalization_stats["mean"][dims], 0.0) - np.testing.assert_array_equal(reader._normalization_stats["std"][dims], 1.0) - np.testing.assert_array_equal(reader._normalization_stats["min"][dims], -1.0) - np.testing.assert_array_equal(reader._normalization_stats["max"][dims], 1.0) - - def test_missing_exclusion_provenance_is_rejected(self, trimmed_stats_case): - root, bucket, trim_csv, result = trimmed_stats_case - payload = json.loads(json.dumps(result)) - payload.pop("excluded_episodes_provenance") - _write_stats(root, payload) - - with pytest.raises(DataContractError, match="excluded_episodes_provenance"): - RoboCOINDataset( - dataset_dir=str(bucket), - normalize_mode="min-max", - trim_csv=str(trim_csv), - ) - - def test_non_train_stats_population_is_rejected(self, trimmed_stats_case): - root, bucket, trim_csv, result = trimmed_stats_case - result["population"]["split"] = "val" - _write_stats(root, result) - - with pytest.raises(DataContractError, match="train-derived"): - RoboCOINDataset( - dataset_dir=str(bucket), - normalize_mode="min-max", - trim_csv=str(trim_csv), - ) - - def test_info_split_drift_after_stats_is_rejected(self, tmp_path): - root = tmp_path / "root" - bucket = _make_bucket(root / "bucket", [10, 10]) - trim_csv = _make_trim_csv(tmp_path) - result = compute_stats_for_robot_type( - "test_robot", - [str(bucket)], - rot6d_identity=False, - trim_csv=str(trim_csv), - split="train", - ) - _write_stats(root, result) - - # The generated stats covered both episodes. Narrowing train afterwards - # must invalidate the stored physical-span digest even though the - # top-level label still truthfully says "train". - _set_info_splits(bucket, {"train": "0:1", "val": "1:2"}) - with pytest.raises(DataContractError, match="effective train population"): - RoboCOINDataset( - dataset_dir=str(bucket), - normalize_mode="min-max", - trim_csv=str(trim_csv), - ) - - def test_sibling_info_split_drift_invalidates_pooled_stats(self, tmp_path): - root = tmp_path / "root" - bucket_a = _make_bucket(root / "bucket-a", [10]) - bucket_b = _make_bucket(root / "bucket-b", [10, 10]) - trim_csv = _write_trim_csv( - tmp_path / "trim.csv", - [_trim_row(dataset="bucket-a")], - ) - result = compute_stats_for_robot_type( - "test_robot", - [str(bucket_a), str(bucket_b)], - rot6d_identity=False, - trim_csv=str(trim_csv), - split="train", - ) - _write_stats(root, result) - - _set_info_splits(bucket_b, {"train": "0:1", "val": "1:2"}) - with pytest.raises(DataContractError, match="contributor 'bucket-b'.*no longer matches"): - RoboCOINDataset( - dataset_dir=str(bucket_a), - normalize_mode="min-max", - trim_csv=str(trim_csv), - ) - - @pytest.mark.parametrize( - ("field", "bad_value"), - [ - ("schema_version", True), - ("policy", "keep"), - ("episode_indices", [True]), - ], - ) - def test_mismatched_exclusion_provenance_is_rejected( - self, - trimmed_stats_case, - field, - bad_value, - ): - root, bucket, trim_csv, result = trimmed_stats_case - payload = json.loads(json.dumps(result)) - if field == "episode_indices": - payload["excluded_episodes_provenance"]["datasets"]["bucket"][field] = bad_value - else: - payload["excluded_episodes_provenance"][field] = bad_value - _write_stats(root, payload) - - with pytest.raises(DataContractError, match="excluded_episodes_provenance"): - RoboCOINDataset( - dataset_dir=str(bucket), - normalize_mode="min-max", - trim_csv=str(trim_csv), - ) - - def test_changed_current_bucket_exclusions_are_rejected(self, trimmed_stats_case): - root, bucket, trim_csv, result = trimmed_stats_case - _write_stats(root, result) - _write_exclusions(bucket, [0]) - - with pytest.raises(DataContractError, match="current exclusion population"): - RoboCOINDataset( - dataset_dir=str(bucket), - normalize_mode="min-max", - trim_csv=str(trim_csv), - ) - - def test_changed_sibling_exclusions_in_pooled_stats_are_rejected(self, tmp_path): - root = tmp_path / "root" - bucket_a = _make_bucket(root / "bucket-a", [10]) - bucket_b = _make_bucket(root / "bucket-b", [10]) - trim_csv = _write_trim_csv( - tmp_path / "trim.csv", - [_trim_row(dataset="bucket-a")], - ) - result = compute_stats_for_robot_type( - "test_robot", - [str(bucket_a), str(bucket_b)], - rot6d_identity=False, - trim_csv=str(trim_csv), - ) - _write_stats(root, result) - _write_exclusions(bucket_b, [0]) - - with pytest.raises(DataContractError, match="contributor 'bucket-b'.*current exclusion population"): - RoboCOINDataset( - dataset_dir=str(bucket_a), - normalize_mode="min-max", - trim_csv=str(trim_csv), - ) - - def test_reader_fails_if_sibling_exclusions_change_during_construction( - self, - tmp_path, - monkeypatch, - ): - root = tmp_path / "root" - bucket_a = _make_bucket(root / "bucket-a", [10]) - bucket_b = _make_bucket(root / "bucket-b", [10]) - trim_csv = _write_trim_csv( - tmp_path / "trim.csv", - [_trim_row(dataset="bucket-a")], - ) - result = compute_stats_for_robot_type( - "test_robot", - [str(bucket_a), str(bucket_b)], - rot6d_identity=False, - trim_csv=str(trim_csv), - ) - _write_stats(root, result) - target = bucket_b / "meta" / "excluded_episodes.json" - replacement = bucket_b / "meta" / "replacement-exclusions.json" - replacement.write_text(json.dumps({"episode_indices": [0]})) - original_post_init = RoboCOINDataset._post_init - - def post_init_then_replace(self, info): - original_post_init(self, info) - replacement.replace(target) - - monkeypatch.setattr(RoboCOINDataset, "_post_init", post_init_then_replace) - with pytest.raises(DataContractError, match="changed while reader construction"): - RoboCOINDataset( - dataset_dir=str(bucket_a), - normalize_mode="min-max", - trim_csv=str(trim_csv), - ) - - def test_reader_fails_if_exclusions_change_during_construction( - self, - trimmed_stats_case, - monkeypatch, - ): - root, bucket, trim_csv, result = trimmed_stats_case - _write_stats(root, result) - target = bucket / "meta" / "excluded_episodes.json" - replacement = bucket / "meta" / "replacement-exclusions.json" - replacement.write_text(json.dumps({"episode_indices": [0]})) - original_load_prompts = RoboCOINDataset._load_prompts - - def load_prompts_then_replace(self): - original_load_prompts(self) - replacement.replace(target) - - monkeypatch.setattr(RoboCOINDataset, "_load_prompts", load_prompts_then_replace) - with pytest.raises(DataContractError, match="changed while reader construction"): - RoboCOINDataset( - dataset_dir=str(bucket), - normalize_mode="min-max", - trim_csv=str(trim_csv), - ) - - def test_missing_provenance_is_rejected_when_trim_is_enabled(self, trimmed_stats_case): - root, bucket, trim_csv, result = trimmed_stats_case - payload = json.loads(json.dumps(result)) - payload.pop("trim_provenance") - _write_stats(root, payload) - - with pytest.raises(DataContractError): - RoboCOINDataset( - dataset_dir=str(bucket), - normalize_mode="min-max", - trim_csv=str(trim_csv), - ) - - @pytest.mark.parametrize( - "missing_field", - ["schema_version", "sha256", "min_len", "zero_span_policy"], - ) - def test_each_missing_provenance_field_is_rejected(self, trimmed_stats_case, missing_field): - root, bucket, trim_csv, result = trimmed_stats_case - payload = json.loads(json.dumps(result)) - payload["trim_provenance"].pop(missing_field) - _write_stats(root, payload) - - with pytest.raises(DataContractError): - RoboCOINDataset( - dataset_dir=str(bucket), - normalize_mode="min-max", - trim_csv=str(trim_csv), - ) - - @pytest.mark.parametrize( - ("field", "bad_value"), - [ - ("schema_version", 2), - ("sha256", "0" * 64), - ("min_len", 2), - ("zero_span_policy", "keep"), - ], - ) - def test_each_mismatched_provenance_field_is_rejected( - self, - trimmed_stats_case, - field, - bad_value, - ): - root, bucket, trim_csv, result = trimmed_stats_case - payload = json.loads(json.dumps(result)) - payload["trim_provenance"][field] = bad_value - _write_stats(root, payload) - - with pytest.raises(DataContractError): - RoboCOINDataset( - dataset_dir=str(bucket), - normalize_mode="min-max", - trim_csv=str(trim_csv), - ) - - def test_trimmed_stats_are_rejected_when_reader_trim_is_disabled(self, trimmed_stats_case): - root, bucket, _, result = trimmed_stats_case - _write_stats(root, result) - - with pytest.raises(DataContractError): - RoboCOINDataset( - dataset_dir=str(bucket), - normalize_mode="min-max", - trim_csv=None, - ) - - def test_legacy_stats_without_provenance_are_accepted_when_trim_is_disabled( - self, - trimmed_stats_case, - ): - root, bucket, _, result = trimmed_stats_case - legacy_payload = json.loads(json.dumps(result)) - legacy_payload.pop("trim_provenance") - legacy_payload.pop("excluded_episodes_provenance") - _write_stats(root, legacy_payload) - - reader = RoboCOINDataset( - dataset_dir=str(bucket), - normalize_mode="min-max", - trim_csv=None, - ) - assert reader._eps_df["length"].tolist() == [10] - assert reader._normalization_stats is not None - - def test_reader_rejects_new_stats_if_csv_changes_after_offsets(self, tmp_path, monkeypatch): - root, bucket = _make_encoded_stats_bucket(tmp_path) - trim_csv = _make_trim_csv(tmp_path) - - replacement = _write_trim_csv( - tmp_path / "replacement-for-b-stats.csv", - [_trim_row(trim_head_to=3, trim_tail_from=7)], - ) - replacement.replace(trim_csv) - stats_b = compute_stats_for_robot_type( - "test_robot", - [str(bucket)], - rot6d_identity=False, - trim_csv=str(trim_csv), - ) - _write_stats(root, stats_b) - - restore_a = _write_trim_csv(tmp_path / "restore-a.csv", [_trim_row()]) - restore_a.replace(trim_csv) - replace_with_b = _write_trim_csv( - tmp_path / "replace-with-b.csv", - [_trim_row(trim_head_to=3, trim_tail_from=7)], - ) - original_add_offsets = RoboCOINDataset._add_data_offsets - - def add_offsets_then_replace(self, eps): - original_add_offsets(self, eps) - replace_with_b.replace(trim_csv) - - monkeypatch.setattr(RoboCOINDataset, "_add_data_offsets", add_offsets_then_replace) - with pytest.raises(DataContractError): - RoboCOINDataset( - dataset_dir=str(bucket), - normalize_mode="min-max", - trim_csv=str(trim_csv), - ) - - def test_reader_rejects_csv_change_without_normalization(self, tmp_path, monkeypatch): - _, bucket = _make_encoded_stats_bucket(tmp_path) - trim_csv = _make_trim_csv(tmp_path) - replacement = _write_trim_csv( - tmp_path / "replacement.csv", - [_trim_row(trim_head_to=3, trim_tail_from=7)], - ) - original_add_offsets = RoboCOINDataset._add_data_offsets - - def add_offsets_then_replace(self, eps): - original_add_offsets(self, eps) - replacement.replace(trim_csv) - - monkeypatch.setattr(RoboCOINDataset, "_add_data_offsets", add_offsets_then_replace) - with pytest.raises(DataContractError, match="changed while reader construction"): - RoboCOINDataset( - dataset_dir=str(bucket), - normalize_mode=None, - trim_csv=str(trim_csv), - ) diff --git a/tests/dataloader/test_robocoin_final.py b/tests/dataloader/test_robocoin_final.py new file mode 100644 index 000000000..c2c056701 --- /dev/null +++ b/tests/dataloader/test_robocoin_final.py @@ -0,0 +1,298 @@ +"""Regression tests for the physically cleaned RoboCOIN training contract.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +import torch +from PIL import Image + +from openwam.dataloader.robocoin import MultiRobotCOINDataset, RoboCOINDataset +from openwam.dataloader.utils.normalization import ROT6D_DIMS_EEF20 + +CAMERA = "observation.images.cam_head_rgb" + + +def _make_bucket(bucket: Path, lengths=(8, 8), *, robot_type="test_robot", dex=False): + meta = bucket / "meta" + (meta / "episodes").mkdir(parents=True) + features = { + CAMERA: {"dtype": "video"}, + "eef_sim_pose_action": {}, + "eef_sim_pose_state": {}, + } + if dex: + motors = ["left_hand_joint_0", "left_hand_joint_1", "right_hand_joint_0"] + features.update({"action": {"names": motors}, "observation.state": {"names": motors}}) + else: + features.update({"gripper_open_scale_action": {}, "gripper_open_scale_state": {}}) + info = { + "fps": 30, + "robot_type": robot_type, + "data_path": "data/chunk-{chunk_index:03d}/file-{file_index:03d}.parquet", + "video_path": "videos/{video_key}/chunk-{chunk_index:03d}/file-{file_index:03d}.mp4", + "features": features, + } + (meta / "info.json").write_text(json.dumps(info), encoding="utf-8") + starts = np.cumsum([0, *lengths[:-1]]) + pd.DataFrame( + { + "episode_index": range(len(lengths)), + "length": lengths, + "dataset_from_index": starts, + "data/chunk_index": 0, + "data/file_index": 0, + f"videos/{CAMERA}/chunk_index": 0, + f"videos/{CAMERA}/file_index": 0, + } + ).to_parquet(meta / "episodes" / "chunk-000.parquet", index=False) + pd.DataFrame({"task_index": [0]}, index=pd.Index(["move the object"], name="task")).to_parquet( + meta / "tasks.parquet" + ) + n = sum(lengths) + pose = np.zeros((n, 12), dtype=np.float32) + pose[:, 0] = np.arange(n) + data = { + "task_index": np.zeros(n, dtype=np.int64), + "episode_index": np.repeat(np.arange(len(lengths)), lengths), + "frame_index": np.concatenate([np.arange(length) for length in lengths]), + "eef_sim_pose_action": list(pose), + "eef_sim_pose_state": list(pose), + } + if dex: + fingers = np.tile([0.1, 0.2, 0.3], (n, 1)).astype(np.float32) + data.update({"action": list(fingers), "observation.state": list(fingers / 2)}) + else: + grip = np.tile([0.25, 0.75], (n, 1)).astype(np.float32) + data.update({"gripper_open_scale_action": list(grip), "gripper_open_scale_state": list(grip)}) + data_dir = bucket / "data" / "chunk-000" + data_dir.mkdir(parents=True) + pd.DataFrame(data).to_parquet(data_dir / "file-000.parquet", index=False) + return bucket + + +def _vectors(dim): + return { + "mean": [0.0] * dim, + "std": [1.0] * dim, + "min": [-1.0] * dim, + "max": [1.0] * dim, + "q01": [-1.0] * dim, + "q99": [1.0] * dim, + } + + +def _write_stats(family, *, robot_type="test_robot", dex=False, payload=None): + meta = family / "meta" + meta.mkdir(exist_ok=True) + raw = payload if payload is not None else {"eef": {**_vectors(20), "robot_type": robot_type}} + if dex: + raw["hand"] = {**_vectors(3), "dof_left": 2, "dof_right": 1} + path = meta / f"stats_{robot_type}.json" + path.write_text(json.dumps(raw), encoding="utf-8") + return path + + +@pytest.fixture +def decode_frame_indices(monkeypatch): + def decode(path, frame_indices, height, width): + del path + return [Image.new("RGB", (width, height), (int(index), 0, 0)) for index in frame_indices] + + monkeypatch.setattr("openwam.dataloader.bases.lerobot_v3_reader._decode_video_frames", decode) + + +def test_family_and_single_bucket_read_existing_stats_after_relocation(tmp_path): + original = tmp_path / "old_family_name" + _make_bucket(original / "bucket-a") + _make_bucket(original / "bucket-b") + path = _write_stats(original) + raw = json.loads(path.read_text(encoding="utf-8")) + # Historical metadata is informational; no old filesystem path is required. + raw.update({"trim_provenance": {"path": "/retired/audit.csv"}, "whole_bucket_exclusions_provenance": {}}) + path.write_text(json.dumps(raw), encoding="utf-8") + family = tmp_path / "relocated_RoboCOIN" + original.rename(family) + + dataset = RoboCOINDataset.from_config({"dataset_dir": str(family), "normalize_mode": "quantile"}) + assert isinstance(dataset, MultiRobotCOINDataset) + assert len(dataset._buckets) == 2 + assert sum(len(bucket._eps_df) for bucket in dataset._buckets) == 4 + for bucket in dataset._buckets: + np.testing.assert_array_equal(bucket._normalization_stats["q99"], np.ones(20)) + leaf = RoboCOINDataset.from_config({"dataset_dir": str(family / "bucket-a"), "normalize_mode": "quantile"}) + assert len(leaf._eps_df) == 2 + + +def test_processed_manifest_is_used_without_additional_exclusions(tmp_path): + family = tmp_path / "family" + former_denied_name = "Airbot_MMK2_storage_peach_pear" + bucket = _make_bucket(family / former_denied_name, lengths=(8, 5, 1)) + # This obsolete artifact must never be opened by a final-data reader. + (bucket / "meta" / "excluded_episodes.json").write_text("invalid old artifact", encoding="utf-8") + dataset = RoboCOINDataset.from_config({"dataset_dir": str(family), "normalize_mode": None}) + leaf = dataset._buckets[0] + assert leaf._eps_df["episode_index"].tolist() == [0, 1, 2] + assert leaf._eps_df["length"].tolist() == [8, 5, 1] + assert leaf._ep_data_row_offset.tolist() == [0, 8, 13] + assert len(leaf) == 14 + assert "trim_csv" not in leaf.CONFIG_KEYS + + +def test_missing_stats_fail_without_writing_or_computing(tmp_path): + bucket = _make_bucket(tmp_path / "family" / "bucket") + files_before = {path.relative_to(tmp_path): path.read_bytes() for path in tmp_path.rglob("*") if path.is_file()} + with pytest.raises(FileNotFoundError, match="stats_test_robot.json"): + RoboCOINDataset(dataset_dir=str(bucket), normalize_mode="quantile") + assert files_before == { + path.relative_to(tmp_path): path.read_bytes() for path in tmp_path.rglob("*") if path.is_file() + } + + +@pytest.mark.parametrize("bad_value", [[0.0] * 19, [float("nan")] * 20]) +def test_invalid_stored_stat_vectors_fail_early(tmp_path, bad_value): + family = tmp_path / "family" + bucket = _make_bucket(family / "bucket") + raw = {"eef": _vectors(20)} + raw["eef"]["q01"] = bad_value + _write_stats(family, payload=raw) + with pytest.raises(ValueError, match="finite vector with shape"): + RoboCOINDataset(dataset_dir=str(bucket), normalize_mode="quantile") + + +def test_stats_robot_type_must_match_bucket(tmp_path): + family = tmp_path / "family" + bucket = _make_bucket(family / "bucket") + _write_stats(family, payload={"eef": {**_vectors(20), "robot_type": "other_robot"}}) + with pytest.raises(ValueError, match="robot_type"): + RoboCOINDataset(dataset_dir=str(bucket), normalize_mode="quantile") + + +def test_active_normalization_fields_are_required(tmp_path): + family = tmp_path / "family" + bucket = _make_bucket(family / "bucket") + _write_stats(family, payload={"eef": {"min": [-1.0] * 20, "max": [1.0] * 20}}) + with pytest.raises(KeyError, match="q01"): + RoboCOINDataset(dataset_dir=str(bucket), normalize_mode="quantile") + + +def test_split_keeps_physical_data_and_video_offsets_aligned(tmp_path, decode_frame_indices): + bucket = _make_bucket(tmp_path / "bucket") + path = bucket / "meta" / "info.json" + info = json.loads(path.read_text(encoding="utf-8")) + info["splits"] = {"train": "1:2"} + path.write_text(json.dumps(info), encoding="utf-8") + dataset = RoboCOINDataset(dataset_dir=str(bucket), normalize_mode=None, num_frames=5, video_stride=1) + assert dataset._ep_data_row_offset.tolist() == [8] + assert dataset._ep_video_frame_offsets[CAMERA].tolist() == [8] + observed = set() + for index in range(len(dataset)): + sample = dataset[index] + action_rows = sample["action"][sample["action_mask"].any(dim=1), 0].to(torch.int64).tolist() + video_rows = [image.getpixel((0, 0))[0] for image, valid in zip(sample["video"], sample["video_mask"]) if valid] + assert action_rows == video_rows[: len(action_rows)] + observed.update(video_rows) + np.testing.assert_allclose(sample["action"][0, [9, 19]], [0.25, 0.75]) + assert observed == set(range(8, 16)) + + +def test_rot6d_remains_identity_under_stored_normalization(tmp_path): + family = tmp_path / "family" + bucket = _make_bucket(family / "bucket") + raw = {"eef": _vectors(20)} + raw["eef"]["q01"] = [-2.0] * 20 + raw["eef"]["q99"] = [2.0] * 20 + _write_stats(family, payload=raw) + dataset = RoboCOINDataset(dataset_dir=str(bucket), normalize_mode="quantile") + np.testing.assert_array_equal(dataset._normalization_stats["q01"][list(ROT6D_DIMS_EEF20)], -1) + np.testing.assert_array_equal(dataset._normalization_stats["q99"][list(ROT6D_DIMS_EEF20)], 1) + assert dataset._normalization_stats["q01"][0] == -2 + + +def test_dexterous_hand_keeps_fingers_and_masks_scalar_grips(tmp_path, decode_frame_indices): + family = tmp_path / "family" + bucket = _make_bucket(family / "dex", dex=True) + _write_stats(family, dex=True) + dataset = RoboCOINDataset( + dataset_dir=str(bucket), normalize_mode="quantile", unify_action=True, num_frames=5, video_stride=1 + ) + sample = dataset[0] + assert sample["action"].shape[-1] == 80 + np.testing.assert_allclose(sample["action"][0, [10, 11, 44]], [0.1, 0.2, 0.3], atol=1e-6) + np.testing.assert_allclose(sample["proprio"][0, [10, 11, 44]], [0.05, 0.1, 0.15], atol=1e-6) + assert sample["action_mask"][0, [10, 11, 44]].all() + assert not sample["action_mask"][0, [9, 43, 12, 45, 79]].any() + + +def test_dexterous_hand_stats_must_match_finger_dof(tmp_path): + family = tmp_path / "family" + bucket = _make_bucket(family / "dex", dex=True) + _write_stats(family, payload={"eef": _vectors(20), "hand": {**_vectors(3), "dof_left": 1, "dof_right": 2}}) + with pytest.raises(ValueError, match="DOF"): + RoboCOINDataset(dataset_dir=str(bucket), normalize_mode="quantile", unify_action=True) + + +def test_video_timestamps_are_independent_of_contiguous_data_rows(tmp_path, decode_frame_indices): + """A native video shard may contain padding absent from the data shard.""" + bucket = _make_bucket(tmp_path / "bucket") + manifest = bucket / "meta" / "episodes" / "chunk-000.parquet" + episodes = pd.read_parquet(manifest) + episodes[f"videos/{CAMERA}/from_timestamp"] = np.array([100, 150]) / 30.0 + episodes[f"videos/{CAMERA}/to_timestamp"] = np.array([108, 158]) / 30.0 + episodes.to_parquet(manifest, index=False) + dataset = RoboCOINDataset(dataset_dir=str(bucket), normalize_mode=None, num_frames=5, video_stride=1) + first = dataset._getitem_impl(0) + second = dataset._getitem_impl(int(dataset._cum_n_starts[1])) + assert first["action"][0, 0].item() == 0 + assert second["action"][0, 0].item() == 8 + assert first["video"][0].getpixel((0, 0))[0] == 100 + assert second["video"][0].getpixel((0, 0))[0] == 150 + + +@pytest.mark.parametrize("bad_start", [-1.0, float("nan"), 0.05]) +def test_invalid_video_timestamps_fail_before_sampling(tmp_path, bad_start): + bucket = _make_bucket(tmp_path / "bucket") + manifest = bucket / "meta" / "episodes" / "chunk-000.parquet" + episodes = pd.read_parquet(manifest) + episodes[f"videos/{CAMERA}/from_timestamp"] = [bad_start, 8 / 30.0] + episodes.to_parquet(manifest, index=False) + with pytest.raises(ValueError, match="timestamp|frame boundar"): + RoboCOINDataset(dataset_dir=str(bucket), normalize_mode=None) + + +def test_family_stats_cannot_fall_back_outside_selected_root(tmp_path): + family = tmp_path / "family" + _make_bucket(family / "bucket") + _write_stats(tmp_path) + with pytest.raises(FileNotFoundError, match="stats_test_robot.json"): + RoboCOINDataset.from_config({"dataset_dir": str(family), "normalize_mode": "quantile"}) + + +def test_float32_video_timestamps_keep_long_shard_frame_offsets(tmp_path): + bucket = _make_bucket(tmp_path / "bucket") + manifest = bucket / "meta" / "episodes" / "chunk-000.parquet" + episodes = pd.read_parquet(manifest) + episodes[f"videos/{CAMERA}/from_timestamp"] = np.array([300000, 300008], dtype=np.float64).astype( + np.float32 + ) / np.float32(30) + episodes[f"videos/{CAMERA}/to_timestamp"] = np.array([300008, 300016], dtype=np.float64).astype( + np.float32 + ) / np.float32(30) + episodes.to_parquet(manifest, index=False) + dataset = RoboCOINDataset(dataset_dir=str(bucket), normalize_mode=None) + assert dataset._ep_video_frame_offsets[CAMERA].tolist() == [300000, 300008] + + +def test_video_timestamp_span_must_match_episode_length(tmp_path): + bucket = _make_bucket(tmp_path / "bucket") + manifest = bucket / "meta" / "episodes" / "chunk-000.parquet" + episodes = pd.read_parquet(manifest) + episodes[f"videos/{CAMERA}/from_timestamp"] = np.array([0, 8]) / 30.0 + episodes[f"videos/{CAMERA}/to_timestamp"] = np.array([9, 16]) / 30.0 + episodes.to_parquet(manifest, index=False) + with pytest.raises(ValueError, match="span.*episode length"): + RoboCOINDataset(dataset_dir=str(bucket), normalize_mode=None) diff --git a/tests/dataloader/test_robocoin_trim.py b/tests/dataloader/test_robocoin_trim.py deleted file mode 100644 index 9400315be..000000000 --- a/tests/dataloader/test_robocoin_trim.py +++ /dev/null @@ -1,715 +0,0 @@ -"""Regression tests for RoboCOIN's CSV-driven dead-frame trimming. - -The end-to-end case uses a synthetic LeRobot v3 bucket whose action payload -and decoded pixels both encode their source row/frame index. This catches a -trim offset being applied to parquet but not video (or vice versa). -""" - -from __future__ import annotations - -import csv -import json -from pathlib import Path - -import numpy as np -import pandas as pd -import pyarrow as pa -import pyarrow.parquet as pq -import pytest -import torch -from PIL import Image - -from openwam.dataloader.robocoin import ( - ROBOCOIN_WHOLE_BUCKET_EXCLUSIONS, - RoboCOINDataset, - _load_trim_spec, -) -from openwam.dataloader.utils.lerobotv3 import DataContractError - -CAM = "observation.images.cam_head_rgb" -TRIM_COLUMNS = ( - "dataset", - "episode_index", - "total_frames", - "trim_head_to", - "trim_tail_from", -) - - -def _write_trim_csv(path: Path, rows: list[dict], columns=TRIM_COLUMNS) -> Path: - with path.open("w", newline="") as fh: - writer = csv.DictWriter(fh, fieldnames=columns, extrasaction="ignore") - writer.writeheader() - writer.writerows(rows) - return path - - -def _trim_row(**overrides) -> dict: - row = { - "dataset": "bucket", - "episode_index": 0, - "total_frames": 10, - "trim_head_to": 2, - "trim_tail_from": 8, - } - row.update(overrides) - return row - - -def _episodes(lengths: list[int]) -> pd.DataFrame: - starts = np.concatenate([[0], np.cumsum(lengths[:-1])]).astype(np.int64) - return pd.DataFrame( - { - "episode_index": np.arange(len(lengths), dtype=np.int64), - "length": lengths, - "_data_row_offset": starts, - f"_video_frame_offset/{CAM}": starts, - "_video_frame_offset/second_camera": starts + 100, - } - ) - - -def _filter(tmp_path: Path, eps: pd.DataFrame, rows: list[dict]) -> pd.DataFrame: - trim_csv = _write_trim_csv(tmp_path / "trim.csv", rows) - reader = RoboCOINDataset.__new__(RoboCOINDataset) - reader._trim_csv = trim_csv - reader._dataset_id = "bucket" - reader._fps = 30.0 - return reader._filter_episodes(eps) - - -@pytest.fixture(autouse=True) -def _isolate_trim_cache(): - # Immutable trim snapshots are shared by many bucket readers. - # Tests reuse tmp_path names across separate pytest processes, so isolate it. - from openwam.dataloader import robocoin - - robocoin._TRIM_SNAPSHOT_CACHE.clear() - yield - robocoin._TRIM_SNAPSHOT_CACHE.clear() - - -class TestLoadTrimSpec: - def test_parses_empty_head_or_tail_and_skips_noop_rows(self, tmp_path): - path = _write_trim_csv( - tmp_path / "trim.csv", - [ - { - "dataset": "bucket-a", - "episode_index": 0, - "total_frames": 10, - "trim_head_to": 2, - "trim_tail_from": "", - }, - { - "dataset": "bucket-a", - "episode_index": 1, - "total_frames": 12, - "trim_head_to": "", - "trim_tail_from": 9, - }, - { - "dataset": "bucket-a", - "episode_index": 2, - "total_frames": 8, - "trim_head_to": "", - "trim_tail_from": "", - }, - { - "dataset": "bucket-b", - "episode_index": 3, - "total_frames": 6, - "trim_head_to": 1, - "trim_tail_from": 5, - }, - ], - ) - - assert _load_trim_spec(path) == { - "bucket-a": {0: (2, None, 10), 1: (0, 9, 12)}, - "bucket-b": {3: (1, 5, 6)}, - } - - def test_explicit_missing_path_fails_fast(self, tmp_path): - with pytest.raises((OSError, ValueError)): - _load_trim_spec(tmp_path / "does-not-exist.csv") - - def test_empty_string_is_not_the_disabled_value(self): - # Only trim_csv=None is an opt-out; a configured empty path is invalid. - with pytest.raises((OSError, ValueError)): - _load_trim_spec("") - - @pytest.mark.parametrize("missing", TRIM_COLUMNS) - def test_missing_required_column_fails_fast(self, tmp_path, missing): - columns = tuple(c for c in TRIM_COLUMNS if c != missing) - row = { - "dataset": "bucket", - "episode_index": 0, - "total_frames": 10, - "trim_head_to": 2, - "trim_tail_from": 8, - } - path = _write_trim_csv(tmp_path / f"missing-{missing}.csv", [row], columns) - with pytest.raises(ValueError): - _load_trim_spec(path) - - def test_empty_total_frames_fails_fast(self, tmp_path): - path = _write_trim_csv( - tmp_path / "empty-total.csv", - [ - { - "dataset": "bucket", - "episode_index": 0, - "total_frames": "", - "trim_head_to": 2, - "trim_tail_from": 8, - } - ], - ) - with pytest.raises(ValueError): - _load_trim_spec(path) - - def test_failed_parse_is_not_cached(self, tmp_path): - from openwam.dataloader import robocoin - - path = _write_trim_csv( - tmp_path / "repairable.csv", - [_trim_row()], - columns=tuple(c for c in TRIM_COLUMNS if c != "total_frames"), - ) - with pytest.raises(ValueError): - _load_trim_spec(path) - assert str(path) not in robocoin._TRIM_SNAPSHOT_CACHE - - _write_trim_csv(path, [_trim_row()]) - assert _load_trim_spec(path) == {"bucket": {0: (2, 8, 10)}} - - def test_cache_refreshes_when_same_path_is_replaced(self, tmp_path): - path = _write_trim_csv(tmp_path / "replaceable.csv", [_trim_row(trim_head_to=2)]) - assert _load_trim_spec(path) == {"bucket": {0: (2, 8, 10)}} - - replacement = _write_trim_csv(tmp_path / "replacement.csv", [_trim_row(trim_head_to=3)]) - replacement.replace(path) - assert _load_trim_spec(path) == {"bucket": {0: (3, 8, 10)}} - - @pytest.mark.parametrize( - "overrides", - [ - {"trim_head_to": -1, "trim_tail_from": ""}, - {"trim_head_to": 11, "trim_tail_from": ""}, - {"trim_head_to": "", "trim_tail_from": -1}, - {"trim_head_to": "", "trim_tail_from": 11}, - {"trim_head_to": 8, "trim_tail_from": 3}, - ], - ids=["negative-head", "head-past-total", "negative-tail", "tail-past-total", "head-after-tail"], - ) - def test_invalid_trim_bounds_fail_fast(self, tmp_path, overrides): - path = _write_trim_csv(tmp_path / "invalid-bounds.csv", [_trim_row(**overrides)]) - with pytest.raises(ValueError): - _load_trim_spec(path) - - def test_duplicate_dataset_episode_key_fails_fast(self, tmp_path): - path = _write_trim_csv( - tmp_path / "duplicate.csv", - [_trim_row(trim_head_to=1), _trim_row(trim_head_to=2)], - ) - with pytest.raises(ValueError): - _load_trim_spec(path) - - -class TestFilterEpisodes: - def test_head_tail_and_both_update_all_offsets(self, tmp_path): - eps = _episodes([10, 10, 10]) - out = _filter( - tmp_path, - eps, - [ - { - "dataset": "bucket", - "episode_index": 0, - "total_frames": 10, - "trim_head_to": 2, - "trim_tail_from": "", - }, - { - "dataset": "bucket", - "episode_index": 1, - "total_frames": 10, - "trim_head_to": "", - "trim_tail_from": 7, - }, - { - "dataset": "bucket", - "episode_index": 2, - "total_frames": 10, - "trim_head_to": 2, - "trim_tail_from": 8, - }, - ], - ) - - assert out["length"].tolist() == [8, 7, 6] - assert out["_data_row_offset"].tolist() == [2, 10, 22] - assert out[f"_video_frame_offset/{CAM}"].tolist() == [2, 10, 22] - assert out["_video_frame_offset/second_camera"].tolist() == [102, 110, 122] - - def test_one_stale_entry_fails_closed_with_data_contract_error(self, tmp_path): - eps = _episodes([10, 10]) - with pytest.raises(DataContractError, match="[Ss][Tt][Aa][Ll][Ee]"): - _filter( - tmp_path, - eps, - [ - { - "dataset": "bucket", - "episode_index": 0, - "total_frames": 999, - "trim_head_to": 1, - "trim_tail_from": 9, - }, - { - # This simulates an undetectable same-length collision. - # Rejecting the bucket prevents this entry being applied. - "dataset": "bucket", - "episode_index": 1, - "total_frames": 10, - "trim_head_to": 2, - "trim_tail_from": 8, - }, - ], - ) - - def test_zero_span_trim_drops_episode(self, tmp_path): - eps = _episodes([10, 10]) - out = _filter( - tmp_path, - eps, - [ - { - "dataset": "bucket", - "episode_index": 0, - "total_frames": 10, - "trim_head_to": 10, - "trim_tail_from": "", - } - ], - ) - - assert out["episode_index"].tolist() == [1] - assert out["_data_row_offset"].tolist() == [10] - assert out[f"_video_frame_offset/{CAM}"].tolist() == [10] - - -def _make_bucket(bucket: Path, lengths: list[int]) -> Path: - meta = bucket / "meta" - (meta / "episodes").mkdir(parents=True) - (meta / "info.json").write_text( - json.dumps( - { - "fps": 30.0, - "robot_type": "test_robot", - "data_path": "data/chunk-{chunk_index:03d}/file-{file_index:03d}.parquet", - "video_path": "videos/{video_key}/chunk-{chunk_index:03d}/file-{file_index:03d}.mp4", - "features": { - CAM: {"dtype": "video"}, - "eef_sim_pose_action": {}, - "gripper_open_scale_action": {}, - "eef_sim_pose_state": {}, - "gripper_open_scale_state": {}, - }, - } - ) - ) - - rows = [] - start = 0 - for ep, length in enumerate(lengths): - rows.append( - { - "episode_index": ep, - "length": length, - "dataset_from_index": start, - "data/chunk_index": 0, - "data/file_index": 0, - f"videos/{CAM}/chunk_index": 0, - f"videos/{CAM}/file_index": 0, - } - ) - start += length - pq.write_table(pa.Table.from_pandas(pd.DataFrame(rows)), meta / "episodes" / "chunk-000.parquet") - pd.DataFrame({"task_index": [0]}, index=pd.Index(["do the thing"], name="task")).to_parquet(meta / "tasks.parquet") - - source_index = np.arange(start, dtype=np.float32) - eef = np.zeros((start, 12), dtype=np.float32) - eef[:, 0] = source_index - episode_index = np.concatenate([np.full(length, ep, dtype=np.int64) for ep, length in enumerate(lengths)]) - frame_index = np.concatenate([np.arange(length, dtype=np.int64) for length in lengths]) - data_dir = bucket / "data" / "chunk-000" - data_dir.mkdir(parents=True) - pq.write_table( - pa.Table.from_pandas( - pd.DataFrame( - { - "task_index": np.zeros(start, dtype=np.int64), - "episode_index": episode_index, - "frame_index": frame_index, - "eef_sim_pose_action": list(eef), - "gripper_open_scale_action": list(np.zeros((start, 2), dtype=np.float32)), - "eef_sim_pose_state": list(eef), - "gripper_open_scale_state": list(np.zeros((start, 2), dtype=np.float32)), - } - ) - ), - data_dir / "file-000.parquet", - ) - return bucket - - -def _set_info_splits(bucket: Path, splits: dict) -> None: - info_path = bucket / "meta" / "info.json" - info = json.loads(info_path.read_text()) - info["splits"] = splits - info_path.write_text(json.dumps(info)) - - -@pytest.fixture -def patch_decode(monkeypatch): - def fake_decode(path, frame_indices, height, width): - del path - frames = [] - for frame_index in frame_indices: - image = Image.new("RGB", (width, height), (0, 0, 0)) - image.putpixel((0, 0), (int(frame_index), 0, 0)) - frames.append(image) - return frames - - monkeypatch.setattr("openwam.dataloader.bases.lerobot_v3_reader._decode_video_frames", fake_decode) - - -def _decoded_index(image: Image.Image) -> int: - return image.getpixel((0, 0))[0] - - -def test_root_from_config_preloads_bad_trim_csv_before_bucket_workers(tmp_path): - root = tmp_path / "root" - _make_bucket(root / "bucket-a", [8]) - _make_bucket(root / "bucket-b", [8]) - bad_csv = _write_trim_csv( - tmp_path / "bad.csv", - [_trim_row()], - columns=tuple(c for c in TRIM_COLUMNS if c != "total_frames"), - ) - - # build_multibucket deliberately tolerates an individual corrupt bucket. - # A shared trim_csv is global configuration, so it must be validated before - # entering those workers and preserve the actionable parse exception. - with pytest.raises(ValueError, match="trim_csv"): - RoboCOINDataset.from_config({"dataset_dir": str(root), "trim_csv": str(bad_csv)}) - - -def test_root_mode_trim_csv_with_no_bucket_key_overlap_fails_fast(tmp_path): - root = tmp_path / "root" - _make_bucket(root / "bucket-a", [8]) - _make_bucket(root / "bucket-b", [8]) - trim_csv = _write_trim_csv( - tmp_path / "wrong-root.csv", - [ - _trim_row( - dataset="bucket-from-another-root", - total_frames=8, - trim_head_to=1, - trim_tail_from=7, - ) - ], - ) - - with pytest.raises(ValueError): - RoboCOINDataset.from_config({"dataset_dir": str(root), "trim_csv": str(trim_csv), "normalize_mode": None}) - - -def test_root_mode_trim_csv_with_partial_bucket_key_overlap_is_allowed(tmp_path): - root = tmp_path / "root" - _make_bucket(root / "bucket-a", [8]) - _make_bucket(root / "bucket-b", [8]) - trim_csv = _write_trim_csv( - tmp_path / "partial-root.csv", - [ - _trim_row(dataset="bucket-a", total_frames=8, trim_head_to=1, trim_tail_from=7), - _trim_row(dataset="retired-bucket", total_frames=8, trim_head_to=2, trim_tail_from=6), - ], - ) - - ds = RoboCOINDataset.from_config({"dataset_dir": str(root), "trim_csv": str(trim_csv), "normalize_mode": None}) - lengths_by_bucket = {bucket._dataset_id: bucket._eps_df["length"].tolist() for bucket in ds._buckets} - assert lengths_by_bucket == {"bucket-a": [6], "bucket-b": [8]} - - -def test_root_mode_drops_whole_bucket_exclusion_before_fanout(tmp_path): - root = tmp_path / "root" - kept = _make_bucket(root / "kept-bucket", [40]) - excluded_name = next(iter(ROBOCOIN_WHOLE_BUCKET_EXCLUSIONS)) - _make_bucket(root / excluded_name, [40]) - - ds = RoboCOINDataset.from_config({"dataset_dir": str(root), "normalize_mode": None}) - - assert [bucket._dataset_dir for bucket in ds._buckets] == [kept] - - -def test_direct_reader_rejects_whole_bucket_exclusion(tmp_path): - excluded_name = next(iter(ROBOCOIN_WHOLE_BUCKET_EXCLUSIONS)) - bucket = _make_bucket(tmp_path / excluded_name, [40]) - - with pytest.raises(DataContractError, match="excluded as a whole"): - RoboCOINDataset(dataset_dir=str(bucket), normalize_mode=None) - - -def test_excluded_bucket_does_not_count_as_trim_csv_overlap(tmp_path): - root = tmp_path / "root" - _make_bucket(root / "kept-bucket", [40]) - excluded_name = next(iter(ROBOCOIN_WHOLE_BUCKET_EXCLUSIONS)) - _make_bucket(root / excluded_name, [40]) - trim_csv = _write_trim_csv( - tmp_path / "excluded-only.csv", - [ - _trim_row( - dataset=excluded_name, - total_frames=40, - trim_head_to=1, - trim_tail_from=39, - ) - ], - ) - - with pytest.raises(ValueError, match="silently no-op"): - RoboCOINDataset.from_config({"dataset_dir": str(root), "trim_csv": str(trim_csv), "normalize_mode": None}) - - -def test_reader_rejects_noncanonical_numeric_data_shard(tmp_path): - bucket = _make_bucket(tmp_path / "bucket", [8]) - canonical = bucket / "data" / "chunk-000" / "file-000.parquet" - canonical.rename(canonical.with_name("file-0.parquet")) - - with pytest.raises(DataContractError, match="non-canonical numeric data shard"): - RoboCOINDataset(dataset_dir=str(bucket), normalize_mode=None) - - -def test_reader_rejects_non_string_data_path_template(tmp_path): - bucket = _make_bucket(tmp_path / "bucket", [8]) - info_path = bucket / "meta" / "info.json" - info = json.loads(info_path.read_text()) - info["data_path"] = None - info_path.write_text(json.dumps(info)) - - with pytest.raises(DataContractError, match="data_path must be a non-empty string"): - RoboCOINDataset(dataset_dir=str(bucket), normalize_mode=None) - - -def test_reader_rejects_malformed_string_data_path_template(tmp_path): - bucket = _make_bucket(tmp_path / "bucket", [8]) - info_path = bucket / "meta" / "info.json" - info = json.loads(info_path.read_text()) - info["data_path"] = "data/chunk-{chunk_index.foo}/file-{file_index}.parquet" - info_path.write_text(json.dumps(info)) - - with pytest.raises(DataContractError, match="invalid info.json data_path template"): - RoboCOINDataset(dataset_dir=str(bucket), normalize_mode=None) - - -def test_root_mode_pins_one_snapshot_for_all_bucket_constructors(tmp_path, monkeypatch): - from openwam.dataloader import robocoin - - root = tmp_path / "root" - _make_bucket(root / "bucket-a", [10]) - _make_bucket(root / "bucket-b", [10]) - trim_csv = _write_trim_csv( - tmp_path / "trim.csv", - [_trim_row(dataset="bucket-a")], - ) - original_load = robocoin._load_trim_snapshot - loads = [] - - def record_load(path): - snapshot = original_load(path) - loads.append(snapshot) - return snapshot - - monkeypatch.setattr(robocoin, "_load_trim_snapshot", record_load) - ds = RoboCOINDataset.from_config({"dataset_dir": str(root), "trim_csv": str(trim_csv), "normalize_mode": None}) - - assert len(loads) == 1 - assert {bucket._dataset_id: bucket._eps_df["length"].tolist() for bucket in ds._buckets} == { - "bucket-a": [6], - "bucket-b": [10], - } - assert all(bucket._trim_snapshot is None for bucket in ds._buckets) - - -def test_root_mode_rejects_csv_change_during_fanout_without_normalization(tmp_path, monkeypatch): - root = tmp_path / "root" - _make_bucket(root / "bucket", [10]) - trim_csv = _write_trim_csv(tmp_path / "trim.csv", [_trim_row()]) - replacement = _write_trim_csv( - tmp_path / "replacement.csv", - [_trim_row(trim_head_to=3, trim_tail_from=7)], - ) - original_add_offsets = RoboCOINDataset._add_data_offsets - - def add_offsets_then_replace(self, eps): - original_add_offsets(self, eps) - replacement.replace(trim_csv) - - monkeypatch.setattr(RoboCOINDataset, "_add_data_offsets", add_offsets_then_replace) - with pytest.raises(DataContractError, match="changed while from_config reader construction"): - RoboCOINDataset.from_config({"dataset_dir": str(root), "trim_csv": str(trim_csv), "normalize_mode": None}) - - -def test_root_mode_unknown_episode_id_fails_closed_against_full_manifest(tmp_path): - root = tmp_path / "root" - bucket = _make_bucket(root / "bucket", [8, 8]) - _set_info_splits(bucket, {"train": "0:1", "val": "1:2"}) - trim_csv = _write_trim_csv( - tmp_path / "unknown-episode.csv", - [ - _trim_row( - dataset="bucket", - episode_index=99, - total_frames=8, - trim_head_to=1, - trim_tail_from=7, - ) - ], - ) - - with pytest.raises(DataContractError, match="[Uu]nknown.*episode"): - RoboCOINDataset.from_config({"dataset_dir": str(root), "trim_csv": str(trim_csv), "normalize_mode": None}) - - -def test_root_mode_trim_entry_outside_selected_split_is_not_unknown(tmp_path): - root = tmp_path / "root" - bucket = _make_bucket(root / "bucket", [8, 8]) - _set_info_splits(bucket, {"train": "0:1", "val": "1:2"}) - trim_csv = _write_trim_csv( - tmp_path / "val-episode.csv", - [ - _trim_row( - dataset="bucket", - episode_index=1, - total_frames=8, - trim_head_to=1, - trim_tail_from=7, - ) - ], - ) - - ds = RoboCOINDataset.from_config({"dataset_dir": str(root), "trim_csv": str(trim_csv), "normalize_mode": None}) - assert [bucket._dataset_id for bucket in ds._buckets] == ["bucket"] - assert ds._buckets[0]._eps_df["episode_index"].tolist() == [0] - assert ds._buckets[0]._eps_df["length"].tolist() == [8] - - -def test_root_mode_stale_trim_entry_outside_selected_split_fails_closed(tmp_path): - root = tmp_path / "root" - bucket = _make_bucket(root / "bucket", [8, 8]) - _set_info_splits(bucket, {"train": "0:1", "val": "1:2"}) - trim_csv = _write_trim_csv( - tmp_path / "stale-val-episode.csv", - [ - _trim_row( - dataset="bucket", - episode_index=1, - total_frames=9, - trim_head_to=1, - trim_tail_from=7, - ) - ], - ) - - with pytest.raises(DataContractError, match="[Ss][Tt][Aa][Ll][Ee]"): - RoboCOINDataset.from_config({"dataset_dir": str(root), "trim_csv": str(trim_csv), "normalize_mode": None}) - - -def test_root_mode_stale_bucket_fails_closed_globally(tmp_path): - root = tmp_path / "root" - _make_bucket(root / "bucket-a", [8]) - _make_bucket(root / "bucket-b", [8]) - trim_csv = _write_trim_csv( - tmp_path / "stale.csv", - [ - { - "dataset": "bucket-a", - "episode_index": 0, - "total_frames": 9, - "trim_head_to": 1, - "trim_tail_from": 7, - } - ], - ) - - with pytest.raises(DataContractError, match="[Ss][Tt][Aa][Ll][Ee]"): - RoboCOINDataset.from_config({"dataset_dir": str(root), "trim_csv": str(trim_csv), "normalize_mode": None}) - - -def test_trimmed_video_and_action_remain_aligned_end_to_end(tmp_path, patch_decode): - lengths = [8, 8, 8, 8] - bucket = _make_bucket(tmp_path / "bucket", lengths) - trim_csv = _write_trim_csv( - tmp_path / "trim.csv", - [ - { - "dataset": bucket.name, - "episode_index": 0, - "total_frames": 8, - "trim_head_to": 2, - "trim_tail_from": "", - }, - { - "dataset": bucket.name, - "episode_index": 1, - "total_frames": 8, - "trim_head_to": "", - "trim_tail_from": 6, - }, - { - "dataset": bucket.name, - "episode_index": 2, - "total_frames": 8, - "trim_head_to": 2, - "trim_tail_from": 6, - }, - { - "dataset": bucket.name, - "episode_index": 3, - "total_frames": 8, - "trim_head_to": 8, - "trim_tail_from": "", - }, - ], - ) - ds = RoboCOINDataset( - dataset_dir=str(bucket), - num_frames=5, - video_stride=1, - height=32, - width=32, - multiview=False, - normalize_mode=None, - trim_csv=trim_csv, - ) - - assert ds._eps_df["episode_index"].tolist() == [0, 1, 2] - assert ds._eps_df["length"].tolist() == [6, 6, 4] - assert ds._ep_data_row_offset.tolist() == [2, 8, 18] - assert ds._ep_video_frame_offsets[CAM].tolist() == [2, 8, 18] - assert len(ds) == 16 - - seen = set() - for idx in range(len(ds)): - sample = ds[idx] - action_rows = sample["action"][sample["action_mask"].any(dim=1), 0].to(torch.int64).tolist() - video_rows = [ - _decoded_index(frame) for frame, valid in zip(sample["video"], sample["video_mask"].tolist()) if valid - ] - assert action_rows == video_rows[: len(action_rows)] - seen.update(video_rows) - - assert seen == set(range(2, 8)) | set(range(8, 14)) | set(range(18, 22)) diff --git a/tests/dataloader/test_stats_accumulator.py b/tests/dataloader/test_stats_accumulator.py new file mode 100644 index 000000000..1e223647e --- /dev/null +++ b/tests/dataloader/test_stats_accumulator.py @@ -0,0 +1,71 @@ +"""Numerical and sampling contracts shared by benchmark statistics tools.""" + +import numpy as np +import pytest + +from openwam.dataloader.utils.stats_computation.accumulator import Accumulator + + +@pytest.mark.parametrize("method", ["update", "update_batch"]) +def test_accumulator_matches_full_population_statistics(method): + rows = np.random.RandomState(17).normal(size=(71, 4)).astype(np.float32) + rows[:, -1] = 7 + accumulator = Accumulator(dim=4, reservoir_cap=100) + for batch in np.array_split(rows, 5): + getattr(accumulator, method)(batch) + stats = accumulator.finalize() + assert accumulator.count == len(rows) + for key, expected in { + "mean": rows.astype(np.float64).mean(axis=0), + "std": np.where(rows.astype(np.float64).std(axis=0) < 1e-8, 1, rows.astype(np.float64).std(axis=0)), + "min": rows.min(axis=0), + "max": rows.max(axis=0), + "q01": np.quantile(rows, 0.01, axis=0), + "q99": np.quantile(rows, 0.99, axis=0), + }.items(): + np.testing.assert_allclose(stats[key], expected, rtol=1e-6, atol=1e-7) + + +def test_reservoir_matches_sequential_algorithm_r_across_batches(): + rows = np.arange(300, dtype=np.float32).reshape(100, 3) + cap = 7 + expected = rows[:cap].copy() + rng = np.random.RandomState(23) + for i in range(cap, len(rows)): + position = rng.randint(0, i + 1) + if position < cap: + expected[position] = rows[i] + accumulator = Accumulator(dim=3, reservoir_cap=cap, seed=23) + for batch in np.array_split(rows, 6): + accumulator.update_batch(batch) + np.testing.assert_array_equal(accumulator._res, expected) + assert accumulator._res_seen == len(rows) + assert accumulator._res_n == cap + + +@pytest.mark.parametrize("method", ["update", "update_batch"]) +@pytest.mark.parametrize( + "rows", [np.ones((3, 1)), np.ones(3), np.array([[np.nan, 0]]), np.array([[np.inf, 0]]), np.array([[1e100, 0]])] +) +def test_invalid_batches_do_not_mutate_statistics(method, rows): + accumulator = Accumulator(dim=2, reservoir_cap=5) + with pytest.raises(ValueError): + getattr(accumulator, method)(rows) + assert accumulator.count == 0 + assert accumulator._res_seen == 0 + np.testing.assert_array_equal(accumulator.mean, [0, 0]) + + +@pytest.mark.parametrize( + "kwargs", [{"dim": 0}, {"dim": -1}, {"reservoir_cap": 0}, {"reservoir_cap": -1}, {"reservoir_cap": 1.5}] +) +def test_accumulator_rejects_invalid_dimensions_or_capacity(kwargs): + with pytest.raises(ValueError, match="positive integers"): + Accumulator(**kwargs) + + +def test_empty_statistics_cannot_produce_infinite_bounds(): + accumulator = Accumulator(dim=2, reservoir_cap=5) + accumulator.update_batch(np.empty((0, 2))) + with pytest.raises(ValueError, match="without any rows"): + accumulator.finalize() diff --git a/tests/test_public_config_defaults.py b/tests/test_public_config_defaults.py index 2b48705ff..23226f375 100644 --- a/tests/test_public_config_defaults.py +++ b/tests/test_public_config_defaults.py @@ -1,6 +1,8 @@ from pathlib import Path +import pytest import yaml +from hydra import compose, initialize_config_dir REPO_ROOT = Path(__file__).resolve().parents[1] @@ -11,9 +13,53 @@ def test_training_does_not_warm_start_from_a_placeholder_by_default(): assert config["training"]["finetune_ckpt_path"] is None -def test_robocoin_trim_manifest_is_opt_in(): - config = yaml.safe_load( - (REPO_ROOT / "configs" / "dataloader" / "pretrain_data" / "robocoin.yaml").read_text(encoding="utf-8") - ) +def test_pretraining_defaults_load_all_four_final_datasets(): + with initialize_config_dir(config_dir=str(REPO_ROOT / "configs"), version_base=None): + config = compose(config_name="train", overrides=["dataloader=mixture"]) - assert config["trim_csv"] is None + expected = { + "agibotworld": "AgiBotWorld-Beta", + "robocoin": "RoboCOIN", + "oxe_droid": "Droid", + "interndata_a1": "InternData-A1", + } + assert config.dataloader.type == "mixture" + assert config.dataloader.weight_strategy == "proportional" + assert set(config.dataloader.datasets) == set(expected) + for name, directory in expected.items(): + source = config.dataloader.datasets[name] + assert source.enabled + assert source.dataset_dir == f"/mnt/data/wangyuran/pretrain_data_final/{directory}" + assert source.normalize_mode == "quantile" + assert source.total_hours is None + assert not {"trim_csv", "stats_root", "a1_stats_root", "segment_start", "segment_end"}.intersection(source) + + assert not (REPO_ROOT / "configs" / "dataloader" / "pretrain_data").exists() + + +@pytest.mark.parametrize("name", ["agibotworld", "robocoin", "oxe_droid", "interndata_a1"]) +def test_flat_pretraining_configs_compose_as_standalone_dataloaders(name): + with initialize_config_dir(config_dir=str(REPO_ROOT / "configs"), version_base=None): + config = compose(config_name="train", overrides=[f"dataloader={name}"]) + assert config.dataloader.type == name + assert config.dataloader.num_frames == 33 + assert config.dataloader.video_stride == 4 + assert config.dataloader.normalize_mode == "quantile" + + +def test_mixture_hydra_source_overrides_remain_independent(): + with initialize_config_dir(config_dir=str(REPO_ROOT / "configs"), version_base=None): + config = compose( + config_name="train", + overrides=[ + "dataloader=mixture", + "dataloader.datasets.robocoin.dataset_dir=/tmp/relocated_robocoin", + "dataloader.datasets.oxe_droid.enabled=false", + "dataloader.datasets.interndata_a1.total_hours=1.5", + ], + ) + assert config.dataloader.datasets.robocoin.dataset_dir == "/tmp/relocated_robocoin" + assert config.dataloader.datasets.agibotworld.dataset_dir.endswith("/AgiBotWorld-Beta") + assert not config.dataloader.datasets.oxe_droid.enabled + assert config.dataloader.datasets.interndata_a1.total_hours == 1.5 + assert config.dataloader.datasets.agibotworld.total_hours is None