diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index b3c59043a8ac..987b3ab42ea7 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -866,6 +866,64 @@ Notes: few large reads); scattered point reads coalesce less. - Blob reads are available only on `scan()`, not on the `search()` queries. +### Contiguous windows for PyTorch + +Install the `torch` extra, then use `to_contiguous_window_dataset` to expose +map-style windows without loading the selected rows or BLOB payloads into Python +memory up front. The Dataset builds a compact index from the group column, order +column, and Paimon row IDs. Each `__getitem__` call fetches only that window from +the snapshot recorded in `dataset.snapshot_id`. + +```shell +pip install pypaimon[torch] +``` + +```python +import torch + + +def float32_window(values): + return torch.tensor(values, dtype=torch.float32) + + +windows = ( + frames.scan() + .where("split = 'train'") + .to_contiguous_window_dataset( + window_size=16, + columns=["state", "action"], + group_key="episode_id", + order_key="step_idx", + tail="pad", + column_transforms={ + "state": float32_window, + "action": float32_window, + }, + ) +) + +sample = windows[0] +assert sample["action"].shape == (16, action_size) +assert sample["is_pad"].shape == (16,) +``` + +The group and order keys in a sample identify the window anchor. Every projected +column contains the whole window. With `tail="drop"`, only full windows are +exposed. With `tail="pad"`, every real row is an anchor; missing suffix values +repeat the last real value by default and `is_pad` is `True` exactly at those +positions. With `tail="error"`, construction fails if any scheduled anchor is +incomplete. Use `pad_values` to override the repeated value for individual +columns. Anchors advance by `stride`, which defaults to one row. + +`column_transforms` receive one padded Python list per projected column. This is +where applications define tensor dtype and shape or decode BLOB bytes. The +optional `adapter` receives the resulting sample mapping and can rename or +combine fields for a model-specific batch contract. The core Dataset does not +know model field names, image formats, or normalization rules. Top-level +functions and callable classes are recommended for transforms and adapters so +the Dataset remains picklable by multi-worker `torch.utils.data.DataLoader` +instances. + ### Distributed BLOB processing with Ray For larger jobs, read descriptors with `to_ray()`, then fetch and process BLOB diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md index af9189bbae5e..f6e7c87617e1 100644 --- a/docs/docs/pypaimon/pytorch.md +++ b/docs/docs/pypaimon/pytorch.md @@ -157,7 +157,49 @@ embedded frame ordinals keep frame mapping out of the normal data file. Use physical video ranges and cache decoder sessions per worker. See [Multimodal API: Video Frame Storage](multimodal-api#video-frame-storage) for the write path and a complete decoder example. +## Contiguous Windows +Use a map-style `ContiguousWindowDataset` when training samples are fixed-size +windows which must not cross a sequence boundary. The dataset builds an index +from only the group column, order column, and Paimon row IDs. Projected values, +including BLOB payloads, are read from the pinned snapshot when a sample is +requested; they are not retained in the index. + +```python +from torch.utils.data import DataLoader + +dataset = ( + frames.scan() + .to_contiguous_window_dataset( + window_size=16, + columns=["state", "image"], + anchor_columns=["image"], + group_key="episode_id", + order_key="step_idx", + tail="pad", + ) +) + +loader = DataLoader(dataset, batch_size=32, num_workers=4, shuffle=True) +``` + +Each item contains the group and order keys, one list for each requested +column, and a boolean `is_pad` tensor where `True` marks padding. Padding +repeats the final real value by default; `pad_values` can override individual +columns. Columns named in `anchor_columns` contain only the first row's value, +which is useful when an observation applies to a full action window. Use +`column_transforms` to convert column lists to tensors and +`adapter` to produce a model-specific sample mapping. Keep these callbacks +picklable when using multiple DataLoader workers. + +Scheduled anchors start at row zero and advance by `stride` (default `1`). +`tail="drop"` omits incomplete windows, `tail="pad"` includes and pads them, +and `tail="error"` rejects a sequence with any scheduled incomplete window. +Rows are sorted by `order_key` inside each `group_key` value. Order values must +be integers which increase by exactly one; duplicates and missing steps are +rejected, and windows never cross groups. The resolved Paimon +snapshot is pinned for the lifetime of the dataset, so later commits cannot +change its index or sample contents. ## File Format Metadata Cache Reusable PyArrow Dataset metadata is cached across reads. Configure its estimated diff --git a/docs/docs/pypaimon/robomind-act-benchmark.md b/docs/docs/pypaimon/robomind-act-benchmark.md new file mode 100644 index 000000000000..37f9e011fbbc --- /dev/null +++ b/docs/docs/pypaimon/robomind-act-benchmark.md @@ -0,0 +1,185 @@ +--- +title: "RoboMIND ACT Storage Benchmark" +sidebar_position: 8 +--- + + + +# RoboMIND ACT Storage Benchmark + +This benchmark measures the same CPU LeRobot ACT training workload over an +original RoboMIND AgileX HDF5 dataset or an already ingested and +canonical-action-backfilled Paimon warehouse. Ingestion and backfill are outside +the timed scope. + +The backends run independently. A resolved experiment document preserves the +shared configuration, normalization, seed, episode selection, Paimon snapshot, +and logical window sequence. Result comparison verifies that contract before it +calculates performance ratios. + +## Install + +Python 3.10 or newer is required. + +```shell +pip install 'pypaimon[act,hdf5]' +``` + +## 1. Prepare the experiment + +```shell +python -m pypaimon.benchmark.act prepare \ + --input /data/RoboMIND/h5_agilex_3rgb \ + --warehouse /data/warehouse \ + --output /data/results/experiment.json +``` + +Preparation is not timed. It verifies that HDF5 discovery matches the Paimon +episodes table, checks versioned action statistics against train-only HDF5 +moments, selects eligible train and validation episodes, pins the frames +snapshot, and materializes deterministic measurement, training, and validation +window indices. + +Without `--experiment`, preparation starts from the packaged +`default_experiment.json`. `--experiment` replaces that definition, so a +custom JSON file must contain every required field. Command-line options then +override individual values: + +```shell +python -m pypaimon.benchmark.act prepare \ + --experiment my-experiment.json \ + --input /data/RoboMIND/h5_agilex_3rgb \ + --warehouse /data/warehouse \ + --action-horizon 32 \ + --batch-size 2 \ + --fetch-batches 8 \ + --rounds 3 \ + --output /data/results/experiment.json +``` + +The resolved experiment embeds the effective parameters as well as: + +- portable source episode metadata and its SHA-256; +- normalization values, scope, version, frame count, and SHA-256; +- selected train and validation episode IDs; +- every logical window index, the window-plan SHA-256, and the + episode-qualified sample-sequence SHA-256; +- the Paimon database, frames table, and pinned snapshot ID. + +## 2. Run each backend + +```shell +python -m pypaimon.benchmark.act run \ + --backend hdf5 \ + --experiment /data/results/experiment.json \ + --input /data/RoboMIND/h5_agilex_3rgb \ + --results-dir /data/results + +python -m pypaimon.benchmark.act run \ + --backend paimon \ + --experiment /data/results/experiment.json \ + --warehouse /data/warehouse \ + --results-dir /data/results +``` + +Use `--output` to choose an exact result path. Otherwise the command writes an +automatically named JSON file below `--results-dir` and prints its absolute +path as a compact JSON object. + +Each result contains the complete resolved experiment and experiment SHA-256, +backend identity, runtime environment, model metadata, planned-sample tensor +fingerprint, three or more raw measurement rounds, and median/minimum/maximum +summary metrics. + +Both adapters produce the same shared sample contract. State and camera images +come from the anchor frame; action covers the complete horizon. HDF5 reads a +window on demand from one episode file. Paimon uses a lazy, snapshot-pinned +`ContiguousWindowDataset`; image columns are anchor-only, and plural +`__getitems__` access coalesces multiple logical batches into a physical +fetch before splitting them back into the unchanged model batch size. + +## 3. Compare results + +Compare explicit files: + +```shell +python -m pypaimon.benchmark.act compare \ + /data/results/robomind-act-hdf5-20260901T010000Z-a1b2c3d4.json \ + /data/results/robomind-act-paimon-20260901T011000Z-e5f6a7b8.json \ + --output /data/results/comparison.json +``` + +Or discover all ACT result documents in a directory: + +```shell +python -m pypaimon.benchmark.act compare \ + --results-dir /data/results \ + --output /data/results/comparison.json +``` + +Directory discovery ignores experiment and prior comparison JSON files. +Results are grouped by experiment SHA-256. Different experiments remain +separate entries in one comparison artifact; only compatible repeated results +for the same experiment and backend are aggregated. + +Within one experiment group, comparison requires identical runtime environment, +model metadata, tensor fingerprint, train-loss trace, and validation-loss trace. +An environment mismatch marks the group `INCOMPATIBLE`; a model, tensor, or +loss mismatch marks it `FAILED`. Neither case produces performance ratios. + +For compatible HDF5 and Paimon results, higher-is-better metrics report +`paimon_over_hdf5`. Lower-is-better latency, time, and memory metrics report +`hdf5_over_paimon`, which is the Paimon speedup or reduction factor. + +## Measurements + +Every backend repeat records: + +- dataset construction time; +- first-batch latency after construction; +- batch-fetch samples per second after warm-up; +- end-to-end fixed ACT optimizer-step time, including dataset fetch; +- per-step loss and compute time after each training batch has been fetched; +- validation loss; +- total measured wall time; +- Python peak allocation from a separate dataset-first-batch replay. + +The shared harness resets Python, NumPy, and Torch random generators before +model construction and enables deterministic Torch algorithms. The logical +window plan is explicit rather than delegated to a streaming reader. + +Python peak allocation uses `tracemalloc` after wall-clock measurement so +tracing overhead does not distort throughput. It does not include every native +Arrow or Torch allocation. The benchmark does not drop the OS page cache. +GPU, multi-worker loading, distributed training, recovery, and policy quality +remain outside this benchmark. + +## Code organization + +- `benchmark.act.harness`: shared ACT tensors, model, trainer, window plan, and + measurement lifecycle; +- `benchmark.act.hdf5`: HDF5 window dataset and train normalization moments; +- `benchmark.act.paimon`: Paimon adapter, snapshot-pinned datasets, and + versioned statistics access; +- `benchmark.act.runner`: experiment preparation and one-backend execution; +- `benchmark.act.compare`: result discovery, compatibility checks, grouping by + experiment, and aggregation of compatible repeated runs; +- `benchmark.act.__main__`: the `prepare`, `run`, and `compare` + command-line interface. diff --git a/paimon-python/pypaimon/benchmark/act/__init__.py b/paimon-python/pypaimon/benchmark/act/__init__.py new file mode 100644 index 000000000000..f6224db90413 --- /dev/null +++ b/paimon-python/pypaimon/benchmark/act/__init__.py @@ -0,0 +1,17 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""ACT training benchmark backends and result comparison.""" diff --git a/paimon-python/pypaimon/benchmark/act/__main__.py b/paimon-python/pypaimon/benchmark/act/__main__.py new file mode 100644 index 000000000000..d1d8560ae275 --- /dev/null +++ b/paimon-python/pypaimon/benchmark/act/__main__.py @@ -0,0 +1,200 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Command-line entry point for ACT benchmark preparation, runs, and reports.""" + +# ruff: noqa: E402 + +import sys + + +def _require_supported_python(version_info): + """Reject runtimes older than the ACT dependencies support.""" + if tuple(version_info[:2]) < (3, 10): + raise RuntimeError("ACT benchmark requires Python 3.10 or newer.") + + +_require_supported_python(sys.version_info) + +import argparse +import copy +import json +import uuid +from datetime import datetime, timezone +from pathlib import Path + +from pypaimon.benchmark.act.compare import ( + compare_results, + load_result_documents, +) +from pypaimon.benchmark.act.experiment import load_experiment +from pypaimon.benchmark.act.runner import prepare_experiment, run_experiment + + +_CONFIG_ARGUMENTS = ( + ("seed", int), + ("action_horizon", int), + ("batch_size", int), + ("optimizer_steps", int), + ("image_height", int), + ("image_width", int), + ("learning_rate", float), + ("weight_decay", float), + ("warmup_batches", int), + ("timed_batches", int), + ("fetch_batches", int), + ("rounds", int), +) + + +def main(argv=None): + """Parse an ACT benchmark subcommand and write its JSON artifact.""" + parser = _parser() + args = parser.parse_args(argv) + if args.command == "prepare": + definition = copy.deepcopy(load_experiment(args.experiment)) + for name, _ in _CONFIG_ARGUMENTS: + value = getattr(args, name) + if value is not None: + definition["config"][name] = value + for name in ( + "statistics_version", "train_episode_id", + "validation_episode_id"): + value = getattr(args, name) + if value is not None: + definition[name] = value + output = Path(args.output) + experiment = prepare_experiment( + args.input, + args.warehouse, + output, + definition=definition, + database=args.database, + ) + _print_artifact("experiment", output, experiment["schema_version"]) + return 0 + if args.command == "run": + experiment = load_experiment(args.experiment) + output = ( + Path(args.output) + if args.output else _artifact_path( + args.results_dir, + "%s-%s" % (experiment["benchmark_id"], args.backend), + ) + ) + result = run_experiment( + args.backend, + args.experiment, + output, + input_root=args.input, + warehouse=args.warehouse, + ) + _print_artifact("result", output, result["status"]) + return 0 + results_dir = args.results_dir + if not args.results and results_dir is None: + results_dir = "act-results" + results = load_result_documents(args.results, results_dir=results_dir) + comparison = compare_results(results) + output = ( + Path(args.output) + if args.output else _artifact_path( + results_dir or "act-results", "comparison") + ) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text( + json.dumps(comparison, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + _print_artifact("comparison", output, comparison["status"]) + return 0 if comparison["status"] == "SUCCEEDED" else 1 + + +def _parser(): + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + commands = parser.add_subparsers(dest="command", required=True) + + prepare = commands.add_parser( + "prepare", + help="Resolve a shared experiment against matching HDF5 and Paimon data.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + prepare.add_argument("--input", required=True, help="RoboMIND HDF5 root.") + prepare.add_argument("--warehouse", required=True, help="Paimon warehouse.") + prepare.add_argument( + "--experiment", + help="Input experiment JSON; packaged defaults are used when omitted.", + ) + prepare.add_argument( + "--output", default="act-results/experiment.json", + help="Resolved experiment JSON path.") + prepare.add_argument("--database", default="robomind") + prepare.add_argument("--statistics-version") + prepare.add_argument("--train-episode-id") + prepare.add_argument("--validation-episode-id") + for name, argument_type in _CONFIG_ARGUMENTS: + prepare.add_argument( + "--" + name.replace("_", "-"), type=argument_type, default=None) + + run = commands.add_parser( + "run", + help="Run one storage backend using a resolved experiment.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + run.add_argument("--backend", required=True, choices=("hdf5", "paimon")) + run.add_argument("--experiment", required=True) + run.add_argument("--input", help="HDF5 root; required for backend=hdf5.") + run.add_argument( + "--warehouse", help="Paimon warehouse; required for backend=paimon.") + run.add_argument("--output", help="Explicit result JSON path.") + run.add_argument( + "--results-dir", default="act-results", + help="Directory for an automatically named result.") + + compare = commands.add_parser( + "compare", + help=( + "Group results by experiment and aggregate compatible repeats." + ), + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + compare.add_argument("results", nargs="*", help="Explicit result JSON files.") + compare.add_argument( + "--results-dir", + help="Also discover ACT result JSON files in this directory.") + compare.add_argument("--output", help="Comparison JSON path.") + return parser + + +def _artifact_path(directory, prefix): + timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + return Path(directory).expanduser() / ( + "%s-%s-%s.json" % (prefix, timestamp, uuid.uuid4().hex[:8])) + + +def _print_artifact(kind, path, status): + print(json.dumps({ + "artifact": str(Path(path).expanduser().resolve()), + "kind": kind, + "status": status, + }, sort_keys=True)) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/paimon-python/pypaimon/benchmark/act/compare.py b/paimon-python/pypaimon/benchmark/act/compare.py new file mode 100644 index 000000000000..ab53f4dd6706 --- /dev/null +++ b/paimon-python/pypaimon/benchmark/act/compare.py @@ -0,0 +1,231 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Validate and aggregate independently produced ACT benchmark results.""" + +import hashlib +import json +from pathlib import Path + + +_METRICS = { + "batch_fetch_samples_per_s": "higher", + "dataset_build_s": "lower", + "first_batch_s": "lower", + "fixed_steps_s": "lower", + "python_peak_allocated_bytes": "lower", + "wall_time_s": "lower", +} + + +def canonical_sha256(value): + """Return the SHA-256 of a JSON value using canonical serialization.""" + payload = json.dumps(value, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def load_result_documents(paths, results_dir=None): + """Load explicit result files plus ACT results discovered in a directory. + + Explicit paths must contain result documents. Directory discovery ignores + experiment and prior comparison JSON files. A path found both ways is read + once, preserving explicit-path order followed by sorted directory entries. + """ + candidates = [Path(path).expanduser().resolve() for path in paths] + explicit = set(candidates) + if results_dir is not None: + directory = Path(results_dir).expanduser().resolve() + candidates.extend(sorted(directory.glob("*.json"))) + seen = set() + results = [] + for path in candidates: + path = path.resolve() + if path in seen: + continue + seen.add(path) + with path.open(encoding="utf-8") as result_file: + document = json.load(result_file) + if document.get("schema_version") != "act-benchmark-result@1": + if path in explicit: + raise ValueError("Not an ACT benchmark result: %s." % path) + continue + results.append(document) + if not results: + raise ValueError("No ACT benchmark result files were found.") + return results + + +def compare_results(results): + """Group result documents by experiment and compare compatible backends. + + Results from different experiment definitions remain separate. Results in + one experiment group must report the same runtime environment; otherwise + the group is marked incompatible and no performance ratios are produced. + + Args: + results: Iterable of decoded ``act-benchmark-result@1`` documents. + + Returns: + A JSON-compatible comparison document with one entry per experiment. + """ + groups = {} + for result in results: + if result.get("schema_version") != "act-benchmark-result@1": + raise ValueError("Unsupported ACT benchmark result schema.") + experiment = result.get("experiment") + if not isinstance(experiment, dict): + raise ValueError("ACT benchmark result has no experiment object.") + experiment_sha256 = canonical_sha256(experiment) + if result.get("experiment_sha256") != experiment_sha256: + raise ValueError("ACT result experiment SHA-256 differs.") + if result.get("status") != "SUCCEEDED": + raise ValueError("ACT comparison requires successful results.") + if result.get("backend") not in ("hdf5", "paimon"): + raise ValueError("ACT result has an unsupported backend.") + groups.setdefault(experiment_sha256, []).append(result) + + experiments = [ + _compare_experiment(experiment_sha256, grouped) + for experiment_sha256, grouped in sorted(groups.items()) + ] + statuses = {item["status"] for item in experiments} + if statuses == {"SUCCEEDED"}: + status = "SUCCEEDED" + elif "FAILED" in statuses: + status = "FAILED" + else: + status = "INCOMPATIBLE" + return { + "schema_version": "act-benchmark-comparison@1", + "status": status, + "experiments": experiments, + } + + +def _compare_experiment(experiment_sha256, results): + environments = { + canonical_sha256(result.get("environment", {})) for result in results + } + by_backend = {} + for result in results: + by_backend.setdefault(result["backend"], []).append(result) + if set(by_backend) != {"hdf5", "paimon"}: + return { + "experiment_sha256": experiment_sha256, + "experiment": results[0]["experiment"], + "status": "INCOMPATIBLE", + "reason": "both hdf5 and paimon results are required", + "backends": sorted(by_backend), + "result_count": len(results), + "metrics": {}, + } + if len(environments) != 1: + return { + "experiment_sha256": experiment_sha256, + "experiment": results[0]["experiment"], + "status": "INCOMPATIBLE", + "reason": "runtime environments differ", + "environment_sha256s": sorted(environments), + "backends": sorted(by_backend), + "metrics": {}, + } + models = {canonical_sha256(result.get("model")) for result in results} + if len(models) != 1: + return _failed_group( + experiment_sha256, by_backend, results, "models differ") + fingerprints = { + result.get("tensor_fingerprint", {}).get("sha256") + for result in results + } + if len(fingerprints) != 1 or None in fingerprints: + return _failed_group( + experiment_sha256, + by_backend, + results, + "tensor fingerprints differ", + ) + loss_traces = {canonical_sha256([{ + "round": run["round"], + "train_loss": run["train_loss"], + "validation_loss": run["validation_loss"], + } for run in result.get("runs", [])]) for result in results} + if len(loss_traces) != 1: + return _failed_group( + experiment_sha256, by_backend, results, "loss traces differ") + + medians = { + backend: _aggregate_backend(items) + for backend, items in by_backend.items() + } + metrics = {} + for name, preferred in _METRICS.items(): + values = { + backend: summary[name] + for backend, summary in medians.items() + if name in summary + } + if values: + metric = dict(values) + metric["preferred"] = preferred + if set(values) == {"hdf5", "paimon"}: + if preferred == "higher" and values["hdf5"]: + metric["paimon_over_hdf5"] = ( + values["paimon"] / values["hdf5"]) + elif preferred == "lower" and values["paimon"]: + metric["hdf5_over_paimon"] = ( + values["hdf5"] / values["paimon"]) + metrics[name] = metric + return { + "experiment_sha256": experiment_sha256, + "experiment": results[0]["experiment"], + "status": "SUCCEEDED", + "environment": results[0]["environment"], + "environment_sha256": next(iter(environments)), + "backends": sorted(by_backend), + "result_count": len(results), + "metrics": metrics, + } + + +def _failed_group(experiment_sha256, by_backend, results, reason): + return { + "experiment_sha256": experiment_sha256, + "experiment": results[0]["experiment"], + "status": "FAILED", + "reason": reason, + "backends": sorted(by_backend), + "result_count": len(results), + "metrics": {}, + } + + +def _aggregate_backend(results): + names = set.intersection(*( + set(result.get("summary", {})) for result in results + )) + aggregated = {} + for name in names: + if name not in _METRICS: + continue + values = [result["summary"][name]["median"] for result in results] + values.sort() + middle = len(values) // 2 + aggregated[name] = ( + values[middle] + if len(values) % 2 + else (values[middle - 1] + values[middle]) / 2.0 + ) + return aggregated diff --git a/paimon-python/pypaimon/benchmark/act/default_experiment.json b/paimon-python/pypaimon/benchmark/act/default_experiment.json new file mode 100644 index 000000000000..91c2ffe291c0 --- /dev/null +++ b/paimon-python/pypaimon/benchmark/act/default_experiment.json @@ -0,0 +1,22 @@ +{ + "benchmark_id": "robomind-act", + "config": { + "action_horizon": 32, + "batch_size": 2, + "fetch_batches": 8, + "image_height": 64, + "image_width": 80, + "learning_rate": 0.0001, + "optimizer_steps": 2, + "rounds": 3, + "seed": 20260825, + "timed_batches": 32, + "warmup_batches": 1, + "weight_decay": 0.0001 + }, + "dataset": "RoboMIND AgileX", + "schema_version": "act-benchmark-experiment@1", + "statistics_version": "robomind-agilex-joint-position@1", + "train_episode_id": null, + "validation_episode_id": null +} diff --git a/paimon-python/pypaimon/benchmark/act/experiment.py b/paimon-python/pypaimon/benchmark/act/experiment.py new file mode 100644 index 000000000000..f1bb675ca298 --- /dev/null +++ b/paimon-python/pypaimon/benchmark/act/experiment.py @@ -0,0 +1,42 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Load the declarative parameters shared by ACT benchmark runs.""" + +import json +from pathlib import Path + + +DEFAULT_EXPERIMENT = Path(__file__).with_name("default_experiment.json") + + +def load_experiment(path=None): + """Load an ACT experiment definition from JSON or the packaged default. + + Args: + path: Optional JSON path. When omitted, the packaged RoboMIND ACT + benchmark defaults are loaded. + + Returns: + A dictionary containing the benchmark identity, normalization version, + episode selection, and shared ACT/measurement configuration. + """ + source = DEFAULT_EXPERIMENT if path is None else Path(path) + with source.expanduser().open(encoding="utf-8") as experiment_file: + experiment = json.load(experiment_file) + if not isinstance(experiment, dict): + raise ValueError("ACT experiment must be a JSON object.") + return experiment diff --git a/paimon-python/pypaimon/benchmark/act/harness.py b/paimon-python/pypaimon/benchmark/act/harness.py new file mode 100644 index 000000000000..ee1aca7a3772 --- /dev/null +++ b/paimon-python/pypaimon/benchmark/act/harness.py @@ -0,0 +1,577 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared deterministic ACT model, trainer, and window plan for benchmarks.""" + +import gc +import hashlib +import json +import math +import random +import time +import tracemalloc +from dataclasses import asdict, dataclass +from io import BytesIO + +import numpy as np +import torch +import torch.nn.functional as functional +from PIL import Image +from torch.utils.data import default_collate + + +CAMERA_KEYS = ( + "observation.images.front", + "observation.images.left_wrist", + "observation.images.right_wrist", +) + + +@dataclass(frozen=True) +class BenchmarkConfig: + """Immutable model, sampling, training, and measurement parameters. + + Every backend reconstructs this configuration from the resolved experiment + so tensor shapes, optimizer behavior, random seeds, and metric boundaries + remain comparable. + """ + + seed: int = 20260825 + action_horizon: int = 32 + batch_size: int = 2 + optimizer_steps: int = 2 + image_height: int = 64 + image_width: int = 80 + learning_rate: float = 1e-4 + weight_decay: float = 1e-4 + warmup_batches: int = 1 + timed_batches: int = 32 + fetch_batches: int = 8 + rounds: int = 3 + + def __post_init__(self): + positive_ints = ( + "action_horizon", + "batch_size", + "optimizer_steps", + "image_height", + "image_width", + "warmup_batches", + "timed_batches", + "fetch_batches", + ) + for name in positive_ints: + value = getattr(self, name) + if ( + isinstance(value, bool) + or not isinstance(value, int) + or value <= 0): + raise ValueError("%s must be a positive int." % name) + if isinstance(self.seed, bool) or not isinstance(self.seed, int): + raise ValueError("seed must be an int.") + if isinstance(self.rounds, bool) or not isinstance(self.rounds, int): + raise ValueError("rounds must be an int.") + if self.rounds < 3: + raise ValueError("rounds must be at least 3.") + if self.learning_rate <= 0: + raise ValueError("learning_rate must be positive.") + if self.weight_decay < 0: + raise ValueError("weight_decay must not be negative.") + + def to_dict(self): + return asdict(self) + + +@dataclass(frozen=True) +class WindowPlan: + """Logical dataset-window indices consumed by one experiment. + + Measurement indices cover warm-up and timed reads, train indices cover + fixed optimizer steps, and validation indices cover the final loss. These + are map-style dataset indices, not Paimon row IDs. ``sha256`` identifies + the exact plan across independent backend processes. + """ + + seed: int + measurement_indices: tuple + train_indices: tuple + validation_indices: tuple + + @property + def sha256(self): + payload = json.dumps( + self.to_dict(), sort_keys=True, separators=(",", ":")) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + def to_dict(self): + return { + "seed": self.seed, + "measurement_indices": list(self.measurement_indices), + "train_indices": list(self.train_indices), + "validation_indices": list(self.validation_indices), + } + + +def build_window_plan(train_window_count, validation_window_count, config): + """Build deterministic measurement, training, and validation indices. + + Args: + train_window_count: Number of complete windows in the train dataset. + validation_window_count: Number of complete validation windows. + config: Shared benchmark configuration supplying counts and the seed. + + Returns: + A :class:`WindowPlan`. When more samples are needed than a dataset + contains, consecutive seeded permutations are concatenated; sampling + does not become independent sampling with replacement. + """ + train_window_count = _positive_int( + train_window_count, "train_window_count") + validation_window_count = _positive_int( + validation_window_count, "validation_window_count") + batch_fetch_count = ( + config.warmup_batches + config.timed_batches) * config.batch_size + train_count = config.optimizer_steps * config.batch_size + return WindowPlan( + seed=config.seed, + measurement_indices=tuple(_repeat_permutations( + train_window_count, batch_fetch_count, config.seed + 1)), + train_indices=tuple(_repeat_permutations( + train_window_count, train_count, config.seed + 2)), + validation_indices=tuple(_repeat_permutations( + validation_window_count, config.batch_size, config.seed + 3)), + ) + + +def decode_rgb_image(payload): + """Decode JPEG/PNG bytes into an ``H x W x 3`` RGB NumPy array. + + Raises: + ValueError: If Pillow cannot decode the payload as an image. + """ + try: + return np.asarray(Image.open(BytesIO(payload)).convert("RGB")) + except Exception as error: + raise ValueError("Cannot decode ACT RGB image bytes.") from error + + +def decode_image_tensor(value): + """Decode bytes or an HDF5 uint8 value into normalized ``C x H x W``. + + The returned NumPy array is float32 with values in ``[0, 1]``. Both + storage backends call this function so image conversion is not part of the + performance difference being measured. + """ + payload = ( + bytes(value) + if isinstance(value, (bytes, bytearray, memoryview)) + else np.asarray(value, dtype=np.uint8).tobytes() + ) + image = decode_rgb_image(payload) + return np.transpose(image, (2, 0, 1)).astype(np.float32) / 255.0 + + +def validate_act_batch(batch, config): + """Validate a collated batch against the shared ACT tensor contract. + + Successful validation returns ``None``. It checks exact fields, tensor + shapes and dtypes, finite values, image range, complete unpadded windows, + and ``sample_id == episode_id#step_idx`` identity. + """ + required = { + "sample_id", "episode_id", "step_idx", "qpos", "action", + "images", "is_pad", + } + if set(batch) != required: + raise ValueError( + "ACT batch fields differ: expected %s, got %s." + % (sorted(required), sorted(batch))) + batch_size = len(batch["sample_id"]) + expected = { + "qpos": ((batch_size, 14), torch.float32), + "action": ((batch_size, config.action_horizon, 14), torch.float32), + "images": ( + (batch_size, len(CAMERA_KEYS), 3) + + tuple(batch["images"].shape[-2:]), + torch.float32, + ), + "is_pad": ((batch_size, config.action_horizon), torch.bool), + "step_idx": ((batch_size,), torch.int64), + } + for name, (shape, dtype) in expected.items(): + value = batch[name] + if not isinstance(value, torch.Tensor): + raise ValueError("%s must be a torch.Tensor." % name) + if tuple(value.shape) != shape: + raise ValueError( + "%s has shape %s; expected %s." + % (name, tuple(value.shape), shape)) + if value.dtype != dtype: + raise ValueError( + "%s has dtype %s; expected %s." % (name, value.dtype, dtype)) + for name in ("qpos", "action", "images"): + if not torch.isfinite(batch[name]).all(): + raise ValueError("%s contains NaN or Inf." % name) + if torch.any(batch["images"] < 0) or torch.any(batch["images"] > 1): + raise ValueError("images must be normalized to [0, 1].") + if batch["is_pad"].any(): + raise ValueError("ACT benchmark windows must be complete and unpadded.") + for sample_id, episode_id, step_idx in zip( + batch["sample_id"], batch["episode_id"], + batch["step_idx"].tolist()): + if sample_id != "%s#%s" % (episode_id, step_idx): + raise ValueError( + "sample_id is not aligned with episode_id and step_idx.") + + +def build_lerobot_batch(batch, config): + """Map a shared ACT batch to LeRobot ``ACTPolicy`` feature names. + + Images are resized bilinearly to the configured height and width when + necessary. State, action, and padding retain their original semantics. + """ + validate_act_batch(batch, config) + images = batch["images"] + target_size = (config.image_height, config.image_width) + if tuple(images.shape[-2:]) != target_size: + flat = images.flatten(0, 1) + flat = functional.interpolate( + flat, size=target_size, mode="bilinear", align_corners=False) + images = flat.reshape(images.shape[:3] + target_size) + result = { + "observation.state": batch["qpos"], + "action": batch["action"], + "action_is_pad": batch["is_pad"], + } + for index, name in enumerate(CAMERA_KEYS): + result[name] = images[:, index] + return result + + +def build_act_policy(config): + """Build the reduced CPU ACT policy used only by this benchmark. + + Returns: + ``(policy, metadata)`` containing the LeRobot policy and a + JSON-compatible description of its architecture and parameter counts. + Pretrained weights are disabled, so this function performs no model + download and does not represent a production training configuration. + """ + try: + import importlib.metadata + from lerobot.configs.types import FeatureType, PolicyFeature + from lerobot.policies.act.configuration_act import ACTConfig + from lerobot.policies.act.modeling_act import ACTPolicy + except ImportError as error: + raise ImportError( + "ACT benchmark requires: " + "pip install -e '.[act]'.") from error + + inputs = { + "observation.state": PolicyFeature(FeatureType.STATE, (14,)), + } + inputs.update({ + name: PolicyFeature( + FeatureType.VISUAL, + (3, config.image_height, config.image_width), + ) + for name in CAMERA_KEYS + }) + act_config = ACTConfig( + input_features=inputs, + output_features={ + "action": PolicyFeature(FeatureType.ACTION, (14,)), + }, + device="cpu", + chunk_size=config.action_horizon, + n_action_steps=config.action_horizon, + vision_backbone="resnet18", + pretrained_backbone_weights=None, + dim_model=64, + n_heads=4, + dim_feedforward=256, + n_encoder_layers=1, + n_decoder_layers=1, + use_vae=True, + latent_dim=16, + n_vae_encoder_layers=1, + kl_weight=10.0, + ) + policy = ACTPolicy(act_config) + return policy, { + "implementation": "lerobot.ACTPolicy", + "lerobot_version": importlib.metadata.version("lerobot"), + "vision_backbone": act_config.vision_backbone, + "pretrained_backbone_weights": act_config.pretrained_backbone_weights, + "chunk_size": act_config.chunk_size, + "dim_model": act_config.dim_model, + "n_heads": act_config.n_heads, + "n_encoder_layers": act_config.n_encoder_layers, + "n_decoder_layers": act_config.n_decoder_layers, + "n_vae_encoder_layers": act_config.n_vae_encoder_layers, + "latent_dim": act_config.latent_dim, + "kl_weight": act_config.kl_weight, + "parameter_count": sum( + parameter.numel() for parameter in policy.parameters()), + "trainable_parameter_count": sum( + parameter.numel() + for parameter in policy.parameters() if parameter.requires_grad), + } + + +def run_backend( + backend, + round_number, + dataset_factory, + plan, + config, + sample_sequence_sha256, + policy_factory=None): + """Measure one backend with the shared plan, model, and trainer. + + ``backend`` is a result label and ``round_number`` identifies the repeat. + ``dataset_factory`` must return ``(train_dataset, validation_dataset)`` and + must be reusable: it is called for the timed run and again by the separate + Python-memory replay. ``policy_factory`` is an optional test hook returning + ``(policy, model_metadata)``. + + Returns: + A JSON-compatible metrics dictionary covering dataset construction, + first batch, timed batch fetch, fixed optimizer steps, validation loss, + and a separate ``tracemalloc`` peak replay. ``fixed_steps_s`` includes + dataset fetch, while each ``train_trace.step_time_s`` starts after its + batch is fetched and covers conversion, forward/backward, and optimizer + update. OS page cache is not controlled and native Arrow/Torch + allocations are outside tracemalloc. + """ + _seed_everything(config.seed) + policy_factory = policy_factory or build_act_policy + started = time.monotonic() + dataset_started = time.monotonic() + train_dataset, validation_dataset = dataset_factory() + dataset_build_s = time.monotonic() - dataset_started + + warmup_sample_count = config.warmup_batches * config.batch_size + warmup_iterator = _iter_logical_batches( + train_dataset, + plan.measurement_indices[:warmup_sample_count], + logical_batch_size=config.batch_size, + fetch_batches=1, + ) + first_batch_started = time.monotonic() + first_batch = next(warmup_iterator) + first_batch_s = time.monotonic() - first_batch_started + validate_act_batch(first_batch, config) + for _ in range(config.warmup_batches - 1): + validate_act_batch(next(warmup_iterator), config) + + batch_fetch_iterator = _iter_logical_batches( + train_dataset, + plan.measurement_indices[warmup_sample_count:], + logical_batch_size=config.batch_size, + fetch_batches=config.fetch_batches, + ) + batch_fetch_seconds = 0.0 + batch_fetch_sample_count = 0 + for _ in range(config.timed_batches): + batch_fetch_started = time.monotonic() + batch = next(batch_fetch_iterator) + batch_fetch_seconds += time.monotonic() - batch_fetch_started + validate_act_batch(batch, config) + batch_fetch_sample_count += len(batch["sample_id"]) + + _seed_everything(config.seed) + policy, model = policy_factory(config) + parameters = ( + policy.get_optim_params() + if hasattr(policy, "get_optim_params") else policy.parameters()) + optimizer = torch.optim.AdamW( + parameters, + lr=config.learning_rate, + weight_decay=config.weight_decay, + ) + policy.train() + train_started = time.monotonic() + losses = [] + for step, batch in enumerate(_iter_logical_batches( + train_dataset, + plan.train_indices, + logical_batch_size=config.batch_size, + fetch_batches=config.fetch_batches, + ), 1): + step_started = time.monotonic() + model_batch = build_lerobot_batch(batch, config) + optimizer.zero_grad(set_to_none=True) + loss, components = policy(model_batch) + if loss.ndim != 0 or not torch.isfinite(loss): + raise FloatingPointError( + "ACT produced a non-finite scalar loss at step %d." % step) + loss.backward() + optimizer.step() + losses.append({ + "step": step, + "total": float(loss.detach()), + "components": { + name: _finite_float(value, name) + for name, value in components.items() + }, + "step_time_s": time.monotonic() - step_started, + }) + fixed_steps_s = time.monotonic() - train_started + if len(losses) != config.optimizer_steps: + raise AssertionError( + "Expected %d optimizer steps, got %d." + % (config.optimizer_steps, len(losses))) + + # ACTPolicy only constructs the VAE posterior needed by its supervised + # loss while the module is in training mode. Keep that mode for validation + # but disable gradients and parameter updates below. + policy.train() + _seed_everything(config.seed + 4) + validation_batch = next(_iter_logical_batches( + validation_dataset, + plan.validation_indices, + logical_batch_size=config.batch_size, + fetch_batches=config.fetch_batches, + )) + with torch.no_grad(): + validation_loss, _ = policy(build_lerobot_batch( + validation_batch, config)) + validation_value = _finite_float(validation_loss, "validation_loss") + wall_time_s = time.monotonic() - started + python_peak = _measure_python_peak(dataset_factory, plan, config) + + return { + "round": round_number, + "backend": backend, + "sample_sequence_sha256": sample_sequence_sha256, + "model": model, + "optimizer": { + "name": "AdamW", + "learning_rate": config.learning_rate, + "weight_decay": config.weight_decay, + }, + "warmup_batches": config.warmup_batches, + "first_batch_s": first_batch_s, + "dataset_build_s": dataset_build_s, + "batch_fetch_samples": batch_fetch_sample_count, + "batch_fetch_s": batch_fetch_seconds, + "batch_fetch_samples_per_s": ( + batch_fetch_sample_count / batch_fetch_seconds), + "fixed_steps_s": fixed_steps_s, + "train_loss": [item["total"] for item in losses], + "train_trace": losses, + "validation_loss": validation_value, + "python_peak_allocated_bytes": python_peak, + "peak_memory_measurement": ( + "python-tracemalloc-separate-dataset-first-batch"), + "wall_time_s": wall_time_s, + } + + +def _measure_python_peak(dataset_factory, plan, config): + """Measure Python allocation peak in a separate dataset-first-batch replay. + + The factory is called again so tracing overhead cannot distort the main + throughput timings. The returned integer is the tracemalloc peak in bytes. + """ + gc.collect() + tracemalloc.start() + try: + train_dataset, _ = dataset_factory() + indices = plan.measurement_indices[ + :config.batch_size * config.fetch_batches + ] + next(_iter_logical_batches( + train_dataset, + indices, + logical_batch_size=config.batch_size, + fetch_batches=config.fetch_batches, + )) + _, peak = tracemalloc.get_traced_memory() + return peak + finally: + tracemalloc.stop() + + +def _repeat_permutations(size, count, seed): + """Return ``count`` indices by concatenating seeded permutations.""" + values = [] + generator = np.random.RandomState(seed) + while len(values) < count: + values.extend(generator.permutation(size).tolist()) + return values[:count] + + +def _seed_everything(seed): + """Reset Python, NumPy, and Torch RNGs and enable deterministic Torch ops.""" + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.use_deterministic_algorithms(True) + + +def _finite_float(value, name): + if isinstance(value, torch.Tensor): + if value.numel() != 1: + raise ValueError("%s must be scalar." % name) + value = float(value.detach()) + else: + value = float(value) + if not math.isfinite(value): + raise FloatingPointError("%s is NaN or Inf." % name) + return value + + +def _positive_int(value, name): + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError("%s must be a positive int." % name) + return value + + +def _iter_logical_batches( + dataset, indices, *, logical_batch_size, fetch_batches): + """Yield collated model batches while coalescing physical dataset reads. + + Args: + dataset: Map-style dataset implementing ``__getitem__`` and optionally + plural ``__getitems__(indices)`` access. + indices: Explicit ordered logical-window indices. Their count must be + divisible by ``logical_batch_size``. + logical_batch_size: Number of samples consumed by one model step. + fetch_batches: Logical batches combined into one physical dataset read. + + Yields: + Collated logical batches in the exact input-index order. A plural + dataset method is preferred when available; otherwise samples are read + individually and split back into the same logical batches. + """ + logical_batch_size = _positive_int( + logical_batch_size, "logical_batch_size") + fetch_batches = _positive_int(fetch_batches, "fetch_batches") + if len(indices) % logical_batch_size: + raise ValueError("indices must contain complete logical batches.") + physical_size = logical_batch_size * fetch_batches + getitems = getattr(dataset, "__getitems__", None) + for offset in range(0, len(indices), physical_size): + physical_indices = list(indices[offset:offset + physical_size]) + if getitems is None: + samples = [dataset[index] for index in physical_indices] + else: + samples = getitems(physical_indices) + for logical_offset in range(0, len(samples), logical_batch_size): + yield default_collate( + samples[logical_offset:logical_offset + logical_batch_size]) diff --git a/paimon-python/pypaimon/benchmark/act/hdf5.py b/paimon-python/pypaimon/benchmark/act/hdf5.py new file mode 100644 index 000000000000..472ead7b4b9d --- /dev/null +++ b/paimon-python/pypaimon/benchmark/act/hdf5.py @@ -0,0 +1,201 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""HDF5 dataset adapter for the RoboMIND ACT benchmark.""" + +import numpy as np +import torch +from torch.utils.data import Dataset + +from pypaimon.benchmark.act.harness import decode_image_tensor + + +QPOS_FIELDS = ( + "puppet/joint_position_left", + "puppet/joint_position_right", +) +ACTION_FIELDS = ( + "master/joint_position_left", + "master/joint_position_right", +) +IMAGE_FIELDS = ( + "observations/rgb_images/camera_front", + "observations/rgb_images/camera_left_wrist", + "observations/rgb_images/camera_right_wrist", +) + + +class Hdf5ACTWindowDataset(Dataset): + """Read complete ACT windows lazily from one HDF5 episode. + + ``episode`` supplies the file path, logical episode ID, and frame count. + For a window anchor, state and three camera images come from the anchor + frame while actions cover ``[anchor, anchor + action_horizon)``. Each + access opens and closes the HDF5 file and returns the shared ACT sample + mapping consumed by :mod:`pypaimon.benchmark.act.harness`. + """ + + def __init__(self, episode, normalization, action_horizon): + self.episode = episode + self.normalization = normalization + self.action_horizon = action_horizon + self.window_count = episode.frame_count - action_horizon + 1 + if self.window_count <= 0: + raise ValueError( + "Episode %s is shorter than action horizon %d." + % (episode.episode_id, action_horizon)) + + def __len__(self): + return self.window_count + + def __getitem__(self, anchor): + """Return the ACT window whose first frame is ``anchor``. + + Negative anchors follow Python sequence semantics. State and images + come from the anchor frame, while action contains the complete horizon. + """ + if anchor < 0: + anchor += self.window_count + if anchor < 0 or anchor >= self.window_count: + raise IndexError(anchor) + import h5py + + with h5py.File(str(self.episode.path), "r") as h5: + qpos = _read_vectors(h5, QPOS_FIELDS, anchor) + action = _read_vectors( + h5, + ACTION_FIELDS, + slice(anchor, anchor + self.action_horizon), + ) + images = np.stack([ + decode_image_tensor(h5[field][anchor]) for field in IMAGE_FIELDS + ]) + qpos = ( + (qpos - self.normalization["qpos_mean"]) + / self.normalization["qpos_std"]) + action = ( + (action - self.normalization["action_mean"]) + / self.normalization["action_std"]) + return { + "sample_id": "%s#%d" % (self.episode.episode_id, anchor), + "episode_id": self.episode.episode_id, + "step_idx": anchor, + "qpos": torch.from_numpy(np.ascontiguousarray(qpos)), + "action": torch.from_numpy(np.ascontiguousarray(action)), + "images": torch.from_numpy(np.ascontiguousarray(images)), + "is_pad": torch.zeros(self.action_horizon, dtype=torch.bool), + } + + +def create_datasets(train_episode, validation_episode, normalization, config): + """Create HDF5 datasets for the experiment's selected episodes. + + Args: + train_episode: Selected training episode with its HDF5 path and frame + count. + validation_episode: Selected validation episode with the same fields. + normalization: Shared state and action normalization arrays. + config: Benchmark configuration containing the action horizon. + + Returns: + ``(train_dataset, validation_dataset)`` in that order. + """ + return ( + Hdf5ACTWindowDataset( + train_episode, normalization, config.action_horizon), + Hdf5ACTWindowDataset( + validation_episode, normalization, config.action_horizon), + ) + + +def compute_normalization(episodes): + """Compute train-only HDF5 state and action normalization. + + Args: + episodes: Discovered episodes carrying ``path``, ``split``, and + ``success`` attributes. + + Returns: + ``(normalization, metadata)`` where normalization contains float32 + arrays used by training. Metadata retains the float64 action moments + and frame count used to validate Paimon statistics without losing + precision. Standard deviations use a ``1e-2`` floor. + """ + train = [ + episode for episode in episodes + if episode.split == "train" and episode.success + ] + if not train: + raise ValueError("No successful train episodes are available.") + qpos = _Moments(14) + action = _Moments(14) + import h5py + + for episode in sorted(train, key=lambda item: item.episode_id): + with h5py.File(str(episode.path), "r") as h5: + qpos.update(_read_vectors( + h5, QPOS_FIELDS, slice(None), dtype=np.float64)) + action.update(_read_vectors( + h5, ACTION_FIELDS, slice(None), dtype=np.float64)) + qpos_mean, qpos_std = qpos.finish() + action_mean, action_std = action.finish() + return ({ + "qpos_mean": qpos_mean.astype(np.float32), + "qpos_std": qpos_std.astype(np.float32), + "action_mean": action_mean.astype(np.float32), + "action_std": action_std.astype(np.float32), + }, { + "action_mean": action_mean, + "action_std": action_std, + "frame_count": action.count, + }) + + +def _read_vectors(h5, fields, selection, dtype=np.float32): + value = np.concatenate([ + np.asarray(h5[field][selection], dtype=dtype) for field in fields + ], axis=-1) + if not np.isfinite(value).all(): + raise ValueError("ACT vector contains NaN or Inf.") + return value + + +class _Moments(object): + """Accumulate float64 population moments with a ``1e-2`` std floor.""" + + def __init__(self, width): + self.count = 0 + self.total = np.zeros(width, dtype=np.float64) + self.total_square = np.zeros(width, dtype=np.float64) + + def update(self, value): + value = np.asarray(value, dtype=np.float64) + if value.ndim != 2 or value.shape[1] != len(self.total): + raise ValueError( + "Unexpected normalization shape %s." % (value.shape,)) + if not np.isfinite(value).all(): + raise ValueError("Normalization input contains NaN or Inf.") + self.count += value.shape[0] + self.total += value.sum(axis=0) + self.total_square += np.square(value).sum(axis=0) + + def finish(self): + if self.count == 0: + raise ValueError("Cannot compute normalization from no frames.") + mean = self.total / self.count + variance = np.maximum( + self.total_square / self.count - np.square(mean), 0.0) + return mean, np.maximum(np.sqrt(variance), 1e-2) diff --git a/paimon-python/pypaimon/benchmark/act/paimon.py b/paimon-python/pypaimon/benchmark/act/paimon.py new file mode 100644 index 000000000000..d99ea4227aae --- /dev/null +++ b/paimon-python/pypaimon/benchmark/act/paimon.py @@ -0,0 +1,152 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Paimon dataset adapter for the RoboMIND ACT benchmark.""" + +import numpy as np +import torch + +from pypaimon.benchmark.act.harness import decode_image_tensor +from pypaimon.sample import robomind_agilex as agilex + + +QPOS_COLUMNS = ( + "state_joint_position_left", + "state_joint_position_right", +) +ACTION_COLUMNS = ("action",) +IMAGE_COLUMNS = ( + "rgb_front", + "rgb_left_wrist", + "rgb_right_wrist", +) + + +class PaimonACTAdapter: + """Convert a contiguous Paimon row window to the shared ACT sample. + + State and camera columns are taken from the anchor row. The action column + covers the full horizon. The returned mapping has the same IDs, tensors, + shapes, and normalization as :class:`Hdf5ACTWindowDataset`. + """ + + def __init__(self, normalization): + self.normalization = normalization + + def __call__(self, sample): + """Convert the generic window mapping into ACT tensors and identity. + + The persisted ``frame_index`` becomes the shared ACT ``step_idx``. + State and image columns are singleton lists; action retains the full + horizon and ``is_pad`` is forwarded unchanged. + """ + qpos = np.concatenate([ + np.asarray(sample[name][0], dtype=np.float32) + for name in QPOS_COLUMNS + ]) + action = np.concatenate([ + np.asarray(sample[name], dtype=np.float32) + for name in ACTION_COLUMNS + ], axis=-1) + images = np.stack([ + decode_image_tensor(sample[name][0]) for name in IMAGE_COLUMNS + ]) + qpos = ( + (qpos - self.normalization["qpos_mean"]) + / self.normalization["qpos_std"]) + action = ( + (action - self.normalization["action_mean"]) + / self.normalization["action_std"]) + episode_id = sample["episode_id"] + step_idx = sample["frame_index"] + return { + "sample_id": "%s#%d" % (episode_id, step_idx), + "episode_id": episode_id, + "step_idx": step_idx, + "qpos": torch.from_numpy(np.ascontiguousarray(qpos)), + "action": torch.from_numpy(np.ascontiguousarray(action)), + "images": torch.from_numpy(np.ascontiguousarray(images)), + "is_pad": sample["is_pad"], + } + + +def create_datasets( + frames, + snapshot_id, + train_episode_id, + validation_episode_id, + normalization, + config): + """Create lazy train and validation windows pinned to one snapshot. + + State and image columns are anchor-only, so one sample reads the initial + joint position and three observation images once rather than once per + action-horizon row. + + Args: + frames: Paimon frames table used to create both scans. + snapshot_id: Snapshot pinned by experiment preparation. Both returned + datasets reject any different resolved snapshot. + train_episode_id: Episode selected for training windows. + validation_episode_id: Episode selected for validation windows. + normalization: Shared state and action normalization arrays. + config: Benchmark configuration containing the action horizon. + + Returns: + ``(train_dataset, validation_dataset)`` in that order, as lazy + ``ContiguousWindowDataset`` instances pinned to ``snapshot_id``. + """ + datasets = tuple( + frames.scan(snapshot_id=snapshot_id).where( + "episode_id = '%s'" % episode_id.replace("'", "''") + ).to_contiguous_window_dataset( + window_size=config.action_horizon, + columns=QPOS_COLUMNS + ACTION_COLUMNS + IMAGE_COLUMNS, + anchor_columns=QPOS_COLUMNS + IMAGE_COLUMNS, + group_key="episode_id", + order_key="frame_index", + stride=1, + tail="drop", + adapter=PaimonACTAdapter(normalization), + ) + for episode_id in (train_episode_id, validation_episode_id) + ) + actual_snapshot_ids = {dataset.snapshot_id for dataset in datasets} + if actual_snapshot_ids != {snapshot_id}: + raise RuntimeError( + "Paimon ACT windows must remain pinned to frames snapshot %s; " + "got %s." % (snapshot_id, sorted(actual_snapshot_ids))) + return datasets + + +def statistics_row(connection, statistics_version): + """Return the unique versioned action-statistics row.""" + escaped = statistics_version.replace("'", "''") + rows = (connection.get_table(agilex.FEATURE_STATS_TABLE).scan() + .where("statistics_version = '%s'" % escaped).to_list()) + if len(rows) != 1: + raise ValueError( + "Expected one normalization row for %r, got %d." + % (statistics_version, len(rows))) + return rows[0] + + +def latest_snapshot_id(table): + """Return the table's latest snapshot ID or fail for an empty table.""" + snapshot = table.raw_table.snapshot_manager().get_latest_snapshot() + if snapshot is None: + raise ValueError("Paimon frames table has no snapshot.") + return snapshot.id diff --git a/paimon-python/pypaimon/benchmark/act/runner.py b/paimon-python/pypaimon/benchmark/act/runner.py new file mode 100644 index 000000000000..d82ca1683954 --- /dev/null +++ b/paimon-python/pypaimon/benchmark/act/runner.py @@ -0,0 +1,739 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Prepare and run RoboMIND ACT benchmarks over HDF5 or Paimon. + +Both adapters consume one immutable :class:`BenchmarkConfig`, one train-only +normalization object, and one explicit window plan. The runner resets the same +seed before constructing the same LeRobot ACT policy and AdamW trainer for each +backend. Each backend runs independently without attempting OS cache control +and writes its tensor fingerprint, loss trace, timing metrics, and Python +allocation metrics to one result JSON document. +Ingestion and canonical-action backfill are deliberately outside the benchmark. +""" + +import gc +import hashlib +import json +import os +import platform +import subprocess +import time +import uuid +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path + +import PIL +import h5py +import numpy as np +import pyarrow as pa +import torch +import pypaimon.multimodal as pmm +from pypaimon import build_info +from pypaimon.benchmark.act.hdf5 import ( + compute_normalization as compute_hdf5_normalization, + create_datasets as create_hdf5_datasets, +) +from pypaimon.benchmark.act.experiment import load_experiment +from pypaimon.benchmark.act.harness import ( + BenchmarkConfig, + WindowPlan, + build_window_plan, + run_backend, +) +from pypaimon.benchmark.act.compare import canonical_sha256 +from pypaimon.benchmark.act.paimon import ( + create_datasets as create_paimon_datasets, + latest_snapshot_id, + statistics_row, +) +from pypaimon.sample import robomind_agilex as agilex + + +@dataclass(frozen=True) +class _BenchmarkEpisode: + path: Path + source_key: str + episode_id: str + split: str + success: bool + frame_count: int + + +def prepare_experiment( + input_root, + warehouse, + output_path, + *, + definition=None, + database=agilex.DEFAULT_DATABASE): + """Resolve a benchmark definition against matching HDF5 and Paimon data. + + Preparation is outside timed benchmark execution. It verifies source + identity and Paimon statistics, selects eligible train/validation episodes, + computes train-only normalization, and fixes every logical window index. + + Args: + input_root: RoboMIND AgileX HDF5 root used as the source of episode + files and raw normalization moments. + warehouse: Existing Paimon warehouse containing the matching ingested + and canonical-action-backfilled dataset. + output_path: Destination for the resolved experiment JSON document. + definition: Optional decoded experiment definition. The packaged + defaults are used when omitted. + database: Paimon database containing the RoboMIND tables. + + Returns: + The resolved, JSON-compatible experiment dictionary written to + ``output_path``. + """ + definition = load_experiment() if definition is None else definition + if definition.get("schema_version") != "act-benchmark-experiment@1": + raise ValueError("Unsupported ACT benchmark experiment schema.") + config = BenchmarkConfig(**definition["config"]) + statistics_version = definition["statistics_version"] + input_root = Path(input_root).expanduser().resolve() + warehouse = Path(warehouse).expanduser().resolve() + output_path = Path(output_path).expanduser().resolve() + + discovered = agilex.discover_episodes(input_root) + connection = pmm.connect( + database=database, options={"warehouse": str(warehouse)}) + source_episodes, source_sha256 = _validate_source_identity( + discovered, _episode_rows(connection)) + source_by_id = {episode.episode_id: episode for episode in source_episodes} + frames = connection.get_table(agilex.FRAMES_TABLE) + frames_snapshot_id = latest_snapshot_id(frames) + normalization, normalization_metadata = _shared_normalization( + source_episodes, + connection, + frames_snapshot_id, + statistics_version, + ) + del normalization + train_episode = _select_episode( + source_by_id, + split="train", + requested=definition.get("train_episode_id"), + action_horizon=config.action_horizon, + ) + validation_episode = _select_episode( + source_by_id, + split="val", + requested=definition.get("validation_episode_id"), + action_horizon=config.action_horizon, + ) + plan = build_window_plan( + train_episode.frame_count - config.action_horizon + 1, + validation_episode.frame_count - config.action_horizon + 1, + config, + ) + sequence_sha256 = _sample_sequence_sha256( + train_episode.episode_id, validation_episode.episode_id, plan) + episodes = sorted(({ + "episode_id": episode.episode_id, + "source_key": episode.source_key, + "split": episode.split, + "success": episode.success, + "frame_count": episode.frame_count, + } for episode in source_episodes), key=lambda item: item["episode_id"]) + experiment = { + "schema_version": "act-benchmark-experiment@1", + "benchmark_id": definition.get("benchmark_id", "robomind-act"), + "dataset": definition.get("dataset", "RoboMIND AgileX"), + "config": config.to_dict(), + "statistics_version": statistics_version, + "train_episode_id": train_episode.episode_id, + "validation_episode_id": validation_episode.episode_id, + "source": { + "sha256": source_sha256, + "episodes": episodes, + }, + "normalization": normalization_metadata, + "window_plan": { + **plan.to_dict(), + "sha256": plan.sha256, + "sample_sequence_sha256": sequence_sha256, + }, + "paimon": { + "database": database, + "frames_table": agilex.FRAMES_TABLE, + "frames_snapshot_id": frames_snapshot_id, + }, + } + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text( + json.dumps(experiment, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return experiment + + +def run_experiment( + backend, + experiment_path, + output_path, + *, + input_root=None, + warehouse=None, + policy_factory=None): + """Run one storage backend against a resolved ACT experiment. + + Args: + backend: Result label and dataset implementation, either ``hdf5`` or + ``paimon``. + experiment_path: Resolved JSON produced by :func:`prepare_experiment`. + output_path: Destination JSON result path. + input_root: Required only for the HDF5 backend. + warehouse: Required only for the Paimon backend. + policy_factory: Optional test hook returning ``(policy, metadata)``. + + Returns: + A JSON-compatible single-backend result containing the resolved + experiment, runtime environment, tensor fingerprint, per-round raw + metrics, and median/min/max summary. + """ + if backend not in ("hdf5", "paimon"): + raise ValueError("backend must be 'hdf5' or 'paimon'.") + experiment = load_experiment(experiment_path) + _validate_resolved_experiment(experiment) + config = BenchmarkConfig(**experiment["config"]) + plan = _window_plan_from_experiment(experiment) + normalization = { + name: np.asarray(value, dtype=np.float32) + for name, value in experiment["normalization"]["values"].items() + } + sequence_sha256 = experiment["window_plan"]["sample_sequence_sha256"] + if backend == "hdf5": + if input_root is None: + raise ValueError("input_root is required for the HDF5 backend.") + episodes = _hdf5_episodes_from_experiment(input_root, experiment) + by_id = {episode.episode_id: episode for episode in episodes} + train_episode = by_id[experiment["train_episode_id"]] + validation_episode = by_id[experiment["validation_episode_id"]] + + def dataset_factory(): + return create_hdf5_datasets( + train_episode, validation_episode, normalization, config) + + source = {"input_root": str(Path(input_root).expanduser().resolve())} + else: + if warehouse is None: + raise ValueError("warehouse is required for the Paimon backend.") + dataset_factory, source = _paimon_factory_from_experiment( + warehouse, experiment, normalization, config) + + started_at = _utc_now() + started = time.monotonic() + fingerprint = _tensor_fingerprint(dataset_factory(), plan) + runs = [] + for round_number in range(1, config.rounds + 1): + runs.append(run_backend( + backend, + round_number, + dataset_factory, + plan, + config, + sequence_sha256, + policy_factory=policy_factory, + )) + gc.collect() + result = { + "schema_version": "act-benchmark-result@1", + "benchmark_id": experiment["benchmark_id"], + "run_id": "%s-%s" % ( + started_at.replace(":", "").replace("-", ""), + uuid.uuid4().hex[:8], + ), + "status": "SUCCEEDED", + "backend": backend, + "experiment": experiment, + "experiment_sha256": canonical_sha256(experiment), + "source": source, + "tensor_fingerprint": fingerprint, + "model": runs[0]["model"], + "runs": runs, + "summary": _summarize(runs), + "environment": _runtime_environment( + Path(__file__).resolve().parents[4]), + "command": _command_argv(), + "timing": {"wall_time_s": time.monotonic() - started}, + "unverified": [ + "OS page cache is uncontrolled; no cache dropping was attempted.", + "CPU fixed-step loss parity proves engineering equivalence, " + "not policy quality.", + "GPU, multi-worker dataset loading, distributed training, and " + "recovery are unverified.", + "Python tracemalloc excludes native Arrow and Torch allocations.", + ], + "started_at": started_at, + "finished_at": _utc_now(), + } + output_path = Path(output_path).expanduser().resolve() + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text( + json.dumps(result, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return result + + +def _validate_resolved_experiment(experiment): + """Reject incomplete or internally inconsistent resolved experiments.""" + required = { + "schema_version", "benchmark_id", "dataset", "config", + "statistics_version", "train_episode_id", "validation_episode_id", + "source", "normalization", "window_plan", "paimon", + } + if experiment.get("schema_version") != "act-benchmark-experiment@1": + raise ValueError("Unsupported ACT benchmark experiment schema.") + missing = required - set(experiment) + if missing: + raise ValueError( + "Resolved ACT experiment is missing: %s." + % ", ".join(sorted(missing))) + source = experiment["source"] + if canonical_sha256(source["episodes"]) != source["sha256"]: + raise ValueError("ACT experiment source-manifest hash differs.") + normalization = experiment["normalization"] + if canonical_sha256(normalization["values"]) != normalization["sha256"]: + raise ValueError("ACT experiment normalization hash differs.") + plan = _window_plan_from_experiment(experiment) + if plan.sha256 != experiment["window_plan"]["sha256"]: + raise ValueError("ACT experiment window-plan hash differs.") + episodes = { + item["episode_id"]: item for item in source["episodes"] + } + try: + train = episodes[experiment["train_episode_id"]] + validation = episodes[experiment["validation_episode_id"]] + except KeyError as error: + raise ValueError( + "ACT experiment selected episode is absent from the source." + ) from error + config = BenchmarkConfig(**experiment["config"]) + expected_plan = build_window_plan( + train["frame_count"] - config.action_horizon + 1, + validation["frame_count"] - config.action_horizon + 1, + config, + ) + if expected_plan.to_dict() != plan.to_dict(): + raise ValueError( + "ACT experiment window plan was not built from its config and " + "selected episodes.") + expected_sequence = _sample_sequence_sha256( + train["episode_id"], validation["episode_id"], plan) + if expected_sequence != experiment["window_plan"][ + "sample_sequence_sha256"]: + raise ValueError("ACT experiment sample-sequence hash differs.") + + +def _window_plan_from_experiment(experiment): + """Reconstruct immutable logical-window indices from JSON values.""" + value = experiment["window_plan"] + return WindowPlan( + seed=value["seed"], + measurement_indices=tuple(value["measurement_indices"]), + train_indices=tuple(value["train_indices"]), + validation_indices=tuple(value["validation_indices"]), + ) + + +def _hdf5_episodes_from_experiment(input_root, experiment): + """Validate HDF5 episode identity and attach manifest frame counts.""" + discovered = agilex.discover_episodes(Path(input_root).expanduser().resolve()) + by_id = {episode.episode_id: episode for episode in discovered} + expected = experiment["source"]["episodes"] + actual_identity = sorted(({ + "episode_id": episode.episode_id, + "source_key": episode.source_key, + "split": episode.split, + "success": episode.success, + } for episode in discovered), key=lambda item: item["episode_id"]) + expected_identity = [{ + "episode_id": item["episode_id"], + "source_key": item["source_key"], + "split": item["split"], + "success": item["success"], + } for item in expected] + if actual_identity != expected_identity: + raise ValueError("HDF5 source differs from the ACT experiment.") + return [ + _BenchmarkEpisode( + path=by_id[item["episode_id"]].path, + source_key=item["source_key"], + episode_id=item["episode_id"], + split=item["split"], + success=item["success"], + frame_count=item["frame_count"], + ) + for item in expected + ] + + +def _paimon_factory_from_experiment( + warehouse, experiment, normalization, config): + """Validate Paimon source/statistics and return a pinned dataset factory.""" + warehouse = Path(warehouse).expanduser().resolve() + paimon = experiment["paimon"] + connection = pmm.connect( + database=paimon["database"], + options={"warehouse": str(warehouse)}, + ) + frames = connection.get_table(paimon["frames_table"]) + snapshot_id = paimon["frames_snapshot_id"] + expected_episodes = experiment["source"]["episodes"] + actual_episodes = sorted(_episode_rows(connection), + key=lambda item: item["episode_id"]) + if actual_episodes != expected_episodes: + raise ValueError("Paimon source differs from the ACT experiment.") + row = statistics_row(connection, experiment["statistics_version"]) + expected_normalization = experiment["normalization"] + action_mean = np.asarray(row["action_mean"], dtype=np.float32) + action_std = np.asarray(row["action_std"], dtype=np.float32) + if ( + row["source_snapshot_id"] != snapshot_id + or row["source_split"] != "train" + or row["frame_count"] != expected_normalization["frame_count"] + or row["feature_name"] != "action" + or row["standard_deviation_floor"] != 1e-2 + or not np.array_equal( + action_mean, normalization["action_mean"]) + or not np.array_equal(action_std, normalization["action_std"])): + raise ValueError( + "Paimon normalization differs from the ACT experiment.") + + def factory(): + return create_paimon_datasets( + frames, + snapshot_id, + experiment["train_episode_id"], + experiment["validation_episode_id"], + normalization, + config, + ) + + return factory, { + "warehouse": str(warehouse), + "database": paimon["database"], + "frames_table": paimon["frames_table"], + "frames_snapshot_id": snapshot_id, + } + + +def _tensor_fingerprint(datasets, plan): + """Hash the exact planned sample IDs and tensors outside timed execution.""" + comparisons = ( + ("train", datasets[0], + sorted(set(plan.measurement_indices + plan.train_indices))), + ("validation", datasets[1], + sorted(set(plan.validation_indices))), + ) + digest = hashlib.sha256() + count = 0 + for split, dataset, indices in comparisons: + for index in indices: + sample = dataset[index] + identity = { + "split": split, + "index": index, + "sample_id": sample["sample_id"], + "episode_id": sample["episode_id"], + "step_idx": sample["step_idx"], + } + digest.update(json.dumps( + identity, sort_keys=True, separators=(",", ":") + ).encode("utf-8")) + for name in ("qpos", "action", "images", "is_pad"): + tensor = sample[name].detach().cpu().contiguous() + digest.update(name.encode("utf-8")) + digest.update(str(tensor.dtype).encode("ascii")) + digest.update(str(tuple(tensor.shape)).encode("ascii")) + digest.update(tensor.numpy().tobytes()) + count += 1 + return { + "sha256": digest.hexdigest(), + "checked_window_count": count, + "fields": [ + "sample_id", "episode_id", "step_idx", "qpos", "action", + "images", "is_pad", + ], + } + + +def _shared_normalization( + episodes, + connection, + frames_snapshot_id, + statistics_version): + """Build one train-only normalization contract for both backends. + + HDF5 supplies state and action moments from successful train episodes. + Versioned Paimon action statistics must match the float64 HDF5 moments, + train scope, frame count, source snapshot, feature name, and ``1e-2`` + standard-deviation floor. + + Returns: + ``(arrays, metadata)`` where arrays are float32 training values and + metadata is JSON-compatible and includes their canonical SHA-256. + """ + normalization, hdf5_metadata = compute_hdf5_normalization(episodes) + action_mean = hdf5_metadata["action_mean"] + action_std = hdf5_metadata["action_std"] + action_count = hdf5_metadata["frame_count"] + row = statistics_row(connection, statistics_version) + if row["source_snapshot_id"] != frames_snapshot_id: + raise ValueError( + "Normalization source snapshot %s differs from frames " + "snapshot %s." + % (row["source_snapshot_id"], frames_snapshot_id)) + if row["source_split"] != "train" or row["frame_count"] != action_count: + raise ValueError( + "Versioned action normalization has the wrong train scope.") + if row["feature_name"] != "action": + raise ValueError("Versioned normalization feature must be action.") + if row["standard_deviation_floor"] != 1e-2: + raise ValueError( + "Versioned normalization must use the 1e-2 std floor.") + stored_mean = np.asarray(row["action_mean"], dtype=np.float64) + stored_std = np.asarray(row["action_std"], dtype=np.float64) + if not ( + np.allclose(stored_mean, action_mean, rtol=1e-10, atol=1e-10) + and np.allclose(stored_std, action_std, rtol=1e-10, atol=1e-10)): + raise ValueError( + "Versioned Paimon action normalization differs from HDF5 source.") + normalization["action_mean"] = stored_mean.astype(np.float32) + normalization["action_std"] = stored_std.astype(np.float32) + serializable = { + name: value.tolist() for name, value in normalization.items() + } + digest = hashlib.sha256(json.dumps( + serializable, sort_keys=True, separators=(",", ":") + ).encode("utf-8")).hexdigest() + return normalization, { + "statistics_version": statistics_version, + "source_split": "train", + "frame_count": action_count, + "standard_deviation_floor": 1e-2, + "values": serializable, + "sha256": digest, + } + + +def _episode_rows(connection): + return connection.get_table(agilex.EPISODES_TABLE).scan().select([ + "episode_id", + "source_key", + "split", + "success", + "frame_count", + ]).to_list() + + +def _validate_source_identity(episodes, rows): + """Match HDF5 discovery to Paimon episodes and return a manifest hash. + + Episode ID, source key, split, and success must match exactly. Paimon's + versioned episode rows contribute frame counts used to build complete + windows. The returned records retain the local HDF5 paths while the hash + covers only portable source metadata. + """ + expected = { + item.episode_id: { + "episode_id": item.episode_id, + "source_key": item.source_key, + "split": item.split, + "success": item.success, + } + for item in episodes + } + actual = { + item["episode_id"]: { + "episode_id": item["episode_id"], + "source_key": item["source_key"], + "split": item["split"], + "success": item["success"], + } + for item in rows + } + if actual != expected or len(actual) != len(rows): + raise ValueError( + "HDF5 and Paimon source identity differ; rebuild or select " + "matching inputs.") + rows_by_id = {item["episode_id"]: item for item in rows} + enriched = [ + _BenchmarkEpisode( + path=item.path, + source_key=item.source_key, + episode_id=item.episode_id, + split=item.split, + success=item.success, + frame_count=rows_by_id[item.episode_id]["frame_count"], + ) + for item in episodes + ] + manifest = sorted([ + { + "episode_id": item.episode_id, + "source_key": item.source_key, + "split": item.split, + "success": item.success, + "frame_count": rows_by_id[item.episode_id]["frame_count"], + } + for item in episodes + ], key=lambda item: item["episode_id"]) + payload = json.dumps(manifest, sort_keys=True, separators=(",", ":")) + return enriched, hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def _select_episode(source_by_id, split, requested, action_horizon): + """Select a successful, split-matching episode long enough for one window. + + An explicit episode is honored when eligible; otherwise the + lexicographically first eligible episode ID is selected. + """ + eligible = { + episode_id: episode + for episode_id, episode in source_by_id.items() + if episode.split == split + and episode.success + and episode.frame_count >= action_horizon + } + if not eligible: + raise ValueError( + "No successful %s episode is long enough for horizon %d." + % (split, action_horizon)) + selected = requested or min(eligible) + if selected not in eligible: + raise ValueError( + "Requested %s episode is missing, unsuccessful, or too short: %s." + % (split, selected)) + return eligible[selected] + + +def _summarize(runs): + """Return median, minimum, and maximum metrics across backend repeats.""" + metrics = ( + "dataset_build_s", + "first_batch_s", + "batch_fetch_samples_per_s", + "fixed_steps_s", + "validation_loss", + "python_peak_allocated_bytes", + "wall_time_s", + ) + result = {"round_count": len(runs)} + for name in metrics: + values = [item[name] for item in runs] + result[name] = { + "median": float(np.median(values)), + "min": float(np.min(values)), + "max": float(np.max(values)), + } + return result + + +def _sample_sequence_sha256(train_episode_id, validation_episode_id, plan): + """Hash episode-qualified sample IDs in measurement/train/validation order.""" + value = { + "batch_fetch": [ + "%s#%d" % (train_episode_id, index) + for index in plan.measurement_indices + ], + "train": [ + "%s#%d" % (train_episode_id, index) + for index in plan.train_indices + ], + "validation": [ + "%s#%d" % (validation_episode_id, index) + for index in plan.validation_indices + ], + } + return hashlib.sha256(json.dumps( + value, sort_keys=True, separators=(",", ":") + ).encode("utf-8")).hexdigest() + + +def _git_head(repository): + try: + return subprocess.check_output( + ["git", "-C", str(repository), "rev-parse", "HEAD"], + stderr=subprocess.DEVNULL, + universal_newlines=True, + ).strip() + except (OSError, subprocess.CalledProcessError): + return "UNKNOWN" + + +def _runtime_environment(repository): + """Return dependency, CPU, thread, and source identity for comparison.""" + source_commit = _git_head(repository) + package_build = build_info.full_version() + if source_commit == "UNKNOWN" and package_build == "UNKNOWN": + raise RuntimeError( + "ACT benchmark cannot determine its source identity.") + return { + "python": platform.python_version(), + "os": platform.platform(), + "machine": platform.machine(), + "cpu_identity": _cpu_identity(), + "cpu_count": os.cpu_count() or 1, + "torch_threads": torch.get_num_threads(), + "torch_interop_threads": torch.get_num_interop_threads(), + "pypaimon_build": package_build, + "numpy": np.__version__, + "pyarrow": pa.__version__, + "h5py": h5py.__version__, + "pillow": PIL.__version__, + "torch": torch.__version__, + "source_commit": source_commit, + } + + +def _cpu_identity(): + """Return the most specific CPU model available from the local OS.""" + if platform.system() == "Darwin": + try: + return subprocess.check_output( + ["sysctl", "-n", "machdep.cpu.brand_string"], + stderr=subprocess.DEVNULL, + universal_newlines=True, + ).strip() + except (OSError, subprocess.CalledProcessError): + pass + identity = platform.processor().strip() + if identity: + return identity + if platform.system() == "Linux": + try: + for line in Path("/proc/cpuinfo").read_text().splitlines(): + if line.startswith(("model name", "Hardware")): + return line.partition(":")[2].strip() + except OSError: + pass + return platform.machine() + + +def _command_argv(): + """Return the invoked Python basename and command-line arguments.""" + import sys + return [os.path.basename(sys.executable)] + list(sys.argv) + + +def _utc_now(): + return datetime.now(timezone.utc).isoformat( + timespec="seconds").replace("+00:00", "Z") diff --git a/paimon-python/pypaimon/multimodal/query.py b/paimon-python/pypaimon/multimodal/query.py index d4491651d8e4..33f0483b519a 100644 --- a/paimon-python/pypaimon/multimodal/query.py +++ b/paimon-python/pypaimon/multimodal/query.py @@ -164,6 +164,70 @@ def to_torch( max_buffer_input_splits=max_buffer_input_splits, ) + def to_contiguous_window_dataset( + self, + *, + window_size, + columns=None, + anchor_columns=None, + group_key="episode_id", + order_key="step_idx", + stride=1, + tail="drop", + column_transforms=None, + pad_values=None, + adapter=None, + blob_parallelism=64): + """Build a snapshot-pinned, map-style Dataset of contiguous rows. + + The Dataset indexes only ``group_key``, ``order_key``, and Paimon row + IDs, then reads projected values on demand. Columns listed in + ``anchor_columns`` are provided to ``column_transforms`` as one-element + lists read from the first row of each window; ``adapter`` receives the + transformed values. ``order_key`` must contain non-null integers that + increase by exactly one within each group. The Dataset sorts rows within + each group and never creates a window across groups. + + Args: + window_size: Number of rows in a complete window. + columns: Value columns to return, excluding the group and order + keys. The scan projection is used when omitted. + anchor_columns: Subset of ``columns`` read only from the window's + first row. + group_key: Column identifying an independent row sequence. + order_key: Integer position column within each group. + stride: Distance between scheduled window starts. + tail: Handling for incomplete final windows: ``drop``, ``pad``, or + ``error``. + column_transforms: Per-column callables applied to value lists. + pad_values: Optional replacement values used by ``tail='pad'``. + adapter: Callable that converts the complete sample mapping. + blob_parallelism: Maximum concurrent BLOB body reads per fetch. + + Returns: + A snapshot-pinned ``ContiguousWindowDataset``. See that class for + padding, mask, transform, and adapter result semantics. + """ + if self._result_factory is not None: + raise TypeError( + "to_contiguous_window_dataset is only supported on scan(), " + "not search queries.") + from pypaimon.multimodal.window_dataset import ContiguousWindowDataset + return ContiguousWindowDataset( + self, + window_size=window_size, + columns=columns, + anchor_columns=anchor_columns, + group_key=group_key, + order_key=order_key, + stride=stride, + tail=tail, + column_transforms=column_transforms, + pad_values=pad_values, + adapter=adapter, + blob_parallelism=blob_parallelism, + ) + def to_ray( self, *, diff --git a/paimon-python/pypaimon/multimodal/window_dataset.py b/paimon-python/pypaimon/multimodal/window_dataset.py new file mode 100644 index 000000000000..6eab515a2cbb --- /dev/null +++ b/paimon-python/pypaimon/multimodal/window_dataset.py @@ -0,0 +1,491 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Snapshot-pinned PyTorch Dataset for contiguous Paimon row windows.""" + +import copy +import operator +from collections import defaultdict +from numbers import Integral + +import torch +from torch.utils.data import Dataset + +from pypaimon.common.options.core_options import CoreOptions +from pypaimon.multimodal.query import ScanQuery +from pypaimon.schema.data_types import is_blob_type, is_map_blob_type +from pypaimon.snapshot.time_travel_util import SCAN_KEYS +from pypaimon.table.special_fields import SpecialFields + + +class ContiguousWindowDataset(Dataset): + """Map-style Dataset which reads fixed row windows on demand. + + The in-memory index contains only group values, order values, and Paimon + row IDs. Each ``__getitem__`` reads the projected rows from the snapshot + resolved while the index was built. Within each group, ``order_key`` must + contain non-null integers that increase by exactly one; rows from different + groups never share a window. ``tail`` controls scheduled anchors whose + remaining rows are shorter than ``window_size``: + + * ``drop`` omits them; + * ``pad`` repeats final values and marks repeats in ``is_pad``; + * ``error`` rejects the dataset. + + The raw result mapping contains scalar group and order values, a + length-``window_size`` Boolean ``is_pad`` tensor, one-element lists for + ``anchor_columns``, and length-``window_size`` lists for other projected + columns. ``anchor_columns`` therefore avoids loading repeated context such + as observation images or initial robot state. ``column_transforms`` then + convert individual column lists before ``adapter`` adapts the complete + mapping to a model-specific contract. + ``blob_parallelism`` controls concurrent BLOB reads for each item or batch. + """ + + _TAIL_POLICIES = ("drop", "pad", "error") + + def __init__( + self, + query, + *, + window_size, + columns=None, + anchor_columns=None, + group_key="episode_id", + order_key="step_idx", + stride=1, + tail="drop", + column_transforms=None, + pad_values=None, + adapter=None, + blob_parallelism=64): + if getattr(query, "_result_factory", None) is not None: + raise TypeError( + "ContiguousWindowDataset is only supported on scan(), " + "not search queries.") + self.window_size = _positive_int(window_size, "window_size") + self.stride = _positive_int(stride, "stride") + if tail not in self._TAIL_POLICIES: + raise ValueError( + "tail must be one of %s; got %r." + % (self._TAIL_POLICIES, tail)) + self.tail = tail + self.group_key = _column(query, group_key, "group_key") + self.order_key = _column(query, order_key, "order_key") + if self.group_key == self.order_key: + raise ValueError("group_key and order_key must name different columns.") + if "is_pad" in (self.group_key, self.order_key): + raise ValueError("group_key and order_key must not be is_pad.") + self.columns = _columns( + query, columns, self.group_key, self.order_key) + self.anchor_columns = _anchor_columns(anchor_columns, self.columns) + anchor_column_set = set(self.anchor_columns) + self._window_columns = [ + name for name in self.columns if name not in anchor_column_set + ] + self.column_transforms = _column_transforms( + column_transforms, self.columns) + self.pad_values = _pad_values(pad_values, self.columns) + if adapter is not None and not callable(adapter): + raise TypeError("adapter must be callable or None.") + self.adapter = adapter + self.blob_parallelism = _positive_int( + blob_parallelism, "blob_parallelism") + + if not query._table.options.row_tracking_enabled(): + raise ValueError( + "ContiguousWindowDataset requires row-tracking.enabled=true.") + + index, snapshot_id = _read_window_index( + query, self.group_key, self.order_key) + self.snapshot_id = snapshot_id + self._table = _pin_table(query._table, snapshot_id) + self._groups, self._anchors = self._build_index(index) + + @classmethod + def from_query(cls, query, **kwargs): + """Build a contiguous-window Dataset from a ``ScanQuery``.""" + return cls(query, **kwargs) + + def __len__(self): + return len(self._anchors) + + def __getitem__(self, index): + """Read one window by map-style Dataset index. + + Negative indices follow Python sequence semantics. The return value is + the pre-adapter mapping described by the class, or the adapter result + when an adapter is configured. + """ + anchor, row_ids = self._resolve_window(index) + rows = self._read_window_rows(row_ids) + anchor_row = ( + self._read_rows(row_ids[:1], self.anchor_columns)[0] + if self.anchor_columns else None + ) + return self._sample(anchor, rows, anchor_row) + + def __getitems__(self, indices): + """Read several Dataset indices while coalescing overlapping row IDs. + + The returned list preserves the requested index order and duplicates. + Coalescing affects only physical reads, not logical sample cardinality. + """ + windows = [self._resolve_window(index) for index in indices] + if not windows: + return [] + row_ids = list(dict.fromkeys( + row_id for _, window_row_ids in windows + for row_id in window_row_ids + )) + rows_by_id = dict(zip(row_ids, self._read_window_rows(row_ids))) + anchor_row_ids = list(dict.fromkeys( + window_row_ids[0] for _, window_row_ids in windows + )) + anchor_rows_by_id = ( + dict(zip( + anchor_row_ids, + self._read_rows(anchor_row_ids, self.anchor_columns), + )) + if self.anchor_columns else {} + ) + return [ + self._sample( + anchor, + [rows_by_id[row_id] for row_id in window_row_ids], + anchor_rows_by_id.get(window_row_ids[0]), + ) + for anchor, window_row_ids in windows + ] + + def _resolve_window(self, index): + index = operator.index(index) + if index < 0: + index += len(self._anchors) + if index < 0 or index >= len(self._anchors): + raise IndexError("window index out of range") + + anchor = self._anchors[index] + group_index, start, valid_count = anchor + row_ids = self._groups[group_index][2] + return anchor, row_ids[start:start + valid_count] + + def _sample(self, anchor, rows, anchor_row=None): + group_index, start, valid_count = anchor + group_key, order_values, _ = self._groups[group_index] + padding_count = self.window_size - valid_count + padding_mask = torch.zeros(self.window_size, dtype=torch.bool) + if padding_count: + padding_mask[valid_count:] = True + sample = { + self.group_key: group_key, + self.order_key: order_values[start], + "is_pad": padding_mask, + } + for name in self.columns: + if name in self.anchor_columns: + values = [copy.deepcopy(anchor_row[name])] + else: + values = [copy.deepcopy(row[name]) for row in rows] + if padding_count and name not in self.anchor_columns: + pad_value = self.pad_values.get(name, values[-1]) + values.extend( + copy.deepcopy(pad_value) for _ in range(padding_count)) + transform = self.column_transforms.get(name) + sample[name] = transform(values) if transform is not None else values + if self.adapter is not None: + return self.adapter(sample) + return sample + + def _build_index(self, index): + """Validate index rows and return grouped row IDs plus window anchors. + + Args: + index: Arrow table containing ``group_key``, ``order_key``, and + Paimon's ``_ROW_ID`` for the resolved snapshot. + + Returns: + ``(groups, anchors)``. Each group stores its key, ordered positions, + and row IDs. Each anchor stores group index, start offset, and the + number of real rows available before optional padding. + """ + group_values = index.column(self.group_key).to_pylist() + order_values = index.column(self.order_key).to_pylist() + row_ids = index.column(SpecialFields.ROW_ID.name).to_pylist() + grouped = defaultdict(list) + for group_key, order_value, row_id in zip( + group_values, order_values, row_ids): + if group_key is None: + raise ValueError("%s must not contain null values." % self.group_key) + if order_value is None: + raise ValueError("%s must not contain null values." % self.order_key) + if isinstance(order_value, bool) or not isinstance(order_value, Integral): + raise ValueError( + "%s must contain integer values." % self.order_key) + try: + grouped[group_key].append((int(order_value), int(row_id))) + except TypeError: + raise ValueError( + "%s values must be hashable." % self.group_key) + + groups = [] + anchors = [] + try: + sorted_groups = sorted(grouped.items(), key=lambda item: item[0]) + except TypeError: + raise ValueError( + "%s values must be mutually orderable." % self.group_key) + for group_key, members in sorted_groups: + try: + members.sort(key=lambda item: item[0]) + except TypeError: + raise ValueError( + "%s values in group %r must be mutually orderable." + % (self.order_key, group_key)) + for previous, current in zip(members, members[1:]): + if previous[0] == current[0]: + raise ValueError( + "Group %s has duplicate order value %r in %s." + % (group_key, current[0], self.order_key)) + if current[0] != previous[0] + 1: + raise ValueError( + "Group %s is not contiguous in %s: %s followed by %s." + % (group_key, self.order_key, + previous[0], current[0])) + + group_index = len(groups) + group_orders = [member[0] for member in members] + group_row_ids = [member[1] for member in members] + groups.append((group_key, group_orders, group_row_ids)) + for start in range(0, len(members), self.stride): + valid_count = min(self.window_size, len(members) - start) + if valid_count < self.window_size: + if self.tail == "drop": + continue + if self.tail == "error": + raise ValueError( + "Group %s has an incomplete window at %s: " + "window_size=%d, available=%d." + % (group_key, group_orders[start], + self.window_size, valid_count)) + anchors.append((group_index, start, valid_count)) + return groups, anchors + + def _read_window_rows(self, row_ids): + if not self._window_columns: + return [{} for _ in row_ids] + return self._read_rows(row_ids, self._window_columns) + + def _read_rows(self, row_ids, columns=None): + """Read projected rows by ID from the pinned snapshot. + + Args: + row_ids: Paimon row IDs to read. Their order and duplicates define + the returned row order. + columns: Projected value columns, or all Dataset columns when + omitted. + + Returns: + A list of row dictionaries aligned one-for-one with ``row_ids``. + The internal ``_ROW_ID`` field is removed, and BLOB descriptors are + resolved to their bodies. + """ + columns = self.columns if columns is None else columns + query = ScanQuery(self._table) + predicate_builder = ( + self._table.new_read_builder() + .with_projection( + [field.name for field in self._table.fields] + + [SpecialFields.ROW_ID.name]) + .new_predicate_builder() + ) + query._predicate = predicate_builder.is_in( + SpecialFields.ROW_ID.name, row_ids) + query._projection = list(columns) + query._include_row_id = True + + blob_columns = [ + field.name for field in self._table.fields + if field.name in columns + and (is_blob_type(field.type) or is_map_blob_type(field.type)) + ] + if blob_columns: + scalar, blobs = query.read_blobs( + blob_columns, parallelism=self.blob_parallelism) + rows = scalar.to_pylist() + for name in blob_columns: + values = blobs[name] + if len(values) != len(rows): + raise RuntimeError( + "BLOB column %s is not row-aligned with a window read." + % name) + for row, value in zip(rows, values): + row[name] = value + else: + rows = query.to_arrow().to_pylist() + + by_row_id = {} + row_id_column = SpecialFields.ROW_ID.name + for row in rows: + row_id = int(row[row_id_column]) + del row[row_id_column] + by_row_id[row_id] = row + missing = [row_id for row_id in row_ids if row_id not in by_row_id] + if missing: + raise RuntimeError( + "Pinned snapshot %s did not return indexed row IDs %s." + % (self.snapshot_id, missing)) + return [by_row_id[row_id] for row_id in row_ids] + + +def _read_window_index(query, group_key, order_key): + index_query = copy.copy(query) + index_query._projection = [group_key, order_key] + index_query._include_row_id = True + read_builder = index_query._configured_read_builder() + plan = read_builder.new_scan().plan() + index = read_builder.new_read().to_arrow(plan.splits()) + if index.num_rows and plan.snapshot_id is None: + raise RuntimeError("Cannot pin the snapshot used to build the window index.") + return index, plan.snapshot_id + + +def _pin_table(table, snapshot_id): + """Pin a table copy to ``snapshot_id``, or reuse it when unresolved.""" + if snapshot_id is None: + return table + scan_keys = set(SCAN_KEYS) + scan_keys.update(option.key() for option in ( + CoreOptions.SCAN_MODE, + CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP, + CoreOptions.SCAN_FILE_CREATION_TIME_MILLIS, + CoreOptions.SCAN_CREATION_TIME_MILLIS, + )) + options = { + key: None for key in scan_keys + if table.options.options.contains_key(key) + } + options[CoreOptions.SCAN_SNAPSHOT_ID.key()] = str(snapshot_id) + return table.copy(options) + + +def _columns(query, columns, group_key, order_key): + available = {field.name for field in query._table.fields} + if columns is None: + if query._projection is None: + columns = [field.name for field in query._table.fields] + else: + columns = list(query._projection) + columns = [name for name in columns + if name not in (group_key, order_key)] + elif isinstance(columns, str): + columns = [columns] + else: + try: + columns = list(columns) + except TypeError: + raise TypeError( + "columns must be a non-empty sequence of column names.") + if not columns: + raise ValueError("columns must contain at least one value column.") + if any(not isinstance(name, str) or not name for name in columns): + raise TypeError("columns must contain only non-empty column names.") + if len(set(columns)) != len(columns): + raise ValueError("columns must not contain duplicates.") + invalid = [name for name in columns if name not in available] + if invalid: + raise ValueError("columns do not exist: %s." % invalid) + reserved = [name for name in columns + if name in (group_key, order_key, "is_pad")] + if reserved: + raise ValueError( + "columns must not include group_key, order_key, or is_pad: %s." + % reserved) + return columns + + +def _anchor_columns(value, columns): + if value is None: + return [] + if isinstance(value, str): + value = [value] + else: + try: + value = list(value) + except TypeError: + raise TypeError( + "anchor_columns must be a sequence of projected column names.") + if any(not isinstance(name, str) or not name for name in value): + raise TypeError( + "anchor_columns must contain only non-empty column names.") + if len(set(value)) != len(value): + raise ValueError("anchor_columns must not contain duplicates.") + invalid = [name for name in value if name not in columns] + if invalid: + raise ValueError( + "anchor_columns must be included in columns: %s." % invalid) + return value + + +def _column_transforms(value, columns): + transforms = _mapping(value, "column_transforms") + _validate_mapping_columns(transforms, columns, "column_transforms") + invalid = [name for name, transform in transforms.items() + if not callable(transform)] + if invalid: + raise TypeError( + "column_transforms values must be callable: %s." % invalid) + return transforms + + +def _pad_values(value, columns): + values = _mapping(value, "pad_values") + _validate_mapping_columns(values, columns, "pad_values") + return values + + +def _mapping(value, name): + if value is None: + return {} + try: + return dict(value) + except (TypeError, ValueError): + raise TypeError("%s must be a mapping or None." % name) + + +def _validate_mapping_columns(value, columns, name): + invalid = [column for column in value if column not in columns] + if invalid: + raise ValueError("%s contains unknown columns: %s." % (name, invalid)) + + +def _column(query, value, name): + if not isinstance(value, str) or not value: + raise TypeError("%s must be a non-empty column name." % name) + available = {field.name for field in query._table.fields} + if value not in available: + raise ValueError("%s column %r does not exist." % (name, value)) + return value + + +def _positive_int(value, name): + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError("%s must be a positive int." % name) + return value + + +__all__ = ["ContiguousWindowDataset"] diff --git a/paimon-python/pypaimon/tests/act_benchmark_test.py b/paimon-python/pypaimon/tests/act_benchmark_test.py new file mode 100644 index 000000000000..9f9fb4357c09 --- /dev/null +++ b/paimon-python/pypaimon/tests/act_benchmark_test.py @@ -0,0 +1,208 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +import pytest + +from pypaimon.benchmark.act.compare import ( + canonical_sha256, + compare_results, + load_result_documents, +) +from pypaimon.benchmark.act.experiment import load_experiment + + +def test_packaged_experiment_contains_a_valid_benchmark_config(): + _require_act_runtime() + from pypaimon.benchmark.act.harness import BenchmarkConfig + + experiment = load_experiment() + + assert experiment["schema_version"] == "act-benchmark-experiment@1" + assert experiment["benchmark_id"] == "robomind-act" + assert BenchmarkConfig(**experiment["config"]).to_dict() == ( + experiment["config"]) + assert experiment["statistics_version"] == ( + "robomind-agilex-joint-position@1") + + +def test_compare_reports_ratio_for_matching_backend_results(): + experiment = {"schema_version": "act-benchmark-experiment@1"} + environment = {"python": "3.10", "machine": "arm64"} + hdf5 = _result("hdf5", experiment, environment, throughput=10.0) + paimon = _result("paimon", experiment, environment, throughput=15.0) + hdf5["summary"]["first_batch_s"] = { + "median": 2.0, "min": 2.0, "max": 2.0} + paimon["summary"]["first_batch_s"] = { + "median": 1.0, "min": 1.0, "max": 1.0} + + comparison = compare_results([hdf5, paimon]) + + assert comparison["status"] == "SUCCEEDED" + assert len(comparison["experiments"]) == 1 + group = comparison["experiments"][0] + assert group["backends"] == ["hdf5", "paimon"] + assert group["metrics"]["batch_fetch_samples_per_s"] == { + "hdf5": 10.0, + "paimon": 15.0, + "paimon_over_hdf5": 1.5, + "preferred": "higher", + } + assert group["metrics"]["first_batch_s"] == { + "hdf5": 2.0, + "paimon": 1.0, + "hdf5_over_paimon": 2.0, + "preferred": "lower", + } + + +def test_compare_rejects_different_tensor_fingerprints(): + experiment = {"schema_version": "act-benchmark-experiment@1"} + environment = {"python": "3.10", "machine": "arm64"} + hdf5 = _result("hdf5", experiment, environment, throughput=10.0) + paimon = _result("paimon", experiment, environment, throughput=15.0) + paimon["tensor_fingerprint"]["sha256"] = "different" + + comparison = compare_results([hdf5, paimon]) + + assert comparison["status"] == "FAILED" + group = comparison["experiments"][0] + assert group["status"] == "FAILED" + assert group["reason"] == "tensor fingerprints differ" + assert group["metrics"] == {} + + +def test_compare_requires_results_from_both_backends(): + experiment = {"schema_version": "act-benchmark-experiment@1"} + environment = {"python": "3.10", "machine": "arm64"} + + comparison = compare_results([ + _result("hdf5", experiment, environment, throughput=10.0), + ]) + + assert comparison["status"] == "INCOMPATIBLE" + group = comparison["experiments"][0] + assert group["status"] == "INCOMPATIBLE" + assert group["reason"] == "both hdf5 and paimon results are required" + assert group["metrics"] == {} + + +def test_load_results_combines_explicit_files_and_directory(tmp_path): + experiment = {"schema_version": "act-benchmark-experiment@1"} + environment = {"python": "3.10", "machine": "arm64"} + hdf5_path = tmp_path / "hdf5.json" + paimon_path = tmp_path / "paimon.json" + ignored_path = tmp_path / "experiment.json" + hdf5_path.write_text(json.dumps( + _result("hdf5", experiment, environment, throughput=10.0))) + paimon_path.write_text(json.dumps( + _result("paimon", experiment, environment, throughput=15.0))) + ignored_path.write_text(json.dumps(experiment)) + + results = load_result_documents( + [hdf5_path], results_dir=tmp_path) + + assert [result["backend"] for result in results] == ["hdf5", "paimon"] + + +def test_compare_groups_multiple_experiments_without_cross_comparing(): + environment = {"python": "3.10", "machine": "arm64"} + results = [] + for seed in (1, 2): + experiment = { + "schema_version": "act-benchmark-experiment@1", + "config": {"seed": seed}, + } + results.extend([ + _result("hdf5", experiment, environment, throughput=10.0), + _result("paimon", experiment, environment, throughput=15.0), + ]) + + comparison = compare_results(results) + + assert comparison["status"] == "SUCCEEDED" + assert len(comparison["experiments"]) == 2 + assert all( + group["backends"] == ["hdf5", "paimon"] + for group in comparison["experiments"] + ) + + +def test_compare_reports_incompatible_runtime_environments(): + experiment = {"schema_version": "act-benchmark-experiment@1"} + hdf5 = _result( + "hdf5", experiment, + {"python": "3.10", "machine": "arm64"}, throughput=10.0) + paimon = _result( + "paimon", experiment, + {"python": "3.11", "machine": "arm64"}, throughput=15.0) + + comparison = compare_results([hdf5, paimon]) + + assert comparison["status"] == "INCOMPATIBLE" + group = comparison["experiments"][0] + assert group["reason"] == "runtime environments differ" + assert len(group["environment_sha256s"]) == 2 + assert group["metrics"] == {} + + +def test_compare_rejects_tampered_result_experiment_hash(): + experiment = {"schema_version": "act-benchmark-experiment@1"} + result = _result( + "hdf5", + experiment, + {"python": "3.10", "machine": "arm64"}, + throughput=10.0, + ) + result["experiment_sha256"] = "tampered" + + with pytest.raises(ValueError, match="experiment SHA-256 differs"): + compare_results([result]) + + +def _require_act_runtime(): + pytest.importorskip("torch") + pytest.importorskip("PIL.Image") + + +def _result(backend, experiment, environment, throughput): + return { + "schema_version": "act-benchmark-result@1", + "status": "SUCCEEDED", + "backend": backend, + "experiment": experiment, + "experiment_sha256": canonical_sha256(experiment), + "environment": environment, + "model": {"implementation": "test-policy", "parameter_count": 1}, + "tensor_fingerprint": { + "sha256": "same-tensors", + "checked_window_count": 2, + }, + "runs": [{ + "round": round_number, + "train_loss": [1.0, 0.5], + "validation_loss": 0.25, + } for round_number in range(1, 4)], + "summary": { + "round_count": 3, + "batch_fetch_samples_per_s": { + "median": throughput, + "min": throughput, + "max": throughput, + }, + }, + } diff --git a/paimon-python/pypaimon/tests/act_runner_test.py b/paimon-python/pypaimon/tests/act_runner_test.py new file mode 100644 index 000000000000..e8f14c674df6 --- /dev/null +++ b/paimon-python/pypaimon/tests/act_runner_test.py @@ -0,0 +1,687 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# ruff: noqa: E402 + +import json +import tracemalloc +from io import BytesIO +from types import SimpleNamespace +from unittest.mock import patch + +import numpy as np +import pytest + + +torch = pytest.importorskip("torch") +Image = pytest.importorskip("PIL.Image") +h5py = pytest.importorskip("h5py") + +import pypaimon.multimodal as pmm +import pypaimon.benchmark.act.harness as act_harness +import pypaimon.benchmark.act.__main__ as act_cli +import pypaimon.benchmark.act.runner as act_runner +from pypaimon.benchmark.act.runner import ( + BenchmarkConfig, + prepare_experiment, + run_experiment, +) +from pypaimon.benchmark.act.experiment import load_experiment +from pypaimon.benchmark.act.compare import canonical_sha256, compare_results +from pypaimon.benchmark.act.harness import build_window_plan, run_backend +from pypaimon.benchmark.act.hdf5 import Hdf5ACTWindowDataset +from pypaimon.benchmark.act.paimon import ( + ACTION_COLUMNS, + IMAGE_COLUMNS, + QPOS_COLUMNS, + create_datasets as create_paimon_datasets, + latest_snapshot_id, +) +from pypaimon.multimodal.query import ScanQuery +from pypaimon.multimodal.window_dataset import ContiguousWindowDataset +from pypaimon.sample import robomind_agilex as agilex + + +def test_logical_batches_coalesce_one_physical_fetch(): + class BatchDataset: + def __init__(self): + self.calls = [] + + def __getitems__(self, indices): + self.calls.append(list(indices)) + return [{"value": torch.tensor(index)} for index in indices] + + dataset = BatchDataset() + + batches = list(act_harness._iter_logical_batches( + dataset, + tuple(range(8)), + logical_batch_size=2, + fetch_batches=4, + )) + + assert dataset.calls == [list(range(8))] + assert [batch["value"].tolist() for batch in batches] == [ + [0, 1], [2, 3], [4, 5], [6, 7], + ] + + +def test_logical_batches_reject_incomplete_batch_tail(): + class BatchDataset: + def __getitems__(self, indices): + return [{"value": torch.tensor(index)} for index in indices] + + with pytest.raises(ValueError, match="complete logical batches"): + list(act_harness._iter_logical_batches( + BatchDataset(), + tuple(range(9)), + logical_batch_size=2, + fetch_batches=4, + )) + + +def test_backend_times_without_tracemalloc_and_measures_memory_separately(): + states = [] + config = BenchmarkConfig( + seed=11, + action_horizon=1, + batch_size=1, + optimizer_steps=1, + image_height=2, + image_width=2, + warmup_batches=1, + timed_batches=1, + rounds=3, + ) + + class TracingDataset(torch.utils.data.Dataset): + def __len__(self): + return 2 + + def __getitem__(self, index): + states.append(tracemalloc.is_tracing()) + return { + "sample_id": "episode-a#%d" % index, + "episode_id": "episode-a", + "step_idx": index, + "qpos": torch.zeros(14), + "action": torch.zeros((1, 14)), + "images": torch.zeros((3, 3, 2, 2)), + "is_pad": torch.zeros(1, dtype=torch.bool), + } + + dataset = TracingDataset() + plan = build_window_plan(len(dataset), len(dataset), config) + result = run_backend( + "test", + 1, + lambda: (dataset, dataset), + plan, + config, + "sequence-sha256", + policy_factory=_policy_factory, + ) + + assert states[0] is False + assert states[-1] is True + assert result["peak_memory_measurement"] == ( + "python-tracemalloc-separate-dataset-first-batch") + + +def test_backend_coalesces_timed_batch_fetches(): + config = BenchmarkConfig( + seed=11, + action_horizon=1, + batch_size=2, + optimizer_steps=1, + image_height=2, + image_width=2, + warmup_batches=1, + timed_batches=4, + fetch_batches=4, + rounds=3, + ) + + clock = SimpleNamespace(value=0.0) + + class BatchDataset(torch.utils.data.Dataset): + def __init__(self): + self.calls = [] + + def __len__(self): + return 16 + + def __getitem__(self, index): + return { + "sample_id": "episode-a#%d" % index, + "episode_id": "episode-a", + "step_idx": index, + "qpos": torch.zeros(14), + "action": torch.zeros((1, 14)), + "images": torch.zeros((3, 3, 2, 2)), + "is_pad": torch.zeros(1, dtype=torch.bool), + } + + def __getitems__(self, indices): + clock.value += 0.25 + self.calls.append(list(indices)) + return [self[index] for index in indices] + + dataset = BatchDataset() + plan = build_window_plan(len(dataset), len(dataset), config) + + def validate(_batch, _config): + clock.value += 1.0 + + with ( + patch.object(act_harness, "_measure_python_peak", return_value=0), + patch.object( + act_harness.time, "monotonic", side_effect=lambda: clock.value), + patch.object(act_harness, "validate_act_batch", side_effect=validate), + ): + result = run_backend( + "test", + 1, + lambda: (dataset, dataset), + plan, + config, + "sequence-sha256", + policy_factory=_policy_factory, + ) + + assert dataset.calls == [ + list(plan.measurement_indices[:2]), + list(plan.measurement_indices[2:10]), + list(plan.train_indices), + list(plan.validation_indices), + ] + assert result["batch_fetch_s"] == 0.25 + + +def _jpeg(value): + buffer = BytesIO() + Image.fromarray(np.full((8, 10, 3), value, dtype=np.uint8)).save( + buffer, format="JPEG") + return np.frombuffer(buffer.getvalue(), dtype=np.uint8) + + +def _write_episode(root, split, name, offset, frames=6): + path = (root / "13_packbowl" / "success_episodes" / split / name + / "data" / "trajectory.hdf5") + path.parent.mkdir(parents=True) + with h5py.File(path, "w") as h5: + h5.create_dataset("language_raw", data=[b"pack the bowl"]) + h5.create_dataset( + "language_distilbert", + data=np.zeros((1, 1, 768), dtype=np.float16), + ) + for index, (_, hdf5_path) in enumerate(agilex.NUMERIC_FIELDS): + values = np.arange(frames * 7, dtype=np.float64).reshape(frames, 7) + h5.create_dataset(hdf5_path, data=values + offset + index * 100) + variable = h5py.vlen_dtype(np.dtype("uint8")) + for image_index, (_, hdf5_path) in enumerate(agilex.IMAGE_FIELDS): + dataset = h5.create_dataset(hdf5_path, (frames,), dtype=variable) + for frame_index in range(frames): + dataset[frame_index] = _jpeg( + offset + image_index + frame_index) + return path + + +@pytest.fixture +def benchmark_input(tmp_path, monkeypatch): + root = tmp_path / "input" + _write_episode(root, "train", "train-a", 1) + _write_episode(root, "train", "train-b", 11) + _write_episode(root, "val", "val-a", 21) + warehouse = tmp_path / "warehouse" + monkeypatch.setattr(agilex, "TABLE_OPTIONS", { + **agilex.TABLE_OPTIONS, + "vector.file.format": "parquet", + }) + agilex.ingest_local(root, warehouse, batch_size=2) + agilex.backfill_canonical_action( + warehouse, statistics_version="act-test@1") + return root, warehouse + + +class _Policy(torch.nn.Module): + + def __init__(self): + super().__init__() + self.scale = torch.nn.Parameter(torch.tensor(0.0)) + + def forward(self, batch): + assert self.training + target = batch["action"].mean() + batch["observation.state"].mean() + loss = (self.scale - target).square() + return loss, { + "l1_loss": loss.detach(), + "kld_loss": torch.tensor(0.0), + } + + +def _policy_factory(config): + return _Policy(), { + "implementation": "test-policy", + "chunk_size": config.action_horizon, + "parameter_count": 1, + } + + +def test_prepare_writes_resolved_experiment(benchmark_input, tmp_path): + input_root, warehouse = benchmark_input + definition = load_experiment() + definition["statistics_version"] = "act-test@1" + definition["config"].update({ + "seed": 17, + "action_horizon": 3, + "batch_size": 2, + "optimizer_steps": 2, + "image_height": 8, + "image_width": 10, + "timed_batches": 2, + }) + output = tmp_path / "experiment.json" + + experiment = prepare_experiment( + input_root, warehouse, output, definition=definition) + + assert json.loads(output.read_text()) == experiment + assert experiment["schema_version"] == "act-benchmark-experiment@1" + assert experiment["train_episode_id"] == "train-a" + assert experiment["validation_episode_id"] == "val-a" + assert experiment["source"]["episodes"][0]["frame_count"] == 6 + assert len(experiment["source"]["sha256"]) == 64 + assert len(experiment["normalization"]["sha256"]) == 64 + assert len(experiment["window_plan"]["sha256"]) == 64 + assert experiment["paimon"]["frames_snapshot_id"] > 0 + + +def test_independent_backend_results_preserve_tensor_and_loss_parity( + benchmark_input, tmp_path): + input_root, warehouse = benchmark_input + definition = load_experiment() + definition["statistics_version"] = "act-test@1" + definition["config"].update({ + "seed": 17, + "action_horizon": 3, + "batch_size": 2, + "optimizer_steps": 2, + "image_height": 8, + "image_width": 10, + "timed_batches": 2, + }) + experiment_path = tmp_path / "experiment.json" + prepare_experiment( + input_root, warehouse, experiment_path, definition=definition) + + hdf5_path = tmp_path / "hdf5-result.json" + hdf5_result = run_experiment( + "hdf5", + experiment_path, + hdf5_path, + input_root=input_root, + policy_factory=_policy_factory, + ) + paimon_path = tmp_path / "paimon-result.json" + paimon_result = run_experiment( + "paimon", + experiment_path, + paimon_path, + warehouse=warehouse, + policy_factory=_policy_factory, + ) + + assert json.loads(hdf5_path.read_text()) == hdf5_result + assert json.loads(paimon_path.read_text()) == paimon_result + assert hdf5_result["schema_version"] == "act-benchmark-result@1" + assert paimon_result["schema_version"] == "act-benchmark-result@1" + assert hdf5_result["experiment"] == paimon_result["experiment"] + assert hdf5_result["tensor_fingerprint"] == ( + paimon_result["tensor_fingerprint"]) + assert [run["train_loss"] for run in hdf5_result["runs"]] == [ + run["train_loss"] for run in paimon_result["runs"] + ] + assert [run["validation_loss"] for run in hdf5_result["runs"]] == [ + run["validation_loss"] for run in paimon_result["runs"] + ] + comparison = compare_results([hdf5_result, paimon_result]) + assert comparison["status"] == "SUCCEEDED" + assert comparison["experiments"][0]["backends"] == ["hdf5", "paimon"] + + +def test_backends_match_the_golden_act_window_contract(benchmark_input): + input_root, warehouse = benchmark_input + normalization = { + "qpos_mean": np.zeros(14, dtype=np.float32), + "qpos_std": np.ones(14, dtype=np.float32), + "action_mean": np.zeros(14, dtype=np.float32), + "action_std": np.ones(14, dtype=np.float32), + } + hdf5 = Hdf5ACTWindowDataset( + SimpleNamespace( + path=(input_root / "13_packbowl" / "success_episodes" / "train" + / "train-a" / "data" / "trajectory.hdf5"), + episode_id="train-a", + frame_count=6, + ), + normalization, + action_horizon=3, + ) + connection = pmm.connect( + database=agilex.DEFAULT_DATABASE, + options={"warehouse": str(warehouse)}, + ) + frames = connection.get_table(agilex.FRAMES_TABLE) + paimon, _ = create_paimon_datasets( + frames, + latest_snapshot_id(frames), + "train-a", + "val-a", + normalization, + BenchmarkConfig(action_horizon=3), + ) + + expected = hdf5[1] + actual = paimon[1] + + assert set(expected) == { + "sample_id", "episode_id", "step_idx", "qpos", "action", + "images", "is_pad", + } + assert expected["sample_id"] == "train-a#1" + assert expected["episode_id"] == "train-a" + assert expected["step_idx"] == 1 + assert torch.equal(expected["qpos"], torch.tensor( + list(range(408, 415)) + list(range(508, 515)), + dtype=torch.float32, + )) + assert torch.equal(expected["action"], torch.tensor([ + list(range(1208, 1215)) + list(range(1308, 1315)), + list(range(1215, 1222)) + list(range(1315, 1322)), + list(range(1222, 1229)) + list(range(1322, 1329)), + ], dtype=torch.float32)) + assert torch.allclose( + expected["images"][:, :, 0, 0], + torch.tensor([[2 / 255] * 3, [3 / 255] * 3, [4 / 255] * 3]), + ) + assert not expected["is_pad"].any() + for name in ("qpos", "action", "images", "is_pad"): + assert torch.equal(expected[name], actual[name]) + for name in ("sample_id", "episode_id", "step_idx"): + assert expected[name] == actual[name] + + +def test_hdf5_window_index_bounds(tmp_path): + path = _write_episode(tmp_path, "train", "train-a", 1) + dataset = Hdf5ACTWindowDataset( + SimpleNamespace( + path=path, + episode_id="train-a", + frame_count=6, + ), + { + "qpos_mean": np.zeros(14, dtype=np.float32), + "qpos_std": np.ones(14, dtype=np.float32), + "action_mean": np.zeros(14, dtype=np.float32), + "action_std": np.ones(14, dtype=np.float32), + }, + action_horizon=3, + ) + + assert dataset[-1]["sample_id"] == "train-a#3" + with pytest.raises(IndexError): + dataset[-len(dataset) - 1] + with pytest.raises(IndexError): + dataset[len(dataset)] + + +def test_paimon_run_rejects_normalization_not_recorded_in_statistics( + benchmark_input, tmp_path): + input_root, warehouse = benchmark_input + definition = load_experiment() + definition["statistics_version"] = "act-test@1" + definition["config"].update({ + "action_horizon": 3, + "batch_size": 1, + "optimizer_steps": 1, + "image_height": 8, + "image_width": 10, + }) + experiment_path = tmp_path / "experiment.json" + experiment = prepare_experiment( + input_root, warehouse, experiment_path, definition=definition) + experiment["normalization"]["values"]["action_mean"][0] += 1 + experiment["normalization"]["sha256"] = canonical_sha256( + experiment["normalization"]["values"]) + experiment_path.write_text(json.dumps(experiment)) + + with pytest.raises(ValueError, match="normalization differs"): + run_experiment( + "paimon", + experiment_path, + tmp_path / "must-not-exist.json", + warehouse=warehouse, + policy_factory=_policy_factory, + ) + + +def test_run_rejects_tampered_source_manifest(benchmark_input, tmp_path): + input_root, warehouse = benchmark_input + definition = load_experiment() + definition["statistics_version"] = "act-test@1" + definition["config"].update({ + "action_horizon": 3, + "batch_size": 1, + "optimizer_steps": 1, + "image_height": 8, + "image_width": 10, + }) + experiment_path = tmp_path / "experiment.json" + experiment = prepare_experiment( + input_root, warehouse, experiment_path, definition=definition) + experiment["source"]["episodes"][0]["frame_count"] += 1 + experiment_path.write_text(json.dumps(experiment)) + + with pytest.raises(ValueError, match="source-manifest hash differs"): + run_experiment( + "hdf5", + experiment_path, + tmp_path / "must-not-exist.json", + input_root=input_root, + policy_factory=_policy_factory, + ) + + +def test_run_rejects_config_not_used_to_build_window_plan( + benchmark_input, tmp_path): + input_root, warehouse = benchmark_input + definition = load_experiment() + definition["statistics_version"] = "act-test@1" + definition["config"].update({ + "action_horizon": 3, + "batch_size": 1, + "optimizer_steps": 1, + "image_height": 8, + "image_width": 10, + }) + experiment_path = tmp_path / "experiment.json" + experiment = prepare_experiment( + input_root, warehouse, experiment_path, definition=definition) + experiment["config"]["seed"] += 1 + experiment_path.write_text(json.dumps(experiment)) + + with pytest.raises(ValueError, match="window plan was not built"): + run_experiment( + "hdf5", + experiment_path, + tmp_path / "must-not-exist.json", + input_root=input_root, + policy_factory=_policy_factory, + ) + + +def test_paimon_windows_are_lazy_snapshot_pinned_and_vortex_independent( + benchmark_input, tmp_path): + input_root, warehouse = benchmark_input + definition = load_experiment() + definition["statistics_version"] = "act-test@1" + definition["config"]["action_horizon"] = 3 + experiment = prepare_experiment( + input_root, + warehouse, + tmp_path / "experiment.json", + definition=definition, + ) + normalization = { + name: np.asarray(value, dtype=np.float32) + for name, value in experiment["normalization"]["values"].items() + } + connection = pmm.connect( + database=agilex.DEFAULT_DATABASE, + options={"warehouse": str(warehouse)}, + ) + frames = connection.get_table(agilex.FRAMES_TABLE) + assert frames.raw_table.table_schema.options["vector.file.format"] == ( + "parquet") + snapshot_id = latest_snapshot_id(frames) + + original = ScanQuery._fetch_bodies + with patch.object( + ScanQuery, "_fetch_bodies", side_effect=original) as fetch: + train, validation = create_paimon_datasets( + frames, + snapshot_id, + "train-a", + "val-a", + normalization, + BenchmarkConfig( + action_horizon=3, + batch_size=1, + optimizer_steps=1, + image_height=8, + image_width=10, + rounds=3, + ), + ) + assert fetch.call_count == 0 + assert isinstance(train, ContiguousWindowDataset) + assert isinstance(validation, ContiguousWindowDataset) + with patch.object( + train, "_read_rows", wraps=train._read_rows) as read_rows: + sample_before_append = train[0] + assert [call.args[1] for call in read_rows.call_args_list] == [ + list(ACTION_COLUMNS), + list(QPOS_COLUMNS + IMAGE_COLUMNS), + ] + assert [len(call.args[0]) for call in read_rows.call_args_list] == [3, 1] + assert fetch.call_count == 1 + assert { + name: len(fetch.call_args.args[1][name]) + for name in IMAGE_COLUMNS + } == {name: 1 for name in IMAGE_COLUMNS} + + scalar, blobs = frames.scan().where( + "episode_id = 'train-a' AND frame_index = 5" + ).read_blobs(IMAGE_COLUMNS) + appended = scalar.to_pylist()[0] + appended["frame_index"] = 6 + for name in IMAGE_COLUMNS: + appended[name] = blobs[name][0] + frames.add([appended]) + + assert train.snapshot_id == snapshot_id + assert validation.snapshot_id == snapshot_id + assert latest_snapshot_id(frames) != snapshot_id + assert len(train) == 4 + sample_after_append = train[0] + for name in ("qpos", "action", "images", "is_pad"): + assert torch.equal( + sample_before_append[name], sample_after_append[name]) + + +def test_requires_at_least_three_measurement_rounds(): + with pytest.raises(ValueError, match="rounds must be at least 3"): + BenchmarkConfig(rounds=2) + + +def test_fetch_batches_must_be_positive(): + assert BenchmarkConfig(fetch_batches=4).fetch_batches == 4 + with pytest.raises(ValueError, match="fetch_batches must be a positive int"): + BenchmarkConfig(fetch_batches=0) + + +def test_cli_exposes_prepare_run_and_compare_contracts(capsys): + with pytest.raises(SystemExit): + act_cli.main(["prepare", "--help"]) + + prepare_help = capsys.readouterr().out + assert "--experiment" in prepare_help + assert "--fetch-batches" in prepare_help + + with pytest.raises(SystemExit): + act_cli.main(["run", "--help"]) + + run_help = capsys.readouterr().out + assert "--backend" in run_help + assert "--experiment" in run_help + assert "--results-dir" in run_help + + with pytest.raises(SystemExit): + act_cli.main(["compare", "--help"]) + + assert "--results-dir" in capsys.readouterr().out + + +def test_cli_requires_python_3_10_or_newer(): + with pytest.raises(RuntimeError, match="Python 3.10 or newer"): + act_cli._require_supported_python((3, 9)) + + act_cli._require_supported_python((3, 10)) + + +def test_automatic_artifact_paths_do_not_overwrite_same_second(tmp_path): + with patch.object(act_cli, "datetime") as now: + now.now.return_value.strftime.return_value = "20260901T120000Z" + + first = act_cli._artifact_path(tmp_path, "robomind-act-hdf5") + second = act_cli._artifact_path(tmp_path, "robomind-act-hdf5") + + assert first != second + assert first.parent == tmp_path + assert second.parent == tmp_path + + +def test_runtime_environment_uses_package_identity_outside_git_checkout( + tmp_path): + with patch.object(act_runner, "_git_head", return_value="UNKNOWN"): + environment = act_runner._runtime_environment(tmp_path) + + assert environment["source_commit"] == "UNKNOWN" + assert environment["pypaimon_build"] != "UNKNOWN" + assert environment["cpu_identity"] + assert environment["cpu_count"] > 0 + assert environment["torch_threads"] > 0 + assert environment["torch_interop_threads"] > 0 + assert all(environment[name] for name in ( + "numpy", "pyarrow", "h5py", "pillow")) + + with ( + patch.object(act_runner, "_git_head", return_value="UNKNOWN"), + patch.object(act_runner.build_info, "full_version", return_value="UNKNOWN"), + pytest.raises(RuntimeError, match="source identity"), + ): + act_runner._runtime_environment(tmp_path) diff --git a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py new file mode 100644 index 000000000000..389d4bb5f36a --- /dev/null +++ b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py @@ -0,0 +1,472 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import os +import pickle +import shutil +import tempfile +import unittest +from unittest.mock import patch + +import pyarrow as pa +import torch + +import pypaimon.multimodal as pmm +from pypaimon.multimodal.query import ScanQuery +from pypaimon.multimodal.window_dataset import ContiguousWindowDataset + + +_TABLE_OPTIONS = { + "row-tracking.enabled": "true", + "data-evolution.enabled": "true", + "deletion-vectors.enabled": "true", + "file.format": "parquet", + "vector.file.format": "parquet", +} + + +class _TensorColumnTransform: + + def __call__(self, values): + return torch.tensor(values, dtype=torch.int64) + + +class _WindowAdapter: + + def __call__(self, sample): + return { + "episode": sample["episode"], + "start": sample["step"], + "values": sample["value"], + "padding_mask": sample["is_pad"], + } + + +class ContiguousWindowDatasetTest(unittest.TestCase): + + def setUp(self): + self.temp_dir = tempfile.mkdtemp(prefix="pypaimon_windows_") + self.conn = pmm.connect(options={ + "warehouse": os.path.join(self.temp_dir, "warehouse"), + }) + + def tearDown(self): + shutil.rmtree(self.temp_dir, ignore_errors=True) + + @staticmethod + def _schema(): + return pa.schema([ + pa.field("episode", pa.string(), nullable=False), + pa.field("step", pa.int32(), nullable=False), + pa.field("value", pa.int32(), nullable=False), + pa.field("payload", pa.large_binary(), nullable=False), + ]) + + @staticmethod + def _row(episode, step): + return { + "episode": episode, + "step": step, + "value": step + (100 if episode == "episode-b" else 0), + "payload": ("%s-%d" % (episode, step)).encode(), + } + + def _table(self, name="frames"): + table = self.conn.create_table( + name, schema=self._schema(), options=_TABLE_OPTIONS) + table.add([ + self._row("episode-b", 2), + self._row("episode-a", 1), + self._row("episode-b", 0), + self._row("episode-a", 0), + self._row("episode-b", 3), + self._row("episode-b", 1), + ]) + return table + + @staticmethod + def _dataset(table, **kwargs): + return ( + table.scan() + .to_contiguous_window_dataset( + window_size=3, + columns=["value", "payload"], + group_key="episode", + order_key="step", + **kwargs, + ) + ) + + def test_sorts_rows_and_never_crosses_episode_boundaries(self): + dataset = self._dataset(self._table()) + + self.assertIsInstance(dataset, torch.utils.data.Dataset) + self.assertEqual(2, len(dataset)) + self.assertIsInstance(dataset.snapshot_id, int) + self.assertNotIn("_episodes", vars(dataset)) + + first = dataset[0] + second = dataset[1] + self.assertEqual("episode-b", first["episode"]) + self.assertEqual(0, first["step"]) + self.assertEqual([100, 101, 102], first["value"]) + self.assertEqual([101, 102, 103], second["value"]) + self.assertFalse(first["is_pad"].any()) + self.assertEqual({"episode-b"}, { + window["episode"] for window in (first, second) + }) + + def test_reads_blob_payloads_only_when_a_window_is_requested(self): + table = self._table() + original = ScanQuery._fetch_bodies + with patch.object(ScanQuery, "_fetch_bodies", side_effect=original) as fetch: + dataset = self._dataset(table) + self.assertEqual(0, fetch.call_count) + + sample = dataset[0] + + self.assertEqual(1, fetch.call_count) + self.assertEqual(3, len(fetch.call_args.args[1]["payload"])) + self.assertEqual( + [b"episode-b-0", b"episode-b-1", b"episode-b-2"], + sample["payload"], + ) + + def test_reads_map_blob_payloads_for_map_only_and_mixed_windows(self): + schema = pa.schema([ + pa.field("episode", pa.string(), nullable=False), + pa.field("step", pa.int32(), nullable=False), + pa.field("payload", pa.large_binary()), + pa.field( + "attachments", + pa.map_(pa.string(), pa.large_binary()), + ), + ]) + table = self.conn.create_table( + "map_blobs", + schema=schema, + options=_TABLE_OPTIONS, + ) + table.add(pa.Table.from_pylist([ + { + "episode": "episode-a", + "step": 0, + "payload": b"scalar-0", + "attachments": { + "body": b"map-0", "empty": b"", "null": None}, + }, + { + "episode": "episode-a", + "step": 1, + "payload": b"scalar-1", + "attachments": None, + }, + ], schema=schema)) + + def window(columns): + return table.scan().to_contiguous_window_dataset( + window_size=2, + columns=columns, + group_key="episode", + order_key="step", + )[0] + + map_only = window(["attachments"]) + mixed = window(["payload", "attachments"]) + + self.assertEqual( + {"body": b"map-0", "empty": b"", "null": None}, + dict(map_only["attachments"][0]), + ) + self.assertIsNone(map_only["attachments"][1]) + self.assertEqual([b"scalar-0", b"scalar-1"], mixed["payload"]) + self.assertEqual(map_only["attachments"], mixed["attachments"]) + + def test_anchor_columns_read_only_the_window_anchor(self): + table = self._table() + original = ScanQuery._fetch_bodies + with patch.object(ScanQuery, "_fetch_bodies", side_effect=original) as fetch: + dataset = self._dataset(table, anchor_columns=["payload"]) + + sample = dataset[0] + + self.assertEqual([100, 101, 102], sample["value"]) + self.assertEqual([b"episode-b-0"], sample["payload"]) + self.assertEqual(1, fetch.call_count) + self.assertEqual(1, len(fetch.call_args.args[1]["payload"])) + + def test_plural_access_coalesces_overlapping_window_reads(self): + dataset = self._dataset( + self._table(), anchor_columns=["payload"]) + + with patch.object( + dataset, "_read_rows", wraps=dataset._read_rows) as read: + actual = dataset.__getitems__([1, 0, 1]) + + self.assertEqual(2, read.call_count) + self.assertEqual(4, len(read.call_args_list[0].args[0])) + self.assertEqual(["value"], read.call_args_list[0].args[1]) + self.assertEqual(2, len(read.call_args_list[1].args[0])) + self.assertEqual(["payload"], read.call_args_list[1].args[1]) + self.assertEqual( + [("episode-b", 1), ("episode-b", 0), ("episode-b", 1)], + [(sample["episode"], sample["step"]) for sample in actual], + ) + self.assertEqual( + [[101, 102, 103], [100, 101, 102], [101, 102, 103]], + [sample["value"] for sample in actual], + ) + self.assertEqual( + [[b"episode-b-1"], [b"episode-b-0"], [b"episode-b-1"]], + [sample["payload"] for sample in actual], + ) + + def test_plural_access_isolates_mutable_cells_between_samples(self): + table = self.conn.create_table( + "mutable_cells", + schema=pa.schema([ + pa.field("episode", pa.string(), nullable=False), + pa.field("step", pa.int32(), nullable=False), + pa.field("values", pa.list_(pa.int32()), nullable=False), + ]), + options=_TABLE_OPTIONS, + ) + table.add([ + {"episode": "episode-a", "step": step, "values": [step]} + for step in range(3) + ]) + + def mutate(values): + for value in values: + value.append(99) + return values + + dataset = table.scan().to_contiguous_window_dataset( + window_size=2, + columns=["values"], + group_key="episode", + order_key="step", + column_transforms={"values": mutate}, + ) + + batched = dataset.__getitems__([0, 1, 0]) + singles = [dataset[index] for index in (0, 1, 0)] + + self.assertEqual( + [sample["values"] for sample in singles], + [sample["values"] for sample in batched], + ) + + def test_pad_tail_repeats_last_row_and_marks_real_padding(self): + dataset = self._dataset( + self._table(), tail="pad", pad_values={"value": -1}) + + self.assertEqual(6, len(dataset)) + short_tail = dataset[1] + long_tail = dataset[-1] + self.assertEqual("episode-a", short_tail["episode"]) + self.assertEqual([1, -1, -1], short_tail["value"]) + self.assertEqual( + [b"episode-a-1"] * 3, short_tail["payload"]) + self.assertEqual([False, True, True], short_tail["is_pad"].tolist()) + self.assertEqual("episode-b", long_tail["episode"]) + self.assertEqual([103, -1, -1], long_tail["value"]) + self.assertEqual([False, True, True], long_tail["is_pad"].tolist()) + + def test_error_tail_rejects_an_incomplete_scheduled_window(self): + with self.assertRaisesRegex( + ValueError, "episode-a.*incomplete.*window_size=3"): + self._dataset(self._table(), tail="error") + + def test_stride_controls_scheduled_window_anchors(self): + dataset = self._dataset(self._table(), stride=2, tail="pad") + + self.assertEqual( + [("episode-a", 0), ("episode-b", 0), ("episode-b", 2)], + [(dataset[index]["episode"], dataset[index]["step"]) + for index in range(len(dataset))], + ) + self.assertEqual( + [False, False, True], dataset[-1]["is_pad"].tolist()) + + def test_rejects_missing_and_duplicate_order_keys_within_a_group(self): + gapped = self.conn.create_table( + "gapped", schema=self._schema(), options=_TABLE_OPTIONS) + gapped.add([ + self._row("episode-a", 0), + self._row("episode-a", 2), + ]) + + with self.assertRaisesRegex( + ValueError, "episode-a.*not contiguous.*0.*2"): + self._dataset(gapped) + + table = self.conn.create_table( + "duplicates", schema=self._schema(), options=_TABLE_OPTIONS) + table.add([ + self._row("episode-a", 0), + self._row("episode-a", 0), + self._row("episode-a", 1), + ]) + + with self.assertRaisesRegex( + ValueError, "episode-a.*duplicate.*order.*0"): + self._dataset(table) + + def test_pins_snapshot_for_later_on_demand_reads(self): + table = self._table() + dataset = self._dataset(table) + snapshot_id = dataset.snapshot_id + + table.add([self._row("episode-b", 4)]) + + self.assertEqual(snapshot_id, dataset.snapshot_id) + self.assertNotEqual( + snapshot_id, table.raw_table.snapshot_manager().get_latest_snapshot().id) + self.assertEqual(2, len(dataset)) + self.assertEqual([101, 102, 103], dataset[-1]["value"]) + + def test_snapshot_pin_clears_scan_mode_before_on_demand_reads(self): + table = self._table() + query = ScanQuery(table.raw_table.copy({"scan.mode": "latest-full"})) + + dataset = query.to_contiguous_window_dataset( + window_size=3, + columns=["value", "payload"], + group_key="episode", + order_key="step", + ) + + self.assertEqual([100, 101, 102], dataset[0]["value"]) + + def test_pickle_round_trip_preserves_snapshot_and_window(self): + dataset = self._dataset( + self._table(), anchor_columns=["payload"]) + expected = dataset[-1] + + restored = pickle.loads(pickle.dumps(dataset)) + + self.assertEqual(dataset.snapshot_id, restored.snapshot_id) + self.assertEqual( + dataset.snapshot_id, + restored._table.options.scan_snapshot_id(), + ) + actual = restored[-1] + self.assertEqual(expected["episode"], actual["episode"]) + self.assertEqual(expected["step"], actual["step"]) + self.assertEqual(expected["value"], actual["value"]) + self.assertEqual(expected["payload"], actual["payload"]) + self.assertTrue(torch.equal(expected["is_pad"], actual["is_pad"])) + + def test_projection_filter_transform_and_dataloader_workers(self): + table = self._table() + dataset = ( + table.scan() + .where("episode = 'episode-b'") + .select(["value"]) + .to_contiguous_window_dataset( + window_size=2, + group_key="episode", + order_key="step", + column_transforms={"value": _TensorColumnTransform()}, + adapter=_WindowAdapter(), + ) + ) + + loader = torch.utils.data.DataLoader( + dataset, batch_size=2, shuffle=False, num_workers=2) + batches = list(loader) + + self.assertEqual(2, len(batches)) + self.assertEqual(torch.int64, batches[0]["values"].dtype) + self.assertEqual((2, 2), tuple(batches[0]["values"].shape)) + self.assertEqual(torch.bool, batches[0]["padding_mask"].dtype) + self.assertEqual([0, 1, 2], [ + start for batch in batches for start in batch["start"].tolist() + ]) + self.assertEqual( + [[100, 101], [101, 102], [102, 103]], + [values for batch in batches for values in batch["values"].tolist()], + ) + self.assertTrue(all( + episode == "episode-b" + for batch in batches for episode in batch["episode"] + )) + + def test_default_keys_and_public_from_query_entry_point(self): + table = self.conn.create_table( + "default_keys", + schema=pa.schema([ + pa.field("episode_id", pa.string(), nullable=False), + pa.field("step_idx", pa.int32(), nullable=False), + pa.field("value", pa.int32(), nullable=False), + ]), + options=_TABLE_OPTIONS, + ) + table.add([ + {"episode_id": "episode-a", "step_idx": 0, "value": 10}, + {"episode_id": "episode-a", "step_idx": 1, "value": 11}, + ]) + + dataset = ContiguousWindowDataset.from_query( + table.scan().select(["value"]), window_size=2) + + self.assertEqual(1, len(dataset)) + self.assertEqual("episode-a", dataset[0]["episode_id"]) + self.assertEqual(0, dataset[0]["step_idx"]) + self.assertEqual([10, 11], dataset[0]["value"]) + + def test_validates_configuration_and_scan_only_contract(self): + table = self._table() + query = table.scan() + for name, value in ( + ("window_size", 0), + ("stride", 0), + ("tail", "unknown"), + ("group_key", "missing"), + ("order_key", "missing")): + kwargs = { + "window_size": 2, + "columns": ["value"], + "stride": 1, + "tail": "drop", + "group_key": "episode", + "order_key": "step", + } + kwargs[name] = value + with self.subTest(name=name), self.assertRaises((TypeError, ValueError)): + query.to_contiguous_window_dataset(**kwargs) + + reserved_table = self.conn.create_table( + "reserved", schema=pa.schema([ + pa.field("is_pad", pa.string(), nullable=False), + pa.field("step", pa.int32(), nullable=False), + pa.field("value", pa.int32(), nullable=False), + ]), options=_TABLE_OPTIONS) + with self.assertRaisesRegex(ValueError, "must not be is_pad"): + reserved_table.scan().to_contiguous_window_dataset( + window_size=2, columns=["value"], + group_key="is_pad", order_key="step") + + with self.assertRaisesRegex(TypeError, "only supported on scan"): + table.search("anything", column="episode").to_contiguous_window_dataset( + window_size=2, columns=["value"], + group_key="episode", order_key="step") + + +if __name__ == "__main__": + unittest.main() diff --git a/paimon-python/setup.py b/paimon-python/setup.py index 12e24632bc3e..e2cfcc33a8af 100644 --- a/paimon-python/setup.py +++ b/paimon-python/setup.py @@ -216,6 +216,15 @@ def read_requirements(): install_requires = read_requirements() +LEROBOT_DEPENDENCIES = [ + # datasets 4.1+ may select PyArrow 21+, while PyPaimon currently + # supports PyArrow <20. Pandas 2.2.2+ supports NumPy 2.x selected + # by LeRobot's media dependencies. + 'datasets>=4,<4.1; python_version>="3.10"', + 'pandas>=2.2.2,<3; python_version>="3.10"', + 'lerobot>=0.4.4,<0.5; python_version>="3.10"', +] + long_description = "See Apache Paimon Python API \ [Doc](https://paimon.apache.org/docs/master/pypaimon/python-api/) for usage." @@ -224,7 +233,12 @@ def read_requirements(): version=VERSION, packages=PACKAGES, include_package_data=True, - package_data={"pypaimon": ["_full_version"]}, + package_data={ + "pypaimon": [ + "_full_version", + "benchmark/act/default_experiment.json", + ], + }, cmdclass={"build_py": PaimonBuildPy, "sdist": PaimonSdist}, install_requires=install_requires, entry_points={ @@ -241,20 +255,16 @@ def read_requirements(): # rosbags is pure Python and does not require a ROS installation. 'rosbags>=0.11.5,<0.12; python_version>="3.10"', ], - 'lerobot': [ - # datasets 4.1+ may select PyArrow 21+, while PyPaimon currently - # supports PyArrow <20. Pandas 2.2.2+ supports NumPy 2.x selected - # by LeRobot's media dependencies. - 'datasets>=4,<4.1; python_version>="3.10"', - 'pandas>=2.2.2,<3; python_version>="3.10"', - 'lerobot>=0.4.4,<0.5; python_version>="3.10"', - ], + 'lerobot': LEROBOT_DEPENDENCIES, 'ray': [ 'ray>=2.10,<3; python_version>="3.8"', ], 'torch': [ 'torch', ], + 'act': LEROBOT_DEPENDENCIES + [ + 'Pillow; python_version>="3.10"', + ], 'daft': [ 'daft>=0.7.6; python_version>="3.10"', ],