diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index b4ac4f4c0..ff878152b 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -272,7 +272,8 @@ topics: title: Differentiable Learning Layers aliases: [learning layers, differentiable IK layer, 可微 IK, 训练层, 学习模块分层] keywords: [DifferentiableIKLayer, KinematicsLayer, SerialChainKinematics, DifferentiableIKCfg, - solve_dls_ik, damped_least_squares_step, Jacobian gradient, dJ/dq, learn_grasp_yaw, grasp yaw, 抓取角度优化, 抓取姿态学习] + solve_dls_ik, damped_least_squares_step, Jacobian gradient, dJ/dq, learn_grasp_yaw, grasp yaw, 抓取角度优化, 抓取姿态学习, + ResNet10Encoder, ResNet10, visual encoder, 视觉编码器] paths: [topics/learning-layers/learning-layers.md] source_of_truth: - embodichain/compute/kinematics/dls.py @@ -442,7 +443,7 @@ topics: status: active - id: data-pipeline title: Data Pipeline - aliases: [online data, online dataset, async recorder, dataset persistence, simple BC demo, ODS documentation, ODS 文档, 数据管线, 在线数据, 数据集持久化, 数据保存, + aliases: [online data, online dataset, async recorder, dataset persistence, simple BC demo, ODS documentation, ODS capability showcase, ODS 能力展示, ODS 文档, 数据管线, 在线数据, 数据集持久化, 数据保存, 深度视频, 数据集校验, 数据集转换, 数据集划分] keywords: [OnlineDataEngine, OnlineDataset, sample_batch, sample_batch_with_context, lock_index, LeRobotRecorder, AsyncLeRobotRecorder, DepthSidecarManager, depth_meta.json, fragment_id, finalize, expert action schema, qpos_slice, qvel_slice, @@ -468,6 +469,7 @@ topics: - docs/source/features/online_data.md - docs/source/features/ods/ - docs/architecture/dataset-loading-design.md + - docs/architecture/ods-data-stream-showcase-design.md - embodichain/lab/gym/envs/managers/dataset_manager.py - embodichain/lab/gym/envs/managers/datasets.py - embodichain/lab/gym/envs/managers/async_datasets.py @@ -479,9 +481,10 @@ topics: - tools/lerobot_export/ - embodichain/lab/scripts/run_task.py - examples/ods/simple_bc_demo/ + - examples/ods/data_stream_demo/ watch_paths: [embodichain/data_pipeline/, embodichain/lab/gym/envs/managers/dataset_manager.py, embodichain/lab/gym/envs/managers/datasets.py, embodichain/lab/gym/envs/managers/async_datasets.py, embodichain/lab/gym/envs/embodied_env.py, embodichain/lab/scripts/run_task.py, - tests/data_pipeline/, embodichain/cli/dataset.py, tools/lerobot_export/, examples/ods/simple_bc_demo/, + tests/data_pipeline/, embodichain/cli/dataset.py, tools/lerobot_export/, examples/ods/simple_bc_demo/, examples/ods/data_stream_demo/, docs/source/features/online_data.md, docs/source/features/ods/, docs/architecture/dataset-loading-design.md] related_topics: [manager-functor, env-framework, task-programs, sensor-system] status: active diff --git a/agent_context/topics/data-pipeline/data-pipeline.md b/agent_context/topics/data-pipeline/data-pipeline.md index f795d24b6..20915d90f 100644 --- a/agent_context/topics/data-pipeline/data-pipeline.md +++ b/agent_context/topics/data-pipeline/data-pipeline.md @@ -16,6 +16,9 @@ belongs to [data-assets](../data-assets/data-assets.md). | Shared sequence contract and sources | `embodichain/data_pipeline/datasets/sequence.py`, `sources.py`, `_sampling.py` | | Bounded numeric LeRobot loading | `embodichain/data_pipeline/datasets/lerobot.py`: `LeRobotDataSource`, explicit field mapping | | Local Franka ChunkBC demo and offline comparisons | `examples/ods/simple_bc_demo/train.py`; local task/config/model/trainer/evaluator and evidence in the same directory | +| Pure delivery capability and standard video decoding comparison | `examples/ods/data_stream_demo/run.py`, `suite.py`, `report.py`; no learner, real SDK/DataLoader decode and parent-owned prefetch quotas | +| Fixed-budget visual training, diversity and total pipeline time | `examples/ods/data_stream_demo/train.py`, `training_metrics.py`, `training_report.py`; single/dual-camera BC, standard dataset recording and completed-training accounting; `training_throughput.py` owns three-stage preflight guidance | +| Dual-camera learning and camera review | `examples/ods/data_stream_demo/converge.py`, `preview.py`; `CONVERGENCE.md` owns the candidate recipe, development selection and one held-out test | | Commit/discard/finalization dispatch | `embodichain/lab/gym/envs/managers/dataset_manager.py` | | Synchronous LeRobot persistence | `embodichain/lab/gym/envs/managers/datasets.py`: `LeRobotRecorder` | | Background persistence | `embodichain/lab/gym/envs/managers/async_datasets.py`: `AsyncLeRobotRecorder` | @@ -38,6 +41,34 @@ Example-local documentation owns its task binding, CLI/artifact names, comparison protocol and historical evidence. Keep generic operational contracts in the core guides. +The [ODS capability-showcase design](../../../docs/architecture/ods-data-stream-showcase-design.md) +owns the fixed-budget single/dual-camera BC comparison and the historical +state/single/dual delivery diagnostic. Its example-local decoded reader is +separate from the public resident numeric reader. Follow it for offline corpus +sizes, source-reuse protection, camera qualification and timing boundaries. +Training comparisons match actual updates and valid supervision; data volume +and trained-state diversity are outcomes. Record only completed optimizer +updates as training coverage, and include recording/finalization and normal +teardown in total pipeline time. Serial trace replay separately attributes +generation/training overlap. +Budget-aware prefetch and three-stage capacity guidance remain example-local; +they do not change public engine/controller scheduling. The +[visual demo README](../../../examples/ods/data_stream_demo/README.md) +owns calibration and pilot commands, including isolated eager controls. +The visual runner can ingest complete producer subbatches before generation +completion. Its raw training/archive/serial owners share one C-episode bound; +serial controls still archive during generation. This example-local handoff +and nested SDK-stage profiling preserve the public sampling contract. +All acquisition controls may overlap recording and generation within the same +storage bounds. Worker SDK service timestamps remain separate from owner +copy/wait spans. `training_audit.py` combines full numeric checks, source-window +SDK reads and `video_integrity.py`'s all-frame software video count/integrity +check; fresh post-timing receipts preserve earlier evidence. +The [experiment evidence index](../../../examples/ods/data_stream_demo/EXPERIMENTS.md) +links immutable historical studies and their workload-specific limitations. +Generated run evidence is excluded from the source tree; it does not establish +training efficiency or policy performance. + `BatchSource.sample_context()` returns owned TensorDicts under a shared `SequenceSpecCfg`; `SequenceDataset` supplies the common iteration interface. The package's `schema.py` owns common semantic validation and schema fingerprints; diff --git a/agent_context/topics/learning-layers/learning-layers.md b/agent_context/topics/learning-layers/learning-layers.md index fd46bc648..73ae1ec67 100644 --- a/agent_context/topics/learning-layers/learning-layers.md +++ b/agent_context/topics/learning-layers/learning-layers.md @@ -8,8 +8,12 @@ These modules do not import learning or simulation. `SerialChainKinematics` returns TCP poses and geometric Jacobians from one FK pass, retaining dJ/dq. `learning/layers/` owns `KinematicsLayer` and `DifferentiableIKLayer` Torch -module interfaces. `learning/models/` owns general MLPs, feature normalization -and configuration-driven MLP construction. These components do not import RL, +module interfaces. `learning/models/` owns general MLPs, feature normalization, +configuration-driven MLP construction, and the trainable `ResNet10Encoder`. +The encoder is imported lazily so torchvision is required only by visual clients; +it consumes float32 NCHW RGB in `[0, 1]` and owns ImageNet normalization. +The ODS visual demo owns the shared-view/history BC head and task losses, not +the general encoder. These components do not import RL, TensorDict or simulation. RL policy interfaces, actor-critic composition and policy registration/builders belong to `learning/rl/policies/`; RL code consumes the shared networks and normalizer. See [RL learning](../rl-learning/rl-learning.md) @@ -57,7 +61,11 @@ part of the training/evaluation entry point. The demo does not execute physical motion; keep task predictors out of production policy registries until a reusable contract exists. -Tests live under `tests/compute/kinematics/` and `tests/learning/layers/`. +Tests live under `tests/compute/kinematics/`, `tests/learning/layers/`, and +`tests/learning/models/`. The visual encoder tests verify reference-structure +forward/gradient equivalence, image ownership and checkpoint round trips; +the ODS demo's `qualify_encoder.py` additionally verifies pinned RLInf classes, +causal multiview inputs, image-only fitting, and explicitly selected CUDA. Validate FK, dJ/dq, full IK and recurrent-state gradients separately; include nonzero TCP offsets, repeated solves, checkpoint migration and import boundaries. GPU validation runs serially with `--run-gpu`. The unified kinematic benchmark's `differentiable` mode measures training and diff --git a/docs/architecture/dataset-loading-design.md b/docs/architecture/dataset-loading-design.md index c2de949cc..09fefc1fb 100644 --- a/docs/architecture/dataset-loading-design.md +++ b/docs/architecture/dataset-loading-design.md @@ -9,6 +9,12 @@ For use and operation, start with the [ODS guide](../source/features/online_data This document owns source/Dataset architecture and implementation boundaries; task parameters and historical run evidence remain with the demo. +The [ODS data-stream showcase](ods-data-stream-showcase-design.md) +uses standard SDK/DataLoader video decoding for its offline comparison. That +example-local reader is separate from the shipped resident numeric backend; its +visual results cannot be obtained by dropping images or preloading decoded +datasets through the current numeric loader. + ## Problem and acceptance criterion `SequenceDataset` supplies one PyTorch context-iteration interface for live and diff --git a/docs/architecture/ods-data-stream-showcase-design.md b/docs/architecture/ods-data-stream-showcase-design.md new file mode 100644 index 000000000..a2581c8fd --- /dev/null +++ b/docs/architecture/ods-data-stream-showcase-design.md @@ -0,0 +1,195 @@ +# ODS visual training and data-stream showcase + +This example measures how asynchronous generation changes **trained-state +coverage and complete pipeline time under a fixed actual training budget**. +It reuses the Franka PickPlace environment and physical expert from +`examples/ods/simple_bc_demo`. The experiment, BC head, recorder and scheduling +remain example-local; public sampling contracts belong to ODS. + +Use the [example README](../../examples/ods/data_stream_demo/README.md) for +commands, configuration guidance and output interpretation. Historical studies, +including failed attempts and negative results, are linked from the compact +[experiment evidence index](../../examples/ods/data_stream_demo/EXPERIMENTS.md). +Generated plots, logs, source snapshots and machine-specific inventories are +kept outside the source change. + +## Fixed-training-budget visual comparison + +The main question is: with identical actual training work and assigned +hardware, how much distinct physical data reaches completed optimizer updates, +and how long does the whole pipeline take? Data volume and final coverage are +outcomes rather than matching constraints. + +| Path | Behavior | +| --- | --- | +| Offline-8 / 32 / 128 | Generate, record and finalize the full corpus, then train through standard LeRobot decoding | +| ODS overlap | Train after initial fill while later generations and recording continue | +| ODS serial | Regenerate the overlap sources and replay its ordered training manifest with generation/model updates serialized | + +Run third-person (`single`) and third-person plus wrist (`dual`) camera +conditions separately. The delivery-only diagnostic also supports state-only +input, but that condition does not participate in the visual training matrix. +Serial may overlap recording with generation or training; it must not acquire +an extra archive-submission barrier that confounds the intended comparison. + +```mermaid +flowchart LR + Expert[Expert generation] --> Shared[Shared cache C] + Shared --> Owned[Owned data at most C] + Owned --> Train[Visual BC updates] + Owned --> Writer[Bounded recorder] + Writer --> Files[Standard LeRobot dataset] + Files --> Decode[SDK video decoding] + Decode --> Train +``` + +The direct owned-data edge is ODS; the finalized-dataset edge is offline. +Every arm saves a standard dataset. The serial control reuses the direct path +while prohibiting generation/model-update overlap. + +## Model and work contract + +The systems default is U=128 updates, B=32 chunks, H=2 causal observations and +F=32 future action targets. Each qualified run completes 4,096 training chunks +and 131,072 valid action targets. Complete windows stay within one episode and +continuity segment. A larger budget is a new frozen recipe; 128 updates do not +imply task convergence. + +The shared lazy `learning/models/ResNet10Encoder` adapts RLInf commit +[`91a302ec50ddacd1fdabc5859b85e3661d768236`](https://github.com/RLinf/RLinf/blob/91a302ec50ddacd1fdabc5859b85e3661d768236/rlinf/models/embodiment/modules/resnet_utils.py): + +- torchvision BasicBlock stages `[1,1,1,1]`, GroupNorm4, epsilon 1e-5; +- eight learned spatial embeddings, 256-dimensional projection, LayerNorm/Tanh; +- scratch initialization with the full encoder trainable; +- one encoder shared across ordered camera views and both history frames. + +Keep attribution and third-party license notices in the module. The compatible +reference stack is PyTorch 2.7/torchvision 0.22. The example requires neither +RLInf runtime nor pretrained weights. + +Adapters supply 256×256 RGB float32 in [0,1]; the encoder applies ImageNet +normalization once. Numeric inputs are two 27-feature robot observations: +qpos, qvel and TCP pose. Cube geometry, success, expert progress and future +observations do not enter policy forward. The 512/512 MLP predicts F×10 +absolute TCP position, rotation and gripper targets. Target construction and +normalized MSE follow ChunkBC. Future data are supervision only. + +Use FP32 AdamW, learning rate 3e-4, weight decay 1e-6 and gradient-norm limit 1, +with augmentation/noise disabled. Match initialization and recipe within each +vision/seed; freeze numeric normalization from the common initial eight +trajectories. Save model, optimizer, normalization, RNG, configuration and step +state. Equal work does not require equal final weights or policy quality. + +## Acquisition, reuse and recording + +Default C=8 shared episodes, P=2 parallel producer environments and lifetime +per-source delivery cap K=128. K is an engineering guard, not an empirically +optimal learning setting. Subsequent acquisition uses publication-granularity +handoff; the initial full C block remains common to normalization. + +The owner polls complete P-episode producer subbatches, validates immutable +source/generation/frame/continuity metadata, and clones under the existing +publication lock. Busy or unpublished rows return no data; overwritten expected +identities and worker failures fail explicitly. Copy all accepted sources in a +generation before requesting its successor. The public engine's immediate +sample/error behavior does not change. + +Raw training, archive-staging and serial-ready owners share a C-episode bound. +Packing creates independent writer packets and releases raw references. The +writer queue holds two packets, plus one active packet and at most one pending +submission. Returned batches are independent of later cache replacement. +Require P*K>=B and enough batch entries to cover every retained source. +`--handoff-mode generation` retains whole-generation acquisition as a control. +Serial archives published subbatches during generation and defers updates until +that generation completes. Trace v2 records retained block boundaries and +recorder/handoff settings; v1 traces remain readable. + +At each batch boundary, budget-aware prefetch requests necessary supply when +remaining credits cannot fund remaining updates. Optional refresh uses measured +generation, recorder-staging, snapshot and recent complete ready-update times. +Cold model startup is excluded from the latter estimate. Unknown timing defers +optional requests. Pending requests coalesce; credits never reset on eviction. +An `eager` diagnostic requests after every switch and reports unused data. + +Each block receives at least one completed batch covering all sources before +replacement. Sample sources evenly, visiting unused legal anchors before +reshuffling. Offline uses the same balancing/anchor rule but may reuse its +fixed corpus to complete the budget. After update U, accept no new requests; +finish and record every already-requested generation. Untrained tails count +as acquisition cost and never as trained diversity. + +Create and finalize datasets through LeRobot 0.4.4, with real Parquet/video +features, 30 FPS and AV1/yuv420p/GOP2/CRF30. Explicit encoder threads default to +one and must match across arms. Camera encoding is sequential to avoid an SDK +process-pool fork from the CUDA learner's writer thread. SDK episode service +has nested validation, frame-ingestion, save/encode/commit and manifest spans; +these partition the outer service and are not additional wall time. + +Offline training uses the official Dataset and CPU TorchCodec 0.5 decoding, +four spawned workers and prefetch factor one. Numeric supervision reads the +required 33 frames; each camera reads only the two causal history frames. +Do not replace standard decoding with a resident fully decoded corpus. + +## Timing, resources and capacity guidance + +Complete pipeline time starts at supervised process launch and ends after +training, all accepted recording, checkpoint writes and normal process closure. +Also record last-update time, first-update latency, waits and recording drain. +Mandatory content audit and requested independent policy evaluation run after +this clock. Zero-scene runs explicitly record evaluation as not requested. +Failures remain visible and do not enter qualified ratios. + +Use separate simulation/training primary GPUs and the same eight-CPU affinity +across compared arms. Record GPU UUIDs, masks, thread budgets, actual context +placement and sampled resources. Hardware admission is not exclusive host +reservation. An explicitly declared passive graphics allowance is recorded and +bounded separately; it grants no extra compute capacity. + +Measure accepted generation episodes/s, recorder busy-service episodes/s, and +complete ready-batch preparation/transfer/optimizer updates/s. Keep pure model +time separate. `min(generation, recording)` estimates bottleneck supply; +`gamma_effective*K/B` estimates its protected update ceiling. Shared-host +contention, startup and cleanup prevent treating either as total-time forecasts. +Overlapping intervals cannot be summed to obtain pipeline time. + +Preflight preserves C/P/B/K and defaults to `max(16,ceil(C*K/B)+8)` updates. +It reports when later generations never reach training. Sweep feasible P +values that divide C, then select the smallest measured P within 90% of the +best effective supply under matched hardware/recipes. A computed K needed to +keep training busy is diagnostic, not permission to raise the reuse guard. +Freeze the workload before comparison; avoid artificial delays or dummy work. + +## Metrics, evidence and acceptance + +Only synchronized completed updates increase training coverage. Preserve +immutable source/frame identities, ordered update manifests and retired-source +counters. Report trained and entropy-effective source counts, occupied initial +XY cells (fixed 8×8 grid), cube XYZ voxels at anchors (2 cm), unique anchors, +history/target frames, target duplication, publish-to-first-training latency, +and lifetime chunks per source. New IDs do not imply new physical states. + +Generate time/coverage, coverage-growth, pipeline-timeline and lifetime-reuse +figures from run logs. Resource metrics distinguish fixed numeric cache size +from actual runtime memory, and finalized disk bytes from unmeasured temporary +peaks. Preserve unused-source counts and all failure reasons. + +Before paired ratios, verify matching initial-model/normalization hashes, +initial-C numeric prefix, camera order, actual U/B/H/F, GPU assignment, CPU +budget and supervisor/process evidence. Serial additionally matches source, +anchor and retained-block order. Neither matching numerical manifests nor +lossy AV1 files establish raw rendered RGB identity. + +Qualification covers reference-aligned encoder outputs/gradients, preprocessing, +view/history order, image-dependent fitting, actual parameter updates and +checkpoint reload. Shared encoder tests own its production contract. Actual +pipeline checks cover causal windows, quotas, bounded ownership, old-output +independence, real SDK decoding, AV1 PSNR/MAE, every saved video frame and normal +worker closure. Add focused production regressions for core defects; do not +build a separate general training-test framework around the example. + +Run qualification, smoke, metric checks, throughput preflight and one-seed pilot +before a formal matrix. The formal recipe is four primary arms×two camera +modes×three seeds, plus six serial replays. Regenerate data, rotate arm order +and retain slow/failed attempts. Evaluate the same 20 development states after +timing as policy-quality context. Report the measured time–coverage tradeoff; +the historical studies do not establish universal speedup or task convergence. diff --git a/docs/source/api_reference/embodichain/embodichain.learning.models.rst b/docs/source/api_reference/embodichain/embodichain.learning.models.rst index 00f2ac0c3..ca480554e 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.models.rst +++ b/docs/source/api_reference/embodichain/embodichain.learning.models.rst @@ -17,6 +17,7 @@ RL policy interfaces and construction belong to MLP EmpiricalNormalizer build_mlp_from_cfg + ResNet10Encoder .. automodule:: embodichain.learning.models :members: @@ -34,3 +35,14 @@ Feature normalization .. automodule:: embodichain.learning.models.normalizer :members: + +Trainable visual features +------------------------- + +``ResNet10Encoder`` maps float32 NCHW RGB images in ``[0, 1]`` at 256 by +256 pixels to 256-dimensional features. It owns ImageNet normalization and +keeps the GroupNorm residual backbone and learned spatial pooling trainable. +Its torchvision dependency is loaded only when the visual encoder is imported. + +.. automodule:: embodichain.learning.models.resnet_encoder + :members: diff --git a/embodichain/learning/models/__init__.py b/embodichain/learning/models/__init__.py index 548450aa4..1b1bf9ac8 100644 --- a/embodichain/learning/models/__init__.py +++ b/embodichain/learning/models/__init__.py @@ -14,19 +14,20 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""General neural networks, normalization and network construction.""" +"""General neural networks, visual features and normalization, imported lazily.""" from __future__ import annotations from importlib import import_module from typing import Any -__all__ = ["MLP", "EmpiricalNormalizer", "build_mlp_from_cfg"] +__all__ = ["MLP", "EmpiricalNormalizer", "build_mlp_from_cfg", "ResNet10Encoder"] _EXPORT_MODULES = { "MLP": "mlp", "EmpiricalNormalizer": "normalizer", "build_mlp_from_cfg": "mlp", + "ResNet10Encoder": "resnet_encoder", } diff --git a/embodichain/learning/models/resnet_encoder.py b/embodichain/learning/models/resnet_encoder.py new file mode 100644 index 000000000..ca81eb933 --- /dev/null +++ b/embodichain/learning/models/resnet_encoder.py @@ -0,0 +1,128 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# Copyright 2025 The RLinf Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Trainable visual features adapted from RLInf's Apache-2.0 ResNet encoder. + +Reference: RLinf/RLinf, commit ``91a302ec50ddacd1fdabc5859b85e3661d768236``, +``rlinf/models/embodiment/modules/resnet_utils.py``. Modifications remove the +classifier, pretrained checkpoint requirement, frozen/detached backbone and +dropout. This module owns ImageNet normalization; callers supply RGB floats in +``[0, 1]``. Spatial pooling contracts directly over pixels, avoiding the +reference implementation's expanded ``[N, C, H, W, 8]`` intermediate. +""" + +from __future__ import annotations + +from functools import partial + +import torch +from torch import nn +from torchvision.models.resnet import BasicBlock, ResNet + +__all__ = ["ResNet10Encoder"] + + +class _SpatialResNet10(ResNet): + def __init__(self) -> None: + super().__init__( + BasicBlock, + [1, 1, 1, 1], + norm_layer=partial(nn.GroupNorm, 4, eps=1e-5), + ) + del self.avgpool + del self.fc + + def _forward_impl(self, images: torch.Tensor) -> torch.Tensor: + features = self.maxpool(self.relu(self.bn1(self.conv1(images)))) + features = self.layer1(features) + features = self.layer2(features) + features = self.layer3(features) + return self.layer4(features) + + +class _SpatialLearnedEmbeddings(nn.Module): + def __init__(self) -> None: + super().__init__() + self.kernel = nn.Parameter(torch.randn(512, 8, 8, 8)) + + def forward(self, features: torch.Tensor) -> torch.Tensor: + # Same channel-major, embedding-minor order as the pinned reference. + return torch.einsum("bchw,chwf->bcf", features, self.kernel).flatten(1) + + +class ResNet10Encoder(nn.Module): + """Encode 256-pixel RGB images into trainable 256-dimensional features. + + The encoder has four single-block residual stages with GroupNorm, eight + learned spatial embeddings, and a Linear/LayerNorm/Tanh projection. It + initializes all parameters locally and never downloads weights. Every + stage remains differentiable in both training and evaluation modes. + + Input is float32 ``[N, 3, 256, 256]`` RGB in ``[0, 1]`` on the model's + device. Value range and finiteness are the data adapter's responsibility; + forward checks tensor metadata without synchronizing CUDA for pixel scans. + Noncontiguous inputs are accepted and are never modified in place. + """ + + output_dim = 256 + image_size = 256 + + def __init__(self) -> None: + super().__init__() + self.register_buffer( + "image_mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) + ) + self.register_buffer( + "image_std", torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) + ) + self.resnet_backbone = _SpatialResNet10() + self.pooling_layer = _SpatialLearnedEmbeddings() + self.mlp = nn.Sequential( + nn.Linear(512 * 8, self.output_dim), + nn.LayerNorm(self.output_dim), + nn.Tanh(), + ) + # Pinned RLInf init_mlp_weights(..., nonlinearity="tanh"). + nn.init.kaiming_normal_(self.mlp[0].weight, mode="fan_out", nonlinearity="tanh") + nn.init.zeros_(self.mlp[0].bias) + + def forward(self, images: torch.Tensor) -> torch.Tensor: + """Return ``[N, 256]`` features from causal RGB float inputs. + + Args: + images: Nonempty float32 NCHW images, already scaled to ``[0, 1]``. + + Returns: + Independently owned feature tensor connected to inputs and weights. + + Raises: + ValueError: Image shape, dtype or device does not match the contract. + """ + if ( + images.ndim != 4 + or images.shape[0] < 1 + or images.shape[1:] != (3, self.image_size, self.image_size) + or images.dtype != torch.float32 + or images.device != self.image_mean.device + ): + raise ValueError( + "ResNet10Encoder requires nonempty float32 [N, 3, 256, 256] " + "RGB images on the model device" + ) + normalized = (images - self.image_mean) / self.image_std + features = self.resnet_backbone(normalized) + return self.mlp(self.pooling_layer(features)) diff --git a/examples/ods/data_stream_demo/.gitignore b/examples/ods/data_stream_demo/.gitignore new file mode 100644 index 000000000..08afe3f11 --- /dev/null +++ b/examples/ods/data_stream_demo/.gitignore @@ -0,0 +1,2 @@ +/results/ +/outputs/ diff --git a/examples/ods/data_stream_demo/CONVERGENCE.md b/examples/ods/data_stream_demo/CONVERGENCE.md new file mode 100644 index 000000000..05d607176 --- /dev/null +++ b/examples/ods/data_stream_demo/CONVERGENCE.md @@ -0,0 +1,127 @@ +# Dual-camera PickPlace learning example + +This recipe adds learning acceptance to the ODS visual example. It reuses the +current single-cycle Franka PickPlace physics, handwritten expert, bounded +publication handoff, visual policy and standard LeRobot recording. Review the +camera preview and resolved recipe before launching a training run. + +## Preview and review + +```bash +python -m examples.ods.data_stream_demo.preview \ + --device cuda:0 --camera-profile pick_place --output outputs/pick-place-preview +python -m examples.ods.data_stream_demo.converge review \ + --output outputs/pick-place-review.json +``` + +Mask an idle GPU before preview, for example `CUDA_VISIBLE_DEVICES=3`. Preview +creates an ODS producer, acquires one physically accepted expert episode, saves +actual RGB images and a side-by-side video, then closes the producer. It does +not construct a policy or optimizer. Preview data are not training or evaluation +data. The resolved `environment-config.json` records camera calibration. + +`pick-place-v2` uses 256×256 third-person RGB with eye[-0.60,-0.45,0.50] and +target[-0.42,0.10,0.08], 60° FOV. The wrist camera is attached to +`FrankaPanda/fr3_hand_tcp`, with local eye[0.12,0,-0.04], target[0,0,0.05], +up[0,1,0] and 90° FOV. This avoids the initial palm occlusion in the historical +mount. The benchmark camera profile remains v3. Training and evaluation load +the same named calibration; mismatched checkpoint calibration is rejected. + +## Candidate training recipe + +The YAML at [configs/pick_place_dual_convergence.yaml](configs/pick_place_dual_convergence.yaml) +is the source of truth. `review` writes a resolved JSON packet and its input SHA. + +| Setting | Candidate | +| --- | --- | +| Task | Current real-contact Franka single PickPlace; unchanged 600-step/30Hz limit and physical success predicate | +| Model | From-scratch shared ResNet10/GroupNorm/spatial embeddings; 256D/view; 512/512 action MLP | +| Input | Two causal frames per view and two 27D robot states; no cube pose or expert phase in policy forward | +| Prediction / execution | 32 absolute TCP/rotation/gripper targets; execute 8 before replanning | +| Work budget | 20,000 updates ×64 chunks =1,280,000 chunks /40,960,000 action targets | +| Optimizer | FP32 AdamW; weight decay1e-6; gradient norm limit1 | +| Learning rate | 500-update linear warmup to3e-4, cosine decay to3e-5 | +| Augmentation | Disabled for the first learning qualification | +| ODS | C32 /P4 /K2048; publication handoff; budget-aware refill; two encoder threads | +| Normalization | Initial32 training episodes only; frozen | +| Development | Initial policy and every2,000 updates; same20 fixed states | +| Development gate | Last two updated checkpoints both reach at least18/20 successes | +| Checkpoint selection | Highest development success, then smallest mean goal error; freeze selection before testing | +| Final test | One independent100-state panel, opened only after the development gate passes; target≥80/100 | +| Runtime bounds | Six-hour whole worker deadline; one-hour bound per evaluation | + +K2048 is a provisional engineering cap, not a learning-qualified optimum. The +declared work requires at least625 trained sources if all deliveries reach +updates. Earlier retirement can require more. Keep all source exposure and +unused-tail records; do not increase K automatically to keep the GPU occupied. + +Training, development, test and preview use distinct seeds20261025,20271025, +20281025 and20291025. Development data never fit normalization or enter a +training batch. The final panel must remain unopened during tuning; after its +outcomes are used to change a recipe, use a new sealed panel for a new claim. + +The success criterion is practical task acceptance, not proof of mathematical +optimizer convergence. The candidate budget is not a promise of80% success. +Retain failed seeds/checkpoints and report an unmet target explicitly. + +## Train after recipe review + +First recheck the unchanged physical expert on the current installed stack. +The existing qualification-only path performs no policy learning and requires +at least38/40 measured successes. It uses no camera input for the expert; the +visual profile changes sensors only, with the same physics and action decoder. + +```bash +CUDA_VISIBLE_DEVICES=3 python -m examples.ods.simple_bc_demo.train \ + --qualification-only --qualification-scenes 40 --no-viewer --no-video \ + --sim-device cuda:0 --train-device cuda:0 \ + --output outputs/pick-place-expert-qualification +``` + +Proceed after that command succeeds and the candidate learning recipe is +reviewed: + +```bash +python -m examples.ods.data_stream_demo.converge train \ + --gpu-ids 3 4 --output outputs/pick-place-dual-learning +tensorboard --logdir outputs/pick-place-dual-learning/training/tensorboard +``` + +GPU indices are examples, not reservations. The supervisor admits two idle +GPUs, masks their UUIDs before model/simulation imports, binds eight CPU cores +and checks owned contexts. An optional `--aux-graphics-gpu-id` retains the +existing bounded passive-graphics allowance. The supervisor retains process +return codes and rejects leftover worker groups. Generation/evaluation use the +simulation device and policy training/inference use the second device. + +The learning driver reuses `_Experiment` for acquisition, protected sampling, +actual updates, recording and failure cleanup. It adds a declared LR schedule, +TensorBoard losses, independent periodic development processes and immutable +checkpoint selection. Evaluation drains lifecycle events and recorder staging +while the producer remains alive. These pauses are part of this learning run; +it is not a fixed-budget speedup measurement or an RLInf adapter. + +The full update budget runs even when an intermediate policy performs well, +so the final-two-development gate can reveal later degradation. If that gate +fails, final testing is skipped and the command exits nonzero. A failed final +test also exits nonzero; the model is never taken over by the expert. + +## Outputs and limits + +`training/` contains `initial.pt`, immutable periodic checkpoints, `final.pt`, +`development.json`, evaluator logs/outcomes, `selection.json`, TensorBoard events, +the existing source/anchor ledger and training trace, and the complete standard +LeRobot dataset. `selection.json` is written before `evaluation/final-test.json`. +`summary.json` keeps engineering completion separate from `learning_target_met`. +Loss, lifted rate, goal error and task success remain separate measures. + +Checkpoints include model, optimizer, normalization, RNG and recipe. This entry +point currently evaluates saved checkpoints through `training_eval`; it does +not implement exact replay/source recovery or a convergence-run resume command. +Do not restart an interrupted run while claiming the original quota or source +cache was restored. Keep the saved outcomes and use a fresh output directory. + +Runtime estimates depend on current generation/recording and evaluation service, +not optimizer time alone. The historical128-update timings cannot be scaled +into a reliable convergence-time estimate. Recheck hardware admission at launch +and retain all waits, recording drain and independent evaluation cost. diff --git a/examples/ods/data_stream_demo/EXPERIMENTS.md b/examples/ods/data_stream_demo/EXPERIMENTS.md new file mode 100644 index 000000000..0e57a0e6b --- /dev/null +++ b/examples/ods/data_stream_demo/EXPERIMENTS.md @@ -0,0 +1,82 @@ +# Experiment evidence and review scope + +The source tree contains the runnable experiment, qualification tools and +report generators. Generated logs, plots, checkpoints, source snapshots and +machine-specific manifests belong in run output directories rather than the +code change. Local `results/` and `outputs/` directories are ignored. + +The complete pre-cleanup tree is preserved at commit +[`3829593c07392431282187ac0d810a8a4a0c5373`](https://github.com/DexForce/EmbodiChain/tree/3829593c07392431282187ac0d810a8a4a0c5373/examples/ods/data_stream_demo/results), +also named by branch `codex/ods-stream-design-backup-20261024`. The links below +use that immutable commit. Removing generated evidence from the PR does not +change experiment qualification or turn historical failures into successes. + +| Study | Scope and retained result | +| --- | --- | +| [Delivery diagnostic](https://github.com/DexForce/EmbodiChain/blob/3829593c07392431282187ac0d810a8a4a0c5373/examples/ods/data_stream_demo/results/2026-10-21/RESULTS.md) | 27 primary runs and nine diagnostics; measures delivery/decoding rather than training efficiency | +| [Copy optimization](https://github.com/DexForce/EmbodiChain/blob/3829593c07392431282187ac0d810a8a4a0c5373/examples/ods/data_stream_demo/results/2026-10-21-copy-optimization/README.md) | CPU copy timings and independent-storage checks; not an end-to-end speedup | +| [Visual training](https://github.com/DexForce/EmbodiChain/blob/3829593c07392431282187ac0d810a8a4a0c5373/examples/ods/data_stream_demo/results/2026-10-21-visual-training/README.md) | 30 attempts, 29 qualified; one audit-process shutdown failure remains failed | +| [Budget-aware prefetch](https://github.com/DexForce/EmbodiChain/blob/3829593c07392431282187ac0d810a8a4a0c5373/examples/ods/data_stream_demo/results/2026-10-22-budget-aware-prefetch/README.md) | Six capacity preflights and eight pilot runs; six primary outcomes/four pairs, with eager controls separate | +| [Publication handoff](https://github.com/DexForce/EmbodiChain/blob/3829593c07392431282187ac0d810a8a4a0c5373/examples/ods/data_stream_demo/results/2026-10-23-publication-handoff/README.md) | Ten qualified pilot/diagnostic runs; earlier six-run smoke is functional qualification only | + +## Measured limits + +The formal visual study uses U128/B32/H2/F32: 4,096 trained chunks and +131,072 valid action targets per attempt. Its complete-time medians are +382.8/478.3 seconds for ODS overlap and 332.7/413.1 seconds for Offline-32 +(single/dual; dual Offline-32 has only two qualified runs). ODS is slower in +that workload. All six ODS policies score 0/20 successes; among qualified +formal evaluations, only Offline-128/dual/seed19001 scores 1/20. These are +systems measurements, not evidence of converged PickPlace learning. The failed +attempt is Offline-32/dual/seed19003; its original evaluation was not run. +Later shutdown repair and supplemental validation do not replace that outcome. + +The later publication pilot keeps the same work budget, uses one seed and +requests no policy evaluation. Both camera modes complete 48 updates before +their generation fully publishes. Post-initial publication-to-first-completed- +update p50/p95 changes from 9.662/19.650 to 0.271/0.310 seconds single, and +24.030/47.079 to 7.698/22.576 seconds dual, relative to generation handoff. +Complete times are 204.842/334.537 seconds for publication overlap, +209.339/335.005 for generation handoff, 200.088/334.321 for serial, and +216.944/350.906 for Offline-32. Serial remains slightly faster than overlap; +state coverage remains similar. This supports earlier usable data, not a +positive isolated overlap speedup or a diversity advantage. + +Serial matches numeric trajectories, initialization, normalization and ordered +source/anchor work. Raw rendered RGB equality is not established, and final +model weights may differ. AV1 is lossy. Shared-host observations and one-seed +pilots do not establish a hardware-independent speedup. + +## Reproduction and artifacts + +Use the [README](README.md) to generate fresh data, run calibration/training, +audit saved standard datasets and render reports. The historical tree retains +exact protocol files, executed source snapshots, qualification outputs, raw +metadata archives and their hash inventories. Large datasets/checkpoints remain +at the explicitly inventoried external local paths; they were never uploaded +to Git and are not portable public artifacts. A rerun must create a fresh output +directory and record its own source, package, hardware and native-library +versions. Compare within that frozen recipe rather than mixing historical runs. + +## Review layers + +[PR #850](https://github.com/DexForce/EmbodiChain/pull/850) owns the standalone +ODS foundation: the flow controller, engine health/status and copy/lifecycle +fixes, their production tests and benchmark, the simple BC caller migration, +and public operational/API documentation. It was merged into `main` as +[`0dcacf739`](https://github.com/DexForce/EmbodiChain/commit/0dcacf7393f30876cd57e65d921f8df5ab1ee6a8) +after the full repository test gate passed, including the runtime-module import +fix. The controller accepts only its engine's live source; resident sources +sample directly through their own interface. + +[PR #828](https://github.com/DexForce/EmbodiChain/pull/828) builds on that +merged foundation and targets `main`. It owns the lazy shared ResNet10 encoder +plus the runnable visual training/data-stream example, qualification, recording, +auditing and reporting. +The delivery diagnostic also supplies helpers used by visual training; moving +individual entry points requires extracting those shared helpers first. + +The 1,040 generated evidence files were removed from the code diff: 174,638,781 +logical bytes and approximately 4.59 million added text lines. Their immutable +backup is linked above. Future VLA/WAM adapters and richer task distributions +require their own implementation changes. diff --git a/examples/ods/data_stream_demo/README.md b/examples/ods/data_stream_demo/README.md new file mode 100644 index 000000000..3b3255c15 --- /dev/null +++ b/examples/ods/data_stream_demo/README.md @@ -0,0 +1,615 @@ +# ODS visual training: diversity and end-to-end cost + +The main experiment asks: **with the same actual training budget and hardware, +how much distinct data reaches the model, and how long does the complete +pipeline take?** It runs real BC training from scratch on the existing Franka +PickPlace task. Only third-person (`single`) and third-person plus wrist +(`dual`) camera policies participate in the training comparison. + +Accepted data volume and trained-state coverage may differ. The common budget +is optimizer updates, batch size and valid supervision, not epochs over +unequal datasets. All paths save a standard LeRobot dataset, checkpoint and +an auditable training manifest. A faster result is interpreted together with +the achieved data coverage and independent policy evaluation. + +The complete contract is in the +[fixed-budget visual design](../../../docs/architecture/ods-data-stream-showcase-design.md#fixed-training-budget-visual-comparison). + +For task learning with periodic student evaluation, start with the separate +[dual-camera convergence recipe](CONVERGENCE.md). It provides an acquisition-only +camera preview and a reviewable configuration before training, then selects a +checkpoint on development states before one held-out test. Its longer learning +budget is separate from the fixed-work system comparison below. + +Historical protocols, negative results and their immutable backup links are +listed in [experiment evidence](EXPERIMENTS.md). The measured pilots establish +earlier usable data, but do not establish a diversity advantage, an isolated +overlap speedup, or a converged PickPlace policy. Generate fresh run outputs +using the commands below; generated evidence is excluded from the source tree. + +## Primary training paths + +| CLI path | Data | Purpose | +| --- | --- | --- | +| `--arm offline --offline-episodes 8` | Finalize 8 episodes, then SDK read/decode and train | Small corpus, low acquisition cost | +| `--arm offline --offline-episodes 32` | Same, with 32 episodes | Intermediate time–coverage point | +| `--arm offline --offline-episodes 128` | Same, with 128 episodes | Larger corpus, higher acquisition cost | +| `--arm ods_overlap` | Start after C episodes, generate and record the next block during training | Main ODS demonstration | +| `--arm ods_serial --trace ...` | Regenerate overlap sources and replay its training manifest without generation/training overlap | Attribution control | + +Defaults are U=128 updates, B=32, H=2, F=32, C=8 shared episodes, P=2 producer +environments and K=128 per-source lifetime ODS deliveries. Qualified runs +complete 4,096 chunks and 131,072 valid future action targets. K is a guarded +engineering setting, not a demonstrated best learning-reuse value. Offline +paths reuse their fixed corpora as required and disclose the resulting reuse. + +ODS switches to a ready new block at a batch boundary after at least one +completed batch has covered every current source. It keeps training the +current block while production runs, within K, and waits when credits are +exhausted. The default `--prefetch-policy budget` requests new data when the +remaining credits cannot fund the remaining updates, or measured generation, +recording backpressure and snapshot times predict a refresh before training +ends. The latter decision uses recent fast, complete batch/update times, not +the first cold CNN step. Unknown timing permits quota-required requests only. +K remains a maximum: a fast producer can replace a block before K is spent. +`--prefetch-policy eager` retains the earlier request-after-every-switch rule +as a diagnostic control; it can accept data that never reaches training. +ODS does not repeat early data without a bound. At update U it stops +requesting new generations, completes any in-flight generation and records +all accepted data, including untrained final sources. Those sources do not +increase diversity metrics. + +The default `--handoff-mode publication` keeps the initial C-episode block for +common normalization, then reads complete accepted producer subbatches of P +episodes as they are published. It does not wait for the rest of that +generation. Raw recording references are released after an independent packet +is packed; the union of training, pending-archive and serial-replay raw owners +stays at most C. The writer queue still has two packets, plus one active writer +and at most one pending submission. Returned training batches remain owned. +Publication mode requires `P*K >= B`; every retained subbatch must fund at least +one full batch. Source quotas are never replenished by replacing a subbatch. + +`--handoff-mode generation` retains the full-generation handoff as a control. +Offline acquisition uses the same chosen handoff and recorder, then waits for +the complete dataset before standard SDK decoding/training. Publication-mode +Serial accepts and archives subbatches during production, holds at most C raw +episodes, and delays model updates until the generation completes. It therefore +removes generation/training overlap without adding an archive-publication +barrier. Trace v2 records retained-source boundaries and recorder/handoff +settings; old v1 traces remain readable. + +`--encoder-threads` exposes the AV1 encoder's CPU thread limit, default 1. +Use the same setting and CPU allocation across compared paths. The SDK writer +remains single-owner and uses sequential camera encoding; it does not fork a +camera encoding pool from the CUDA learner's writer thread. Encoding settings +are recorded and checked against the finalized standard dataset. + +The model follows the pinned RLInf ResNet10/GroupNorm/spatial-embedding design, +with a shared trainable CNN, eight spatial embeddings and 256-dimensional +projection, followed by a 512/512 BC head. Both camera histories share encoder +weights. It starts from scratch, uses only two causal RGB images per camera +and two 27-feature robot states, and predicts 32 future TCP/rotation/gripper +targets. Cube geometry is retained only for supervision/auditing, not policy +input. ImageNet normalization runs once on RGB float32 [0,1]. Reference +alignment, gradients, image-dependent small-data fitting and reload checks +must pass before formal measurements. + +## Run the visual training experiment + +Use the py311 environment with matching DexSim/Franka assets, PyTorch 2.7, +torchvision 0.22, LeRobot 0.4.4 and TorchCodec 0.5. The example requires two idle +primary GPUs with sufficient memory: one for simulation, one for CNN training. +Use the qualified native device mapping and disclose any explicitly reserved +auxiliary graphics device. Physical IDs below are examples, not reservations; +replace them with idle devices on the current machine. + +```bash +conda activate py311 +python -m pip install 'torch==2.7.0' 'torchvision==0.22.0' \ + --index-url https://download.pytorch.org/whl/cu128 +python -m pip install -e '.[vision]' 'lerobot==0.4.4' 'torchcodec==0.5' +python -m examples.ods.data_stream_demo.train --help + +python -m examples.ods.data_stream_demo.train \ + --suite smoke --gpu-ids 2 3 --output outputs/ods-visual-smoke + +python -m examples.ods.data_stream_demo.train \ + --suite calibrate --gpu-ids 2 3 --output outputs/ods-visual-calibration + +python -m examples.ods.data_stream_demo.train \ + --suite pilot --gpu-ids 2 3 --evaluation-scenes 0 \ + --output outputs/ods-visual-pilot + +python -m examples.ods.data_stream_demo.train \ + --suite report --gpu-ids 2 3 --output outputs/ods-visual-formal +``` + +Freeze model, camera profile, hardware and training budget after preflight; +do not increase epochs or insert delays to manufacture overlap. Formal +execution has 24 primary runs: four paths×two camera modes×three fixed seeds, +plus six matching serial replays, for 30 runs total. Primary runs regenerate +their data, rotate path order and retain failed/slower attempts. Preflight +measures real producer and learner throughput, not a sleep consumer. + +### Select a configuration from measured capacity + +Preflight preserves the requested C/P/B/K instead of substituting a fixed +producer configuration. Its default update count is +`max(16, ceil(C*K/B)+8)`, so it can observe new-generation use after the first +block's credits run out. `--calibration-updates` overrides that diagnostic +budget; a short run without later-generation training is explicitly marked +insufficient for steady-supply conclusions. + +Measure three capacities in each visual condition: accepted generation +`gamma_gen` (episodes/s), SDK write/encode busy-service `gamma_record` +(episodes/s), and ready batch preparation + transfer + synchronized training +`mu_ready` (updates/s). The pure model update rate is retained separately. +Use `gamma_effective = min(gamma_gen, gamma_record)` as a bottleneck estimate, +then compare `gamma_effective*K/B` with `mu_ready`. This is a service-capacity +diagnostic, not an end-to-end speed prediction: staging, startup, shutdown and +shared resource contention still cost time. Never add overlapping stage +durations and call the sum pipeline time. + +Sweep only feasible P values that divide C, holding camera settings, hardware, +CPU affinity, B, K and model fixed. Select the smallest measured P reaching +90% of the best measured effective supply for that condition. Incompatible +hardware/recipes or unqualified runs cannot enter that selection. Increasing P +cannot solve a recorder bottleneck. If one producer setting must serve both +camera conditions, use the smallest tested P reaching that threshold in both, +and retain each condition's individual choice. The implied K for full learner +utilization +is reported as a diagnostic, **not** permission to raise the reuse guard or a +recommended learning setting. Keep K bounded and report unavoidable waiting +when production cannot support the actual model workload. + +For example, measure P=2/4/8 on the same admitted GPU pair and compare the +saved calibration packets. Recheck device availability before each command: + +```bash +python -m examples.ods.data_stream_demo.train --suite calibrate \ + --producer-envs 2 --gpu-ids 2 3 --output outputs/ods-capacity/p2 +python -m examples.ods.data_stream_demo.train --suite calibrate \ + --producer-envs 4 --gpu-ids 2 3 --output outputs/ods-capacity/p4 +python -m examples.ods.data_stream_demo.train --suite calibrate \ + --producer-envs 8 --gpu-ids 2 3 --output outputs/ods-capacity/p8 +python -m examples.ods.data_stream_demo.training_throughput \ + outputs/ods-capacity/p2 outputs/ods-capacity/p4 outputs/ods-capacity/p8 \ + --output outputs/ods-capacity/selection +``` + +The one-seed `pilot` runs Offline-32, budget-aware ODS overlap and its exact +numeric/source-anchor serial replay, plus a separately reported eager control, +in both camera modes. Each mode keeps the same GPU pair and eight CPU cores +across all four paths. It holds the actual training budget fixed, retains +full-corpus audits, and writes its primary report from `main/` so the eager +control cannot contaminate baseline aggregates. It is a scoped engineering +comparison, not the three-seed formal study. Zero evaluation scenes explicitly +means no new policy-quality evidence. Run the full matrix only after this +pilot establishes the intended resource and scheduling behavior. + +The supervisor also owns device selection for single runs. It masks the +selected UUID pair, assigns logical `cuda:0` to generation and `cuda:1` to +training, and audits the owned process group: + +```bash +python -m examples.ods.data_stream_demo.train \ + --suite single --arm ods_overlap --vision single --gpu-ids 2 3 \ + --output outputs/ods-visual/overlap-single + +python -m examples.ods.data_stream_demo.train \ + --suite single --arm offline --offline-episodes 32 --vision dual \ + --gpu-ids 2 3 \ + --output outputs/ods-visual/offline32-dual +``` + +Do not reuse already finalized files as the cold offline baseline. Both +pipelines include fresh acquisition, failed expert attempts, encoding, +initialization, actual optimizer work, checkpoint writing, finalization and +normal shutdown. Last-update time is recorded separately from total time, +so pending recording work cannot disappear from the comparison. +All paths permit recording to overlap the next generation within the same +bounded storage. Request the next shared-cache fill before waiting for old +owned episodes to enter the writer queue; release those old snapshots before +taking new ones. Offline still starts training only after its full corpus is +finalized. Serial still starts each generation after the preceding optimizer +updates finish, and starts its next updates only after generation completes. +An extra recorder-submission barrier must not make either control slower. + +The supervisor then launches a separate CPU corpus/decoder audit and the +independent policy evaluation. Their durations are saved and excluded from +pipeline time. A failed audit or evaluation retains the original measured +pipeline time and disqualifies that run for comparative claims. +After the stage leader exits, the supervisor allows up to two seconds for its +owned process group to disappear. Remaining members, including zombies, still +fail the stage; PID, parent, state and termination evidence are retained in +`process_group_shutdown`. This wait belongs to the stage, including pipeline +time for the training stage. Historical studies and failures are preserved separately in +[experiment evidence](EXPERIMENTS.md). +Actual runs remain unqualified while validation is pending. The supervisor +persists the requested evaluation scene count; formal reports require all +requested scenes and normal evaluation closure, or an explicit zero-scene +`not_requested` status. + +On a qualified driver, the new training supervisor can accept an explicit +`--aux-graphics-gpu-id` for the native driver's passive graphics footprint. +Outside each pipeline's primary pair it permits only owned pure-G contexts of +at most 8 MiB per PID, requires 16 GiB free and records other board activity. +The flag alone grants no compute capacity. If the same board is explicitly +listed in `--gpu-ids`, it may also serve as another pipeline's primary after +passing all three idle readings. Every process group remains limited to its +own two primary GPUs for computation. This exception belongs to the training supervisor. +The historical delivery supervisor below retains its stricter requirement +that an auxiliary board also pass three idle readings. + +## Training outputs and interpretation + +Render only retained training logs after workers exit: + +```bash +python -m examples.ods.data_stream_demo.training_report \ + outputs/ods-visual-formal --output outputs/ods-visual-report +``` + +Per-run `summary.json` contains the fixed work ledger and `training_metrics`. +`training-events.jsonl` records immutable completed-update source/anchor IDs, +losses and monotonic stage timestamps. `training-sources.json` retains exact +per-source frame counters, including retired sources. The report writes +JSON/Markdown and four measured figures, each in PNG/SVG/PDF: + +1. Total time versus initial-state coverage and effective trained sources. +2. Coverage growth against wall time and completed optimizer updates. +3. Generation, recording, data preparation, training and cleanup timeline. +4. Lifetime source use, including retired/unused sources and the ODS K line. + +`writer-sdk-events.jsonl` streams the recording thread's actual SDK service +intervals. They include writing, encoding and recording bookkeeping, while +prefetch-worker SDK fetch spans include temporal reads, decoding and validation. +The timeline distinguishes both from owner copying and waiting; overlapping +worker durations are not additive pipeline time or isolated codec CPU time. +Each episode's existing SDK service event now contains four nested stage +intervals: provenance validation, `add_frame`, `save_episode`, and reference/ +manifest bookkeeping. Their durations are also aggregated in writer stats. +`save_episode` includes encoding and SDK commit work; it is not a pure codec +kernel measurement. Nested intervals must not be added again to their parent. +`updates_before_generation_fully_published` counts completed updates before +the last episode publication of their generation. Copying a source alone does +not increment this metric. `peak_owned_raw_episodes` audits the C-slot owner +union; packed writer packets, training batches and simulator/runtime memory +are separately bounded or measured. `recording_capacity_wait` indicates lack +of owned raw slots and can overlap ongoing production; it is not an isolated +causal attribution of all waiting to the recorder. +`dataset_bytes` and `checkpoint_bytes` record finalized logical file sizes, +excluding temporary disk peaks and later audit/evaluation artifacts. + +The post-timing audit checks every numeric frame, an SDK H2 window from every +source, sparse RGB PSNR/MAE references, and sequential software decoding of +**every saved AV1 frame**, including untrained tails. Full-video frame hashes +cover decoded YUV420p planes; they do not imply lossless RGB or exhaustive SDK +random-seek equivalence. Fresh audit receipt/artifact paths are required, so +rechecking a historical run preserves the original evidence: + +```bash +python -m examples.ods.data_stream_demo.training_audit \ + --run-dir outputs/ods-visual/overlap-single \ + --output outputs/ods-visual/overlap-single/post-audit-new.json +``` + +Only completed optimizer updates increase training diversity. Metrics include +unique sources, entropy-based effective sources, fixed 1 cm initial XY bins, +2 cm cube-state voxels at trained anchors, unique anchor/history/target frame +IDs and duplicate-target fraction. New source IDs alone do not demonstrate +physical-state novelty; overlapping prediction windows are deduplicated. +A registered-but-untrained source contributes to acquisition cost, not coverage. + +Failed, unsupported or unqualified runs remain in the report but are excluded +from ratios. Formal same-vision/seed ratios require matching initial-model and +normalization hashes, the common initial-C numeric prefix and camera order, +the exact declared/completed U/B/H/F budget, the same assigned generation and +training GPU UUIDs, and the same eight CPU cores. Supervisor launch masks, +thread budgets, placement audits, normal-exit records and total times must +agree; missing evidence excludes the comparison. Smoke/calibration and legacy +CPU metric fixtures do not produce formal speedup claims. AV1 is lossy; +retain PSNR/MAE audits and the independent +20-state development evaluation beside timing. A smaller offline corpus may +finish sooner or a larger one may cover more states. No universal speedup or +policy-learning improvement is assumed. + +## Historical delivery diagnostic + +The original no-model entry point below remains available to isolate cache, +delivery, conversion and standard video-decode costs. Its state-only condition +belongs to that diagnostic, not the visual training matrix above. It does not +import a learner or produce policy-performance evidence. + +The historical delivery and CPU-copy studies remain available through +[experiment evidence](EXPERIMENTS.md). Their measurements retain their original +scope and are not relabeled as visual-training results. + +## Compared paths + +| Arm | Source availability | Per-source chunk budget | +| --- | --- | --- | +| `lerobot` | Finalized standard Parquet/MP4 corpus; SDK/DataLoader/TorchCodec access | K | +| `streaming` | ODS initial fill, followed by asynchronous cache replacement | Unlimited | +| `protected` | Same ODS refill recipe | K | +| `lerobot_unbounded` | Full finalized LeRobot corpus, decoded on demand; appendix diagnostic | Unlimited | +| `frozen` | ODS initial cache only; appendix diagnostic | K | +| `lerobot_same_cache` | First C LeRobot episodes; appendix diagnostic | K | + +Each arm supports `state`, `single` and `dual`: no camera, fixed third-person +camera, and the same camera plus a wrist-mounted camera. Visual output is +256x256 RGB float32 in [0,1], with all context frames delivered. ODS native +RGBA-to-RGB conversion is timed, as are LeRobot decoding and collation. + +The formal defaults are C=8 cached episodes, P=2 producer environments, +N=32 accepted episodes, B=32 chunks, history=2, future=32, K=128. Each chunk +contains 33 context frames and costs one credit, independent of its camera +count. K is a demonstration value, not an optimal learning-reuse setting. +An evicted source's unused credits are retired, never transferred to another +source. The main suite uses seeds 19001, 19002 and 19003; appendix diagnostics +use the first seed only. + +## Requirements + +Use a source installation of EmbodiChain with a matching DexSim runtime, +Franka assets and a supported CUDA simulation device. State-only acquisition +uses `no-render`; visual acquisition still needs native headless rendering. +The standard reader requires the pinned SDK/decoder and an FFmpeg build +supporting the selected AV1 encoder: + +Validation used PyTorch 2.7.0+cu128 with the SDK/decoder versions below. The +isolated DataLoader lifecycle hook is tested against that PyTorch version; +use a matching simulator and Torch/TorchCodec build. + +```bash +python -m pip install -e . +python -m pip install 'lerobot==0.4.4' 'torchcodec==0.5' +python -m examples.ods.data_stream_demo.run --help +``` + +Select idle physical GPUs and CPU thread settings before starting Python: + +```bash +export CUDA_VISIBLE_DEVICES="$(nvidia-smi --id=7 --query-gpu=uuid --format=csv,noheader)" +export ODS_GPU_UUID="$CUDA_VISIBLE_DEVICES" +export SDL_VIDEODRIVER=dummy +unset DISPLAY WAYLAND_DISPLAY +export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 +``` + +This selects physical GPU7 by its NVML UUID. Pass logical +`--sim-device cuda:0` for the single run, or use `--gpu-ids 7` for a suite. +The matching native runtime maps the primary CUDA UUID to the primary Vulkan +UUID. Qualify compute **and graphics** placement with the actual ODS subprocess +path before running the matrix. On the validation host, Vulkan instance/driver +initialization additionally creates a small graphics footprint on GPU7, so +GPU7 is the qualified strict single-GPU example there. Other machines require +their own device qualification; unmasked physical ordinals are not a fallback. +GPU7 is an example selection, not a reservation. Check that the selected +device is idle and has enough memory. Do not silently shrink an unsupported +workload. The default dual-camera native image cache alone occupies about +2.34 GiB; a returned RGB float32 batch contains about 1.55 GiB of image +payload, before decoder, normalization and prefetch temporaries. Shared +memory, RAM and VRAM are checked separately. + +The acquisition reuses the existing BC demo's physical task and expert but +does not import its trainer or modify its sensor-off defaults. This example +owns its additional camera profiles and decoded data adapter. + +The qualified camera profile is `v3`: fixed third-person and side-mounted +wrist views, `fast-rt`, SPP=1, offscreen OptiX spatial denoising, no window +denoising or super-resolution. V1's self-occluded wrist view and noisy raw +rendering were rejected before timing. V3 also separates the inherited +cube/goal material UIDs so their original orange/green colors remain distinct; +it changes no physical property or success criterion. + +## Run and inspect + +Use a fresh output directory for every run. The stable single-run entry is: + +```bash +python -m examples.ods.data_stream_demo.run \ + --suite single --arm protected --vision state --sim-device cuda:0 \ + --output outputs/ods-stream/protected-state + +python -m examples.ods.data_stream_demo.run \ + --suite single --arm lerobot --vision single --sim-device cuda:0 \ + --output outputs/ods-stream/lerobot-single +``` + +These commands select paths; an engineering run alone is not a completed +27-run report. Freeze calibration, rates, hardware and profile before comparing +formal runs. The matrix is three main arms by three visual conditions by +three seeds, with nine separately identified appendix diagnostics. Do not mix +different profiles in one report input directory. + +Start with the small engineering suite before allocating the formal budget: + +```bash +python -m examples.ods.data_stream_demo.run \ + --suite smoke --gpu-ids 7 --output outputs/ods-stream-smoke +``` + +`--gpu-ids` accepts physical GPU IDs. If omitted, the launcher selects devices +with three zero-utilization readings and at least 16 GiB free, then assigns at +most one child per selected GPU and eight disjoint CPUs per child. It sets +both `CUDA_VISIBLE_DEVICES` and `ODS_GPU_UUID` to the selected physical UUID, +then passes logical `cuda:0`. Reproduction does not require several GPUs. + +The parent periodically audits compute and graphics contexts belonging to +its child PID tree or process group. An owned context on another GPU fails +the cell and terminates only that owned process group. Context samples are +retained below `_supervision/`; unrelated processes are never terminated. +There is no global GPU0-idle prerequisite. Periodic checks can miss short-lived +contexts, so they supplement qualification of the matching native UUID mapping +and do not establish an exclusive-device reservation. Previously blocked +calibration attempts remain failed evidence and do not become a usable protocol. + +For the qualified validation-host driver, an optional auxiliary reservation +allows another idle primary GPU while accounting for that graphics footprint: + +```bash +python -m examples.ods.data_stream_demo.run \ + --suite calibrate --gpu-ids 1 3 4 5 6 --aux-graphics-gpu-id 7 \ + --output outputs/ods-stream-calibration-aux +``` + +The auxiliary GPU must pass the same three idle readings and 16 GiB admission, +is excluded from every primary assignment, and is shared by this suite's jobs. +Only its exact UUID with pure `G` type and known per-PID memory at most 8 MiB is +allowed. Unknown memory, compute contexts, larger graphics footprints or any +other GPU fail the job. Primary and auxiliary owned memory are both recorded, +deduplicating `(PID, UUID)` records before summing across devices. This scope +uses primary GPUs **plus a shared reserved auxiliary GPU**, and is not a +single-GPU result. Low memory and `G` type alone do not prove no computation; +the permission relies on the separate controlled native-runtime qualification. + +An independent minimal libvulkan probe, without CUDA/SDL/DexSim or device/queue submission, +reproduced the 6 MiB GPU7 footprint after `vkCreateInstance` and its disappearance +after instance destruction. A real dual-camera ODS episode then qualified +primary GPU3 plus the explicit GPU7 footprint with graceful process cleanup. +These observations support a bounded **Vulkan instance/driver initialization** +scope, not a claim of zero overhead or a general graphics exception. + +The default remains strict. Calibrate, report and appendix must each receive +the same explicit auxiliary flag when opting in; the frozen protocol records +its UUID, model, memory limit and the offscreen environment. The launcher uses +`SDL_VIDEODRIVER=dummy` with `DISPLAY`/`WAYLAND_DISPLAY` unset. This removes the +X11 dependency; it does not remove the Vulkan graphics footprint. + +Smoke checks all nine main cells with one seed, C=P=2, N=8, B=4, K=8, +30 seconds of consumption and a 16 chunks/s offered load. Images and context +length remain unchanged. Smoke verifies paths and cleanup, not the formal load +or generation coverage. + +Calibrate capacity, then freeze the resulting protocol for the main suite: + +```bash +python -m examples.ods.data_stream_demo.run \ + --suite calibrate --gpu-ids 7 \ + --output outputs/ods-stream-calibration + +python -m examples.ods.data_stream_demo.run \ + --suite report --gpu-ids 7 \ + --protocol outputs/ods-stream-calibration/calibration.json \ + --output outputs/ods-stream-main +``` + +The protocol fixes the common offered rate and observation duration based on +measured generation and ready-source throughput; it does not select a +learning-optimal K. Formal runs do not silently shorten an expensive visual +condition. Calibration and smoke results remain separate from main estimates. +Use a custom `--suite single --duration-seconds ...` run to explore a different +window. The smoke recipe and formal protocol have fixed windows. + +Calibration performs six cells: protected ODS and unbounded standard decoded +LeRobot for each visual condition. It requires all 32 accepted episodes and +four completed generations. The packet records generation rates, observed +batch-service estimates, actual decoded-reader throughput, source hashes and +camera profile. It is not a throughput guarantee or a learning-budget search. +Code/profile changes invalidate the packet and require recalibration. + +After the main suite, run the nine appendix diagnostics using the same packet +and its seed 19001 corpora: + +```bash +python -m examples.ods.data_stream_demo.run \ + --suite appendix --gpu-ids 7 \ + --protocol outputs/ods-stream-calibration/calibration.json \ + --dataset-root outputs/ods-stream-main \ + --output outputs/ods-stream-appendix +``` + +For appendix, `--dataset-root` is the **main report root**. The launcher looks +for `lerobot//seed19001/lerobot_dataset/` below it. Frozen starts a new +physical producer; the two SDK diagnostics reuse those finalized files. +For `--suite single`, `--dataset-root` instead names one concrete finalized +LeRobot dataset directory. + +Each child has a finite watchdog and its own process group. The parent +passes the same-host monotonic launch timestamp so cold first-batch timing +includes Python/library startup and instrumentation setup. Exact commands, +device/CPU assignments and stdout are stored below `_supervision/`; final +per-run summaries remain below `//seed/`. Failures, +timeouts, cancellation and unsupported cells remain in `suite-summary.json`. +The launcher kills only a process group it created, and never retries a +failed experimental cell automatically. + +The LeRobot writer uses the official SDK and real MP4 video features. The +reader calls the public Dataset temporal API in spawned DataLoader workers, +with actual CPU TorchCodec decoding. It does not load a decoded corpus into +resident tensors. The parent owns delivered and outstanding K reservations; +recreating an iterator does not replenish them. Normal OS file caches, +decoder instances and bounded prefetch are retained and disclosed. The +experiment does not clear system-wide caches. + +Closing first stops new dispatch and drains existing prefetch without counting +it as delivered. Every reader and quality-audit worker must exit with code zero; +joining an aborted worker or forcing termination fails the run. Summaries retain +discarded reservations and shutdown time. Iterator restart keeps already +delivered K usage. + +## Render results + +Rendering reads retained logs and performs no acquisition, decoding or new +timing experiment: + +```bash +python -m examples.ods.data_stream_demo.report outputs/ods-stream-main \ + --output outputs/ods-stream-report --animation +``` + +The input must be the exact experiment root containing run directories, or +one run's `summary.json`. Do not point it at an unrelated results tree. It +reads `summary.json`, `lifecycle.jsonl` and `batches.jsonl`; malformed records +raise an error instead of being silently dropped. + +Every suite renders its report automatically after all children stop. A manual +render can add the optional animation or combine retained results deliberately; +keep smoke, calibration and different protocol versions clearly identified. + +Output includes one Markdown report, `report.json`, and four figures in PNG, +SVG and PDF: + +1. Cache residency, generation and delivered-source timeline. +2. Per-source delivered chunks, including evicted identities and the K line. +3. First-batch latency, consumer p95 and delivered rate versus offered demand. +4. State/single/dual throughput, logical bandwidth, memory and storage costs. + +`--animation` additionally replays logged protected-run deliveries as a GIF. +It is generated after timing and is not a simulation recording. Late complete +batches remain in lifetime K accounting and are marked with crosses in the +timeline; they do not enter observation-window throughput or latency estimates. +Raw camera references and decoded image quality remain part of the run artifacts. + +Filled points and median/range aggregates use only complete, explicitly +qualified `report`/`appendix` runs. Failed, unsupported, unqualified, smoke +and calibration runs remain visible; measured points are hollow and missing +metrics are unknown, never zero. One-seed appendix results remain diagnostics. +The report does not select favorable seeds or rank a performance winner. + +## Read the evidence correctly + +Launch-to-first-batch includes generation and the chosen availability barrier: +LeRobot waits for N finalized episodes, whereas ODS starts after C episodes. +Source-ready first-batch starts from available data, but LeRobot still reads +and decodes. These two timing scopes overlap and must not be added. Their +entire difference cannot be attributed to shared-memory transport. + +Wait counters overlap only the source-ready observation window; cold generation +waits remain in `lifetime_data_wait_seconds`. Request throttling and terminal +supply-exhausted idle time are separate. Controller recovery includes the final +successful retry service, so wait and service counters are not mutually +exclusive phases and must not be summed as a wall-time breakdown. + +Protected streaming can deliberately wait for new data after exhausting the +current budget. Report this separately from slow copying or decoding. A +source ID is not a measure of state coverage; delivered chunks are not +optimizer updates. AV1 is lossy, so throughput comparisons retain quality +and encoding settings. Fixed ODS cache bytes do not imply constant total +process memory or lower memory than a bounded DataLoader. + +This demo can establish measured availability, reuse protection, provenance +and resource tradeoffs for these workloads. It does not establish policy +quality, training acceleration, a universal ODS speedup, or that the broader +LeRobot ecosystem cannot stream data. diff --git a/examples/ods/data_stream_demo/__init__.py b/examples/ods/data_stream_demo/__init__.py new file mode 100644 index 000000000..52ec4a8f7 --- /dev/null +++ b/examples/ods/data_stream_demo/__init__.py @@ -0,0 +1,21 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Real ODS generation and standard-decoded LeRobot capability experiments.""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/examples/ods/data_stream_demo/acquisition.py b/examples/ods/data_stream_demo/acquisition.py new file mode 100644 index 000000000..dde8d25b3 --- /dev/null +++ b/examples/ods/data_stream_demo/acquisition.py @@ -0,0 +1,448 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Real expert acquisition shared by the stream and file-backed comparisons.""" + +from __future__ import annotations + +import math +from collections.abc import Iterator +from typing import TYPE_CHECKING, Any, Literal + +from examples.ods.simple_bc_demo.config import load_task_config + +if TYPE_CHECKING: + from tensordict import TensorDict + + from embodichain.data_pipeline.engine import OnlineDataEngine + +__all__ = [ + "CAMERA_PROFILE_VERSION", + "VisionMode", + "camera_uids", + "camera_configs", + "camera_profile_version", + "load_acquisition_config", + "make_engine", + "iter_generation_episodes", + "poll_generation_episodes", +] + +VisionMode = Literal["state", "single", "dual"] + +# The original +Y wrist mount self-occluded. Version 2 qualified a fixed +# +X-side mount with SPP=1/OptiX, without temporal reconstruction or +# super-resolution. Version 3 keeps those settings and separates the inherited +# cube/goal material identities so their existing colors remain independent. +CAMERA_PROFILE_VERSION = "v3" + + +def camera_profile_version(profile: str = "benchmark") -> str: + """Identify the benchmark or close-up PickPlace camera calibration.""" + if profile not in ("benchmark", "pick_place"): + raise ValueError(f"Unknown camera profile: {profile!r}") + return CAMERA_PROFILE_VERSION if profile == "benchmark" else "pick-place-v2" + + +def camera_uids(vision: VisionMode) -> tuple[str, ...]: + """Return the frozen camera order, rejecting unknown profiles.""" + if vision == "state": + return () + if vision == "single": + return ("cam_high",) + if vision == "dual": + return ("cam_high", "cam_wrist") + raise ValueError(f"Unknown vision profile: {vision!r}") + + +def _camera( + uid: str, extrinsics: dict[str, Any], *, fov_degrees: float = 60.0 +) -> dict[str, Any]: + # Square pinhole images use the same horizontal and vertical FOV. + focal = 128.0 / math.tan(math.radians(fov_degrees / 2)) + return { + "sensor_type": "Camera", + "uid": uid, + "width": 256, + "height": 256, + "intrinsics": [focal, focal, 128.0, 128.0], + "near": 0.005, + "far": 3.0, + "enable_color": True, + "enable_depth": False, + "enable_mask": False, + "enable_normal": False, + "enable_position": False, + "extrinsics": extrinsics, + } + + +def camera_configs( + vision: VisionMode, *, camera_profile: str = "benchmark" +) -> list[dict[str, Any]]: + """Return ordered camera calibration without importing simulation bindings.""" + camera_profile_version(camera_profile) + cameras = camera_uids(vision) + sensors = [] + if cameras: + eye, target = [-0.50, -0.90, 0.88], [-0.50, 0.10, 0.20] + if camera_profile == "pick_place": + eye, target = [-0.60, -0.45, 0.50], [-0.42, 0.10, 0.08] + sensors.append( + _camera( + "cam_high", + { + "eye": eye, + "target": target, + "up": [0.0, 0.0, 1.0], + }, + ) + ) + if vision == "dual": + sensors.append( + _camera( + "cam_wrist", + { + "parent": "FrankaPanda/fr3_hand_tcp", + "eye": ( + [0.12, 0.0, -0.04] + if camera_profile == "pick_place" + else [0.10, 0.0, -0.10] + ), + "target": ( + [0.0, 0.0, 0.05] + if camera_profile == "pick_place" + else [0.0, 0.0, 0.03] + ), + "up": [0.0, 1.0, 0.0], + }, + fov_degrees=90.0 if camera_profile == "pick_place" else 60.0, + ) + ) + return sensors + + +def load_acquisition_config( + vision: VisionMode, + device: str, + producer_envs: int = 2, + seed: int = 19001, + *, + camera_profile: str = "benchmark", +) -> dict[str, Any]: + """Expand the existing physical task with the showcase's sensor profile. + + The expert, reset distribution, physical success checks and action decoder + are the existing BC task's implementations. Only the independent camera + suite, scene material identities and run-local production settings change. + This function never imports a learner, policy model or trainer. + + Args: + vision: State-only, third-person, or third-person and wrist acquisition. + device: Explicit simulation tensor device, such as ``cuda:0``. + producer_envs: Number of parallel physical environments. + seed: Expert/reset random seed. + camera_profile: Historical benchmark or close-up PickPlace calibration. + + Returns: + Expanded Gym configuration; callers may serialize it as a run artifact. + """ + cameras = camera_uids(vision) + version = camera_profile_version(camera_profile) + if type(producer_envs) is not int or producer_envs < 1: + raise ValueError("producer_envs must be a positive integer") + if type(seed) is not int or seed < 0: + raise ValueError("seed must be a nonnegative integer") + config = load_task_config(device) + # Native materials are registered by UID. Both inherited shapes otherwise + # use "default_mat", causing the goal's green to replace the cube's orange. + # Keep their original RGBA values and every physical property unchanged. + for collection, object_uid in ( + ("rigid_object", "cube"), + ("background", "goal_marker"), + ): + obj = next(obj for obj in config[collection] if obj["uid"] == object_uid) + obj["shape"]["visual_material"]["uid"] = f"ods_stream_{object_uid}_material" + config.update( + num_envs=producer_envs, + seed=seed, + headless=True, + enable_sensor=bool(cameras), + camera_profile_version=version, + renderer="fast-rt" if cameras else "no-render", + render_cfg={ + "spp": 1, + "denoising": { + "window": "off", + "offscreen": "optix" if cameras else "off", + }, + }, + ) + # Replace the inherited preview camera; state-only keeps no hidden sensor. + config["sensor"] = camera_configs(vision, camera_profile=camera_profile) + return config + + +def make_engine( + vision: VisionMode, + device: str, + *, + producer_envs: int = 2, + seed: int = 19001, + buffer_size: int = 8, + budget: int | None = None, + max_episode_steps: int = 600, + initialization_timeout: float = 1800.0, + shutdown_timeout: float = 30.0, + camera_profile: str = "benchmark", +) -> OnlineDataEngine: + """Allocate, but do not start, the manual CPU-shared-memory producer. + + Args: + vision: Frozen camera profile. + device: Simulation device; shared data remains on CPU. + producer_envs: Parallel expert environments, dividing ``buffer_size``. + seed: Acquisition seed, independent of the consumer RNG. + buffer_size: Number of rolling trajectory slots. + budget: Optional per-episode chunk budget and equal hard ceiling. + max_episode_steps: Physical episode and shared-row capacity. + initialization_timeout: Initial accepted-data fill deadline in seconds. + shutdown_timeout: Graceful producer shutdown bound in seconds. + camera_profile: Camera calibration shared with policy evaluation. + + Returns: + Unstarted engine with source provenance and bounded lifecycle logging. + """ + from embodichain.data_pipeline.engine import OnlineDataEngine, OnlineDataEngineCfg + + config = load_acquisition_config( + vision, device, producer_envs, seed, camera_profile=camera_profile + ) + if type(buffer_size) is not int or buffer_size < 1 or buffer_size % producer_envs: + raise ValueError("buffer_size must be positive and divisible by producer_envs") + if type(max_episode_steps) is not int or max_episode_steps < 1: + raise ValueError("max_episode_steps must be a positive integer") + config["max_episode_steps"] = max_episode_steps + return OnlineDataEngine( + OnlineDataEngineCfg( + gym_config=config, + state_dim=9, + action_dim=7, + action_semantics="policy_request", + action_contracts=("eef_pose.relative@1", "gripper.parallel@1"), + buffer_size=buffer_size, + max_episode_steps=max_episode_steps, + buffer_device="cpu", + refill_policy="manual", + include_provenance=True, + episode_sample_budget=budget, + episode_sample_hard_limit=budget, + lifecycle_event_capacity=2 * buffer_size, + initialization_timeout=initialization_timeout, + shutdown_timeout=shutdown_timeout, + ) + ) + + +def poll_generation_episodes( + engine: OnlineDataEngine, + generation: int, + next_episode_ordinal: int, + *, + max_count: int, +) -> list[TensorDict]: + """Copy an available ordered prefix without waiting for a generation. + + The owner advances the global ``next_episode_ordinal`` by the number + returned, and copies all rows before requesting another generation. + The producer publishes parallel episodes together, so a completed subbatch + is readable while it generates subsequent subbatches. A busy publication + lock or an unpublished next episode returns an empty list; this function + never waits for a writer or consumes the engine's sampling quota. + + Args: + engine: Started owner-side manual engine with provenance enabled. + generation: Requested generation whose rows have not been replaced. + next_episode_ordinal: Next global accepted episode ordinal. Generation + ``g`` contains ordinals ``(g - 1) * buffer_size`` through + ``g * buffer_size - 1``; its end ordinal is also accepted. + max_count: Maximum owned episode copies to return, bounded by the + caller's available retention slots. Zero performs only health and + generation checks. + + Returns: + Independent valid-prefix TensorDicts in consecutive slot order. Advance + the ordinal only after taking ownership of these returned episodes. + + Raises: + RuntimeError: The worker failed, a later generation was requested, or + an expected publication has been lost. + ValueError: Arguments or published validity/provenance are inconsistent. + """ + import torch + + if type(generation) is not int or generation < 1: + raise ValueError("generation must be a positive integer") + generation_start = (generation - 1) * engine.buffer_size + generation_end = generation * engine.buffer_size + if ( + type(next_episode_ordinal) is not int + or not generation_start <= next_episode_ordinal <= generation_end + ): + raise ValueError("next_episode_ordinal must belong to the requested generation") + if type(max_count) is not int or max_count < 0: + raise ValueError("max_count must be a nonnegative integer") + engine._require_owner_process("poll_generation_episodes") + status = engine.get_flow_status() + if status["requested_generation"] < generation: + raise ValueError("generation has not been requested") + if status["requested_generation"] != generation: + raise RuntimeError("Cannot acquire a generation after a newer refill request") + if not max_count or next_episode_ordinal == generation_end: + return [] + + publication_lock = engine._lock_index.get_lock() + if not publication_lock.acquire(False): + return [] + episodes = [] + try: + engine.check_health() + write_start, write_end = engine._lock_index[:] + next_slot = next_episode_ordinal - generation_start + for slot in range(next_slot, min(engine.buffer_size, next_slot + max_count)): + if write_start <= slot < write_end: + break + row = engine.shared_buffer[slot] + published_generation = row.get("ods_generation", None) + if published_generation is None: + raise ValueError("Incremental acquisition requires source provenance") + first_generation = int(published_generation[0]) + if first_generation != generation: + if ( + first_generation > generation + or status["completed_generation"] >= generation + or 0 <= slot < write_start + ): + raise RuntimeError( + f"Expected generation {generation} slot {slot} was not retained" + ) + break + valid = row["valid"].bool() + length = int(valid.sum()) + if length < 1 or not torch.equal( + valid, torch.arange(valid.numel(), device=valid.device) < length + ): + raise ValueError("Acquisition requires a nonempty valid prefix") + prefix = row[:length] + continuity = prefix.get("continuity_id", None) + if continuity is not None and ( + (continuity < 0).any() or (continuity != continuity[0]).any() + ): + raise ValueError("Acquisition cannot cross continuity boundaries") + for key, expected in ( + ("ods_run_id", engine.run_id), + ("ods_generation", generation), + ("ods_episode_id", generation_start + slot), + ): + if (prefix[key] != expected).any(): + raise ValueError(f"Episode {key} differs from publication identity") + if not torch.equal( + prefix["ods_frame_index"], + torch.arange(length, device=valid.device), + ): + raise ValueError("Episode frame ordering is inconsistent") + # Payload and identity must be copied under the same lock that + # publishes the next write window. The caller may retain the copy + # across subsequent fills without pinning shared cache slots. + episodes.append(prefix.clone()) + engine.check_health() + finally: + publication_lock.release() + return episodes + + +def iter_generation_episodes( + engine: OnlineDataEngine, generation: int +) -> Iterator[TensorDict]: + """Copy completed episodes one at a time for bounded file-writer staging. + + The caller must finish this iterator before requesting the next refill. + Each yielded TensorDict owns an independent, valid-prefix-only copy of one + episode, including its native images and immutable provenance. Keeping only + a bounded writer queue avoids duplicating the whole visual cache. + + Args: + engine: Started owner-side manual-refill engine. + generation: Completed generation that must still occupy every slot. + + Yields: + An independent ``[valid_steps]`` episode in cache-slot order. + + Raises: + RuntimeError: A refill is pending or data no longer matches generation. + ValueError: Episode validity/provenance/continuity is inconsistent. + """ + import torch + + if type(generation) is not int or generation < 1: + raise ValueError("generation must be a positive integer") + engine._require_owner_process("iter_generation_episodes") + for slot in range(engine.buffer_size): + status = engine.get_flow_status() + if ( + status["requested_generation"] != generation + or status["completed_generation"] != generation + ): + raise RuntimeError( + "Snapshot requires a completed, non-refilling generation" + ) + # The core has no full-episode snapshot API. Use the same publication + # lock as sampling and the existing BC offline collector. Never return + # a shared-buffer view, even when the producer is nominally idle. + with engine._lock_index.get_lock(): + if tuple(engine._lock_index[:]) != (-1, -1): + raise RuntimeError("Snapshot cannot copy an active write window") + row = engine.shared_buffer[slot] + valid = row["valid"].bool() + length = int(valid.sum()) + if length < 1 or not torch.equal( + valid, torch.arange(valid.numel(), device=valid.device) < length + ): + raise ValueError("Offline acquisition requires a nonempty valid prefix") + episode = row[:length].clone() + continuity = episode.get("continuity_id", None) + if continuity is not None and ( + (continuity < 0).any() or (continuity != continuity[0]).any() + ): + raise ValueError("Offline acquisition cannot cross continuity boundaries") + for key, expected in ( + ("ods_run_id", engine.run_id), + ("ods_generation", generation), + ): + if (episode[key] != expected).any(): + raise ValueError(f"Episode {key} differs from snapshot identity") + if ( + (episode["ods_episode_id"] < 0).any() + or (episode["ods_episode_id"] != episode["ods_episode_id"][0]).any() + or not torch.equal( + episode["ods_frame_index"], + torch.arange(length, device=episode.device), + ) + ): + raise ValueError( + "Episode source identity or frame ordering is inconsistent" + ) + yield episode diff --git a/examples/ods/data_stream_demo/configs/pick_place_dual_convergence.yaml b/examples/ods/data_stream_demo/configs/pick_place_dual_convergence.yaml new file mode 100644 index 000000000..875fb2bad --- /dev/null +++ b/examples/ods/data_stream_demo/configs/pick_place_dual_convergence.yaml @@ -0,0 +1,22 @@ +# Candidate learning recipe. Review before running the train command. +updates: 20000 +batch_size: 64 +buffer_size: 32 +producer_envs: 4 +episode_budget: 2048 +learning_rate: 0.0003 +final_learning_rate: 0.00003 +warmup_updates: 500 +prediction_horizon: 32 +action_horizon: 8 +encoder_threads: 2 +evaluate_every: 2000 +development_scenes: 20 +development_success_rate: 0.9 +test_scenes: 100 +test_success_rate: 0.8 +training_seed: 20261025 +development_seed: 20271025 +test_seed: 20281025 +run_timeout: 21600 +evaluation_timeout: 3600 diff --git a/examples/ods/data_stream_demo/converge.py b/examples/ods/data_stream_demo/converge.py new file mode 100644 index 000000000..5689fb2f2 --- /dev/null +++ b/examples/ods/data_stream_demo/converge.py @@ -0,0 +1,569 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Train the dual-camera PickPlace policy with explicit learning acceptance. + +Review the resolved recipe first; preview.py acquires expert images without +training. The train command reuses the existing bounded ODS visual pipeline. +Development evaluation selects a checkpoint before a single held-out test. +""" + +from __future__ import annotations + +import argparse +from dataclasses import asdict +import hashlib +import json +import math +import os +from pathlib import Path +import subprocess +import sys +import time +import traceback +from typing import Any + +from embodichain.utils import configclass + +__all__ = ["ConvergenceCfg", "load_recipe", "review_recipe"] + +_ROOT = Path(__file__).resolve().parents[3] +_DEFAULT = Path(__file__).with_name("configs") / "pick_place_dual_convergence.yaml" + + +@configclass +class ConvergenceCfg: + """Fixed useful training work, protected supply and evaluation thresholds.""" + + updates: int = 20000 + batch_size: int = 64 + buffer_size: int = 32 + producer_envs: int = 4 + episode_budget: int = 2048 + learning_rate: float = 3e-4 + final_learning_rate: float = 3e-5 + warmup_updates: int = 500 + prediction_horizon: int = 32 + action_horizon: int = 8 + encoder_threads: int = 2 + evaluate_every: int = 2000 + development_scenes: int = 20 + development_success_rate: float = 0.9 + test_scenes: int = 100 + test_success_rate: float = 0.8 + training_seed: int = 20261025 + development_seed: int = 20271025 + test_seed: int = 20281025 + run_timeout: int = 21600 + evaluation_timeout: int = 3600 + + def __post_init__(self) -> None: + positive = ( + "updates", + "batch_size", + "buffer_size", + "producer_envs", + "episode_budget", + "warmup_updates", + "prediction_horizon", + "action_horizon", + "encoder_threads", + "evaluate_every", + "development_scenes", + "test_scenes", + "run_timeout", + "evaluation_timeout", + ) + for name in positive: + if type(getattr(self, name)) is not int or getattr(self, name) < 1: + raise ValueError(f"{name} must be a positive integer") + if ( + self.buffer_size % self.producer_envs + or self.batch_size % self.buffer_size + or self.producer_envs * self.episode_budget < self.batch_size + ): + raise ValueError("Require C divisible by P, B divisible by C, and P*K>=B") + if ( + self.updates % self.evaluate_every + or self.updates < 2 * self.evaluate_every + or self.warmup_updates >= self.updates + or self.action_horizon > self.prediction_horizon + ): + raise ValueError( + "Budget must fund two evaluation intervals and fit warmup/horizon" + ) + for name in ( + "learning_rate", + "final_learning_rate", + "development_success_rate", + "test_success_rate", + ): + value = getattr(self, name) + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + ): + raise ValueError(f"{name} must be a finite number") + if not 0 < self.final_learning_rate <= self.learning_rate: + raise ValueError("Learning rates must be finite and 0 float: + """Return the declared linear-warmup/cosine learning rate for an update.""" + if step <= self.warmup_updates: + return self.learning_rate * max(1, step) / self.warmup_updates + progress = min( + 1.0, (step - self.warmup_updates) / (self.updates - self.warmup_updates) + ) + return ( + self.final_learning_rate + + (self.learning_rate - self.final_learning_rate) + * (1 + math.cos(math.pi * progress)) + / 2 + ) + + +def load_recipe(path: Path) -> ConvergenceCfg: + """Load an explicit YAML recipe, rejecting unknown or invalid fields.""" + import yaml + + path = Path(path) + text = path.read_text() + values = json.loads(text) if path.suffix == ".json" else yaml.safe_load(text) + if not isinstance(values, dict): + raise ValueError("Convergence YAML must be a mapping") + cfg = ConvergenceCfg(**values) + cfg.validate() + return cfg + + +def review_recipe(cfg: ConvergenceCfg) -> dict[str, Any]: + """Describe the complete candidate without creating an engine or optimizer.""" + from .acquisition import camera_profile_version, camera_configs + + cameras = camera_configs("dual", camera_profile="pick_place") + return { + "recipe": asdict(cfg), + "task": "SimpleBCPickPlace-Franka-v1", + "model": "scratch ResNet10/GroupNorm + spatial embeddings + MLP512/512", + "policy_input": "two causal RGB frames per camera + two 27D robot states", + "camera_profile": camera_profile_version("pick_place"), + "cameras": cameras, + "history": 2, + "precision": "FP32", + "optimizer": "AdamW", + "weight_decay": 1e-6, + "gradient_norm_limit": 1.0, + "normalization": "initial training cache only; frozen", + "trained_chunks": cfg.updates * cfg.batch_size, + "valid_action_targets": cfg.updates * cfg.batch_size * cfg.prediction_horizon, + "minimum_trained_sources_from_cap": math.ceil( + cfg.updates * cfg.batch_size / cfg.episode_budget + ), + "reuse_profile_qualified": False, + "development_gate": "last two updated checkpoints both meet development threshold", + "selection": "highest dev success, then lowest mean goal error, before final test", + "final_test": "one fixed panel, opened only after development gate passes", + "training_started": False, + } + + +def _worker(output: Path, cfg: ConvergenceCfg) -> int: + from .train import _Experiment, _validate, build_parser + from torch.utils.tensorboard import SummaryWriter + + args = build_parser().parse_args( + [ + "--output", + str(output), + "--arm", + "ods_overlap", + "--vision", + "dual", + "--updates", + str(cfg.updates), + "--batch-size", + str(cfg.batch_size), + "--buffer-size", + str(cfg.buffer_size), + "--producer-envs", + str(cfg.producer_envs), + "--offline-episodes", + str(cfg.buffer_size), + "--episode-budget", + str(cfg.episode_budget), + "--prediction-horizon", + str(cfg.prediction_horizon), + "--action-horizon", + str(cfg.action_horizon), + "--encoder-threads", + str(cfg.encoder_threads), + "--seed", + str(cfg.training_seed), + "--run-timeout", + str(cfg.run_timeout), + "--evaluation-scenes", + "0", + "--phase", + "convergence", + ] + ) + args.camera_profile = "pick_place" + from .acquisition import camera_profile_version + + args.camera_profile_version = camera_profile_version(args.camera_profile) + args.convergence_recipe = asdict(cfg) + _validate(args) + + class LearningExperiment(_Experiment): + def __init__(self) -> None: + super().__init__(args) + self.development: list[dict[str, Any]] = [] + self.board = SummaryWriter(str(output / "tensorboard")) + + def save_checkpoint(self, name: str, step: int) -> None: + if (self.output / name).exists(): + raise FileExistsError("Learning checkpoints are immutable") + if step == 0: + for group in self.optimizer.param_groups: + group["lr"] = cfg.rate(0) + super().save_checkpoint(name, step) + + def evaluate( + self, checkpoint: Path, scenes: int, seed: int, label: str + ) -> dict[str, Any]: + receipt = self.output / "evaluation" / f"{label}.json" + receipt.parent.mkdir(exist_ok=True) + command = [ + sys.executable, + "-m", + "examples.ods.data_stream_demo.training_eval", + "--checkpoint", + str(checkpoint), + "--vision", + "dual", + "--scenes", + str(scenes), + "--seed", + str(seed), + "--output", + str(receipt), + ] + started = self.now() + with receipt.with_suffix(".log").open("w") as log: + child = subprocess.Popen( + command, cwd=_ROOT, stdout=log, stderr=subprocess.STDOUT + ) + try: + deadline = time.monotonic() + cfg.evaluation_timeout + while child.poll() is None: + if self.flow is None: + self.poll() + else: + self.flow.poll() + if time.monotonic() >= deadline: + raise TimeoutError( + f"{label} evaluation exceeded its deadline" + ) + time.sleep(0.1) + if child.returncode: + raise RuntimeError( + f"{label} evaluator failed; see {receipt.with_suffix('.log')}" + ) + finally: + if child.poll() is None: + child.terminate() + try: + child.wait(timeout=10) + except subprocess.TimeoutExpired: + child.kill() + child.wait(timeout=10) + report = json.loads(receipt.read_text()) + if ( + not report.get("passed") + or not report.get("closed") + or report.get("expert_initialized") + ): + raise RuntimeError( + "Independent evaluator failed its student-only/closure contract" + ) + if self.flow is not None: + self.metrics.interval( + "learning_development_evaluation", started, self.now() + ) + return report + + def develop(self, checkpoint: Path, step: int) -> None: + report = self.evaluate( + checkpoint, + cfg.development_scenes, + cfg.development_seed, + f"dev-{step:08d}", + ) + row = {"step": step, "checkpoint": str(checkpoint), **report} + self.development.append(row) + (output / "development.json").write_text( + json.dumps(self.development, indent=2) + ) + for name in ("success_rate", "lift_rate", "goal_error", "episode_length"): + self.board.add_scalar("development/" + name, report[name], step) + self.board.flush() + + def bootstrap(self, episodes: list[Any]) -> None: + super().bootstrap(episodes) + self.develop(output / "initial.pt", 0) + + def update(self, *positional: Any, **keywords: Any) -> None: + step = self.completed_updates + 1 + for group in self.optimizer.param_groups: + group["lr"] = cfg.rate(step) + super().update(*positional, **keywords) + if step % 10 == 0: + for name, value in self.metrics._last_losses.items(): + self.board.add_scalar("training/" + name, value, step) + self.board.add_scalar("training/learning_rate", cfg.rate(step), step) + self.board.add_scalar("ods/generation", self.generation, step) + if step % cfg.evaluate_every == 0: + name = f"checkpoint-{step:08d}.pt" + self.save_checkpoint(name, step) + self.develop(output / name, step) + + experiment = LearningExperiment() + result: dict[str, Any] = {} + try: + result = experiment.execute() + candidates = [row for row in experiment.development if row["step"] > 0] + selected = max( + candidates, key=lambda row: (row["success_rate"], -row["goal_error"]) + ) + stable = all( + row["success_rate"] >= cfg.development_success_rate + for row in candidates[-2:] + ) + selection = { + "checkpoint": selected["checkpoint"], + "step": selected["step"], + "sha256": selected["checkpoint_sha256"], + "development_gate_passed": stable, + } + (output / "selection.json").write_text(json.dumps(selection, indent=2)) + final = None + if stable: + final = experiment.evaluate( + Path(selected["checkpoint"]), + cfg.test_scenes, + cfg.test_seed, + "final-test", + ) + met = ( + stable + and final is not None + and final["success_rate"] >= cfg.test_success_rate + ) + result.update( + learning_target_met=met, + development_gate_passed=stable, + selected_checkpoint=selection, + final_test=final, + learning_status="qualified" if met else "budget_exhausted_not_converged", + ) + except BaseException as error: + result = ( + experiment.failure(error) + if not experiment.closed + else { + **result, + "status": "failed", + "error": repr(error), + "learning_target_met": False, + "traceback": traceback.format_exc(), + } + ) + finally: + experiment.board.close() + (output / "summary.json").write_text(json.dumps(result, indent=2)) + return 0 if result.get("learning_target_met") else 2 + + +def _supervise(args: argparse.Namespace, cfg: ConvergenceCfg) -> int: + import psutil + from .suite import _gpu_contexts, _partition_contexts, _terminate + from .training_suite import _devices + + args.output = args.output.resolve() + args.output.mkdir(parents=True, exist_ok=False) + (args.output / "review.json").write_text(json.dumps(review_recipe(cfg), indent=2)) + frozen_recipe = args.output / "recipe.json" + frozen_recipe.write_text(json.dumps(asdict(cfg), indent=2) + "\n") + pairs, audit = _devices(args) + if len(pairs) != 1: + raise ValueError("Convergence training requires exactly two admitted GPUs") + pair = pairs[0] + devices = [pair["simulation"], pair["training"]] + auxiliary = audit["aux_graphics_gpu"] + allowed = {gpu["uuid"] for gpu in devices} + affinity = pair["cpu_affinity"] + environment = os.environ.copy() + for key in ("DISPLAY", "WAYLAND_DISPLAY", "ODS_AUX_GRAPHICS_GPU_UUID"): + environment.pop(key, None) + environment.update( + CUDA_VISIBLE_DEVICES=",".join(gpu["uuid"] for gpu in devices), + ODS_GPU_UUID=devices[0]["uuid"], + OMP_NUM_THREADS="1", + MKL_NUM_THREADS="1", + OPENBLAS_NUM_THREADS="1", + SDL_VIDEODRIVER="dummy", + ) + if auxiliary is not None: + environment["ODS_AUX_GRAPHICS_GPU_UUID"] = auxiliary["uuid"] + command = [ + sys.executable, + "-m", + "examples.ods.data_stream_demo.converge", + "train", + "--config", + str(frozen_recipe), + "--output", + str(args.output / "training"), + "--worker", + ] + started = time.monotonic() + failure: BaseException | None = None + environment["ODS_PROCESS_LAUNCH_MONOTONIC_NS"] = str(time.monotonic_ns()) + with (args.output / "worker.log").open("w") as stream: + child = subprocess.Popen( + command, + cwd=_ROOT, + env=environment, + stdout=stream, + stderr=subprocess.STDOUT, + start_new_session=True, + preexec_fn=lambda: os.sched_setaffinity(0, affinity), + ) + try: + while child.poll() is None: + pids = {child.pid} + try: + pids.update( + p.pid + for p in psutil.Process(child.pid).children(recursive=True) + ) + except psutil.NoSuchProcess: + pass + contexts = [] + for context in _gpu_contexts(): + try: + owned = ( + context["pid"] in pids + or os.getpgid(context["pid"]) == child.pid + ) + except ProcessLookupError: + continue + if owned: + contexts.append(context) + outside = [c for c in contexts if c["gpu_uuid"] not in allowed] + _, unexpected = _partition_contexts( + outside, + "outside-primary-pair", + None if auxiliary is None else auxiliary["uuid"], + ) + if unexpected: + raise RuntimeError( + f"Owned work used unassigned GPU contexts: {unexpected}" + ) + if time.monotonic() - started > cfg.run_timeout + 120: + raise TimeoutError("Convergence worker exceeded total runtime") + time.sleep(1) + except BaseException as error: + failure = error + finally: + if child.poll() is None: + _terminate(child) + members = [] + deadline = time.monotonic() + 2 + while time.monotonic() < deadline: + members = [] + for process in psutil.process_iter(["pid"]): + try: + if os.getpgid(process.pid) == child.pid: + members.append(process.pid) + except ProcessLookupError: + pass + if not members: + break + time.sleep(0.05) + (args.output / "process.json").write_text( + json.dumps( + { + "returncode": child.returncode, + "remaining_group_members": members, + "gpu_pair": devices, + "cpu_affinity": affinity, + "command": command, + "wall_seconds": time.monotonic() - started, + "error": None if failure is None else repr(failure), + }, + indent=2, + ) + ) + if members: + _terminate(child) + raise RuntimeError("Convergence worker left owned processes") + if failure is not None: + raise failure + return child.returncode + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("command", choices=("review", "train")) + parser.add_argument("--config", type=Path, default=_DEFAULT) + parser.add_argument("--output", type=Path) + parser.add_argument("--gpu-ids", type=int, nargs=2) + parser.add_argument("--aux-graphics-gpu-id", type=int) + parser.add_argument("--worker", action="store_true", help=argparse.SUPPRESS) + args = parser.parse_args() + cfg = load_recipe(args.config) + if args.command == "review": + review = review_recipe(cfg) + review["config_sha256"] = hashlib.sha256(args.config.read_bytes()).hexdigest() + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(review, indent=2) + "\n") + print(json.dumps(review, indent=2)) + return + if args.output is None or (not args.worker and args.gpu_ids is None): + parser.error( + "train requires --output and an explicit generation/training --gpu-ids pair" + ) + raise SystemExit( + _worker(args.output, cfg) if args.worker else _supervise(args, cfg) + ) + + +if __name__ == "__main__": + main() diff --git a/examples/ods/data_stream_demo/lerobot_io.py b/examples/ods/data_stream_demo/lerobot_io.py new file mode 100644 index 000000000..cc6f02e5e --- /dev/null +++ b/examples/ods/data_stream_demo/lerobot_io.py @@ -0,0 +1,1421 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Standard SDK recording and genuinely decoded, bounded LeRobot consumption. + +The parent owns exposure reservations. Spawned DataLoader workers own their +SDK Dataset/CPU decoder and never receive a copy of the exposure ledger. +""" + +from __future__ import annotations + +import copy +import json +import math +import os +import queue +import threading +import time +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager +from importlib.metadata import version +from pathlib import Path +from typing import Any, TextIO + +import numpy as np +import torch +from torch.utils.data import DataLoader, Dataset, default_collate + +from embodichain.utils import configclass + +from .payload import numeric_episode_digest + +__all__ = [ + "LeRobotReaderCfg", + "StandardLeRobotWriter", + "StandardLeRobotReader", + "TrainingLeRobotReader", + "read_lerobot_manifest", + "audit_video_quality", +] + +_SIDECAR = "meta/ods_data_stream.json" +_PROVENANCE = "ods.provenance" + + +def _check_versions() -> dict[str, str]: + found = {name: version(name) for name in ("lerobot", "torchcodec")} + if found != {"lerobot": "0.4.4", "torchcodec": "0.5"}: + raise RuntimeError( + f"This protocol requires LeRobot0.4.4/TorchCodec0.5: {found}" + ) + return found + + +def _feature_names(metadata: Mapping[str, Any]) -> list[str]: + return [ + f"{'/'.join(feature['path'])}/{index}" + for feature in metadata["state_features"] + for index in range(feature["width"]) + ] + + +class StandardLeRobotWriter: + """Write accepted CPU episodes through one bounded SDK-owner thread. + + Args: + destination: New output directory; existing destinations are rejected. + metadata: Canonical payload description. + episode_count: Exact accepted count, or ``None`` for a dynamic corpus. + fps: Simulation observation frequency. + queue_size: Queued raw episodes, excluding the active writer/submission. + operation_timeout: Bound on handoff/finalization and worker join. + events_path: Optional new JSONL file outside the corpus. The SDK-owner + thread streams chronological service intervals without callbacks; + log failures are terminal just like recording failures. + encoder_threads: Positive thread limit for each SDK AV1 encoder. + + Real video features use the pinned SDK's ordinary AV1 per-episode encoding + path. A failed/partial corpus remains marked INCOMPLETE and is not readable + through this adapter. The calling run must treat an unjoined native writer + as a failure; Python cannot forcibly interrupt an in-process encoder. + """ + + def __init__( + self, + destination: Path | str, + metadata: Mapping[str, Any], + *, + episode_count: int | None = None, + fps: int = 30, + queue_size: int = 2, + operation_timeout: float = 1800.0, + events_path: Path | str | None = None, + encoder_threads: int = 1, + ) -> None: + if episode_count is not None and ( + type(episode_count) is not int or episode_count < 1 + ): + raise ValueError("episode_count must be a positive integer or None") + for name, value in ( + ("fps", fps), + ("queue_size", queue_size), + ("encoder_threads", encoder_threads), + ): + if type(value) is not int or value < 1: + raise ValueError(f"{name} must be a positive integer") + if not math.isfinite(operation_timeout) or operation_timeout <= 0: + raise ValueError("operation_timeout must be finite and positive") + _check_versions() + self.root = Path(destination) + self._events_path = None if events_path is None else Path(events_path) + if self._events_path is not None: + resolved = self._events_path.resolve() + if ( + self.root.resolve() == resolved + or self.root.resolve() in resolved.parents + ): + raise ValueError("Writer event log must be outside the recorded corpus") + if self.root.exists(): + raise FileExistsError(self.root) + self.root.mkdir(parents=True) + (self.root / "INCOMPLETE").write_text("Accepted episodes are not finalized.\n") + self.metadata = copy.deepcopy(dict(metadata)) + self.episode_count, self.fps = episode_count, fps + self.encoder_threads = encoder_threads + self.operation_timeout = float(operation_timeout) + self._owner = (os.getpid(), threading.get_ident()) + self._queue: queue.Queue[dict[str, Any] | None] = queue.Queue(queue_size) + self._abort = threading.Event() + self._error: BaseException | None = None + self._closed = False + self._submitted = 0 + self._final_episode_count = episode_count + self._identities: set[tuple[int, int]] = set() + self._stats: dict[str, Any] = { + "submitted_episodes": 0, + "written_episodes": 0, + "valid_frames": 0, + "queue_capacity": queue_size, + "peak_queued_episodes": 0, + "queue_wait_seconds": 0.0, + "sdk_create_seconds": 0.0, + "sdk_write_encode_seconds": 0.0, + "sdk_episode_stage_seconds": { + name: 0.0 + for name in ( + "validation_provenance", + "add_frame", + "save_episode", + "references_manifest", + ) + }, + "sdk_finalize_seconds": 0.0, + "finalized": False, + "worker_joined": False, + } + if self._events_path is not None: + self._stats["event_log_rows"] = 0 + self._thread = threading.Thread( + target=self._run, name="stream-lerobot-writer", daemon=True + ) + self._thread.start() + + def _check(self) -> None: + if self._owner != (os.getpid(), threading.get_ident()): + raise RuntimeError("Writer operations require the creator process/thread") + if self._error is not None: + raise RuntimeError("Standard LeRobot recording failed") from self._error + if self._closed: + raise RuntimeError("Standard LeRobot writer is closed") + + @property + def stats(self) -> dict[str, Any]: + """Return detached counters; component wall times can overlap generation.""" + return copy.deepcopy(self._stats) + + def check_health(self) -> None: + """Surface asynchronous recording errors from the creator thread.""" + self._check() + + def _put( + self, item: dict[str, Any] | None, progress: Callable[[], None] | None + ) -> None: + started = time.monotonic() + deadline = started + self.operation_timeout + while True: + self._check() + if progress is not None: + progress() + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("LeRobot writer queue handoff timed out") + try: + self._queue.put(item, timeout=min(remaining, 0.05)) + break + except queue.Full: + continue + self._stats["queue_wait_seconds"] += time.monotonic() - started + self._stats["peak_queued_episodes"] = max( + self._stats["peak_queued_episodes"], self._queue.qsize() + ) + + def submit( + self, episode: dict[str, Any], *, progress: Callable[[], None] | None = None + ) -> None: + """Transfer an owned packed episode to the bounded writer queue. + + Args: + episode: ``pack_episode`` output; caller must not mutate it afterwards. + progress: Creator-thread cancellation/ODS-event-draining callback. + """ + identity = self._validate_submission(episode) + self._put(episode, progress) + self._record_submission(identity) + + def _validate_submission(self, episode: Mapping[str, Any]) -> tuple[int, int]: + self._check() + if self.episode_count is not None and self._submitted >= self.episode_count: + raise ValueError("More accepted episodes than the declared corpus size") + identity = tuple(int(v) for v in episode["frame_ids"][0, :2]) + if identity in self._identities: + raise ValueError("Duplicate immutable source episode") + return identity + + def _record_submission(self, identity: tuple[int, int]) -> None: + self._identities.add(identity) + self._submitted += 1 + self._stats["submitted_episodes"] = self._submitted + self._stats["peak_queued_episodes"] = max( + self._stats["peak_queued_episodes"], self._queue.qsize() + ) + + def try_submit(self, episode: dict[str, Any]) -> bool: + """Transfer an owned packed episode only if the bounded queue has room. + + A ``False`` result transfers no ownership and changes neither accepted + identities nor submission counts. Caller retains the same episode for a + later retry. Worker failures remain terminal and are never hidden as a + full-queue result. On success, caller must not mutate the episode. + """ + identity = self._validate_submission(episode) + try: + self._queue.put_nowait(episode) + except queue.Full: + self._check() + return False + self._record_submission(identity) + return True + + def _write_event( + self, + stream: TextIO | None, + event: Mapping[str, Any] | None, + *, + error: BaseException | None = None, + ) -> None: + if stream is None or event is None: + return + row = { + "schema": "ods-sdk-writer-service/v1", + "kind": "interval", + **event, + "end_ns": time.monotonic_ns(), + "worker_pid": os.getpid(), + "worker_thread_id": threading.get_ident(), + "status": "complete" if error is None else "failed", + } + if error is not None: + row["error"] = repr(error) + stream.write(json.dumps(row, allow_nan=False) + "\n") + self._stats["event_log_rows"] += 1 + + @contextmanager + def _episode_stage(self, event: dict[str, Any], name: str) -> Iterator[None]: + """Attach disjoint substage timings without adding service-log rows.""" + stage: dict[str, Any] = { + "start_ns": time.monotonic_ns(), + "status": "failed", + } + event["stage_intervals"][name] = stage + try: + yield + stage["status"] = "complete" + finally: + stage["end_ns"] = time.monotonic_ns() + stage["duration_ns"] = stage["end_ns"] - stage["start_ns"] + + def _run(self) -> None: + dataset = None + events = None + active_event = None + try: + if self._events_path is not None: + self._events_path.parent.mkdir(parents=True, exist_ok=True) + events = self._events_path.open("x", buffering=1, encoding="utf-8") + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + features = { + "observation.state": { + "dtype": "float32", + "shape": (self.metadata["state_width"],), + "names": _feature_names(self.metadata), + }, + "action": { + "dtype": "float32", + "shape": (self.metadata["action_width"],), + "names": [ + f"action/{i}" for i in range(self.metadata["action_width"]) + ], + }, + _PROVENANCE: { + "dtype": "int64", + "shape": (6,), + "names": [ + "run_id", + "episode_id", + "frame_index", + "generation", + "continuity_id", + "published_at_monotonic_ns", + ], + }, + } + for name, feature in self.metadata["cameras"].items(): + h, w, _ = feature["shape"] + features[f"observation.images.{name}"] = { + "dtype": "video", + "shape": (h, w, 3), + "names": ["height", "width", "channels"], + } + started = time.monotonic() + # create() requires a nonexistent root; retain the explicit partial marker. + (self.root / "INCOMPLETE").unlink() + self.root.rmdir() + active_event = { + "name": "sdk_create", + "start_ns": time.monotonic_ns(), + "source_id": None, + "episode_ordinal": None, + "scope": "LeRobotDataset.create; excludes SDK import, prior directory setup and event-log serialization", + } + dataset = LeRobotDataset.create( + "local/ods-data-stream", + fps=self.fps, + features=features, + root=self.root, + robot_type="franka_panda", + use_videos=bool(self.metadata["cameras"]), + tolerance_s=1e-4, + video_backend="torchcodec", + batch_encoding_size=1, + vcodec="libsvtav1", + encoder_threads=self.encoder_threads, + ) + self._write_event(events, active_event) + active_event = None + (self.root / "INCOMPLETE").write_text( + "Accepted episodes are not finalized.\n" + ) + self._stats["sdk_create_seconds"] = time.monotonic() - started + manifest: dict[str, Any] = { + "format": "ods-data-stream-lerobot-v1", + "payload": self.metadata, + "fps": self.fps, + "versions": _check_versions(), + "episodes": [], + "references": [], + "video": { + "codec": "libsvtav1", + "pix_fmt": "yuv420p", + "gop": 2, + "crf": 30, + "encoder_threads": self.encoder_threads, + "parallel_encoding": False, + }, + } + offset = 0 + while not self._abort.is_set(): + try: + episode = self._queue.get(timeout=0.05) + except queue.Empty: + continue + if episode is None: + break + started = time.monotonic() + active_event = { + "name": "sdk_episode_write_encode", + "start_ns": time.monotonic_ns(), + "source_id": episode["frame_ids"][0, :2].tolist(), + "episode_ordinal": len(manifest["episodes"]), + "scope": "accepted-frame validation/provenance, add_frame, synchronous save_episode/AV1 encoding and raw-reference/hash/manifest bookkeeping; excludes queue wait and event-log serialization", + "stage_intervals": {}, + "stage_scope": "disjoint substages within this interval; add_frame includes frame assembly and synchronous image writes; save_episode includes SDK validation, statistics, numeric commit and AV1 encoding; aggregate stages only across successful episodes", + } + with self._episode_stage(active_event, "validation_provenance"): + length = len(episode["valid"]) + if not episode["valid"].all(): + raise ValueError("Writer requires trimmed accepted frames") + provenance = torch.cat( + [ + episode["frame_ids"], + episode["generations"][:, None], + episode["continuity_id"][:, None], + episode["published_at_monotonic_ns"][:, None], + ], + -1, + ) + with self._episode_stage(active_event, "add_frame"): + for step in range(length): + if self._abort.is_set(): + raise RuntimeError("LeRobot recording aborted") + frame = { + "observation.state": episode["state"][step], + "action": episode["actions"][step], + _PROVENANCE: provenance[step], + "task": "Pick up the cube and place it in the fixed goal region.", + } + frame.update( + { + f"observation.images.{name}": image[step] + for name, image in episode["images"].items() + } + ) + dataset.add_frame(frame) + with self._episode_stage(active_event, "save_episode"): + # SDK camera parallelism forks subprocesses from this thread; + # keep it disabled in a learner that may initialize CUDA. + dataset.save_episode(parallel_encoding=False) + with self._episode_stage(active_event, "references_manifest"): + if not manifest["episodes"]: + refs = self.root / "meta" / "ods_raw_references" + refs.mkdir(exist_ok=True) + for name, images in episode["images"].items(): + for step in sorted({0, length // 2, length - 1}): + filename = f"{name}-{step}.npy" + np.save( + refs / filename, + images[step].numpy(), + allow_pickle=False, + ) + manifest["references"].append( + { + "camera": name, + "index": offset + step, + "path": f"meta/ods_raw_references/{filename}", + } + ) + manifest["episodes"].append( + { + "index": len(manifest["episodes"]), + "offset": offset, + "length": length, + "source_id": episode["frame_ids"][0, :2].tolist(), + "generation": int(episode["generations"][0]), + "numeric_sha256": numeric_episode_digest(episode), + } + ) + offset += length + self._stats["written_episodes"] += 1 + self._stats["valid_frames"] += length + self._stats["sdk_write_encode_seconds"] += time.monotonic() - started + for name, stage in active_event["stage_intervals"].items(): + self._stats["sdk_episode_stage_seconds"][name] += ( + stage["duration_ns"] / 1e9 + ) + self._write_event(events, active_event) + active_event = None + del episode + if self._abort.is_set(): + raise RuntimeError("LeRobot recording aborted") + if len(manifest["episodes"]) != self._final_episode_count: + raise ValueError("Final accepted corpus episode count mismatch") + manifest["accepted_episode_count"] = self._final_episode_count + started = time.monotonic() + active_event = { + "name": "sdk_finalize", + "start_ns": time.monotonic_ns(), + "source_id": None, + "episode_ordinal": None, + "scope": "LeRobotDataset.finalize; excludes final provenance-sidecar/marker writes and event-log serialization", + } + dataset.finalize() + dataset = None + self._write_event(events, active_event) + active_event = None + self._stats["sdk_finalize_seconds"] = time.monotonic() - started + manifest["required_files"] = [ + str(path.relative_to(self.root)) + for path in sorted(self.root.rglob("*")) + if path.is_file() and path.name != "INCOMPLETE" + ] + (self.root / _SIDECAR).write_text( + json.dumps(manifest, indent=2, allow_nan=False) + "\n" + ) + (self.root / "INCOMPLETE").unlink() + self._stats["finalized"] = True + except BaseException as error: + self._error = error + if active_event is not None: + try: + self._write_event(events, active_event, error=error) + except BaseException as event_error: + self._stats["event_log_error"] = repr(event_error) + self.root.mkdir(parents=True, exist_ok=True) + (self.root / "INCOMPLETE").write_text(repr(error) + "\n") + finally: + if dataset is not None: + cleanup_event = { + "name": "sdk_finalize", + "start_ns": time.monotonic_ns(), + "source_id": None, + "episode_ordinal": None, + "cleanup": True, + "scope": "cleanup LeRobotDataset.finalize after a failed/aborted recording; excludes event-log serialization", + } + try: + dataset.finalize() + self._write_event(events, cleanup_event) + except BaseException as error: + self._stats["cleanup_error"] = repr(error) + if events is not None: + try: + events.close() + except BaseException as error: + self._error = self._error or error + self._stats["event_log_error"] = repr(error) + self._stats["finalized"] = False + (self.root / "INCOMPLETE").write_text(repr(self._error) + "\n") + + def _join(self, progress: Callable[[], None] | None, timeout: float) -> None: + deadline = time.monotonic() + timeout + while self._thread.is_alive(): + if progress is not None: + progress() + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError( + "Native LeRobot writer did not join within its deadline" + ) + self._thread.join(min(0.05, remaining)) + self._stats["worker_joined"] = True + + def finalize( + self, + *, + expected_episode_count: int | None = None, + progress: Callable[[], None] | None = None, + ) -> dict[str, Any]: + """Drain encoding/finalization and return detached component counters. + + Args: + expected_episode_count: Required actual accepted count for a dynamic + corpus. A fixed-count writer permits only its declared count. + progress: Owner-thread event/cancellation callback while waiting. + + Returns: + Recording counters. A failed corpus is never exposed as finalized. + """ + self._check() + if expected_episode_count is None: + expected_episode_count = self.episode_count + if type(expected_episode_count) is not int or expected_episode_count < 1: + raise ValueError("Dynamic finalization requires a positive episode count") + if ( + self.episode_count is not None + and expected_episode_count != self.episode_count + ): + raise ValueError("Final count differs from the declared corpus size") + if self._submitted != expected_episode_count: + raise ValueError( + "Cannot finalize before all accepted episodes are submitted" + ) + self._final_episode_count = expected_episode_count + self._put(None, progress) + self._join(progress, self.operation_timeout) + self._check() + self._closed = True + return self.stats + + def abort(self) -> None: + """Cancel pending episodes and bound native writer shutdown to 30 seconds.""" + if self._owner != (os.getpid(), threading.get_ident()): + raise RuntimeError("Abort requires the creator process/thread") + if self._closed: + return + self._abort.set() + self._join(None, min(30.0, self.operation_timeout)) + while True: + try: + self._queue.get_nowait() + except queue.Empty: + break + self._closed = True + + +@configclass +class LeRobotReaderCfg: + """Bounded, standard SDK reader settings used by all visual conditions.""" + + batch_size: int = 32 + history_length: int = 2 + prediction_horizon: int = 32 + episode_budget: int | None = 128 + seed: int = 29001 + num_workers: int = 4 + prefetch_factor: int = 1 + timeout: float = 120.0 + + def __post_init__(self) -> None: + for name in ( + "batch_size", + "history_length", + "prediction_horizon", + "num_workers", + "prefetch_factor", + ): + if type(getattr(self, name)) is not int or getattr(self, name) < 1: + raise ValueError(f"{name} must be a positive integer") + if self.episode_budget is not None and ( + type(self.episode_budget) is not int or self.episode_budget < 1 + ): + raise ValueError("episode_budget must be a positive integer or None") + if not math.isfinite(self.timeout) or self.timeout <= 0: + raise ValueError("timeout must be finite and positive") + + +class _DecodedContexts(Dataset): + def __init__( + self, + root: Path, + manifest: Mapping[str, Any], + cfg: LeRobotReaderCfg, + *, + history_only_images: bool = False, + require_complete: bool = False, + ) -> None: + self.root, self.manifest, self.cfg = root, manifest, cfg + self.history_only_images = history_only_images + self.require_complete = require_complete + self._dataset: Any = None + + def __len__(self) -> int: + return sum(episode["length"] for episode in self.manifest["episodes"]) + + def __getitem__(self, key: tuple[int, int]) -> dict[str, Any]: + reservation, index = key + started_ns = time.monotonic_ns() + if self._dataset is None: + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + offsets = [ + step / self.manifest["fps"] + for step in range( + 1 - self.cfg.history_length, self.cfg.prediction_horizon + ) + ] + delta_timestamps = { + field: offsets for field in ("observation.state", "action", _PROVENANCE) + } + image_offsets = ( + offsets[: self.cfg.history_length] + if self.history_only_images + else offsets + ) + delta_timestamps.update( + { + f"observation.images.{name}": image_offsets + for name in self.manifest["payload"]["cameras"] + } + ) + self._dataset = LeRobotDataset( + "local/ods-data-stream", + root=self.root, + delta_timestamps=delta_timestamps, + tolerance_s=1e-4, + video_backend="torchcodec", + ) + item = self._dataset[index] # Public SDK read, including real video decoding. + provenance = item[_PROVENANCE].to(dtype=torch.int64) + anchor = self.cfg.history_length - 1 + padding = item[f"{_PROVENANCE}_is_pad"].bool() + continuity = provenance[:, 4] + compatible = (~padding) & (continuity == continuity[anchor]) + # Stop at the nearest continuity break; do not reconnect matching IDs. + positions = torch.arange(len(compatible)) + lower = int( + torch.where((~compatible) & (positions < anchor), positions + 1, 0).max() + ) + upper = int( + torch.where( + (~compatible) & (positions > anchor), positions - 1, len(compatible) - 1 + ).min() + ) + valid = compatible & (positions >= lower) & (positions <= upper) + if self.require_complete: + if not valid.all(): + raise ValueError( + "Training plan requested an incomplete or discontinuous context" + ) + for field in ("observation.state", "action"): + if item[f"{field}_is_pad"].bool().any(): + raise ValueError( + "Training numeric context unexpectedly requires padding" + ) + if (provenance[:, :2] != provenance[anchor, :2]).any() or ( + provenance[1:, 2] != provenance[:-1, 2] + 1 + ).any(): + raise ValueError("Training context frame/source order is inconsistent") + if (provenance[:, 3] != provenance[anchor, 3]).any(): + raise ValueError("Training context crosses generation boundaries") + if int(item["index"]) != index: + raise ValueError( + "SDK anchor index differs from the fixed training plan" + ) + result: dict[str, Any] = { + "state": item["observation.state"].float(), + "actions": item["action"].float(), + "frame_ids": provenance[:, :3].clone(), + "generations": provenance[:, 3].clone(), + "continuity_id": continuity.clone(), + "published_at_monotonic_ns": provenance[:, 5].clone(), + "valid": valid, + "loss_mask": valid & (positions >= anchor), + "images": {}, + "_reservation_id": reservation, + } + if self.require_complete: + result["_global_anchor_index"] = index + for name in self.manifest["payload"]["cameras"]: + field = f"observation.images.{name}" + image_valid = ( + valid[: self.cfg.history_length] if self.history_only_images else valid + ) + image_padding = ( + padding[: self.cfg.history_length] + if self.history_only_images + else padding + ) + if not torch.equal(item[f"{field}_is_pad"].bool(), image_padding): + raise ValueError("SDK image/numeric timestamp padding differs") + image = item[field] + if image.ndim == 3 and len(image_valid) == 1: + image = image.unsqueeze(0) + if ( + image.dtype != torch.float32 + or image.ndim != 4 + or image.shape[0] != len(image_valid) + or image.shape[1] != 3 + ): + raise ValueError("SDK decoder must return temporal RGB float32 CHW") + result["images"][name] = ( + image + if self.require_complete + else image.masked_fill(~image_valid[:, None, None, None], 0) + ) + if not self.require_complete: + for name in ("state", "actions"): + result[name] = result[name].masked_fill(~valid[:, None], 0) + for name in ( + "frame_ids", + "generations", + "continuity_id", + "published_at_monotonic_ns", + ): + mask = valid[:, None] if name == "frame_ids" else valid + result[name].masked_fill_(~mask, -1) + ended_ns = time.monotonic_ns() + result["_service_seconds"] = (ended_ns - started_ns) / 1e9 + result["_service_start_ns"] = started_ns + result["_service_end_ns"] = ended_ns + result["_worker_pid"] = os.getpid() + return result + + +def _initialize_worker(worker_id: int) -> None: + torch.set_num_threads(1) + os.environ["HF_HUB_OFFLINE"] = "1" + os.environ["HF_DATASETS_OFFLINE"] = "1" + + +def _finish_loader( + iterator: Any, diagnostics: dict[str, Any], *, timeout: float +) -> None: + """Drain a finite dispatch stream before requiring clean worker shutdown.""" + workers = list(getattr(iterator, "_workers", ())) + started = time.monotonic() + drained = 0 + error: BaseException | None = None + cleanup_errors: list[str] = [] + forced = False + try: + while True: + remaining = started + timeout - time.monotonic() + if remaining <= 0: + raise TimeoutError( + "LeRobot prefetch drain missed its shutdown deadline" + ) + iterator._timeout = min(timeout, remaining) + try: + batch = next(iterator) + except StopIteration: + break + drained += len(batch["valid"]) + del batch + except BaseException as failure: + error = failure + finally: + try: + # PyTorch has no public DataLoader.close(). Reading remains through + # the public SDK; this private call is isolated to lifecycle cleanup. + iterator._shutdown_workers() + except BaseException as failure: + cleanup_errors.append(repr(failure)) + if error is None: + error = failure + for worker in workers: + for stop in (None, worker.terminate, worker.kill): + try: + if stop is not None and worker.is_alive(): + forced = True + stop() + worker.join(timeout=5) + if not worker.is_alive(): + break + except BaseException as failure: + cleanup_errors.append(repr(failure)) + if error is None: + error = failure + codes = [worker.exitcode for worker in workers] + forced |= any(code in (-15, -9) for code in codes) + diagnostics.update( + worker_pids=[worker.pid for worker in workers], + worker_exitcodes=codes, + workers_joined=all(not worker.is_alive() for worker in workers), + forced_termination=forced, + discarded_prefetch_chunks=drained, + shutdown_seconds=time.monotonic() - started, + cleanup_errors=cleanup_errors, + ) + if error is None and ( + forced or not diagnostics["workers_joined"] or any(code != 0 for code in codes) + ): + error = RuntimeError(f"LeRobot worker shutdown was not clean: {diagnostics}") + if error is not None: + diagnostics["error"] = repr(error) + raise error + + +class _ReservationSampler: + """Parent-only episode-uniform index selection and shared exposure ledger.""" + + def __init__(self, episodes: list[dict[str, Any]], cfg: LeRobotReaderCfg) -> None: + self.episodes, self.cfg = episodes, cfg + self.delivered = [0] * len(episodes) + self.reserved = [0] * len(episodes) + self.pending: dict[int, list[int]] = {} + self.generator = torch.Generator().manual_seed(cfg.seed) + self.next_reservation = 0 + self.cancelled_chunks = 0 + self.dispatch_enabled = True + + def __iter__(self) -> Iterator[list[tuple[int, int]]]: + while self.dispatch_enabled: + available = [ + index + for index in range(len(self.episodes)) + if self.cfg.episode_budget is None + or self.delivered[index] + self.reserved[index] + < self.cfg.episode_budget + ] + if not available: + return + if self.cfg.episode_budget is not None: + remaining = sum( + self.cfg.episode_budget - self.delivered[i] - self.reserved[i] + for i in available + ) + if remaining < self.cfg.batch_size: + return + selected: list[int] = [] + indices: list[int] = [] + for _ in range(self.cfg.batch_size): + row = available[ + int(torch.randint(len(available), (), generator=self.generator)) + ] + selected.append(row) + self.reserved[row] += 1 + episode = self.episodes[row] + anchor = int( + torch.randint(episode["length"], (), generator=self.generator) + ) + indices.append(episode["offset"] + anchor) + if ( + self.cfg.episode_budget is not None + and self.delivered[row] + self.reserved[row] + == self.cfg.episode_budget + ): + available.remove(row) + reservation = self.next_reservation + self.next_reservation += 1 + self.pending[reservation] = selected + yield [(reservation, index) for index in indices] + + def commit(self, reservation: int) -> None: + for row in self.pending.pop(reservation): + self.reserved[row] -= 1 + self.delivered[row] += 1 + + def cancel_outstanding(self) -> None: + self.cancelled_chunks += sum(self.reserved) + self.pending.clear() + self.reserved = [0] * len(self.episodes) + + +class _PlannedReservationSampler(_ReservationSampler): + """Finite parent-only reservations with no plan generation in workers.""" + + def __init__( + self, + episodes: list[dict[str, Any]], + cfg: LeRobotReaderCfg, + batch_plan: list[list[tuple[int, int]]], + ) -> None: + super().__init__(episodes, cfg) + self.batch_plan = copy.deepcopy(batch_plan) + self.planned_indices: dict[int, list[int]] = {} + self.cursor = 0 + for batch in self.batch_plan: + if len(batch) != cfg.batch_size: + raise ValueError("All planned training batches must have the same size") + for pair in batch: + if not isinstance(pair, (tuple, list)) or len(pair) != 2: + raise ValueError( + "Plan entries must be (episode_ordinal, anchor_frame)" + ) + episode, anchor = pair + if type(episode) is not int or episode < 0 or episode >= len(episodes): + raise ValueError("Planned episode ordinal is outside the corpus") + if ( + type(anchor) is not int + or anchor < cfg.history_length - 1 + or anchor + cfg.prediction_horizon > episodes[episode]["length"] + ): + raise ValueError("Planned anchor does not have a complete context") + + def __iter__(self) -> Iterator[list[tuple[int, int]]]: + while self.dispatch_enabled and self.cursor < len(self.batch_plan): + batch = self.batch_plan[self.cursor] + self.cursor += 1 + reservation = self.next_reservation + self.next_reservation += 1 + rows = [row for row, _ in batch] + indices = [self.episodes[row]["offset"] + anchor for row, anchor in batch] + for row in rows: + self.reserved[row] += 1 + self.pending[reservation] = rows + self.planned_indices[reservation] = indices + yield [(reservation, index) for index in indices] + + def commit(self, reservation: int) -> None: + super().commit(reservation) + del self.planned_indices[reservation] + + def cancel_outstanding(self) -> None: + super().cancel_outstanding() + self.planned_indices.clear() + + +def read_lerobot_manifest(root: Path | str) -> dict[str, Any]: + """Inspect a finalized corpus without loading decoded frames or the SDK. + + Args: + root: Standard corpus directory with the showcase sidecar. + + Returns: + Detached payload metadata and episode ordinals, offsets, lengths, + immutable sources, generations and numeric payload digests. + """ + root = Path(root) + if (root / "INCOMPLETE").exists(): + raise RuntimeError("Cannot read an incomplete LeRobot corpus") + manifest = json.loads((root / _SIDECAR).read_text()) + for filename in manifest["required_files"]: + if not (root / filename).is_file(): + raise FileNotFoundError(f"Incomplete local corpus: {filename}") + return manifest + + +class StandardLeRobotReader: + """Deliver decoded contexts using spawned SDK workers and parent reservations. + + Args: + root: Finalized SDK corpus with the showcase sidecar. + episodes: Optional logical episode indices, e.g. the first C diagnostic. + **kwargs: :class:`LeRobotReaderCfg` fields. + + K accounts for successfully returned complete chunks, with outstanding + prefetch separately reserved. No decoded corpus is cached by this wrapper. + A decoder/worker error is terminal. ``restart`` preserves delivered credits + and cancels only prefetched work after all old workers have stopped. + """ + + def __init__( + self, root: Path | str, *, episodes: list[int] | None = None, **kwargs: Any + ) -> None: + _check_versions() + self.root = Path(root) + self.manifest = read_lerobot_manifest(self.root) + self.metadata = copy.deepcopy(self.manifest["payload"]) + self.cfg = LeRobotReaderCfg(**kwargs) + self._owner = (os.getpid(), threading.get_ident()) + all_episodes = self.manifest["episodes"] + if episodes is not None and ( + not episodes + or len(set(episodes)) != len(episodes) + or any( + type(i) is not int or i < 0 or i >= len(all_episodes) for i in episodes + ) + ): + raise ValueError("episodes must contain unique logical episode indices") + selected = ( + all_episodes if episodes is None else [all_episodes[i] for i in episodes] + ) + self._sampler = self._make_sampler(selected) + self._dataset = self._make_dataset() + self._loader: DataLoader | None = None + self._iterator: Any = None + self._failed = False + self._closed = False + self._shutdown: dict[str, Any] = {} + self._shutdown_error: BaseException | None = None + self._stats: dict[str, Any] = { + "delivered_batches": 0, + "sdk_getitem_calls": 0, + "returned_valid_camera_frames": 0, + "requested_camera_timestamps": 0, + "worker_fetch_seconds": 0.0, + "worker_pids": [], + "reader": "official LeRobotDataset.__getitem__ + delta_timestamps", + "video_backend": "torchcodec CPU", + "decoded_corpus_preloaded": False, + "prefetch_batches": self.cfg.num_workers * self.cfg.prefetch_factor, + "persistent_workers": True, + "ordered_delivery": True, + "pinned_memory": False, + } + self._start() + + def _make_sampler(self, episodes: list[dict[str, Any]]) -> _ReservationSampler: + return _ReservationSampler(episodes, self.cfg) + + def _make_dataset(self) -> _DecodedContexts: + return _DecodedContexts(self.root, self.manifest, self.cfg) + + def _check(self) -> None: + if self._owner != (os.getpid(), threading.get_ident()): + raise RuntimeError("Reader operations require the creator process/thread") + if self._closed or self._failed: + raise RuntimeError("Reader is closed or terminally failed") + + def _start(self) -> None: + self._sampler.dispatch_enabled = True + self._loader = DataLoader( + self._dataset, + batch_sampler=self._sampler, + num_workers=self.cfg.num_workers, + multiprocessing_context="spawn", + prefetch_factor=self.cfg.prefetch_factor, + persistent_workers=True, + in_order=True, + pin_memory=False, + worker_init_fn=_initialize_worker, + timeout=self.cfg.timeout, + generator=torch.Generator().manual_seed(self.cfg.seed), + collate_fn=default_collate, + ) + self._iterator = iter(self._loader) + + def next_batch(self) -> dict[str, Any]: + """Return one decoded batch; StopIteration denotes a depleted fixed corpus. + + Returns: + Canonical CPU payload with ``_service`` worker timing diagnostics. + """ + self._check() + try: + batch = next(self._iterator) + reservations = batch.pop("_reservation_id") + if not (reservations == reservations[0]).all(): + raise RuntimeError("Mixed DataLoader exposure reservations") + reservation = int(reservations[0]) + service = batch.pop("_service_seconds") + started_ns = batch.pop("_service_start_ns") + ended_ns = batch.pop("_service_end_ns") + item_pids = batch.pop("_worker_pid") + worker_pids = item_pids.unique().tolist() + if (started_ns < 0).any() or (ended_ns < started_ns).any(): + raise RuntimeError("Invalid SDK worker host-monotonic service interval") + self._validate_batch(batch, reservation) + # Collation, SDK decode and normalization all succeeded before debit. + self._sampler.commit(reservation) + self._stats["delivered_batches"] += 1 + self._stats["sdk_getitem_calls"] += self.cfg.batch_size + self._stats["worker_fetch_seconds"] += service.sum().item() + cameras = len(self.metadata["cameras"]) + image_length = ( + next(iter(batch["images"].values())).shape[1] + if cameras + else batch["valid"].shape[1] + ) + self._stats["returned_valid_camera_frames"] += ( + int(batch["valid"][:, :image_length].sum()) * cameras + ) + self._stats["requested_camera_timestamps"] += ( + len(batch["valid"]) * image_length * cameras + ) + self._stats["worker_pids"] = sorted( + set(self._stats["worker_pids"]) | set(worker_pids) + ) + batch["_service"] = { + "worker_fetch_seconds": service.sum().item(), + "worker_item_seconds": service.tolist(), + "worker_fetch_intervals": [ + { + "worker_pid": pid, + "start_ns": int(started_ns[item_pids == pid].min()), + "end_ns": int(ended_ns[item_pids == pid].max()), + "item_count": int((item_pids == pid).sum()), + "scope": "contiguous batch item-fetch envelope including SDK Dataset startup/read, CPU video decode and temporal conversion/validation; excludes worker interpreter startup, collation, IPC and parent waiting; not pure decode, GPU busy time or critical-path time", + } + for pid in worker_pids + ], + "scope": "worker SDK fetch including decode, temporal conversion and padding; first item includes Dataset startup", + } + return batch + except StopIteration: + raise + except BaseException: + self._failed = True + raise + + def _validate_batch(self, batch: Mapping[str, Any], reservation: int) -> None: + selected = self._sampler.pending[reservation] + expected_ids = torch.tensor( + [self._sampler.episodes[row]["source_id"] for row in selected], + dtype=torch.int64, + ) + anchor = self.cfg.history_length - 1 + if not batch["valid"][:, anchor].all() or not torch.equal( + batch["frame_ids"][:, anchor, :2], expected_ids + ): + raise RuntimeError( + "Decoded batch source identities differ from its reservation" + ) + + @property + def stats(self) -> dict[str, Any]: + """Return delivery-local decode evidence; it excludes unused prefetch. + + Worker service is aggregate worker time, not additive critical-path wall + time. Requested timestamps include clamped boundary duplicates. + """ + return copy.deepcopy(self._stats) + + def snapshot_budget(self) -> dict[str, Any]: + """Return delivery and live-prefetch counts for every immutable source.""" + return { + "episode_budget": self.cfg.episode_budget, + "delivered_chunks": sum(self._sampler.delivered), + "reserved_chunks": sum(self._sampler.reserved), + "cancelled_prefetch_chunks": self._sampler.cancelled_chunks, + "pending_batches": len(self._sampler.pending), + "remaining_chunks": ( + None + if self.cfg.episode_budget is None + else len(self._sampler.episodes) * self.cfg.episode_budget + - sum(self._sampler.delivered) + - sum(self._sampler.reserved) + ), + "episodes": [ + { + "source_id": episode["source_id"], + "assigned": self.cfg.episode_budget, + "delivered": delivered, + "reserved": reserved, + "remaining": ( + None + if self.cfg.episode_budget is None + else self.cfg.episode_budget - delivered - reserved + ), + } + for episode, delivered, reserved in zip( + self._sampler.episodes, + self._sampler.delivered, + self._sampler.reserved, + ) + ], + } + + def _stop_workers(self) -> None: + if self._iterator is None: + if self._shutdown_error is not None: + raise self._shutdown_error + return + self._sampler.dispatch_enabled = False + self._shutdown = {"unconsumed_prefetch": self.snapshot_budget()} + try: + # Discard already reserved outputs without committing exposure or + # dispatching replacements. Workers become idle before teardown. + _finish_loader( + self._iterator, self._shutdown, timeout=min(30.0, self.cfg.timeout) + ) + except BaseException as error: + self._failed = True + self._shutdown_error = error + raise + finally: + self._iterator, self._loader = None, None + + def restart(self) -> None: + """Restart worker iteration without resetting delivered exposure credits.""" + self._check() + try: + self._stop_workers() + finally: + self._sampler.cancel_outstanding() + self._start() + + def close(self) -> dict[str, Any]: + """Drain unused prefetch without delivery, then require clean worker exits. + + Returns: + Shutdown status and the unconsumed reservation snapshot. + """ + if self._owner != (os.getpid(), threading.get_ident()): + raise RuntimeError("Reader close requires the creator process/thread") + if not self._closed: + try: + self._stop_workers() + finally: + self._sampler.cancel_outstanding() + self._closed = True + if self._shutdown_error is not None: + raise self._shutdown_error + return copy.deepcopy(self._shutdown) + + def __enter__(self) -> StandardLeRobotReader: + return self + + def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> None: + self.close() + + +class TrainingLeRobotReader(StandardLeRobotReader): + """Read an exact training plan with history-only standard video decoding. + + Args: + root: Finalized corpus. :func:`read_lerobot_manifest` inspects its + episode offsets/lengths before constructing a plan. + batch_plan: Ordered batches of ``(episode_ordinal, anchor_frame)``. + Anchors are zero-based positions within the saved episode, not + immutable source frame IDs. Every context must be complete and + remain within one continuity segment. Batches must be equal-sized. + history_length: Numeric and image causal history, including the anchor. + prediction_horizon: Numeric supervised horizon, including the anchor. + num_workers: Spawned official SDK/TorchCodec CPU reader workers. + prefetch_factor: Maximum prefetched batches per worker. + timeout: Worker fetch and bounded shutdown timeout in seconds. + seed: DataLoader worker seeding only; the caller owns all plan sampling. + + Numeric tensors cover ``history_length + prediction_horizon - 1`` frames. + Image tensors cover only ``history_length`` frames; future image timestamps + are never requested from the SDK. Plans are single-use, and ``restart`` is + explicitly unsupported so consumed training work cannot silently repeat. + """ + + def __init__( + self, + root: Path | str, + batch_plan: list[list[tuple[int, int]]], + *, + history_length: int = 2, + prediction_horizon: int = 32, + num_workers: int = 4, + prefetch_factor: int = 1, + timeout: float = 120.0, + seed: int = 29001, + ) -> None: + if not isinstance(batch_plan, list) or not batch_plan or not batch_plan[0]: + raise ValueError("Training requires a nonempty finite batch plan") + self._training_batch_plan = copy.deepcopy(batch_plan) + super().__init__( + root, + batch_size=len(batch_plan[0]), + history_length=history_length, + prediction_horizon=prediction_horizon, + episode_budget=None, + num_workers=num_workers, + prefetch_factor=prefetch_factor, + timeout=timeout, + seed=seed, + ) + self._stats.update( + plan_batches=len(batch_plan), + image_timestamps_per_chunk=history_length, + numeric_timestamps_per_chunk=history_length + prediction_horizon - 1, + complete_contexts_required=True, + restart_supported=False, + ) + + @property + def episodes(self) -> list[dict[str, Any]]: + """Return detached ordinal, offset, length and immutable-source metadata.""" + return copy.deepcopy(self.manifest["episodes"]) + + def _make_sampler( + self, episodes: list[dict[str, Any]] + ) -> _PlannedReservationSampler: + return _PlannedReservationSampler(episodes, self.cfg, self._training_batch_plan) + + def _make_dataset(self) -> _DecodedContexts: + if len(self.metadata["cameras"]) not in (1, 2): + raise ValueError("Visual training requires one or two configured cameras") + return _DecodedContexts( + self.root, + self.manifest, + self.cfg, + history_only_images=True, + require_complete=True, + ) + + def _validate_batch(self, batch: dict[str, Any], reservation: int) -> None: + super()._validate_batch(batch, reservation) + indices = batch.pop("_global_anchor_index") + expected = torch.tensor( + self._sampler.planned_indices[reservation], dtype=torch.int64 + ) + if not torch.equal(indices, expected) or not batch["valid"].all(): + raise RuntimeError("Decoded training anchors differ from the fixed plan") + + def restart(self) -> None: + """Reject replay of a single-use training plan; create a new reader instead.""" + raise NotImplementedError( + "Fixed training plans are single-use and cannot restart" + ) + + +def audit_video_quality(root: Path | str) -> dict[str, Any]: + """Decode small raw-reference samples outside timed consumption. + + Args: + root: Finalized corpus containing raw first-episode RGB samples. + + Returns: + Per-camera MAE and PSNR; lossless matches use null PSNR plus a flag. + No raw images beyond the fixed references are retained. + """ + root = Path(root) + manifest = json.loads((root / _SIDECAR).read_text()) + refs = manifest["references"] + if not refs: + return {"samples": [], "scope": "state-only; no video"} + cfg = LeRobotReaderCfg( + batch_size=1, + history_length=1, + prediction_horizon=1, + episode_budget=None, + num_workers=1, + ) + dataset = _DecodedContexts(root, manifest, cfg) + loader = DataLoader( + dataset, + batch_size=1, + sampler=[(i, ref["index"]) for i, ref in enumerate(refs)], + num_workers=1, + multiprocessing_context="spawn", + prefetch_factor=1, + worker_init_fn=_initialize_worker, + timeout=cfg.timeout, + ) + iterator = iter(loader) + samples = [] + shutdown: dict[str, Any] = {} + try: + for index, batch in enumerate(iterator): + ref = refs[index] + decoded = batch["images"][ref["camera"]][0, 0].movedim(0, -1).double() * 255 + raw = torch.from_numpy( + np.load(root / ref["path"], allow_pickle=False) + ).double() + diff = decoded - raw + mse = float(diff.square().mean()) + samples.append( + { + "camera": ref["camera"], + "index": ref["index"], + "mae_uint8": float(diff.abs().mean()), + "psnr_db": 10 * math.log10(255**2 / mse) if mse else None, + "exact": mse == 0, + } + ) + finally: + _finish_loader(iterator, shutdown, timeout=min(30.0, cfg.timeout)) + return { + "samples": samples, + "worker_cleanup": shutdown, + "scope": "fixed raw RGB references versus standard CPU video decoding; excluded from timed consumption", + } diff --git a/examples/ods/data_stream_demo/metrics.py b/examples/ods/data_stream_demo/metrics.py new file mode 100644 index 000000000..b390f7a9d --- /dev/null +++ b/examples/ods/data_stream_demo/metrics.py @@ -0,0 +1,727 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Source-identity accounting and bounded instrumentation for data delivery.""" + +from __future__ import annotations + +import json +import os +import threading +import time +from collections import Counter +from pathlib import Path +from typing import Any + +import numpy as np +import psutil +import torch + +from .payload import payload_bytes + +__all__ = ["StreamMetrics", "ResourceMonitor", "write_json", "validate_batch"] + + +def write_json(path: Path, value: Any) -> None: + """Write a complete finite JSON artifact atomically.""" + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(value, indent=2, allow_nan=False) + "\n") + temporary.replace(path) + + +def validate_batch( + batch: dict[str, Any], + *, + batch_size: int, + history: int, + future: int, + cameras: tuple[str, ...], + full_images: bool = False, +) -> None: + """Check causal/source boundaries and bounded image samples on owned data. + + Args: + batch: Common CPU payload returned by either data path. + batch_size: Requested context count. + history: Causal history extent including now. + future: Prediction extent including now. + cameras: Required camera UIDs. + full_images: Inspect every pixel in a separate engineering run. Normal + instrumented runs check a fixed spatial grid instead. + """ + length = history + future - 1 + valid, ids = batch["valid"], batch["frame_ids"] + if valid.shape != (batch_size, length) or valid.dtype != torch.bool: + raise ValueError("Context validity shape/dtype mismatch") + if ids.shape != (batch_size, length, 3) or ids.dtype != torch.int64: + raise ValueError("Context source ID shape/dtype mismatch") + if not valid[:, history - 1].all(): + raise ValueError("Every sampled current anchor must be valid") + anchor = ids[:, history - 1] + offsets = torch.arange(1 - history, future) + if (ids[~valid] != -1).any(): + raise ValueError("Invalid source padding must be -1") + if ((ids[..., :2] != anchor[:, None, :2]) & valid[..., None]).any(): + raise ValueError("A context crosses episode/run identity") + if ((ids[..., 2] != anchor[:, None, 2] + offsets) & valid).any(): + raise ValueError("Context frame indices are not causal and contiguous") + loss = batch["loss_mask"] + if not torch.equal(loss, valid & (offsets >= 0)): + raise ValueError("Loss mask includes history or invalid targets") + continuity = batch["continuity_id"] + if ((continuity < 0) & valid).any() or ( + (continuity != continuity[:, history - 1, None]) & valid + ).any(): + raise ValueError("Context crosses a continuity boundary") + for name in ("state", "actions"): + value = batch[name] + if ( + value.shape[:2] != valid.shape + or value.dtype != torch.float32 + or value.device.type != "cpu" + or not torch.isfinite(value).all() + or (value[~valid] != 0).any() + ): + raise ValueError(f"Invalid numeric payload: {name}") + if set(batch["images"]) != set(cameras): + raise ValueError("Missing or unexpected camera data") + for name, image in batch["images"].items(): + if ( + image.shape != (batch_size, length, 3, 256, 256) + or image.dtype != torch.float32 + or image.device.type != "cpu" + ): + raise ValueError(f"Image contract mismatch: {name}") + sample = image if full_images else image[..., ::32, ::32] + if ( + not torch.isfinite(sample).all() + or (sample < 0).any() + or (sample > 1).any() + or (sample[~valid] != 0).any() + ): + raise ValueError(f"Image range/padding mismatch: {name}") + + +class StreamMetrics: + """Count immutable source deliveries, retaining evictions and failed runs.""" + + def __init__(self, output: Path, *, history: int, quota: int | None) -> None: + self.output, self.history, self.quota = output, history, quota + launch = os.environ.get("ODS_PROCESS_LAUNCH_MONOTONIC_NS") + self.origin_ns = ( + int(launch) + if launch is not None + else int( + (time.monotonic() - (time.time() - psutil.Process().create_time())) + * 1e9 + ) + ) + self.origin = self.origin_ns / 1e9 + self.episodes: dict[str, dict[str, Any]] = {} + self.delivered = 0 + self.window_delivered = 0 + self.post_window_chunks = 0 + self.generations: Counter[int] = Counter() + self.window_generations: Counter[int] = Counter() + self.first_sample_by_generation: dict[int, float] = {} + self.publications: dict[int, dict[str, float]] = {} + self.latencies: list[float] = [] + self.services: list[float] = [] + self.worker_services: list[float] = [] + self.waits: Counter[str] = Counter() + self.window_waits: Counter[str] = Counter() + self.first_batch: float | None = None + self.reader_started: float | None = None + self.measurement_started: float | None = None + self.measurement_deadline: float | None = None + self.last_batch: float | None = None + self.overlap_generations: set[int] = set() + self.publication_completed: dict[int, float] = {} + self.image_samples = 0 + self.bytes = 0 + self._lifecycle = (output / "lifecycle.jsonl").open("w", buffering=1) + self._batches = (output / "batches.jsonl").open("w", buffering=1) + + def elapsed(self) -> float: + """Return launch-relative host monotonic seconds.""" + return time.monotonic() - self.origin + + def events(self, events: list[dict[str, Any]]) -> None: + """Record publications/evictions without keeping their frame arrays.""" + for event in events: + source = f"{event['run_id']}:{event['episode_id']}" + at = (event["monotonic_ns"] - self.origin_ns) / 1e9 + record = {**event, "source_id": source, "time_seconds": at} + self._lifecycle.write(json.dumps(record) + "\n") + row = self.episodes.setdefault( + source, + { + "source_id": source, + "generation": int(event["generation"]), + "delivered_chunks": 0, + "first_consumed_seconds": None, + "evicted": False, + }, + ) + if event["kind"] == "published": + row["published_seconds"] = at + row["slot"] = event["slot"] + self.publications.setdefault(int(event["generation"]), {})[source] = at + else: + row["evicted"] = True + row["evicted_seconds"] = at + + def interval(self, kind: str, started: float, ended: float) -> None: + """Record cooperative flow-control waits using absolute host times.""" + self.waits[kind] += ended - started + if ( + self.measurement_started is not None + and self.measurement_deadline is not None + ): + self.window_waits[kind] += max( + 0.0, + min(ended, self.measurement_deadline) + - max(started, self.origin + self.measurement_started), + ) + + def generation_complete(self, generation: int) -> None: + """Observe a completed fill without equating it with first publication.""" + self.publication_completed.setdefault(generation, self.elapsed()) + + def batch( + self, + batch: dict[str, Any], + *, + requested_at: float, + returned_at: float, + service_seconds: float, + sampled_at: float | None = None, + within_window: bool = True, + ) -> None: + """Record one successfully returned batch, counting each anchor once.""" + at = returned_at - self.origin + if self.first_batch is None: + self.first_batch = at + self.last_batch = at + ids = batch["frame_ids"][:, self.history - 1] + versions = batch["generations"][:, self.history - 1] + sources = Counter(f"{int(row[0])}:{int(row[1])}" for row in ids) + for source, count in sources.items(): + index = next( + i + for i, row in enumerate(ids) + if f"{int(row[0])}:{int(row[1])}" == source + ) + version = int(versions[index]) + row = self.episodes.setdefault( + source, + { + "source_id": source, + "generation": version, + "delivered_chunks": 0, + "first_consumed_seconds": None, + "evicted": False, + }, + ) + row["delivered_chunks"] += count + row["generation"] = version + if row["first_consumed_seconds"] is None: + row["first_consumed_seconds"] = at + if self.quota is not None and row["delivered_chunks"] > self.quota: + raise ValueError(f"Source delivery exceeded K: {source}") + counts = Counter(map(int, versions.tolist())) + self.generations.update(counts) + chunks = len(ids) + latency = returned_at - requested_at + valid_images = int(batch["valid"].sum()) * len(batch["images"]) + byte_count = payload_bytes(batch) + self.delivered += chunks + if within_window: + self.window_delivered += chunks + self.window_generations.update(counts) + self.image_samples += valid_images + self.bytes += byte_count + self.latencies.append(latency) + self.services.append(service_seconds) + for generation in counts: + self.first_sample_by_generation.setdefault( + generation, (sampled_at or returned_at) - self.origin + ) + if "_service" in batch: + values = batch["_service"]["worker_item_seconds"] + self.worker_services.append(sum(values) / len(values)) + else: + self.post_window_chunks += chunks + self._batches.write( + json.dumps( + { + "time_seconds": at, + "chunks": chunks, + "generations": counts, + "episode_counts": sources, + "request_latency_seconds": latency, + "service_seconds": service_seconds, + "valid_image_samples": valid_images, + "payload_bytes": byte_count, + "within_window": within_window, + "sampled_seconds": (sampled_at or returned_at) - self.origin, + } + ) + + "\n" + ) + + def summary(self, duration: float, *, cache_episodes: int = 0) -> dict[str, Any]: + """Return measured counters; empty timing populations remain unknown.""" + + def quantile(values: list[float], q: float) -> float | None: + return float(np.quantile(values, q)) if values else None + + self.overlap_generations = { + generation + for generation, sampled in self.first_sample_by_generation.items() + if generation > 1 + and len(self.publications.get(generation, {})) == cache_episodes + and sampled < max(self.publications[generation].values()) + } + return { + "first_batch_from_launch_seconds": self.first_batch, + "first_batch_from_reader_seconds": ( + self.first_batch - self.reader_started + if self.first_batch is not None and self.reader_started is not None + else None + ), + "source_ready_seconds": self.reader_started, + "delivered_chunks": self.delivered, + "window_delivered_chunks": self.window_delivered, + "post_window_chunks": self.post_window_chunks, + "delivered_chunks_per_second": ( + self.window_delivered / duration if duration else None + ), + "request_latency_p50_seconds": quantile(self.latencies, 0.5), + "request_latency_p95_seconds": quantile(self.latencies, 0.95), + "service_p50_seconds": quantile(self.services, 0.5), + "service_p95_seconds": quantile(self.services, 0.95), + "worker_item_service_p50_seconds": quantile(self.worker_services, 0.5), + "consumed_generations": sorted(self.window_generations), + "generation_chunk_counts": dict(self.generations), + "overlap_generations": sorted(self.overlap_generations), + "data_wait_seconds": self.window_waits["data_wait"], + "lifetime_data_wait_seconds": self.waits["data_wait"], + "throttle_seconds": self.window_waits["usage_throttle"] + + self.window_waits["request_throttle"], + "usage_throttle_seconds": self.window_waits["usage_throttle"], + "request_throttle_seconds": self.window_waits["request_throttle"], + "supply_exhausted_idle_seconds": self.window_waits["supply_exhausted_idle"], + "wait_scope": "measurement-window overlap; controller recovery includes its successful retry service; not a disjoint time decomposition; cold generation waits retained only in lifetime_data_wait_seconds", + "valid_image_samples": self.image_samples, + "payload_bytes": self.bytes, + "logical_payload_bytes_per_second": ( + self.bytes / duration if duration else None + ), + "cap_violations": sum( + self.quota is not None and row["delivered_chunks"] > self.quota + for row in self.episodes.values() + ), + "observation_scope": "in-window complete-batch CPU delivery; late completions retained in lifetime ledger only; no optimizer", + } + + def close(self) -> None: + """Close event logs after final worker draining.""" + self._lifecycle.close() + self._batches.close() + + +class ResourceMonitor: + """Stream sampled owned memory/CPU and explicitly whole-board GPU activity. + + CPU deltas count continuing processes in the owner's descendant tree. New + and disappeared processes are reported instead of imputing unknown work. + NVML utilization is board-wide, including foreign work; only context-memory + accounting is filtered to owned PIDs. No CUDA context is initialized. + """ + + def __init__(self, output: Path, interval: float = 2.0) -> None: + self._output, self._interval = output, interval + self._stop = threading.Event() + self._thread: threading.Thread | None = None + self.peak_pss: int | None = None + self.peak_vram: int | None = None + self.primary_device_peak: int | None = None + self.context_peaks: dict[str, int] = {} + self.samples = 0 + self.errors: list[str] = [] + self._error_count = 0 + self._previous_cpu: dict[tuple[int, float], tuple[float, float]] = {} + self._previous_cpu_at: float | None = None + self._affinity: list[int] = [] + self._cpu_count = 0 + self._cpu = { + "valid_interval_count": 0, + "observed_cpu_seconds": 0.0, + "observed_wall_seconds": 0.0, + "peak_core_equivalents": None, + "peak_assigned_capacity_percent": None, + "new_process_observations": 0, + "missing_process_observations": 0, + } + self._gpu: dict[str, dict[str, Any]] = {} + self._observer_failed = False + + def _error(self, message: str) -> None: + self._error_count += 1 + if message not in self.errors and len(self.errors) < 32: + self.errors.append(message) + + def start(self) -> None: + """Start one bounded-lifetime CPU observer thread.""" + self._thread = threading.Thread(target=self._run, daemon=True) + self._thread.start() + + def _owned_processes(self, owner: psutil.Process) -> list[psutil.Process]: + # psutil.children(recursive=True) builds a host-wide PID/PPID map on + # Linux. Follow only owned thread child files instead; include children + # created by native/Python threads, not only the main thread. + if not Path("/proc/self/task").is_dir(): + return [owner, *owner.children(recursive=True)] + pending: list[tuple[int, int | None]] = [(owner.pid, None)] + seen: set[int] = set() + processes = [] + while pending: + pid, parent = pending.pop() + if pid in seen: + continue + seen.add(pid) + try: + process = owner if parent is None else psutil.Process(pid) + if parent is not None and process.ppid() != parent: + continue + processes.append(process) + for task in Path(f"/proc/{pid}/task").iterdir(): + try: + pending.extend( + (int(child), pid) + for child in (task / "children").read_text().split() + ) + except (FileNotFoundError, ProcessLookupError): + continue + except (FileNotFoundError, ProcessLookupError, psutil.NoSuchProcess): + continue + except (PermissionError, psutil.AccessDenied) as error: + self._error(f"Owned child discovery unavailable: {error}") + return processes + + def _cpu_sample(self, processes: list[psutil.Process], at: float) -> dict[str, Any]: + current = {} + for process in processes: + try: + identity = (process.pid, process.create_time()) + values = process.cpu_times() + current[identity] = (float(values.user), float(values.system)) + except (psutil.NoSuchProcess, psutil.AccessDenied): + continue + previous = self._previous_cpu + continuing = current.keys() & previous.keys() + added, missing = ( + current.keys() - previous.keys(), + previous.keys() - current.keys(), + ) + elapsed = None if self._previous_cpu_at is None else at - self._previous_cpu_at + deltas = [ + (current[key][0] - previous[key][0], current[key][1] - previous[key][1]) + for key in continuing + ] + valid = ( + elapsed is not None + and elapsed > 0 + and bool(deltas) + and all(user >= 0 and system >= 0 for user, system in deltas) + ) + user = sum(value[0] for value in deltas) if valid else None + system = sum(value[1] for value in deltas) if valid else None + cores = (user + system) / elapsed if valid else None + capacity = 100 * cores / self._cpu_count if valid and self._cpu_count else None + if valid: + self._cpu["valid_interval_count"] += 1 + self._cpu["observed_cpu_seconds"] += user + system + self._cpu["observed_wall_seconds"] += elapsed + self._cpu["peak_core_equivalents"] = max( + self._cpu["peak_core_equivalents"] or 0, cores + ) + if capacity is not None: + self._cpu["peak_assigned_capacity_percent"] = max( + self._cpu["peak_assigned_capacity_percent"] or 0, capacity + ) + self._cpu["new_process_observations"] += len(added) + self._cpu["missing_process_observations"] += len(missing) + self._previous_cpu, self._previous_cpu_at = current, at + return { + "interval_seconds": elapsed, + "user_cpu_seconds": user, + "system_cpu_seconds": system, + "core_equivalents": cores, + "assigned_capacity_percent": capacity, + "observed_processes": len(current), + "continuing_processes": len(continuing), + "new_pids": sorted(key[0] for key in added), + "missing_pids": sorted(key[0] for key in missing), + "assigned_cpu_count": self._cpu_count, + "assigned_cpu_affinity": self._affinity, + "scope": "user+system CPU deltas for continuing owned PIDs; excludes unknown first/new/disappeared-process intervals; 100% capacity equals the owner's assigned affinity", + } + + def _primary_devices( + self, nvml: Any + ) -> tuple[Any, dict[str, tuple[Any, list[str]]]]: + visible = [ + value.strip() + for value in os.environ.get("CUDA_VISIBLE_DEVICES", "").split(",") + if value.strip() + ] + override = os.environ.get("ODS_GPU_UUID") + if not visible and override: + visible = [override] + config_file = self._output / "config.json" + config = json.loads(config_file.read_text()) if config_file.exists() else {} + roles: dict[int, list[str]] = {} + for field, role in (("sim_device", "simulation"), ("train_device", "training")): + device = config.get(field, "") + if device.startswith("cuda:"): + roles.setdefault(int(device.split(":")[-1]), []).append(role) + devices = {} + legacy = None + for index, value in enumerate(visible): + try: + handle = ( + nvml.nvmlDeviceGetHandleByIndex(int(value)) + if value.isdigit() + else nvml.nvmlDeviceGetHandleByUUID(value) + ) + uuid = nvml.nvmlDeviceGetUUID(handle) + if isinstance(uuid, bytes): + uuid = uuid.decode() + assigned = roles.get( + index, ["primary" if len(visible) == 1 else f"visible_cuda:{index}"] + ) + devices[uuid] = (handle, assigned) + self._gpu[uuid] = { + "roles": assigned, + "sample_count": 0, + "gpu_sum": 0.0, + "gpu_peak": None, + "memory_sum": 0.0, + "memory_peak": None, + } + if legacy is None or value == override or uuid == override: + legacy = handle + except Exception as error: + self._error(f"Primary GPU {value} unavailable: {error}") + return legacy, devices + + def _run(self) -> None: + handle = None + device_handles = [] + primary_devices = {} + try: + import pynvml + + pynvml.nvmlInit() + device_handles = [ + pynvml.nvmlDeviceGetHandleByIndex(i) + for i in range(pynvml.nvmlDeviceGetCount()) + ] + handle, primary_devices = self._primary_devices(pynvml) + except Exception as error: + self._error(f"NVML unavailable: {error}") + owner = psutil.Process() + try: + self._affinity = owner.cpu_affinity() + except (AttributeError, psutil.AccessDenied, NotImplementedError): + self._affinity = ( + sorted(os.sched_getaffinity(0)) + if hasattr(os, "sched_getaffinity") + else [] + ) + self._cpu_count = len(self._affinity) or os.cpu_count() or 1 + try: + with (self._output / "resources.jsonl").open("w", buffering=1) as stream: + while not self._stop.is_set(): + sampled_at = time.monotonic() + pss = 0 + available = False + processes = self._owned_processes(owner) + owned_pids = {process.pid for process in processes} + cpu = self._cpu_sample(processes, sampled_at) + for process in processes: + try: + value = getattr(process.memory_full_info(), "pss", None) + if value is not None: + available = True + pss += value + except (psutil.NoSuchProcess, psutil.AccessDenied): + continue + vram = None + if handle is not None: + try: + vram = int(pynvml.nvmlDeviceGetMemoryInfo(handle).used) + except Exception as error: + self._error(str(error)) + handle = None + gpu = {} + for uuid, (device_handle, roles) in primary_devices.items(): + try: + utilization = pynvml.nvmlDeviceGetUtilizationRates( + device_handle + ) + active, memory = float(utilization.gpu), float( + utilization.memory + ) + if not (0 <= active <= 100 and 0 <= memory <= 100): + raise ValueError("NVML utilization is outside [0,100]") + gpu[uuid] = { + "roles": roles, + "gpu_percent": active, + "memory_controller_percent": memory, + "scope": "whole board including foreign work; not owned kernel attribution", + } + aggregate = self._gpu[uuid] + aggregate["sample_count"] += 1 + aggregate["gpu_sum"] += active + aggregate["memory_sum"] += memory + aggregate["gpu_peak"] = max( + aggregate["gpu_peak"] or 0, active + ) + aggregate["memory_peak"] = max( + aggregate["memory_peak"] or 0, memory + ) + except Exception as error: + self._error(f"GPU utilization {uuid} unavailable: {error}") + gpu[uuid] = { + "roles": roles, + "gpu_percent": None, + "memory_controller_percent": None, + "scope": "whole board including foreign work; unavailable sample", + } + contexts = {} + context_supported = bool(device_handles) + for device_handle in device_handles: + try: + uuid = pynvml.nvmlDeviceGetUUID(device_handle) + if isinstance(uuid, bytes): + uuid = uuid.decode() + per_pid = {} + total = pynvml.nvmlDeviceGetMemoryInfo(device_handle).total + for query in ( + pynvml.nvmlDeviceGetComputeRunningProcesses, + pynvml.nvmlDeviceGetGraphicsRunningProcesses, + ): + for item in query(device_handle): + value = item.usedGpuMemory + if ( + item.pid in owned_pids + and isinstance(value, int) + and 0 <= value <= total + ): + per_pid[item.pid] = max( + per_pid.get(item.pid, 0), value + ) + if per_pid: + contexts[uuid] = sum(per_pid.values()) + except Exception: + context_supported = False + if context_supported: + self.peak_vram = max( + self.peak_vram or 0, sum(contexts.values()) + ) + for uuid, value in contexts.items(): + self.context_peaks[uuid] = max( + self.context_peaks.get(uuid, 0), value + ) + if available: + self.peak_pss = max(self.peak_pss or 0, pss) + if vram is not None: + self.primary_device_peak = max( + self.primary_device_peak or 0, vram + ) + self.samples += 1 + stream.write( + json.dumps( + { + "monotonic": sampled_at, + "pss_bytes": pss if available else None, + "device_used_bytes": vram, + "owned_context_bytes": ( + contexts if context_supported else None + ), + "cpu_utilization": cpu, + "primary_gpu_utilization": gpu, + }, + allow_nan=False, + ) + + "\n" + ) + self._stop.wait(self._interval) + except Exception as error: + self._observer_failed = True + self._error(f"Resource observer failed: {error}") + + def close(self) -> dict[str, Any]: + """Join instrumentation and return sampled maxima/utilization aggregates.""" + self._stop.set() + if self._thread is not None: + self._thread.join(timeout=10) + if self._thread.is_alive(): + raise RuntimeError("Resource observer did not join") + cpu = dict(self._cpu) + elapsed = cpu["observed_wall_seconds"] + cores = cpu["observed_cpu_seconds"] / elapsed if elapsed else None + cpu.update( + mean_core_equivalents=cores, + mean_assigned_capacity_percent=( + 100 * cores / self._cpu_count if cores is not None else None + ), + assigned_cpu_affinity=self._affinity, + assigned_cpu_count=self._cpu_count, + scope="continuing owned PID-tree CPU time; short-lived/new/missing PID intervals may be omitted; capacity denominator is owner's assigned affinity", + ) + gpu = {} + for uuid, aggregate in self._gpu.items(): + count = aggregate["sample_count"] + gpu[uuid] = { + "roles": aggregate["roles"], + "sample_count": count, + "mean_sampled_gpu_percent": ( + aggregate["gpu_sum"] / count if count else None + ), + "peak_sampled_gpu_percent": aggregate["gpu_peak"], + "mean_sampled_memory_controller_percent": ( + aggregate["memory_sum"] / count if count else None + ), + "peak_sampled_memory_controller_percent": aggregate["memory_peak"], + "scope": "whole-board NVML recent-window utilization including foreign work; sample mean, not exclusive owned-kernel busy time", + } + return { + "peak_pss_bytes": self.peak_pss, + "peak_vram_bytes": self.peak_vram, + "primary_device_sampled_peak_bytes": self.primary_device_peak, + "owned_device_context_sampled_peaks": self.context_peaks, + "resource_sample_count": self.samples, + "resource_errors": self.errors, + "resource_error_count": self._error_count, + "observer_complete": not self._observer_failed, + "cpu_utilization": cpu, + "primary_gpu_utilization": gpu, + "memory_scope": "sampled owned-process PSS and NVML context memory across devices (max per PID/device before sum); primary device total separately diagnostic; not guaranteed peaks", + } diff --git a/examples/ods/data_stream_demo/payload.py b/examples/ods/data_stream_demo/payload.py new file mode 100644 index 000000000..3fd4c6dbe --- /dev/null +++ b/examples/ods/data_stream_demo/payload.py @@ -0,0 +1,249 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Owned CPU payloads shared by the live and standard-decoded showcase arms.""" + +from __future__ import annotations + +import hashlib +import math +from collections.abc import Mapping +from typing import Any + +import torch +from tensordict import TensorDictBase + +__all__ = [ + "describe_raw_buffer", + "normalize_online_batch", + "pack_episode", + "numeric_episode_digest", + "payload_bytes", +] + + +def _path(key: str | tuple[str, ...]) -> tuple[str, ...]: + return (key,) if isinstance(key, str) else key + + +def describe_raw_buffer(raw: TensorDictBase) -> dict[str, Any]: + """Describe numeric observation order and native image layout without copying. + + Args: + raw: ODS rollout buffer or sampled context with nested ``obs``. + + Returns: + JSON-compatible metadata used by both transports. State leaves are + sorted by their complete path; image features retain camera UIDs. + """ + batch_dims = raw.batch_dims + state_features: list[dict[str, Any]] = [] + cameras: dict[str, Any] = {} + for key, tensor in sorted( + raw["obs"].items(include_nested=True, leaves_only=True), + key=lambda pair: _path(pair[0]), + ): + path = _path(key) + shape = list(tensor.shape[batch_dims:]) + if len(path) == 3 and path[0] == "sensor" and path[-1] == "color": + if tensor.dtype != torch.uint8 or len(shape) != 3 or shape[-1] != 4: + raise ValueError("Native camera color must be HWC RGBA uint8") + cameras[path[1]] = {"path": list(path), "shape": shape} + elif path[0] == "sensor": + raise ValueError(f"Unsupported sensor feature: {path}") + else: + if tensor.dtype == torch.bool or not tensor.is_floating_point(): + raise ValueError(f"State observation must be floating point: {path}") + state_features.append( + {"path": list(path), "shape": shape, "width": math.prod(shape)} + ) + action_shape = list(raw["actions"].shape[batch_dims:]) + if not state_features or len(action_shape) != 1: + raise ValueError("Require numeric state features and flat action vectors") + return { + "format": "ods-data-stream-payload-v1", + "state_features": state_features, + "state_width": sum(feature["width"] for feature in state_features), + "action_width": action_shape[0], + "cameras": cameras, + "image_output": "RGB float32 channel-first [0,1]", + } + + +def _numeric(raw: TensorDictBase, metadata: Mapping[str, Any]) -> torch.Tensor: + shape = tuple(raw.batch_size) + return torch.cat( + [ + raw[("obs", *feature["path"])].reshape(*shape, feature["width"]) + for feature in metadata["state_features"] + ], + -1, + ).to(device="cpu", dtype=torch.float32) + + +def _provenance(raw: TensorDictBase) -> dict[str, torch.Tensor]: + keys = raw.keys() + if "frame_ids" in keys: + ids = raw["frame_ids"] + else: + ids = torch.stack( + [raw[key] for key in ("ods_run_id", "ods_episode_id", "ods_frame_index")], + -1, + ) + result = {"frame_ids": ids.to(device="cpu", dtype=torch.int64).clone()} + for canonical, native in ( + ("generations", "ods_generation"), + ("published_at_monotonic_ns", "ods_published_at_monotonic_ns"), + ): + key = canonical if canonical in keys else native + result[canonical] = raw[key].to(device="cpu", dtype=torch.int64).clone() + if "continuity_id" not in keys: + raise ValueError("The showcase requires explicit continuity_id") + result["continuity_id"] = ( + raw["continuity_id"].to(device="cpu", dtype=torch.int64).clone() + ) + return result + + +def normalize_online_batch( + raw: TensorDictBase, metadata: Mapping[str, Any] +) -> dict[str, Any]: + """Convert one owned ODS context to the common RGB float32 contract. + + Args: + raw: Owned, zero-padded Engine or OnlineDataSource context with batch + axes ``[B,L]``. + metadata: Result of :func:`describe_raw_buffer`. + + Returns: + Owned CPU tensors, including zero images/numerics and -1 source padding. + The caller includes this conversion in measured delivery time. + """ + if raw.batch_dims != 2: + raise ValueError("A sampled context must have [batch,time] axes") + valid = raw["valid"].to(device="cpu", dtype=torch.bool).clone() + result: dict[str, Any] = { + "state": _numeric(raw, metadata), + "actions": raw["actions"].to(device="cpu", dtype=torch.float32).clone(), + "valid": valid, + "loss_mask": raw["loss_mask"].to(device="cpu", dtype=torch.bool).clone(), + **_provenance(raw), + "images": {}, + } + divisor = torch.tensor(255.0, device="cpu", dtype=torch.float32) + for name, feature in metadata["cameras"].items(): + rgb = raw[("obs", *feature["path"])][..., :3].movedim(-1, -3) + image = torch.empty(rgb.shape, device="cpu", dtype=torch.float32) + if rgb.device.type == "cpu": + # Cast and normalize in one pass. Engine padding is already zero, + # and division preserves it without scanning the image again. + torch.div(rgb, divisor, out=image) + else: + image.copy_(rgb).div_(255) + result["images"][name] = image + for key in ("state", "actions"): + result[key].masked_fill_(~valid[..., None], 0) + for key in ( + "frame_ids", + "generations", + "continuity_id", + "published_at_monotonic_ns", + ): + mask = valid[..., None] if key == "frame_ids" else valid + result[key].masked_fill_(~mask, -1) + return result + + +def pack_episode(rawrow: TensorDictBase, metadata: Mapping[str, Any]) -> dict[str, Any]: + """Copy one accepted cache row into a bounded writer-owned CPU episode. + + Args: + rawrow: One row with a single time batch axis and immutable provenance. + metadata: Shared ordered observation and camera description. + + Returns: + An owned episode trimmed to real frames, with RGB uint8 HWC images. + Interior invalid gaps are rejected instead of silently retiming data. + """ + if rawrow.batch_dims != 1: + raise ValueError("pack_episode requires one time-axis row") + valid = rawrow["valid"].bool() + positions = valid.nonzero().flatten() + if not positions.numel(): + raise ValueError("Cannot persist an empty accepted episode") + start, stop = int(positions[0]), int(positions[-1]) + 1 + if not valid[start:stop].all(): + raise ValueError("Accepted episode contains interior invalid gaps") + row = rawrow[start:stop] + episode: dict[str, Any] = { + "state": _numeric(row, metadata), + "actions": row["actions"].to(device="cpu", dtype=torch.float32).clone(), + "valid": torch.ones(stop - start, dtype=torch.bool), + **_provenance(row), + "images": {}, + } + for name, feature in metadata["cameras"].items(): + episode["images"][name] = row[("obs", *feature["path"])][..., :3].to( + device="cpu", copy=True, memory_format=torch.contiguous_format + ) + for key in ("state", "actions"): + if not torch.isfinite(episode[key]).all(): + raise ValueError(f"Nonfinite accepted {key}") + ids = episode["frame_ids"] + if (ids[:, :2] != ids[0, :2]).any() or (ids[1:, 2] != ids[:-1, 2] + 1).any(): + raise ValueError("Accepted episode must retain one contiguous source identity") + if (ids < 0).any() or (episode["continuity_id"] < 0).any(): + raise ValueError("Accepted frames require valid source and continuity IDs") + generations = episode["generations"] + if (generations < 1).any() or (generations != generations[0]).any(): + raise ValueError("An accepted episode must belong to one positive generation") + return episode + + +def numeric_episode_digest(episode: Mapping[str, Any]) -> str: + """Hash accepted numeric payload independently of run IDs and clocks. + + Args: + episode: Packed accepted episode. + + Returns: + SHA256 covering ordered state/actions, lengths and continuity boundaries. + """ + digest = hashlib.sha256() + for name in ("state", "actions", "continuity_id"): + value = episode[name].contiguous() + digest.update(str((name, tuple(value.shape), str(value.dtype))).encode()) + digest.update(value.numpy().tobytes()) + return digest.hexdigest() + + +def payload_bytes(payload: Mapping[str, Any]) -> int: + """Count tensor payload bytes, including image padding. + + Args: + payload: Canonical batch or packed episode. + + Returns: + Logical tensor bytes; views and temporaries are not a peak-memory model. + """ + return sum( + ( + value.numel() * value.element_size() + if isinstance(value, torch.Tensor) + else payload_bytes(value) if isinstance(value, Mapping) else 0 + ) + for value in payload.values() + ) diff --git a/examples/ods/data_stream_demo/preview.py b/examples/ods/data_stream_demo/preview.py new file mode 100644 index 000000000..a968d2503 --- /dev/null +++ b/examples/ods/data_stream_demo/preview.py @@ -0,0 +1,187 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Preview the measured dual-camera expert trajectory before visual BC training. + +This entry point performs acquisition only. It never constructs a policy, +optimizer or learner. Use a fresh output directory and mask an idle simulation +GPU before launch. Camera poses and task physics come from acquisition.py. +""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path +import time +from typing import Any + +from .acquisition import ( + camera_profile_version, + camera_uids, + iter_generation_episodes, + load_acquisition_config, + make_engine, +) + +__all__ = ["capture_preview"] + + +def capture_preview( + output: Path, *, device: str, seed: int, camera_profile: str = "pick_place" +) -> dict[str, Any]: + """Save actual camera RGB frames/video from one accepted physical episode. + + Args: + output: New directory for the preview and resolved camera configuration. + device: Explicit simulation device within the launcher's GPU mask. + seed: Preview-only generation seed, separate from training/evaluation. + camera_profile: Camera calibration to review before training. + + Returns: + Camera/frame metadata and clean producer shutdown evidence. + """ + import imageio.v2 as imageio + import matplotlib + import numpy as np + import torch + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + output = Path(output).resolve() + output.mkdir(parents=True, exist_ok=False) + torch.set_num_threads(1) + os.environ["EMBODICHAIN_SIM_EXIT_PROCESS"] = "0" + config = load_acquisition_config( + "dual", device, producer_envs=1, seed=seed, camera_profile=camera_profile + ) + (output / "environment-config.json").write_text(json.dumps(config, indent=2)) + engine = make_engine( + "dual", + device, + producer_envs=1, + seed=seed, + buffer_size=1, + initialization_timeout=300, + camera_profile=camera_profile, + ) + started = time.monotonic() + report: dict[str, Any] = { + "status": "acquiring", + "seed": seed, + "device": device, + "camera_profile_version": camera_profile_version(camera_profile), + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "training_updates": 0, + } + try: + engine.start() + episodes = list(iter_generation_episodes(engine, 1)) + if len(episodes) != 1: + raise RuntimeError("Preview requires exactly one accepted episode") + raw = episodes[0] + report.update( + frames=len(raw), + engine_run_id=engine.run_id, + producer_pid=engine._sim_process.pid, + engine_stats=engine.get_stats(), + acquisition_seconds=time.monotonic() - started, + ) + except BaseException as error: + report.update(status="failed", error=repr(error)) + raise + finally: + process = engine._sim_process + popen = None if process is None else process._popen + engine.stop() + report["producer_exitcode"] = None if popen is None else popen.returncode + report["producer_closed"] = engine._sim_process is None + (output / "preview.json").write_text(json.dumps(report, indent=2) + "\n") + if report["producer_exitcode"] != 0: + raise RuntimeError("Preview producer failed to close normally") + + cameras = camera_uids("dual") + images = { + uid: raw["obs", "sensor", uid, "color"][..., :3].cpu().numpy() + for uid in cameras + } + if any(value.shape != (len(raw), 256, 256, 3) for value in images.values()): + raise ValueError("Preview camera shape differs from the training contract") + indices = np.linspace(0, len(raw) - 1, 6).astype(int) + fig, axes = plt.subplots(3, 4, figsize=(12, 9), constrained_layout=True) + for ordinal, frame in enumerate(indices): + row, pair = divmod(ordinal, 2) + for view, uid in enumerate(cameras): + ax = axes[row, pair * 2 + view] + ax.imshow(images[uid][frame]) + ax.set_title(f"{uid} | frame {frame} | {frame / 30:.1f} s", fontsize=9) + ax.axis("off") + fig.suptitle("Franka PickPlace: measured expert cameras; no policy training") + fig.savefig(output / "camera-preview.png", dpi=150) + plt.close(fig) + imageio.imwrite( + output / "initial-cameras.png", + np.concatenate([images[uid][0] for uid in cameras], axis=1), + ) + with imageio.get_writer( + output / "dual-view-expert.mp4", + fps=30, + codec="libx264", + quality=8, + ffmpeg_params=["-threads", "1"], + ) as writer: + for frame in range(len(raw)): + writer.append_data( + np.concatenate([images[uid][frame] for uid in cameras], axis=1) + ) + report.update( + status="complete", + cameras=list(cameras), + fps=30, + observation_shape=[256, 256, 3], + initial_cube_xy=raw["obs", "cube_pose"][0, :2].tolist(), + maximum_cube_height=float(raw["obs", "cube_pose"][:, 2].max()), + sensor_preview="Actual uint8 RGB camera observations; no generated imagery", + ) + (output / "preview.json").write_text(json.dumps(report, indent=2) + "\n") + return report + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--seed", type=int, default=20291025) + parser.add_argument( + "--camera-profile", choices=("benchmark", "pick_place"), default="pick_place" + ) + args = parser.parse_args() + print( + json.dumps( + capture_preview( + args.output, + device=args.device, + seed=args.seed, + camera_profile=args.camera_profile, + ) + ) + ) + + +if __name__ == "__main__": + main() diff --git a/examples/ods/data_stream_demo/qualify_encoder.py b/examples/ods/data_stream_demo/qualify_encoder.py new file mode 100644 index 000000000..0411927cf --- /dev/null +++ b/examples/ods/data_stream_demo/qualify_encoder.py @@ -0,0 +1,405 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Executable CNN qualification against pinned RLInf and image-only targets. + +This is an engineering qualification entry point, not a timing benchmark or a +policy-performance result. CUDA use is explicit; no GPU is selected implicitly. +""" + +from __future__ import annotations + +import argparse +import ast +from copy import deepcopy +from functools import partial +import hashlib +import json +from pathlib import Path +import time +from typing import Any +from urllib.request import urlopen + +import torch +from torch import nn +from torchvision.models.resnet import BasicBlock, ResNet + +from embodichain.learning.models import ResNet10Encoder +from examples.ods.data_stream_demo.visual_policy import VisualChunkBC + +__all__ = [] + +_REFERENCE_COMMIT = "91a302ec50ddacd1fdabc5859b85e3661d768236" +_REFERENCE_SHA256 = "5776e8a65d1d14dbcfeae070cb0ef23773d2a00a21a59f4f5f5b1933dae4cd47" +_REFERENCE_URL = ( + "https://raw.githubusercontent.com/RLinf/RLinf/" + f"{_REFERENCE_COMMIT}/rlinf/models/embodiment/modules/resnet_utils.py" +) + + +def _upstream_reference(source_path: Path | None) -> tuple[nn.Module, nn.Module]: + source = ( + source_path.read_bytes() + if source_path is not None + else urlopen(_REFERENCE_URL, timeout=30).read() + ) + digest = hashlib.sha256(source).hexdigest() + if digest != _REFERENCE_SHA256: + raise ValueError(f"Pinned RLInf source SHA256 mismatch: {digest}") + selected = {"MyGroupNorm", "ResNet10", "SpatialLearnedEmbeddings"} + tree = ast.Module( + body=[ + node + for node in ast.parse(source).body + if isinstance(node, ast.ClassDef) and node.name in selected + ], + type_ignores=[], + ) + namespace = { + "torch": torch, + "nn": nn, + "partial": partial, + "BasicBlock": BasicBlock, + "ResNet": ResNet, + } + # Only hash-verified, explicitly selected upstream class bodies are executed. + exec(compile(tree, _REFERENCE_URL, "exec"), namespace) + backbone = namespace["ResNet10"](pre_pooling=True) + del backbone.fc + del backbone.avgpool + pooling = namespace["SpatialLearnedEmbeddings"](8, 8, 512, 8) + return backbone, pooling + + +def _reference_check(source_path: Path | None) -> dict[str, Any]: + torch.manual_seed(481) + actual = ResNet10Encoder() + backbone, pooling = _upstream_reference(source_path) + backbone.load_state_dict(actual.resnet_backbone.state_dict()) + pooling.load_state_dict(actual.pooling_layer.state_dict()) + projection = deepcopy(actual.mlp) + images = torch.rand(1, 3, 256, 256, requires_grad=True) + upstream_images = images.detach().clone().requires_grad_() + intermediates: dict[str, list[torch.Tensor]] = {} + handles = [] + for name in ("conv1", "layer1", "layer2", "layer3", "layer4"): + intermediates[name] = [] + for module in (getattr(actual.resnet_backbone, name), getattr(backbone, name)): + handles.append( + module.register_forward_hook( + lambda _module, _args, output, name=name: intermediates[ + name + ].append(output.detach().clone()) + ) + ) + result = actual(images) + upstream_features = backbone( + (upstream_images - actual.image_mean) / actual.image_std + ) + upstream_embedding = pooling(upstream_features) + expected = projection(upstream_embedding) + for handle in handles: + handle.remove() + for values in intermediates.values(): + torch.testing.assert_close(values[0], values[1], rtol=0, atol=0) + actual_embedding = actual.pooling_layer(upstream_features) + torch.testing.assert_close( + actual_embedding, upstream_embedding, rtol=2e-5, atol=5e-5 + ) + torch.testing.assert_close(result, expected, rtol=5e-5, atol=5e-6) + objective = torch.linspace(-1, 1, 256) + (result * objective).sum().backward() + (expected * objective).sum().backward() + torch.testing.assert_close(images.grad, upstream_images.grad, rtol=5e-4, atol=2e-5) + differences = {} + for section, current, reference in ( + ("backbone", actual.resnet_backbone, backbone), + ("spatial", actual.pooling_layer, pooling), + ("projection", actual.mlp, projection), + ): + for (name, parameter), (_, upstream_parameter) in zip( + current.named_parameters(), reference.named_parameters(), strict=True + ): + assert parameter.grad is not None and upstream_parameter.grad is not None + torch.testing.assert_close( + parameter.grad, upstream_parameter.grad, rtol=5e-4, atol=3e-4 + ) + differences[f"{section}.{name}"] = float( + (parameter.grad - upstream_parameter.grad).abs().max() + ) + return { + "passed": True, + "commit": _REFERENCE_COMMIT, + "source_sha256": _REFERENCE_SHA256, + "intermediates": list(intermediates), + "output_max_abs_error": float((result - expected).abs().max().detach()), + "embedding_max_abs_error": float( + (actual_embedding - upstream_embedding).abs().max().detach() + ), + "parameter_gradient_max_abs_errors": differences, + "scope": "encoder structure/weights with freeze, detach and dropout removed", + } + + +def _policy(cameras: tuple[str, ...], horizon: int = 2) -> VisualChunkBC: + return VisualChunkBC( + torch.zeros(27), + torch.ones(27), + torch.zeros(10), + torch.ones(10), + cameras, + horizon, + ) + + +def _policy_contract_check() -> dict[str, Any]: + torch.manual_seed(712) + policy = _policy(("third_person", "wrist")) + proprio = torch.rand(2, 2, 27) + images = {camera: torch.rand(2, 2, 3, 256, 256) for camera in policy.cameras} + observed = [] + handle = policy.encoder.register_forward_pre_hook( + lambda _module, arguments: observed.append(arguments[0].detach().clone()) + ) + prediction = policy(proprio, images) + handle.remove() + expected_images = torch.stack([images[c] for c in policy.cameras], 2).flatten(0, 2) + torch.testing.assert_close(observed[0], expected_images, rtol=0, atol=0) + with torch.no_grad(): + batch_features = policy.encoder(expected_images) + individual = torch.cat( + [policy.encoder(frame[None]) for frame in expected_images] + ) + torch.testing.assert_close(batch_features, individual, rtol=5e-5, atol=8e-6) + for camera in policy.cameras: + changed = dict(images) + changed[camera] = 1 - images[camera] + assert not torch.allclose(prediction, policy(proprio, changed)) + geometry = torch.rand(2, 3, 51) + actions = torch.rand(2, 3, 7) * 0.1 + valid = torch.ones(2, 3, dtype=torch.bool) + mask = valid.clone() + mask[:, 0] = False + original_inputs = [] + hook = policy.register_forward_pre_hook( + lambda _module, args: original_inputs.append(args[0].detach().clone()) + ) + first_loss = policy.loss( + geometry, actions, images, loss_mask=mask, valid_mask=valid + ) + altered = geometry.clone() + altered[:, 2:] += 0.3 + policy.loss(altered, actions, images, loss_mask=mask, valid_mask=valid) + hook.remove() + torch.testing.assert_close(original_inputs[0], original_inputs[1], rtol=0, atol=0) + assert original_inputs[0].shape == (2, 2, 27) + (first_loss[0] + first_loss[1]).backward() + gradients = {} + for name, parameter in policy.named_parameters(): + assert parameter.grad is not None and torch.isfinite(parameter.grad).all(), name + assert parameter.grad.norm() > 0, name + gradients[name] = float(parameter.grad.norm()) + restored = _policy(policy.cameras) + restored.load_state_dict(deepcopy(policy.state_dict())) + with torch.no_grad(): + torch.testing.assert_close( + restored(proprio, images), prediction, rtol=0, atol=0 + ) + rejected = 0 + for invalid in ( + {"third_person": images["third_person"]}, + {**images, "extra": images["wrist"]}, + ): + try: + policy(proprio, invalid) + except ValueError: + rejected += 1 + try: + _policy(tuple(reversed(policy.cameras))).load_state_dict(policy.state_dict()) + except ValueError: + rejected += 1 + assert rejected == 3 + return { + "passed": True, + "flatten_order": "batch, history, configured camera", + "future_geometry_does_not_enter_policy": True, + "camera_and_checkpoint_rejections": rejected, + "parameter_gradient_norms": gradients, + } + + +def _fit_check(device: torch.device, steps: int) -> dict[str, Any]: + torch.manual_seed(918) + results = {} + for cameras in (("third_person",), ("third_person", "wrist")): + policy = _policy(cameras, horizon=2).to(device) + count = 2 ** len(cameras) + bits = torch.tensor( + [ + [float((sample >> view) & 1) for view in range(len(cameras))] + for sample in range(count) + ], + device=device, + ) + proprio = torch.zeros(count, 2, 27, device=device) + images = {} + for view, camera in enumerate(cameras): + image = torch.full((count, 2, 3, 256, 256), 0.25, device=device) + # Spatially localized RGB information is supplied solely by this view. + image[:, :, view, 32:224, 32:224] = ( + 0.1 + bits[:, view, None, None, None] * 0.8 + ) + images[camera] = image + targets = torch.zeros(count, 2, 10, device=device) + for view in range(len(cameras)): + targets[:, :, view] = (bits[:, view, None] * 2 - 1) * 0.5 + optimizer = torch.optim.AdamW(policy.parameters(), lr=3e-4, weight_decay=1e-6) + before = {name: p.detach().clone() for name, p in policy.named_parameters()} + with torch.no_grad(): + initial = float((policy(proprio, images) - targets).square().mean()) + losses = [] + for _ in range(steps): + optimizer.zero_grad(set_to_none=True) + loss = (policy(proprio, images) - targets).square().mean() + assert torch.isfinite(loss) + loss.backward() + nn.utils.clip_grad_norm_(policy.parameters(), 1) + optimizer.step() + losses.append(float(loss.detach())) + with torch.no_grad(): + fitted = policy(proprio, images) + final = float((fitted - targets).square().mean()) + max_error = float((fitted - targets).abs().max()) + assert final < initial * 0.05 and max_error < 0.15, (initial, final, max_error) + changed_stages = {} + for stage in ( + "encoder.resnet_backbone.conv1", + "encoder.resnet_backbone.layer1", + "encoder.resnet_backbone.layer2", + "encoder.resnet_backbone.layer3", + "encoder.resnet_backbone.layer4", + "encoder.pooling_layer", + "encoder.mlp", + "head", + ): + changed_stages[stage] = any( + not torch.equal(parameter.detach(), before[name]) + for name, parameter in policy.named_parameters() + if name.startswith(stage) + ) + assert all(changed_stages.values()) + ablations = {} + for camera in cameras: + ablated = dict(images) + ablated[camera] = images[camera][:1].expand_as(images[camera]) + with torch.no_grad(): + ablated_loss = float( + (policy(proprio, ablated) - targets).square().mean() + ) + assert ablated_loss > max(final * 5, 0.01) + ablations[camera] = ablated_loss + results["single" if len(cameras) == 1 else "dual"] = { + "initial_mse": initial, + "final_mse": final, + "max_abs_error": max_error, + "steps": steps, + "samples": count, + "proprioception_constant": True, + "view_ablation_mse": ablations, + "updated_stages": changed_stages, + "losses": losses, + } + return {"passed": True, "cases": results} + + +def _cuda_check(device: torch.device) -> dict[str, Any]: + torch.manual_seed(44) + cpu = ResNet10Encoder().eval() + image = torch.rand(1, 3, 256, 256) + with torch.no_grad(): + expected = cpu(image) + gpu = deepcopy(cpu).to(device) + with torch.no_grad(): + actual = gpu(image.to(device)).cpu() + torch.testing.assert_close(actual, expected, atol=1e-4, rtol=2e-4) + del gpu + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats(device) + model = _policy(("third_person", "wrist"), horizon=32).to(device) + optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-6) + images = { + camera: torch.rand(32, 2, 3, 256, 256, device=device) + for camera in model.cameras + } + optimizer.zero_grad(set_to_none=True) + prediction = model(torch.zeros(32, 2, 27, device=device), images) + prediction.square().mean().backward() + optimizer.step() + torch.cuda.synchronize(device) + return { + "passed": True, + "device": str(device), + "name": torch.cuda.get_device_name(device), + "cpu_cuda_max_abs_error": float((actual - expected).abs().max()), + "b32_dual_h2_f32_peak_allocated_bytes": torch.cuda.max_memory_allocated(device), + "b32_dual_h2_f32_peak_reserved_bytes": torch.cuda.max_memory_reserved(device), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--reference-source", type=Path) + parser.add_argument("--device", default="cpu") + parser.add_argument("--fit-steps", type=int, default=160) + parser.add_argument("--skip-fit", action="store_true") + parser.add_argument("--cpu-threads", type=int, default=4) + args = parser.parse_args() + if args.fit_steps < 1 or args.cpu_threads < 1: + parser.error("fit steps and CPU threads must be positive") + torch.set_num_threads(args.cpu_threads) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + started = time.monotonic() + result: dict[str, Any] = { + "passed": False, + "torch": torch.__version__, + "device": args.device, + "reference": _reference_check(args.reference_source), + "policy_contract": _policy_contract_check(), + } + device = torch.device(args.device) + if not args.skip_fit: + result["image_only_fit"] = _fit_check(device, args.fit_steps) + if device.type == "cuda": + result["cuda"] = _cuda_check(device) + result["passed"] = True + result["wall_seconds"] = time.monotonic() - started + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(result, indent=2) + "\n") + print( + json.dumps( + { + "passed": result["passed"], + "output": str(args.output), + "seconds": result["wall_seconds"], + } + ) + ) + + +if __name__ == "__main__": + main() diff --git a/examples/ods/data_stream_demo/report.py b/examples/ods/data_stream_demo/report.py new file mode 100644 index 000000000..71cdb0ebf --- /dev/null +++ b/examples/ods/data_stream_demo/report.py @@ -0,0 +1,787 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Render measured data-stream results without rerunning acquisition or decoding. + +Run: python -m examples.ods.data_stream_demo.report OUTPUT_ROOT --animation +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from statistics import median +from typing import Any + +__all__ = ["write_report"] + +_ARMS = { + "lerobot": "Decoded LeRobot / K", + "streaming": "ODS streaming / unlimited", + "protected": "ODS protected / K", + "lerobot_unbounded": "Decoded LeRobot / unlimited", + "frozen": "ODS frozen / C / K", + "lerobot_same_cache": "Decoded LeRobot / C / K", +} +_MAIN = ("lerobot", "streaming", "protected") +_VISIONS = ("state", "single", "dual") +_COLORS = ("#4676ac", "#dc8b36", "#23977c", "#9657a4", "#c36b72", "#667d3c") +_METRICS = ( + "first_batch_from_launch_seconds", + "first_batch_from_reader_seconds", + "delivered_chunks_per_second", + "request_latency_p50_seconds", + "request_latency_p95_seconds", + "service_p50_seconds", + "service_p95_seconds", + "accepted_episodes_per_second", + "data_wait_seconds", + "throttle_seconds", + "cache_bytes", + "peak_pss_bytes", + "peak_vram_bytes", + "final_disk_bytes", + "peak_temporary_disk_bytes", + "logical_payload_bytes_per_second", + "quality_psnr_db", + "quality_mae", + "cap_violations", + "delivered_chunks", + "window_delivered_chunks", + "post_window_chunks", + "unconsumed_prefetch_chunks", +) + + +def _number(value: Any) -> float | None: + return ( + float(value) + if isinstance(value, (int, float)) + and not isinstance(value, bool) + and math.isfinite(value) + else None + ) + + +def _json(path: Path) -> dict[str, Any]: + value = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, dict): + raise ValueError(f"Expected an object in {path}") + return value + + +def _events(path: Path) -> list[dict[str, Any]]: + if not path.exists(): + return [] + rows = [] + for line_number, line in enumerate( + path.read_text(encoding="utf-8").splitlines(), 1 + ): + if not line.strip(): + continue + try: + row = json.loads(line) + except ValueError as error: + raise ValueError(f"Invalid JSON at {path}:{line_number}") from error + if not isinstance(row, dict): + raise ValueError(f"Expected an object at {path}:{line_number}") + rows.append(row) + return rows + + +def _load(root: Path) -> list[dict[str, Any]]: + paths = sorted(root.rglob("summary.json")) if root.is_dir() else [root] + runs = [] + for path in paths: + summary = _json(path) + config = summary.get("config", {}) + if config.get("arm") not in _ARMS or config.get("vision") not in _VISIONS: + raise ValueError(f"Unknown experiment arm/vision in {path}") + if summary.get("status") not in {"complete", "failed", "unsupported"}: + raise ValueError(f"Unknown run status in {path}") + runs.append( + { + "path": str(path.resolve()), + "summary": summary, + "config": config, + "batches": _events(path.parent / "batches.jsonl"), + "lifecycle": _events(path.parent / "lifecycle.jsonl"), + } + ) + if not runs: + raise ValueError(f"No run summaries found in {root}") + return sorted( + runs, + key=lambda r: ( + r["config"]["arm"], + r["config"]["vision"], + str(r["config"].get("seed")), + ), + ) + + +def _qualified(run: dict[str, Any]) -> bool: + return ( + run["summary"]["status"] == "complete" + and run["config"].get("suite") in {"report", "appendix"} + and run["summary"].get("metrics", {}).get("qualification_passed") is True + ) + + +def _metric(run: dict[str, Any], key: str) -> float | None: + return _number(run["summary"].get("metrics", {}).get(key)) + + +def _group(runs: list[dict[str, Any]], arm: str, vision: str) -> list[dict[str, Any]]: + return [ + r for r in runs if r["config"]["arm"] == arm and r["config"]["vision"] == vision + ] + + +def _stats(runs: list[dict[str, Any]], metric: str) -> dict[str, Any]: + values = [_metric(r, metric) for r in runs if _qualified(r)] + values = [v for v in values if v is not None] + return { + "n": len(values), + "median": median(values) if values else None, + "min": min(values) if values else None, + "max": max(values) if values else None, + } + + +def _formatted(stats: dict[str, Any], scale: float = 1.0) -> str: + if not stats["n"]: + return "unknown (n=0)" + values = [stats[k] / scale for k in ("median", "min", "max")] + return f"{values[0]:.3g} [{values[1]:.3g}, {values[2]:.3g}]; n={stats['n']}" + + +def _cell(value: Any) -> str: + return str(value).replace("|", "\\|").replace("\n", " ") + + +def _table(headers: list[str], rows: list[list[Any]]) -> list[str]: + return [ + "| " + " | ".join(headers) + " |", + "| " + " | ".join("---" for _ in headers) + " |", + *("| " + " | ".join(_cell(v) for v in row) + " |" for row in rows), + ] + + +def _save(figure: Any, output: Path, name: str) -> None: + for extension in ("png", "svg", "pdf"): + figure.savefig(output / f"{name}.{extension}", dpi=170, bbox_inches="tight") + + +def _plot_points( + axis: Any, + runs: list[dict[str, Any]], + metric: str, + x: float, + color: str, + *, + scale: float = 1.0, +) -> None: + for i, run in enumerate(runs): + value = _metric(run, metric) + if value is not None: + axis.scatter( + x + (i - (len(runs) - 1) / 2) * 0.045, + value / scale, + s=23, + facecolors=color if _qualified(run) else "none", + edgecolors=color, + linewidths=1, + zorder=3, + ) + stats = _stats(runs, metric) + if stats["n"]: + center = stats["median"] / scale + axis.errorbar( + x, + center, + yerr=[[center - stats["min"] / scale], [stats["max"] / scale - center]], + fmt="_", + color=color, + markersize=14, + capsize=3, + zorder=4, + ) + + +def _timeline(runs: list[dict[str, Any]], plt: Any, output: Path) -> Any: + figure, axes = plt.subplots(3, 1, figsize=(13, 11), squeeze=False) + for vision, axis in zip(_VISIONS, axes[:, 0]): + selected = [ + r + for r in runs + if r["config"]["vision"] == vision and r["config"]["arm"] in _MAIN + ] + for row, run in enumerate(selected): + cache = run["config"].get("cache_episodes", 1) + cache = cache if isinstance(cache, int) and cache > 0 else 1 + end_times = [ + _number(e.get("time_seconds")) + for e in run["lifecycle"] + run["batches"] + ] + end = max((t for t in end_times if t is not None), default=0.0) + first = _metric(run, "first_batch_from_launch_seconds") + ready = _metric(run, "source_ready_seconds") + offline = run["config"]["arm"].startswith("lerobot") + cache_end = min(end, ready) if offline and ready is not None else end + if offline and ready is not None: + duration = _metric(run, "actual_measurement_seconds") + files_end = max(end, ready + duration) if duration is not None else end + axis.hlines( + row, ready, files_end, colors="#596874", linewidth=3, alpha=0.65 + ) + if ready is not None: + axis.plot(ready, row - 0.46, marker="|", color="#777777", markersize=9) + if first is not None: + axis.plot(first, row - 0.46, marker="*", color="#333333", markersize=6) + active = {} + for event in sorted( + run["lifecycle"], key=lambda e: _number(e.get("time_seconds")) or 0 + ): + time_value, slot = _number(event.get("time_seconds")), event.get("slot") + if time_value is None or not isinstance(slot, int): + continue + if slot in active: + start, generation = active.pop(slot) + axis.broken_barh( + [(start, max(0, min(time_value, cache_end) - start))], + (row - 0.37 + 0.6 * slot / cache, 0.6 / cache), + facecolors=plt.get_cmap("tab10")(int(generation) % 10), + alpha=0.4, + ) + if event.get("kind") == "published": + active[slot] = (time_value, event.get("generation", 0)) + for slot, (start, generation) in active.items(): + axis.broken_barh( + [(start, max(0, cache_end - start))], + (row - 0.37 + 0.6 * slot / cache, 0.6 / cache), + facecolors=plt.get_cmap("tab10")(int(generation) % 10), + alpha=0.4, + ) + for batch in run["batches"]: + time_value = _number(batch.get("time_seconds")) + if time_value is None: + continue + for generation, count in batch.get("generations", {}).items(): + if _number(count) is not None and count > 0: + axis.scatter( + time_value, + row + 0.3, + s=min(24, 2 + count), + color=plt.get_cmap("tab10")(int(generation) % 10), + marker="x" if batch.get("within_window") is False else "o", + linewidths=( + 0.8 if batch.get("within_window") is False else 0 + ), + ) + if not run["lifecycle"]: + axis.text( + 0.01, + row, + "no lifecycle records", + transform=axis.get_yaxis_transform(), + fontsize=7, + va="center", + ) + axis.set_yticks( + range(len(selected)), + [ + f"{_ARMS[r['config']['arm']]} / {r['config'].get('seed', '?')} / {r['summary']['status']}" + for r in selected + ], + fontsize=7, + ) + axis.set_title( + f"{vision}: cache residency (bands) and delivered generations (dots; x = late)" + ) + axis.set_xlabel( + "Seconds from process launch; | = reader start, * = first batch" + ) + axis.grid(axis="x", alpha=0.2) + from matplotlib.patches import Patch + from matplotlib.lines import Line2D + + generations = sorted( + { + int(e["generation"]) + for r in runs + for e in r["lifecycle"] + if _number(e.get("generation")) is not None + } + | {int(g) for r in runs for b in r["batches"] for g in b.get("generations", {})} + ) + if generations: + figure.legend( + handles=[ + Patch(color=plt.get_cmap("tab10")(g % 10), label=f"Generation {g}") + for g in generations + ] + + [ + Line2D( + [0], + [0], + color="#596874", + linewidth=3, + label="Finalized corpus on disk", + ) + ], + loc="upper center", + bbox_to_anchor=(0.5, 1.02), + ncol=min(6, len(generations) + 1), + ) + figure.tight_layout() + _save(figure, output, "01_timeline") + return figure + + +def _reuse(runs: list[dict[str, Any]], plt: Any, output: Path) -> None: + figure, axes = plt.subplots(1, 3, figsize=(14, 4), squeeze=False) + for vision, axis in zip(_VISIONS, axes[0]): + caps = set() + for a, arm in enumerate(_MAIN): + selected = _group(runs, arm, vision) + for s, run in enumerate(selected): + usage = run["summary"].get("episode_usage", []) + for i, episode in enumerate(usage): + value = _number(episode.get("delivered_chunks")) + if value is None: + continue + evicted = episode.get("evicted") is True + axis.scatter( + a + (s - 1) * 0.16 + ((i % 7) - 3) * 0.012, + value, + marker="D" if evicted else "o", + s=12, + facecolors=( + "none" if evicted or not _qualified(run) else _COLORS[a] + ), + edgecolors=_COLORS[a], + alpha=0.7, + linewidths=0.7, + ) + cap = _number(run["config"].get("quota")) + if cap is not None: + caps.add(cap) + for cap in caps: + axis.axhline( + cap, color="#555555", linestyle="--", linewidth=1, label=f"K={cap:g}" + ) + axis.set_xticks(range(3), ["LeRobot / K", "ODS / no K", "ODS / K"], rotation=15) + axis.set_title(vision) + axis.set_ylabel("Delivered chunks per immutable source") + axis.set_yscale("symlog", linthresh=10) + axis.set_ylim(bottom=0) + axis.grid(axis="y", alpha=0.2) + if caps: + axis.legend(fontsize=8) + figure.suptitle("All identities; hollow diamonds mark collector-cache eviction") + figure.tight_layout() + _save(figure, output, "02_reuse") + plt.close(figure) + + +def _latency(runs: list[dict[str, Any]], plt: Any, output: Path) -> None: + figure, axes = plt.subplots(3, 3, figsize=(13, 10), squeeze=False) + for v, vision in enumerate(_VISIONS): + for a, arm in enumerate(_MAIN): + selected = _group(runs, arm, vision) + _plot_points( + axes[v, 0], + selected, + "first_batch_from_launch_seconds", + a - 0.12, + _COLORS[a], + ) + _plot_points( + axes[v, 0], + selected, + "first_batch_from_reader_seconds", + a + 0.12, + "#777777", + ) + _plot_points( + axes[v, 1], + selected, + "request_latency_p95_seconds", + a, + _COLORS[a], + scale=0.001, + ) + for s, run in enumerate(selected): + demand = _number(run["config"].get("requested_chunks_per_second")) + delivered = _metric(run, "delivered_chunks_per_second") + if demand is not None and demand > 0 and delivered is not None: + axes[v, 2].scatter( + a + (s - 1) * 0.045, + delivered / demand, + facecolors=_COLORS[a] if _qualified(run) else "none", + edgecolors=_COLORS[a], + s=23, + ) + ratios = [ + _metric(r, "delivered_chunks_per_second") + / r["config"]["requested_chunks_per_second"] + for r in selected + if _qualified(r) + and _metric(r, "delivered_chunks_per_second") is not None + and (_number(r["config"].get("requested_chunks_per_second")) or 0) > 0 + ] + if ratios: + middle = median(ratios) + axes[v, 2].errorbar( + a, + middle, + yerr=[[middle - min(ratios)], [max(ratios) - middle]], + fmt="_", + markersize=14, + capsize=3, + color=_COLORS[a], + ) + for axis, title, label in zip( + axes[v], + ( + "First batch: launch (color) / reader (gray)", + "Consumer-visible p95 batch latency", + "Delivered rate / nominal demand", + ), + ("Seconds", "Milliseconds", "Ratio"), + ): + axis.set_title(f"{vision}: {title}", fontsize=9) + axis.set_ylabel(label) + axis.set_xticks( + range(3), ["LeRobot / K", "ODS / no K", "ODS / K"], rotation=15 + ) + axis.grid(axis="y", alpha=0.2) + axes[v, 2].axhline(1, color="#555555", linestyle="--", linewidth=1) + figure.suptitle("Seed points; median/range use complete, qualified runs only") + figure.tight_layout() + _save(figure, output, "03_latency_demand") + plt.close(figure) + + +def _visual_load(runs: list[dict[str, Any]], plt: Any, output: Path) -> None: + fields = ( + ("delivered_chunks_per_second", "Delivered chunks/s", 1), + ("accepted_episodes_per_second", "Accepted episodes/s", 1), + ("logical_payload_bytes_per_second", "Logical returned MiB/s", 1024**2), + ("peak_pss_bytes", "Sampled process-tree peak PSS (GiB)", 1024**3), + ("peak_vram_bytes", "Sampled VRAM (GiB; declared scope)", 1024**3), + ("final_disk_bytes", "Final dataset storage (GiB)", 1024**3), + ) + figure, axes = plt.subplots(2, 3, figsize=(13, 8), squeeze=False) + for axis, (field, title, scale) in zip(axes.flat, fields): + for a, arm in enumerate(_MAIN): + medians, positions = [], [] + for v, vision in enumerate(_VISIONS): + selected = _group(runs, arm, vision) + x = v + (a - 1) * 0.17 + _plot_points(axis, selected, field, x, _COLORS[a], scale=scale) + stats = _stats(selected, field) + if stats["n"]: + positions.append(x) + medians.append(stats["median"] / scale) + axis.plot( + positions, + medians, + color=_COLORS[a], + linewidth=1, + alpha=0.65, + label=_ARMS[arm], + ) + axis.set_xticks(range(3), _VISIONS) + axis.set_title(title, fontsize=10) + axis.grid(axis="y", alpha=0.2) + axes[0, 0].legend(fontsize=7) + figure.suptitle( + "Visual workload: filled=qualified; hollow=incomplete/unqualified; absent=unknown" + ) + figure.tight_layout() + _save(figure, output, "04_visual_load") + plt.close(figure) + + +def _animate(runs: list[dict[str, Any]], plt: Any, output: Path) -> None: + from matplotlib.animation import FuncAnimation, PillowWriter + + selected = [r for r in runs if r["config"]["arm"] == "protected"] + if not selected: + return + seed = min(str(r["config"].get("seed")) for r in selected) + selected = [r for r in selected if str(r["config"].get("seed")) == seed] + figure, axes = plt.subplots( + len(selected), 1, figsize=(9, 3 * len(selected)), squeeze=False + ) + cursors = [] + end = 0.0 + for run, axis in zip(selected, axes[:, 0]): + points = sorted( + ( + (_number(b.get("time_seconds")), _number(b.get("chunks"))) + for b in run["batches"] + if _number(b.get("time_seconds")) is not None + and _number(b.get("chunks")) is not None + ) + ) + total = 0.0 + times, counts = [0.0], [0.0] + for time_value, chunks in points: + total += chunks + times.append(time_value) + counts.append(total) + end = max(end, max(times)) + axis.step(times, counts, where="post", color="#23977c") + axis.set( + title=f"Protected / {run['config']['vision']} / seed {seed}", + xlabel="Seconds from launch", + ylabel="Lifetime delivered chunks", + ) + cursors.append(axis.axvline(0, color="#dc8b36")) + figure.tight_layout() + + def update(frame: int) -> list[Any]: + for cursor in cursors: + cursor.set_xdata([end * frame / 119] * 2) + return cursors + + animation = FuncAnimation(figure, update, frames=120, blit=False) + animation.save(output / "data_flow.gif", writer=PillowWriter(fps=6)) + plt.close(figure) + + +def write_report( + input_root: Path, output_dir: Path | None = None, *, animation: bool = False +) -> Path: + """Write one report and four figure sets from retained experiment logs. + + Unknown measurements remain unknown. Aggregates exclude failed, unsupported + and unqualified runs, whose summaries and any measurements remain visible. + No performance winner or learning conclusion is inferred. + + Args: + input_root: Root containing run directories, or one summary JSON file. + output_dir: Artifact directory, defaulting to ``input_root/report``. + animation: Also render a posthoc GIF from recorded protected-run batches. + + Returns: + Path to the generated Markdown report. + """ + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + runs = _load(Path(input_root)) + output = ( + Path(output_dir) + if output_dir is not None + else ( + Path(input_root).parent / "report" + if Path(input_root).is_file() + else Path(input_root) / "report" + ) + ) + output.mkdir(parents=True, exist_ok=True) + groups = [] + for arm in _ARMS: + for vision in _VISIONS: + selected = _group(runs, arm, vision) + if selected: + groups.append( + { + "arm": arm, + "vision": vision, + "attempted": len(selected), + "qualified": sum(_qualified(r) for r in selected), + "metrics": {key: _stats(selected, key) for key in _METRICS}, + } + ) + serialized = { + "aggregation": "suite=report|appendix, complete and qualification_passed=true only; median [min,max]; unknown is null", + "groups": groups, + "runs": [{"summary_path": r["path"], **r["summary"]} for r in runs], + } + (output / "report.json").write_text( + json.dumps(serialized, indent=2, allow_nan=False) + "\n", encoding="utf-8" + ) + figure = _timeline(runs, plt, output) + plt.close(figure) + _reuse(runs, plt, output) + _latency(runs, plt, output) + _visual_load(runs, plt, output) + if animation: + _animate(runs, plt, output) + lines = [ + "# ODS data-stream capability report", + "", + "This report measures data delivery, not training or policy quality. It does not assume ODS is faster than decoded LeRobot.", + "", + "All attempted runs are retained below. Filled plot markers and aggregate median [min,max] include only complete, explicitly qualified report/appendix runs. Hollow markers preserve measurements from incomplete, unqualified, smoke or calibration runs; missing measurements are unknown, never zero. Appendix results are one-seed diagnostics when only one seed is present.", + "", + "## Time and memory", + "", + ] + lines += _table( + [ + "Group / vision", + "Attempted / qualified", + "First batch from launch (s)", + "First batch from reader (s)", + "Batch p95 (s)", + "Chunks/s", + "Peak PSS (GiB)", + "Final disk (GiB)", + ], + [ + [ + f"{'Main' if g['arm'] in _MAIN else 'Appendix'}: {_ARMS[g['arm']]} / {g['vision']}", + f"{g['attempted']} / {g['qualified']}", + *( + _formatted(g["metrics"][key], scale) + for key, scale in ( + ("first_batch_from_launch_seconds", 1), + ("first_batch_from_reader_seconds", 1), + ("request_latency_p95_seconds", 1), + ("delivered_chunks_per_second", 1), + ("peak_pss_bytes", 1024**3), + ("final_disk_bytes", 1024**3), + ) + ), + ] + for g in groups + ], + ) + lines += [ + "", + "Launch latency includes the collect-N LeRobot barrier versus the initial-C ODS barrier. Reader latency begins with source-ready access; LeRobot still performs real video decoding. These scopes overlap and must not be added. Final-batch latency includes waiting and output conversion; worker item-service time is a separate component measurement, not decoder kernel time.", + "", + "## Integrity and exposure", + "", + ] + lines += _table( + [ + "Group / vision", + "Delivered chunks", + "Within window", + "After window", + "K violations", + "Unconsumed prefetch", + "PSNR (dB)", + "MAE", + ], + [ + [ + f"{_ARMS[g['arm']]} / {g['vision']}", + *( + _formatted(g["metrics"][key]) + for key in ( + "delivered_chunks", + "window_delivered_chunks", + "post_window_chunks", + "cap_violations", + "unconsumed_prefetch_chunks", + "quality_psnr_db", + "quality_mae", + ) + ), + ] + for g in groups + ], + ) + lines += [ + "", + "One credit is one anchor-source chunk, independent of camera count or context length. Delivered chunks and reuse are lifetime counts, including any complete batch returned after the observation deadline; late batches are excluded from window throughput and latency metrics. Source counts include evicted identities. A cap bounds delivery rather than optimizer updates. AV1 video is lossy; speed must be considered alongside the reported raw-reference quality, codec and recording configuration. Infinite PSNR for exact equality requires an explicit equality count in the run metadata; it is not replaced with a finite score.", + "", + "## Run qualification", + "", + ] + lines += _table( + ["Run", "Suite", "Seed", "Status", "Qualified", "Reason / error"], + [ + [ + f"{_ARMS[r['config']['arm']]} / {r['config']['vision']}", + r["config"].get("suite", "unknown"), + r["config"].get("seed", "unknown"), + r["summary"]["status"], + r["summary"].get("metrics", {}).get("qualification_passed", "unknown"), + r["summary"].get("error") + or r["summary"].get("qualification") + or "See retained summary", + ] + for r in runs + ], + ) + lines += ["", "## Figures", ""] + for name, caption in ( + ( + "01_timeline", + "Availability, cache residency and delivered-generation timelines. Offline collector bands stop at reader start (after engine shutdown); gray lines show finalized files on disk independently of remaining K. Unrecorded write locks are not inferred.", + ), + ( + "02_reuse", + "Per-source reuse; hollow diamonds mean collector-cache eviction. For LeRobot, all finalized corpus files remain available despite collector eviction. The K line is a protection setting, not an optimal learning budget.", + ), + ( + "03_latency_demand", + "Cold versus source-ready first batch, consumer p95 and served demand. Quota waits and service bottlenecks must be interpreted using the retained counters.", + ), + ( + "04_visual_load", + "State/single/dual workload costs. Logical MiB/s is returned tensor payload, not measured bus or disk bandwidth. Memory coverage follows each run's instrumentation declaration.", + ), + ): + lines += [ + f"![{caption}]({name}.png)", + "", + caption + f" [SVG]({name}.svg) · [PDF]({name}.pdf)", + "", + ] + if (output / "data_flow.gif").exists(): + lines += [ + "![Posthoc recorded delivery animation](data_flow.gif)", + "", + "Animation replays logged deliveries; it is not a recording of simulation or a separate measurement.", + "", + ] + lines += [ + "## Interpretation limits", + "", + "A new source identity does not establish broader state coverage. Fixed cache tensor bytes do not establish constant total process memory. Protected waiting can intentionally reduce delivered throughput; it must not be presented as a copy/decoder bottleneck. This finalized local LeRobot path does not establish that the LeRobot ecosystem cannot stream. No universal speedup, learning-efficiency or policy-performance claim follows from these measurements.", + "", + "Raw summaries and input paths, including configuration, qualification failures and cleanup records, are preserved in [report.json](report.json). Re-run this renderer after interrupted experiments have written their final summaries; partial JSONL lines fail visibly rather than being discarded.", + ] + report_path = output / "report.md" + report_path.write_text("\n".join(lines) + "\n", encoding="utf-8") + return report_path + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("input_root", type=Path) + parser.add_argument("--output", type=Path) + parser.add_argument("--animation", action="store_true") + arguments = parser.parse_args() + print( + write_report( + arguments.input_root, arguments.output, animation=arguments.animation + ) + ) diff --git a/examples/ods/data_stream_demo/run.py b/examples/ods/data_stream_demo/run.py new file mode 100644 index 000000000..c8dbf320e --- /dev/null +++ b/examples/ods/data_stream_demo/run.py @@ -0,0 +1,924 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Run real ODS and standard-decoded LeRobot delivery without a learner. + +Run ``python -m examples.ods.data_stream_demo.run --help`` for the isolated +single-run interface and supervised smoke/calibration/report suites. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import os +import subprocess +import sys +import time +import traceback +from importlib.metadata import PackageNotFoundError, version +from pathlib import Path +from typing import Any + +import torch +from tensordict import TensorDict + +from embodichain.data_pipeline.datasets import ( + DataBudgetExhaustedError, + OnlineDataSource, + SequenceSpecCfg, +) +from embodichain.data_pipeline.runtime import ODSFlowController, ODSFlowControllerCfg + +from . import acquisition +from .lerobot_io import ( + StandardLeRobotReader, + StandardLeRobotWriter, + audit_video_quality, +) +from .metrics import ResourceMonitor, StreamMetrics, validate_batch, write_json +from .payload import ( + describe_raw_buffer, + normalize_online_batch, + numeric_episode_digest, + pack_episode, +) + +__all__ = ["build_parser", "run_one", "main"] + +ARMS = ( + "lerobot", + "streaming", + "protected", + "frozen", + "lerobot_unbounded", + "lerobot_same_cache", +) + + +class _ObservationComplete(Exception): + pass + + +class _SupplyExhausted(Exception): + pass + + +class _MeasuredOnlineSource(OnlineDataSource): + """Time successful immediate sampling separately from controller retries.""" + + last_service_seconds: float = 0.0 + + def sample_context( + self, + batch_size: int, + spec: SequenceSpecCfg, + *, + generator: torch.Generator | None = None, + ) -> TensorDict: + started = time.monotonic() + result = super().sample_context(batch_size, spec, generator=generator) + self.last_service_seconds = time.monotonic() - started + return result + + +def build_parser() -> argparse.ArgumentParser: + """Build explicit acquisition, decoding and finite-run settings.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--suite", + choices=("single", "smoke", "calibrate", "report", "appendix"), + default="single", + ) + parser.add_argument("--arm", choices=ARMS, default="protected") + parser.add_argument( + "--vision", choices=("state", "single", "dual"), default="state" + ) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--dataset-root", type=Path) + parser.add_argument("--protocol", type=Path) + parser.add_argument("--gpu-ids", nargs="+", type=int) + parser.add_argument( + "--aux-graphics-gpu-id", + type=int, + help="Suite-only reservation for a separately qualified tiny Vulkan graphics context", + ) + parser.add_argument("--sim-device", default="cuda:0") + parser.add_argument("--seed", type=int, default=19001) + parser.add_argument("--buffer-size", type=int, default=8) + parser.add_argument("--producer-envs", type=int, default=2) + parser.add_argument("--accepted-episodes", type=int, default=32) + parser.add_argument("--max-episode-steps", type=int, default=600) + parser.add_argument("--batch-size", type=int, default=32) + parser.add_argument("--history-length", type=int, default=2) + parser.add_argument("--prediction-horizon", type=int, default=32) + parser.add_argument("--episode-budget", type=int, default=128) + parser.add_argument("--request-rate", type=float, default=64.0) + parser.add_argument("--duration-seconds", type=float, default=180.0) + parser.add_argument("--num-workers", type=int, default=4) + parser.add_argument("--prefetch-factor", type=int, default=1) + parser.add_argument("--data-timeout", type=float, default=120.0) + parser.add_argument("--generation-timeout", type=float, default=1800.0) + parser.add_argument("--initialization-timeout", type=float, default=1800.0) + parser.add_argument("--max-batches", type=int) + parser.add_argument("--full-image-check", action="store_true") + parser.add_argument( + "--calibration-run", action="store_true", help=argparse.SUPPRESS + ) + parser.add_argument( + "--suite-tag", + choices=("single", "smoke", "calibration", "report", "appendix"), + default="single", + help=argparse.SUPPRESS, + ) + return parser + + +def _environment() -> dict[str, Any]: + root = Path(__file__).resolve().parents[3] + sources = [ + *Path(__file__).parent.glob("*.py"), + root / "examples/ods/simple_bc_demo/environment.py", + root / "embodichain/data_pipeline/engine/data.py", + root / "embodichain/data_pipeline/runtime.py", + ] + packages = {} + for package in ("torch", "lerobot", "torchcodec", "av", "dexsim"): + try: + packages[package] = version(package) + except PackageNotFoundError: + packages[package] = None + return { + "python": sys.version, + "packages": packages, + "commit": subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=root, text=True + ).strip(), + "source_sha256": { + str(path.relative_to(root)): hashlib.sha256(path.read_bytes()).hexdigest() + for path in sources + }, + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "auxiliary_graphics_gpu_uuid": os.environ.get("ODS_AUX_GRAPHICS_GPU_UUID"), + "headless_environment": { + key: os.environ.get(key) + for key in ("SDL_VIDEODRIVER", "DISPLAY", "WAYLAND_DISPLAY") + }, + "cpu_affinity": sorted(os.sched_getaffinity(0)), + "torch_threads": torch.get_num_threads(), + "cache_policy": "fresh processes; warm software/filesystem caches; no global cache purge", + } + + +def _schema() -> dict[str, Any]: + return { + "observation": {"version": "data-stream-raw-v1", "timing": "pre_action"}, + "action": { + "semantics": "policy_request", + "names": ["dx", "dy", "dz", "drx", "dry", "drz", "gripper"], + "step_dt": 1 / 30, + "terms": [ + {"term": {"contract": "eef_pose.relative@1"}}, + {"term": {"contract": "gripper.parallel@1"}}, + ], + }, + } + + +def _disk_bytes(path: Path) -> int: + return sum(item.stat().st_size for item in path.rglob("*") if item.is_file()) + + +def _initial_manifest(engine: Any, metadata: dict[str, Any]) -> list[dict[str, Any]]: + """Hash initial numeric episodes under the copy lock without cloning images.""" + records = [] + with engine._lock_index.get_lock(): + if tuple(engine._lock_index[:]) != (-1, -1): + raise RuntimeError("Initial manifest requires an idle completed fill") + for slot in range(engine.buffer_size): + row = engine.shared_buffer[slot] + valid = row["valid"].bool() + state = torch.cat( + [ + row[("obs", *feature["path"])][valid].reshape(-1, feature["width"]) + for feature in metadata["state_features"] + ], + -1, + ) + numeric = { + "state": state, + "actions": row["actions"][valid].clone(), + "continuity_id": row["continuity_id"][valid].clone(), + } + records.append( + { + "source_id": [ + int(row["ods_run_id"][valid][0]), + int(row["ods_episode_id"][valid][0]), + ], + "generation": int(row["ods_generation"][valid][0]), + "steps": int(valid.sum()), + "numeric_sha256": numeric_episode_digest(numeric), + } + ) + return records + + +def _validate_args(args: argparse.Namespace) -> None: + if getattr(args, "aux_graphics_gpu_id", None) is not None: + raise ValueError("--aux-graphics-gpu-id requires a supervised suite") + for key in ( + "buffer_size", + "producer_envs", + "accepted_episodes", + "batch_size", + "episode_budget", + "history_length", + "prediction_horizon", + "max_episode_steps", + "num_workers", + "prefetch_factor", + ): + if getattr(args, key) < 1: + raise ValueError(f"{key} must be positive") + if ( + args.buffer_size % args.producer_envs + or args.accepted_episodes % args.buffer_size + ): + raise ValueError("C must be divisible by P and N by C") + if args.accepted_episodes < args.buffer_size: + raise ValueError("N cannot be smaller than C") + for key in ( + "request_rate", + "duration_seconds", + "data_timeout", + "generation_timeout", + "initialization_timeout", + ): + if not math.isfinite(getattr(args, key)) or getattr(args, key) <= 0: + raise ValueError(f"{key} must be finite and positive") + if args.max_batches is not None and args.max_batches < 1: + raise ValueError("max_batches must be positive") + if args.dataset_root is not None and not args.arm.startswith("lerobot"): + raise ValueError("An existing corpus requires a LeRobot arm") + if args.output.exists() and any(args.output.iterdir()): + raise FileExistsError("Use a fresh output directory") + + +def run_one(args: argparse.Namespace) -> dict[str, Any]: + """Run and audit one finite real-physics or standard-reader experiment. + + Args: + args: Parsed, explicit single-run settings. + + Returns: + Complete or failed summary, also persisted after every cleanup attempt. + """ + _validate_args(args) + args.output.mkdir(parents=True, exist_ok=True) + quota = ( + None if args.arm in ("streaming", "lerobot_unbounded") else args.episode_budget + ) + offline = args.arm.startswith("lerobot") + cfg = { + key: str(value) if isinstance(value, Path) else value + for key, value in vars(args).items() + } + cfg.update( + suite="calibration" if args.calibration_run else args.suite_tag, + cache_episodes=args.buffer_size, + quota=quota, + requested_chunks_per_second=args.request_rate, + measurement_seconds=args.duration_seconds, + history=args.history_length, + future=args.prediction_horizon, + camera_profile_version=getattr(acquisition, "CAMERA_PROFILE_VERSION", "v1"), + ) + write_json(args.output / "config.json", cfg) + write_json(args.output / "environment.json", _environment()) + tracker = StreamMetrics(args.output, history=args.history_length, quota=quota) + resources = ResourceMonitor(args.output) + resources.start() + engine = flow = reader = writer = None + stats: dict[str, Any] = {} + budget: dict[str, Any] = {} + reader_stats: dict[str, Any] = {} + cleanup: dict[str, Any] = {"errors": []} + metrics: dict[str, Any] = {} + manifest: list[dict[str, Any]] = [] + metadata: dict[str, Any] = {} + status, error, end_reason = "complete", None, "observation_window" + corpus = args.dataset_root + measurement_start = None + measured_duration = 0.0 + deadline = None + in_sampling = False + + def sampling_progress() -> None: + if not in_sampling: + return + if deadline is not None and time.monotonic() >= deadline: + raise _ObservationComplete + if engine is not None and quota is not None: + state = engine.get_flow_status() + pending = state["requested_generation"] > state["completed_generation"] + final = ( + args.arm == "frozen" + or state["successful_episodes"] >= args.accepted_episodes + ) + if ( + not pending + and final + and state["unlocked_remaining_chunks"] < args.batch_size + ): + raise _SupplyExhausted + + try: + if corpus is None: + engine = acquisition.make_engine( + args.vision, + args.sim_device, + producer_envs=args.producer_envs, + seed=args.seed, + buffer_size=args.buffer_size, + budget=None if offline else quota, + max_episode_steps=args.max_episode_steps, + initialization_timeout=args.initialization_timeout, + ) + write_json(args.output / "acquisition_config.json", engine.cfg.gym_config) + flow = ODSFlowController( + engine, + ODSFlowControllerCfg( + data_wait_timeout=args.data_timeout, + generation_timeout=args.generation_timeout, + max_generated_episodes=( + args.buffer_size + if args.arm == "frozen" + else args.accepted_episodes + ), + ), + on_events=tracker.events, + on_interval=tracker.interval, + progress=sampling_progress, + ) + engine.start() + flow.drain_events() + tracker.generation_complete(1) + metadata = describe_raw_buffer(engine.shared_buffer) + metadata.update(action_schema=_schema()["action"], step_dt=1 / 30) + # Explicit recording configuration supplements vector order; these + # are policy requests, never relabeled joint/controller commands. + metadata["action_parameters"] = engine.cfg.gym_config["env"]["actions"] + write_json(args.output / "payload_schema.json", metadata) + write_json( + args.output / "initial_numeric_manifest.json", + _initial_manifest(engine, metadata), + ) + if offline: + corpus = args.output / "lerobot_dataset" + writer = StandardLeRobotWriter( + corpus, metadata, episode_count=args.accepted_episodes + ) + for generation in range( + 1, args.accepted_episodes // args.buffer_size + 1 + ): + flow.wait_for_generation(generation) + tracker.generation_complete(generation) + for rawrow in acquisition.iter_generation_episodes( + engine, generation + ): + episode = pack_episode(rawrow, metadata) + manifest.append( + { + "source_id": episode["frame_ids"][0, :2].tolist(), + "steps": len(episode["valid"]), + "generation": generation, + "numeric_sha256": numeric_episode_digest(episode), + } + ) + writer.submit(episode, progress=flow.poll) + del episode, rawrow + if generation < args.accepted_episodes // args.buffer_size: + requested = flow.request_refill() + if requested != generation + 1: + raise RuntimeError("Offline collection generation mismatch") + metrics["writer"] = writer.finalize(progress=flow.poll) + writer = None + stats = engine.get_stats() + engine.stop() + flow.drain_events() + engine = flow = None + write_json(args.output / "numeric_manifest.json", manifest) + tracker.reader_started = tracker.elapsed() + measurement_start = tracker.origin + tracker.reader_started + tracker.measurement_started = tracker.reader_started + minimum_finish = measurement_start + args.duration_seconds + deadline = measurement_start + ( + max(args.duration_seconds, 900.0) + if args.calibration_run and not offline + else args.duration_seconds + ) + tracker.measurement_deadline = deadline + if offline: + reader = StandardLeRobotReader( + corpus, + batch_size=args.batch_size, + history_length=args.history_length, + prediction_horizon=args.prediction_horizon, + episode_budget=quota, + seed=args.seed + 10000, + num_workers=args.num_workers, + prefetch_factor=args.prefetch_factor, + timeout=args.data_timeout, + episodes=( + list(range(args.buffer_size)) + if args.arm == "lerobot_same_cache" + else None + ), + ) + metadata = reader.metadata + if args.suite_tag in ("report", "appendix"): + expected = ( + args.buffer_size + if args.arm == "lerobot_same_cache" + else args.accepted_episodes + ) + if len(reader.snapshot_budget()["episodes"]) != expected: + raise ValueError( + "Formal reader corpus count differs from the accepted-data budget" + ) + if ( + metadata.get("action_schema") != _schema()["action"] + or metadata.get("step_dt") != 1 / 30 + ): + raise ValueError( + "Formal corpus action semantics/control period differ" + ) + # Keep never-selected offline identities visible in the audit too. + for row in reader.snapshot_budget()["episodes"]: + source = ":".join(map(str, row["source_id"])) + tracker.episodes.setdefault( + source, + { + "source_id": source, + "delivered_chunks": 0, + "first_consumed_seconds": None, + "evicted": False, + }, + ) + else: + source = _MeasuredOnlineSource(engine, schema=_schema()) + sequence = SequenceSpecCfg( + history_length=args.history_length, + prediction_horizon=args.prediction_horizon, + ) + rng = torch.Generator().manual_seed(args.seed + 10000) + cameras = acquisition.camera_uids(args.vision) + if set(metadata["cameras"]) != set(cameras): + raise ValueError( + "Corpus camera profile does not match the requested visual condition" + ) + next_request = measurement_start + exhausted = False + batches = 0 + retained = None + retained_digest = None + while time.monotonic() < deadline: + if flow is not None: + flow.poll() + state = engine.get_flow_status() + completed = int(state["completed_generation"]) + tracker.generation_complete(completed) + pending = state["requested_generation"] > completed + if ( + not pending + and args.arm != "frozen" + and state["successful_episodes"] < args.accepted_episodes + ): + pending = flow.request_refill() is not None + if ( + quota is not None + and not pending + and state["unlocked_remaining_chunks"] < args.batch_size + ): + exhausted = True + end_reason = "usage_budget" + if ( + args.calibration_run + and time.monotonic() >= minimum_finish + and completed >= args.accepted_episodes // args.buffer_size + ): + end_reason = "calibration_complete" + break + if exhausted: + idle_started = min(time.monotonic(), deadline) + time.sleep(min(0.02, max(0.0, deadline - time.monotonic()))) + tracker.interval( + "supply_exhausted_idle", + idle_started, + min(time.monotonic(), deadline), + ) + continue + throttle_started = min(time.monotonic(), deadline) + while time.monotonic() < min(next_request, deadline): + if flow is not None: + flow.poll() + time.sleep( + min(0.02, max(0.0, min(next_request, deadline) - time.monotonic())) + ) + if next_request > throttle_started: + tracker.interval( + "request_throttle", + throttle_started, + min(time.monotonic(), deadline), + ) + if time.monotonic() >= deadline: + break + requested_at = time.monotonic() + next_request = requested_at + args.batch_size / args.request_rate + try: + in_sampling = True + if offline: + batch = reader.next_batch() + sampled_at = time.monotonic() + else: + raw = flow.sample_context( + source, args.batch_size, sequence, generator=rng + ) + sampled_at = time.monotonic() + batch = normalize_online_batch(raw, metadata) + del raw + except (StopIteration, _SupplyExhausted): + exhausted = True + end_reason = "usage_budget" + continue + except _ObservationComplete: + break + finally: + in_sampling = False + returned_at = time.monotonic() + service_seconds = ( + float(batch["_service"]["worker_fetch_seconds"]) + if offline + else source.last_service_seconds + returned_at - sampled_at + ) + validate_batch( + batch, + batch_size=args.batch_size, + history=args.history_length, + future=args.prediction_horizon, + cameras=cameras, + full_images=args.full_image_check, + ) + tracker.batch( + batch, + requested_at=requested_at, + returned_at=returned_at, + service_seconds=service_seconds, + sampled_at=sampled_at, + within_window=returned_at <= deadline, + ) + if retained is None: + # Views intentionally retain the original returned allocation. + # This is a real ownership audit, with at most one extra batch + # (1.55 GiB dual RGB backing storage) included in memory samples. + retained = { + "state": batch["state"][0], + "images": { + name: image[0, args.history_length - 1] + for name, image in batch["images"].items() + }, + } + retained_digest = _retained_digest(retained) + batches += 1 + del batch + if returned_at > deadline: + end_reason = "observation_window" + break + if args.max_batches is not None and batches >= args.max_batches: + end_reason = "batch_budget" + break + measured_duration = min(time.monotonic(), deadline) - measurement_start + if retained is not None and _retained_digest(retained) != retained_digest: + raise ValueError("An earlier owned output changed after source replacement") + metrics["retained_output_check"] = retained is not None + if engine is not None: + stats = engine.get_stats() + budget = engine.get_sampling_budget_snapshot() + flow.drain_events() + if reader is not None: + reader_stats = reader.stats + budget = reader.snapshot_budget() + metrics.update( + tracker.summary(measured_duration, cache_episodes=args.buffer_size) + ) + metrics.update( + cache_bytes=( + sum( + value.numel() * value.element_size() + for value in engine.shared_buffer.values( + include_nested=True, leaves_only=True + ) + ) + if engine is not None + else None + ), + accepted_episodes=stats.get( + "successful_episodes", + len(getattr(reader, "manifest", {}).get("episodes", [])), + ), + accepted_episodes_per_second=( + stats.get("successful_episodes", 0) / stats["production_seconds"] + if stats.get("production_seconds", 0) > 0 + else None + ), + actual_measurement_seconds=measured_duration, + final_disk_bytes=_disk_bytes(corpus) if corpus is not None else 0, + peak_temporary_disk_bytes=None, + unconsumed_prefetch_chunks=( + budget.get("reserved_chunks", 0) if offline else 0 + ), + corpus_episodes=( + len(reader.manifest["episodes"]) + if offline + else stats.get("successful_episodes") + ), + visible_source_episodes=( + len(budget.get("episodes", [])) if offline else None + ), + retained_output_backing_scope="one original final batch retained for ownership audit, included in PSS samples", + component_service_scope=( + "aggregate worker SDK item fetch/decode/padding before DataLoader collation; not consumer next() latency" + if offline + else "successful immediate Source sampling/copy plus RGB conversion; controller retry waits excluded" + ), + ) + if quota is not None: + delivered = ( + budget.get("delivered_chunks") + if offline + else budget.get("total_delivered_chunks") + ) + if delivered != tracker.delivered: + raise ValueError("Source ledger and consumer deliveries disagree") + required_generations = min(3, args.accepted_episodes // args.buffer_size) + require_streaming = args.suite_tag == "report" and not offline + actual_count = ( + metrics["visible_source_episodes"] + if offline + else stats.get("successful_episodes") + ) + expected_count = ( + args.buffer_size + if args.arm in ("frozen", "lerobot_same_cache") + else args.accepted_episodes + ) + expected_generations = expected_count // args.buffer_size + # Acceptance is counted before the final row is published/unlocked. + # Formal runs must prove the complete finite supply, not only N accepts. + if stats: + complete_supply = ( + stats.get("completed_generation") == expected_generations + and stats.get("requested_generation") == expected_generations + and set(tracker.publications) == set(range(1, expected_generations + 1)) + and all( + len(sources) == args.buffer_size + for sources in tracker.publications.values() + ) + and len( + { + source + for sources in tracker.publications.values() + for source in sources + } + ) + == expected_count + ) + else: + visible_ids = { + tuple(row["source_id"]) for row in budget.get("episodes", []) + } + corpus_rows = [ + row + for row in reader.manifest["episodes"] + if tuple(row["source_id"]) in visible_ids + ] + complete_supply = ( + len(corpus_rows) == expected_count + and len({tuple(row["source_id"]) for row in corpus_rows}) + == expected_count + and all( + sum(row["generation"] == generation for row in corpus_rows) + == args.buffer_size + for generation in range(1, expected_generations + 1) + ) + ) + checks = { + "valid_batches": bool(tracker.window_delivered), + "cap": metrics["cap_violations"] == 0, + "retained_output": metrics["retained_output_check"], + "generation_coverage": ( + len(tracker.window_generations) >= required_generations + if require_streaming + else True + ), + "partial_publication_delivery": ( + bool(tracker.overlap_generations) if require_streaming else True + ), + "standard_reader": ( + reader_stats.get("decoded_corpus_preloaded") is False + if offline + else True + ), + "camera_delivery": ( + ( + reader_stats.get("requested_camera_timestamps", 0) > 0 + if cameras + else True + ) + if offline + else True + ), + "accepted_data_budget": ( + actual_count == expected_count + if args.suite_tag in ("report", "appendix") + else True + ), + "complete_publication": ( + complete_supply if args.suite_tag in ("report", "appendix") else True + ), + } + metrics["qualification_checks"] = checks + metrics["qualification_passed"] = all(checks.values()) + except BaseException as failure: + status, error = "failed", { + "type": type(failure).__name__, + "message": str(failure), + "traceback": traceback.format_exc(), + } + elapsed = ( + min(time.monotonic(), deadline) - measurement_start + if measurement_start is not None and deadline is not None + else 0 + ) + metrics.update( + tracker.summary( + max(measured_duration, elapsed), cache_episodes=args.buffer_size + ) + ) + metrics["qualification_passed"] = False + finally: + # Failure audits precede shutdown and preserve already committed chunks. + # They must not replace the original exception or prevent other cleanup. + for name, action in ( + ( + "engine_stats", + lambda: engine.get_stats() if engine is not None else None, + ), + ( + "sampling_budget", + lambda: ( + reader.snapshot_budget() + if reader is not None + else ( + engine.get_sampling_budget_snapshot() + if engine is not None + else None + ) + ), + ), + ): + try: + value = action() + if value is not None: + if name == "engine_stats": + stats = value + else: + budget = value + except BaseException as failure: + cleanup.setdefault("audit_errors", []).append( + {"component": name, "error": repr(failure)} + ) + for name, action in ( + ("reader", lambda: reader.close() if reader is not None else None), + ("writer", lambda: writer.abort() if writer is not None else None), + ("engine", lambda: engine.stop() if engine is not None else None), + ("events", lambda: flow.drain_events() if flow is not None else None), + ): + try: + result = action() + if result is not None: + cleanup[name] = result + except BaseException as failure: + cleanup["errors"].append({"component": name, "error": repr(failure)}) + try: + metrics.update(resources.close()) + except BaseException as failure: + cleanup["errors"].append({"component": "resources", "error": repr(failure)}) + tracker.close() + if cleanup["errors"]: + status = "failed" + metrics["qualification_passed"] = False + if reader is not None: + reader_stats = reader.stats + charged = budget.get("delivered_chunks", budget.get("total_delivered_chunks")) + if charged is not None: + metrics["unvalidated_delivered_chunks"] = max( + 0, charged - tracker.delivered + ) + metrics.update( + tracker.summary( + measured_duration + or ( + min(time.monotonic(), deadline) - measurement_start + if deadline is not None and measurement_start is not None + else 0 + ), + cache_episodes=args.buffer_size, + ) + ) + if status == "complete" and offline and corpus is not None: + try: + quality = audit_video_quality(corpus) + write_json(args.output / "video_quality.json", quality) + metrics["video_quality"] = quality + samples = quality.get("samples", []) + psnr = [row["psnr_db"] for row in samples if row["psnr_db"] is not None] + metrics["quality_psnr_db"] = sum(psnr) / len(psnr) if psnr else None + metrics["quality_mae"] = ( + sum(row["mae_uint8"] for row in samples) / len(samples) + if samples + else None + ) + except BaseException as failure: + status = "failed" + error = { + "type": type(failure).__name__, + "message": str(failure), + "phase": "posthoc_quality", + } + metrics["qualification_passed"] = False + summary = { + "config": cfg, + "status": status, + "error": error, + "end_reason": end_reason, + "metrics": metrics, + "episode_usage": list(tracker.episodes.values()), + "engine_stats": stats, + "sampling_budget": budget, + "reader_stats": reader_stats, + "cleanup": cleanup, + } + write_json(args.output / "summary.json", summary) + return summary + + +def _retained_digest(value: dict[str, Any]) -> str: + digest = hashlib.sha256(value["state"].contiguous().numpy().tobytes()) + for name, image in sorted(value["images"].items()): + digest.update(name.encode()) + digest.update(image.contiguous().numpy().tobytes()) + return digest.hexdigest() + + +def main() -> None: + """Run one experiment or delegate an isolated finite suite.""" + args = build_parser().parse_args() + if args.suite != "single": + from .suite import run_suite + + run_suite(args) + return + result = run_one(args) + print( + json.dumps( + { + "status": result["status"], + "qualified": result["metrics"].get("qualification_passed"), + "output": str(args.output), + } + ), + flush=True, + ) + if result["status"] != "complete": + raise SystemExit(2) + + +if __name__ == "__main__": + main() diff --git a/examples/ods/data_stream_demo/suite.py b/examples/ods/data_stream_demo/suite.py new file mode 100644 index 000000000..00faa128f --- /dev/null +++ b/examples/ods/data_stream_demo/suite.py @@ -0,0 +1,1020 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Supervise finite real-data matrices with disjoint CPU/GPU assignments.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import os +import random +import re +import shutil +import signal +import subprocess +import sys +import time +import xml.etree.ElementTree as ET +from pathlib import Path +from typing import Any + +import psutil + +from .acquisition import CAMERA_PROFILE_VERSION +from .report import write_report + +__all__ = ["run_suite"] + +_ROOT = Path(__file__).resolve().parents[3] +_VISIONS = ("state", "single", "dual") +_MAIN = ("lerobot", "streaming", "protected") +_SEEDS = (19001, 19002, 19003) +_ORDER_SEED = 47001 +_AUX_MAX_BYTES = 8 * 1024**2 +_HEADLESS_ENV = {"SDL_VIDEODRIVER": "dummy", "DISPLAY": None, "WAYLAND_DISPLAY": None} +_BASE = { + "buffer_size": 8, + "producer_envs": 2, + "accepted_episodes": 32, + "max_episode_steps": 600, + "batch_size": 32, + "history_length": 2, + "prediction_horizon": 32, + "episode_budget": 128, + "num_workers": 4, + "prefetch_factor": 1, + "data_timeout": 120.0, + "generation_timeout": 1800.0, + "initialization_timeout": 1800.0, +} +_LIMITATIONS = [ + "Select a physical NVML GPU UUID, set CUDA_VISIBLE_DEVICES and ODS_GPU_UUID to that UUID, and use logical cuda:0.", + "The matching native runtime must qualify CUDA UUID to Vulkan UUID mapping; no fallback to an unmasked or unreserved GPU is permitted.", + "Selected GPUs have three idle readings and at least 16GiB free. This is not an exclusive-device reservation.", + "An optional auxiliary graphics GPU is separately reserved and excluded from primary assignments; only known per-PID pure-G footprints up to 8 MiB are allowed there.", + "The parent samples compute and graphics contexts attributed to its child PID tree or process group; short-lived contexts can fall between samples.", + "No system-wide file cache is purged. Calibration is a preflight observation, not a guaranteed concurrent capacity.", +] + + +def _write(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(value, indent=2, allow_nan=False) + "\n", encoding="utf-8" + ) + + +def _hashes() -> dict[str, str]: + paths = [ + *Path(__file__).parent.glob("*.py"), + _ROOT / "examples/ods/simple_bc_demo/config.py", + _ROOT / "examples/ods/simple_bc_demo/environment.py", + *_ROOT.joinpath("examples/ods/simple_bc_demo/configs").rglob("*.yml"), + *_ROOT.joinpath("examples/ods/simple_bc_demo/configs").rglob("*.yaml"), + _ROOT / "embodichain/data_pipeline/engine/data.py", + _ROOT / "embodichain/data_pipeline/runtime.py", + ] + return { + str(path.relative_to(_ROOT)): hashlib.sha256(path.read_bytes()).hexdigest() + for path in sorted(set(paths)) + } + + +def _gpu_reading() -> list[dict[str, Any]]: + output = subprocess.check_output( + [ + "nvidia-smi", + "--query-gpu=index,uuid,name,utilization.gpu,memory.free,memory.total", + "--format=csv,noheader,nounits", + ], + text=True, + timeout=15, + ) + rows = [] + for line in output.splitlines(): + index, uuid, name, utilization, free, total = ( + part.strip() for part in line.split(",") + ) + rows.append( + { + "index": int(index), + "uuid": uuid, + "name": name, + "utilization_percent": int(utilization), + "free_mib": int(free), + "total_mib": int(total), + } + ) + return rows + + +def _gpu_contexts() -> list[dict[str, Any]]: + """Read compute and graphics contexts without acting on foreign processes.""" + output = subprocess.check_output(["nvidia-smi", "-q", "-x"], text=True, timeout=10) + document = ET.fromstring(output) + contexts = [] + gpus = document.findall("gpu") + if not gpus: + raise RuntimeError("nvidia-smi XML did not contain GPU identity records") + for gpu in gpus: + uuid = gpu.findtext("uuid") + if not uuid: + raise RuntimeError("nvidia-smi XML omitted a GPU UUID") + for process in gpu.findall("./processes/process_info"): + pid = process.findtext("pid", "") + if pid.isdigit(): + contexts.append( + { + "pid": int(pid), + "gpu_uuid": uuid, + "type": process.findtext("type"), + "used_memory": process.findtext("used_memory"), + } + ) + return contexts + + +def _memory_bytes(value: str | None) -> int | None: + match = re.fullmatch(r"([0-9]+(?:\.[0-9]+)?)\s*(B|KiB|MiB|GiB)", value or "") + if match is None: + return None + scale = {"B": 1, "KiB": 1024, "MiB": 1024**2, "GiB": 1024**3}[match[2]] + size = float(match[1]) * scale + return math.ceil(size) if math.isfinite(size) else None + + +def _aux_scope(auxiliary: dict[str, Any] | None) -> dict[str, Any] | None: + if auxiliary is None: + return None + return { + "physical_index": auxiliary["index"], + "uuid": auxiliary["uuid"], + "model": auxiliary["name"], + "allowed_type": "G", + "max_per_pid_bytes": _AUX_MAX_BYTES, + "shared_by_suite_jobs": True, + } + + +def _partition_contexts( + owned: list[dict[str, Any]], selected: str, auxiliary: str | None +) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + totals: dict[int, int] = {} + invalid = set() + for row in owned: + if row["gpu_uuid"] != auxiliary: + continue + memory = _memory_bytes(row.get("used_memory")) + if row.get("type") != "G" or memory is None: + invalid.add(row["pid"]) + else: + totals[row["pid"]] = max(totals.get(row["pid"], 0), memory) + allowed, unexpected = [], [] + for row in owned: + if row["gpu_uuid"] == selected: + continue + if ( + auxiliary is not None + and row["gpu_uuid"] == auxiliary + and row["pid"] not in invalid + and totals.get(row["pid"], _AUX_MAX_BYTES + 1) <= _AUX_MAX_BYTES + ): + allowed.append(row) + else: + unexpected.append(row) + return allowed, unexpected + + +def _guard_devices(active: dict[int, dict[str, Any]]) -> None: + """Enforce primary placement and the optional bounded reserved graphics scope.""" + contexts = _gpu_contexts() + owners = {} + groups = {} + for index, item in active.items(): + process = item["process"] + groups[process.pid] = index + owners[process.pid] = index + try: + for child in psutil.Process(process.pid).children(recursive=True): + owners[child.pid] = index + except (psutil.NoSuchProcess, psutil.AccessDenied): + pass + attributed: dict[int, list[dict[str, Any]]] = {index: [] for index in active} + for context in contexts: + owner = owners.get(context["pid"]) + if owner is None: + try: + owner = groups.get(os.getpgid(context["pid"])) + except ProcessLookupError: + continue + if owner is not None: + attributed[owner].append(context) + now = time.monotonic() + for index, item in active.items(): + owned = attributed[index] + allowed, unexpected = _partition_contexts( + owned, item["gpu_uuid"], item.get("aux_gpu_uuid") + ) + for row in allowed: + if row not in item["device_guard"]["allowed_aux_contexts"]: + item["device_guard"]["allowed_aux_contexts"].append(row) + memory_by_process: dict[tuple[int, str], int | None] = {} + for row in owned: + key = (row["pid"], row["gpu_uuid"]) + memory_by_process.setdefault(key, None) + value = _memory_bytes(row.get("used_memory")) + if value is not None: + memory_by_process[key] = max(memory_by_process[key] or 0, value) + memories = list(memory_by_process.values()) + total_memory = ( + sum(memories) if all(value is not None for value in memories) else None + ) + if total_memory is not None: + item["device_guard"]["peak_known_owned_context_bytes"] = max( + item["device_guard"].get("peak_known_owned_context_bytes", 0), + total_memory, + ) + item["device_guard"]["samples"] += 1 + item["device_guard"]["selected_context_seen"] |= any( + row["gpu_uuid"] == item["gpu_uuid"] for row in owned + ) + if unexpected: + item["device_guard"]["unexpected_contexts"].extend(unexpected) + item["device_error"] = ( + f"Owned GPU contexts violate selected UUID {item['gpu_uuid']}: {unexpected}" + ) + with (item["supervision"] / "device-contexts.jsonl").open( + "a", encoding="utf-8" + ) as stream: + stream.write( + json.dumps( + { + "seconds_from_launch": now - item["started"], + "selected_gpu_uuid": item["gpu_uuid"], + "aux_graphics_gpu_uuid": item.get("aux_gpu_uuid"), + "owned_contexts": owned, + "owned_context_memory_bytes": total_memory, + "allowed_aux_contexts": allowed, + "unexpected_contexts": unexpected, + } + ) + + "\n" + ) + + +def _select_devices( + requested: list[int] | None, output: Path, *, auxiliary_id: int | None = None +) -> tuple[list[dict[str, Any]], dict[str, Any]]: + readings = [] + for index in range(3): + readings.append({"monotonic_seconds": time.monotonic(), "gpus": _gpu_reading()}) + if index < 2: + time.sleep(2) + by_id = [{row["index"]: row for row in reading["gpus"]} for reading in readings] + _write( + output / "gpu-preflight.json", + {"readings": readings, "limitations": _LIMITATIONS}, + ) + auxiliary = None + if auxiliary_id is not None: + if auxiliary_id < 0 or requested is not None and auxiliary_id in requested: + raise ValueError( + "The reserved auxiliary GPU must be nonnegative and distinct from every primary GPU" + ) + if not all( + auxiliary_id in reading + and reading[auxiliary_id]["utilization_percent"] == 0 + and reading[auxiliary_id]["free_mib"] >= 16 * 1024 + for reading in by_id + ): + raise RuntimeError( + "Auxiliary GPU requires three idle readings and at least 16 GiB free" + ) + auxiliary = by_id[-1][auxiliary_id] + ids = ( + requested + if requested is not None + else [index for index in sorted(by_id[-1]) if index != auxiliary_id] + ) + if len(set(ids)) != len(ids) or any(index < 0 for index in ids): + raise ValueError("GPU IDs must be unique nonnegative physical indices") + selected, rejected = [], [] + for index in ids: + if all( + index in reading + and reading[index]["utilization_percent"] == 0 + and reading[index]["free_mib"] >= 16 * 1024 + for reading in by_id + ): + selected.append(by_id[-1][index]) + else: + rejected.append( + { + "index": index, + "reason": "not present, not idle in all readings, or less than 16GiB free", + } + ) + cpus = sorted(os.sched_getaffinity(0)) + for gpu in selected[len(cpus) // 8 :]: + rejected.append( + {"index": gpu["index"], "reason": "not enough disjoint eight-CPU slots"} + ) + selected = selected[: len(cpus) // 8] + if not selected: + raise RuntimeError("No idle GPU with 16GiB free and eight available CPUs") + for index, gpu in enumerate(selected): + gpu["cpu_affinity"] = cpus[index * 8 : (index + 1) * 8] + audit = { + "readings": readings, + "selected": selected, + "aux_graphics_gpu": auxiliary, + "aux_graphics_scope": _aux_scope(auxiliary), + "rejected": rejected, + "cpu_budget_per_job": 8, + "limitations": _LIMITATIONS, + } + _write(output / "gpu-preflight.json", audit) + return selected, audit + + +def _protocol(path: Path | None) -> dict[str, Any]: + if path is None: + raise ValueError("Report and appendix require --protocol calibration.json") + protocol = json.loads(path.read_text(encoding="utf-8")) + if ( + protocol.get("status") != "complete" + or protocol.get("format") != "ods-stream-protocol-v1" + ): + raise ValueError("A complete, qualified calibration protocol is required") + if ( + protocol.get("source_sha256") != _hashes() + or protocol.get("camera_profile_version") != CAMERA_PROFILE_VERSION + ): + raise ValueError( + "Code/profile differs from calibration; recalibrate under the new profile" + ) + for name in ("request_rate", "duration_seconds"): + value = protocol.get(name) + if ( + isinstance(value, bool) + or not isinstance(value, (float, int)) + or not math.isfinite(value) + or value <= 0 + ): + raise ValueError(f"Invalid calibrated {name}") + if protocol["duration_seconds"] > 900: + raise ValueError( + "Calibrated observation window exceeds the 900-second protocol ceiling" + ) + if protocol.get("headless_environment") != _HEADLESS_ENV: + raise ValueError("Headless environment differs from the qualified protocol") + return protocol + + +def _jobs( + args: argparse.Namespace, protocol: dict[str, Any] | None +) -> list[dict[str, Any]]: + base = dict(_BASE if protocol is None else protocol["base_settings"]) + suite = args.suite + if suite == "smoke": + base.update( + buffer_size=2, + producer_envs=2, + accepted_episodes=8, + batch_size=4, + episode_budget=8, + ) + rate, duration = 16.0, 30.0 + combinations = [ + (arm, vision, _SEEDS[0]) for arm in _MAIN for vision in _VISIONS + ] + elif suite == "calibrate": + rate, duration = 1e9, 60.0 + combinations = [ + (arm, vision, _SEEDS[0]) + for arm in ("protected", "lerobot_unbounded") + for vision in _VISIONS + ] + else: + rate, duration = protocol["request_rate"], protocol["duration_seconds"] + combinations = ( + [ + (arm, vision, seed) + for seed in _SEEDS + for arm in _MAIN + for vision in _VISIONS + ] + if suite == "report" + else [ + (arm, vision, _SEEDS[0]) + for arm in ("frozen", "lerobot_unbounded", "lerobot_same_cache") + for vision in _VISIONS + ] + ) + jobs = [] + for arm, vision, seed in combinations: + destination = args.output / arm / vision / f"seed{seed}" + corpus = None + if suite == "appendix" and arm.startswith("lerobot"): + corpus = ( + args.dataset_root + / "lerobot" + / vision + / f"seed{seed}" + / "lerobot_dataset" + ) + jobs.append( + { + "arm": arm, + "vision": vision, + "seed": seed, + "output": destination, + "dataset_root": corpus, + "settings": base, + "rate": rate, + "duration": duration, + "suite": "calibration" if suite == "calibrate" else suite, + "calibration": suite == "calibrate", + } + ) + random.Random(_ORDER_SEED).shuffle(jobs) + return jobs + + +def _command(job: dict[str, Any], gpu: dict[str, Any]) -> list[str]: + command = [ + "taskset", + "--cpu-list", + ",".join(map(str, gpu["cpu_affinity"])), + sys.executable, + "-m", + "examples.ods.data_stream_demo.run", + "--suite", + "single", + "--suite-tag", + job["suite"], + "--arm", + job["arm"], + "--vision", + job["vision"], + "--seed", + str(job["seed"]), + "--output", + str(job["output"]), + "--sim-device", + "cuda:0", + "--request-rate", + str(job["rate"]), + "--duration-seconds", + str(job["duration"]), + ] + for name, value in job["settings"].items(): + command.extend(["--" + name.replace("_", "-"), str(value)]) + if job["dataset_root"] is not None: + command.extend(["--dataset-root", str(job["dataset_root"])]) + if job["calibration"]: + command.append("--calibration-run") + return command + + +def _deadline(job: dict[str, Any]) -> float: + settings = job["settings"] + generations = settings["accepted_episodes"] // settings["buffer_size"] + return float( + settings["initialization_timeout"] + + generations * settings["generation_timeout"] + + job["duration"] + + 1800 + + 2 * settings["data_timeout"] + + 180 + ) + + +def _mark(job: dict[str, Any], status: str, reason: str) -> None: + path = job["output"] / "summary.json" + if path.exists(): + try: + summary = json.loads(path.read_text(encoding="utf-8")) + except ValueError: + (path.parent / "summary.partial.json").write_bytes(path.read_bytes()) + summary = {} + else: + summary = {} + summary.setdefault( + "config", + { + **job["settings"], + "arm": job["arm"], + "vision": job["vision"], + "seed": job["seed"], + "suite": job["suite"], + "cache_episodes": job["settings"]["buffer_size"], + "quota": ( + None + if job["arm"] in {"streaming", "lerobot_unbounded"} + else job["settings"]["episode_budget"] + ), + "requested_chunks_per_second": job["rate"], + "measurement_seconds": job["duration"], + }, + ) + summary["status"] = status + if not summary.get("error"): + summary["error"] = {"phase": "supervisor", "message": reason} + summary["supervisor_error"] = reason + summary.setdefault("metrics", {})["qualification_passed"] = False + summary.setdefault("episode_usage", []) + _write(path, summary) + + +def _terminate(process: subprocess.Popen[Any]) -> None: + for sig, timeout in ((signal.SIGTERM, 5), (signal.SIGKILL, 10)): + try: + os.killpg(process.pid, sig) + except ProcessLookupError: + pass + try: + process.wait(timeout=timeout) + except subprocess.TimeoutExpired: + continue + if process.poll() is None: + raise TimeoutError(f"Owned process group {process.pid} did not terminate") + + +def _supervise( + jobs: list[dict[str, Any]], + devices: list[dict[str, Any]], + auxiliary: dict[str, Any] | None = None, +) -> list[dict[str, Any]]: + pending = list(jobs) + active: dict[int, dict[str, Any]] = {} + results = [] + cancelled = False + last_guard = 0.0 + + def cancel(signum: int, frame: Any) -> None: + nonlocal cancelled + cancelled = True + + previous = { + sig: signal.signal(sig, cancel) for sig in (signal.SIGINT, signal.SIGTERM) + } + try: + while pending or active: + for gpu in devices: + if not pending or gpu["index"] in active or cancelled: + continue + job = pending.pop(0) + if job["dataset_root"] is not None and not job["dataset_root"].is_dir(): + reason = f"Required finalized corpus missing: {job['dataset_root']}" + _mark(job, "unsupported", reason) + results.append( + { + "output": str(job["output"]), + "status": "unsupported", + "reason": reason, + } + ) + continue + # Keep supervision files outside the child's initially empty + # output directory; run_one deliberately rejects nonempty runs. + supervision = ( + job["output"].parents[2] + / "_supervision" + / job["output"].relative_to(job["output"].parents[2]) + ) + supervision.mkdir(parents=True) + command = _command(job, gpu) + environment = os.environ.copy() + for name, value in _HEADLESS_ENV.items(): + if value is None: + environment.pop(name, None) + else: + environment[name] = value + environment.pop("ODS_AUX_GRAPHICS_GPU_UUID", None) + if auxiliary is not None: + environment["ODS_AUX_GRAPHICS_GPU_UUID"] = auxiliary["uuid"] + environment.update( + CUDA_VISIBLE_DEVICES=gpu["uuid"], + ODS_GPU_UUID=gpu["uuid"], + OMP_NUM_THREADS="1", + MKL_NUM_THREADS="1", + OPENBLAS_NUM_THREADS="1", + ODS_CPU_BUDGET="8", + ) + launch = { + "command": command, + "cwd": str(_ROOT), + "gpu": gpu, + "aux_graphics_scope": _aux_scope(auxiliary), + "environment": { + name: environment.get(name) + for name in ( + "CUDA_VISIBLE_DEVICES", + "ODS_GPU_UUID", + "ODS_AUX_GRAPHICS_GPU_UUID", + "OMP_NUM_THREADS", + "MKL_NUM_THREADS", + "OPENBLAS_NUM_THREADS", + "SDL_VIDEODRIVER", + "DISPLAY", + "WAYLAND_DISPLAY", + ) + }, + "timeout_seconds": _deadline(job), + "limitations": _LIMITATIONS, + } + _write(supervision / "launch.json", launch) + stream = (supervision / "process.log").open("w", encoding="utf-8") + started_ns = time.monotonic_ns() + started = started_ns / 1e9 + environment["ODS_PROCESS_LAUNCH_MONOTONIC_NS"] = str(started_ns) + try: + process = subprocess.Popen( + command, + cwd=_ROOT, + env=environment, + stdout=stream, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + except BaseException as error: + stream.close() + _mark(job, "failed", repr(error)) + results.append( + { + "output": str(job["output"]), + "status": "failed", + "error": repr(error), + } + ) + continue + active[gpu["index"]] = { + "job": job, + "process": process, + "stream": stream, + "started": started, + "deadline": started + _deadline(job), + "supervision": supervision, + "started_ns": started_ns, + "gpu_uuid": gpu["uuid"], + "aux_gpu_uuid": ( + auxiliary["uuid"] if auxiliary is not None else None + ), + "device_error": None, + "device_guard": { + "samples": 0, + "selected_context_seen": False, + "unexpected_contexts": [], + "allowed_aux_contexts": [], + }, + } + launch["process_launch_monotonic_ns"] = started_ns + launch["environment"]["ODS_PROCESS_LAUNCH_MONOTONIC_NS"] = str( + started_ns + ) + _write(supervision / "launch.json", launch) + print( + f"Started {job['suite']}: {job['arm']}/{job['vision']}/seed{job['seed']} on GPU{gpu['index']}", + flush=True, + ) + if active and time.monotonic() - last_guard >= 2: + try: + _guard_devices(active) + except ( + OSError, + ValueError, + RuntimeError, + ET.ParseError, + subprocess.SubprocessError, + ) as error: + for item in active.values(): + item["device_error"] = ( + f"Cannot audit owned GPU placement: {error!r}" + ) + last_guard = time.monotonic() + for index, item in list(active.items()): + process, job = item["process"], item["job"] + timeout = time.monotonic() >= item["deadline"] + if ( + process.poll() is None + and not timeout + and not cancelled + and item["device_error"] is None + ): + continue + if timeout or cancelled or item["device_error"] is not None: + _terminate(process) + leaked_group = False + if not timeout and not cancelled and item["device_error"] is None: + try: + os.killpg(process.pid, 0) + except ProcessLookupError: + pass + else: + leaked_group = True + _terminate(process) + item["stream"].close() + summary_path = job["output"] / "summary.json" + reason = ( + item["device_error"] + if item["device_error"] is not None + else ( + "suite cancellation" + if cancelled + else ( + "process deadline exceeded" + if timeout + else ( + "Child left an owned process group alive" + if leaked_group + else "" + ) + ) + ) + ) + if reason or process.returncode != 0 or not summary_path.exists(): + _mark( + job, + "failed", + reason + or f"Child exit={process.returncode}, summary exists={summary_path.exists()}", + ) + summary = json.loads(summary_path.read_text(encoding="utf-8")) + summary["device_supervision"] = { + "primary_gpu_uuid": item["gpu_uuid"], + "aux_graphics_scope": _aux_scope(auxiliary), + "guard": item["device_guard"], + } + _write(summary_path, summary) + result = { + "output": str(job["output"]), + "supervision": str(item["supervision"]), + "status": summary["status"], + "qualified": summary.get("metrics", {}).get("qualification_passed"), + "returncode": process.returncode, + "elapsed_seconds": time.monotonic() - item["started"], + "timeout": timeout, + "cancelled": cancelled, + "leaked_process_group": leaked_group, + "pgid": process.pid, + "process_launch_monotonic_ns": item["started_ns"], + "device_guard": item["device_guard"], + "aux_graphics_scope": _aux_scope(auxiliary), + } + _write(job["output"] / "process.json", result) + results.append(result) + print( + f"Finished {job['arm']}/{job['vision']}/seed{job['seed']}: {result['status']}, qualified={result['qualified']}", + flush=True, + ) + del active[index] + if cancelled: + for job in pending: + job["output"].mkdir(parents=True) + _mark(job, "failed", "Suite cancelled before launch") + results.append( + { + "output": str(job["output"]), + "status": "failed", + "cancelled": True, + } + ) + pending.clear() + if pending or active: + time.sleep(0.25) + finally: + for item in active.values(): + try: + _terminate(item["process"]) + finally: + item["stream"].close() + for sig, handler in previous.items(): + signal.signal(sig, handler) + return results + + +def _calibration( + output: Path, + jobs: list[dict[str, Any]], + devices: list[dict[str, Any]], + auxiliary: dict[str, Any] | None = None, +) -> dict[str, Any]: + gamma, capacities, issues = {}, {}, [] + for job in jobs: + summary = json.loads( + (job["output"] / "summary.json").read_text(encoding="utf-8") + ) + key = f"{job['arm']}/{job['vision']}" + metrics, stats = summary.get("metrics", {}), summary.get("engine_stats", {}) + if ( + summary["status"] != "complete" + or metrics.get("qualification_passed") is not True + ): + issues.append( + f"{key}: calibration execution or engineering qualification failed" + ) + continue + if ( + stats.get("successful_episodes", 0) < _BASE["accepted_episodes"] + or stats.get("completed_generation", 0) < 4 + ): + issues.append( + f"{key}: all 32 accepted episodes/four generations were not completed" + ) + continue + if job["arm"] == "protected": + seconds = stats.get("production_seconds", 0) + service = metrics.get("service_p50_seconds") + if seconds <= 0 or service is None or service <= 0: + issues.append(f"{key}: missing production or batch-service timing") + continue + gamma[job["vision"]] = stats["successful_episodes"] / seconds + capacities[key] = { + "chunks_per_second": _BASE["batch_size"] / service, + "scope": "successful immediate Source sampling/copy and RGB conversion p50 with concurrent production; excludes Flow retry waits; not an isolated maximum", + } + else: + rate = metrics.get("delivered_chunks_per_second") + if rate is None or rate <= 0: + issues.append(f"{key}: no measured standard-reader throughput") + continue + capacities[key] = { + "chunks_per_second": rate, + "scope": "source-ready, high-demand unbounded standard Dataset/DataLoader/TorchCodec delivery", + } + rate = ( + 4 * _BASE["episode_budget"] * max(gamma.values()) if len(gamma) == 3 else None + ) + duration = ( + max(180.0, max(4 * _BASE["buffer_size"] / value for value in gamma.values())) + if len(gamma) == 3 + else None + ) + if duration is not None and duration > 900: + issues.append( + "Measured production requires an observation window beyond 900 seconds" + ) + return { + "format": "ods-stream-protocol-v1", + "status": "complete" if not issues and len(capacities) == 6 else "failed", + "camera_profile_version": CAMERA_PROFILE_VERSION, + "source_sha256": _hashes(), + "base_settings": _BASE, + "request_rate": rate, + "duration_seconds": duration, + "generation_episodes_per_second": gamma, + "ready_source_capacity": capacities, + "gpu_models": sorted({gpu["name"] for gpu in devices}), + "aux_graphics_scope": _aux_scope(auxiliary), + "headless_environment": _HEADLESS_ENV, + "order_seed": _ORDER_SEED, + "seeds": list(_SEEDS), + "issues": issues, + "limitations": _LIMITATIONS, + } + + +def run_suite(args: argparse.Namespace) -> dict[str, Any]: + """Execute a supervised matrix and retain every attempted or unsupported cell. + + Args: + args: ``run.build_parser`` arguments with a non-single suite selection. + + Returns: + Machine-readable suite status, also saved beside the per-cell artifacts. + Failed/unsupported execution exits with code 2 after rendering its report. + """ + args = argparse.Namespace(**vars(args)) + args.output = args.output.resolve() + if args.dataset_root is not None: + args.dataset_root = args.dataset_root.resolve() + if args.output.exists() and any(args.output.iterdir()): + raise FileExistsError("Use a fresh suite output directory") + if args.suite == "appendix" and args.dataset_root is None: + raise ValueError( + "Appendix requires --dataset-root pointing to the main REPORT ROOT" + ) + protocol = ( + _protocol(args.protocol) if args.suite in {"report", "appendix"} else None + ) + args.output.mkdir(parents=True, exist_ok=True) + jobs = _jobs(args, protocol) + _write( + args.output / "suite-plan.json", + { + "suite": args.suite, + "jobs": [ + { + **job, + "output": str(job["output"]), + "dataset_root": ( + str(job["dataset_root"]) + if job["dataset_root"] is not None + else None + ), + } + for job in jobs + ], + "order_seed": _ORDER_SEED, + "aux_graphics_gpu_id": args.aux_graphics_gpu_id, + "headless_environment": _HEADLESS_ENV, + "source_sha256": _hashes(), + "limitations": _LIMITATIONS, + }, + ) + preflight_error = None + devices = [] + auxiliary = None + try: + if shutil.which("taskset") is None: + raise RuntimeError( + "taskset is required to assign eight disjoint CPUs per child" + ) + devices, device_audit = _select_devices( + args.gpu_ids, args.output, auxiliary_id=args.aux_graphics_gpu_id + ) + auxiliary = device_audit["aux_graphics_gpu"] + _gpu_contexts() # Qualification requires compute and graphics visibility. + if ( + protocol is not None + and sorted({gpu["name"] for gpu in devices}) != protocol["gpu_models"] + ): + raise RuntimeError( + "Selected GPU models differ from the frozen calibration hardware" + ) + if protocol is not None and protocol.get("aux_graphics_scope") != _aux_scope( + auxiliary + ): + raise RuntimeError( + "Auxiliary graphics reservation differs from the frozen protocol; pass the same explicit auxiliary flag/scope" + ) + except ( + OSError, + ValueError, + RuntimeError, + ET.ParseError, + subprocess.SubprocessError, + ) as error: + preflight_error = repr(error) + for job in jobs: + job["output"].mkdir(parents=True) + _mark(job, "unsupported", preflight_error) + results = [ + { + "output": str(job["output"]), + "status": "unsupported", + "reason": preflight_error, + } + for job in jobs + ] + else: + results = _supervise(jobs, devices, auxiliary) + summary = { + "suite": args.suite, + "status": ( + "complete" + if all(result["status"] == "complete" for result in results) + else "failed" + ), + "qualification_passed": all( + result.get("qualified") is True for result in results + ), + "results": results, + "preflight_error": preflight_error, + "limitations": _LIMITATIONS, + "aux_graphics_scope": _aux_scope(auxiliary), + "requested_aux_graphics_gpu_id": args.aux_graphics_gpu_id, + } + if args.suite == "calibrate": + calibration = _calibration(args.output, jobs, devices, auxiliary) + _write(args.output / "calibration.json", calibration) + summary["calibration_status"] = calibration["status"] + if calibration["status"] != "complete": + summary["status"] = "failed" + _write(args.output / "suite-summary.json", summary) + try: + summary["report"] = str(write_report(args.output)) + except Exception as error: + summary["report_error"] = repr(error) + summary["status"] = "failed" + _write(args.output / "suite-summary.json", summary) + if summary["status"] != "complete": + raise SystemExit(2) + return summary diff --git a/examples/ods/data_stream_demo/train.py b/examples/ods/data_stream_demo/train.py new file mode 100644 index 000000000..1b5dba309 --- /dev/null +++ b/examples/ods/data_stream_demo/train.py @@ -0,0 +1,1544 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Compare visual BC diversity and complete-pipeline time at a fixed update budget. + +Run ``python -m examples.ods.data_stream_demo.train --help``. The supervisor +selects physical GPU UUIDs before launching the real acquisition/learner worker. +Every arm records standard LeRobot data; offline training really decodes it. +""" + +from __future__ import annotations + +import argparse +from collections import Counter +import hashlib +import json +import math +import os +from pathlib import Path +import random +import time +import traceback +from typing import Any + +__all__ = ["build_parser", "run_one", "main"] + + +def build_parser() -> argparse.ArgumentParser: + """Expose explicit resource roles and real, fixed training-work budgets.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--suite", + choices=("single", "smoke", "calibrate", "pilot", "report"), + default="single", + ) + parser.add_argument( + "--arm", choices=("offline", "ods_overlap", "ods_serial"), default="ods_overlap" + ) + parser.add_argument("--vision", choices=("single", "dual"), default="single") + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--gpu-ids", nargs="+", type=int) + parser.add_argument("--aux-graphics-gpu-id", type=int) + parser.add_argument("--sim-device", default="cuda:0") + parser.add_argument("--train-device", default="cuda:1") + parser.add_argument("--seed", type=int, default=19001) + parser.add_argument("--updates", type=int, default=128) + parser.add_argument("--batch-size", type=int, default=32) + parser.add_argument("--buffer-size", type=int, default=8) + parser.add_argument("--producer-envs", type=int, default=2) + parser.add_argument("--episode-budget", type=int, default=128) + parser.add_argument( + "--handoff-mode", choices=("generation", "publication"), default="publication" + ) + parser.add_argument("--encoder-threads", type=int, default=1) + parser.add_argument( + "--prefetch-policy", choices=("budget", "eager"), default="budget" + ) + parser.add_argument("--calibration-updates", type=int) + parser.add_argument("--offline-episodes", type=int, default=32) + parser.add_argument("--history-length", type=int, default=2) + parser.add_argument("--prediction-horizon", type=int, default=32) + parser.add_argument("--action-horizon", type=int, default=16) + parser.add_argument("--max-episode-steps", type=int, default=600) + parser.add_argument("--evaluation-scenes", type=int, default=20) + parser.add_argument("--initialization-timeout", type=float, default=1800) + parser.add_argument("--generation-timeout", type=float, default=1800) + parser.add_argument("--data-timeout", type=float, default=120) + parser.add_argument("--run-timeout", type=float, default=3600) + parser.add_argument("--trace", type=Path) + parser.add_argument("--worker", action="store_true", help=argparse.SUPPRESS) + parser.add_argument("--phase", default="single", help=argparse.SUPPRESS) + return parser + + +def _config(args: argparse.Namespace) -> dict[str, Any]: + return { + key: str(value) if isinstance(value, Path) else value + for key, value in vars(args).items() + } + + +def _digest(values: dict[str, Any]) -> str: + import torch + + digest = hashlib.sha256() + for key, value in sorted(values.items()): + digest.update(key.encode()) + if isinstance(value, torch.Tensor): + value = value.detach().cpu().contiguous() + digest.update(str((tuple(value.shape), value.dtype)).encode()) + digest.update(memoryview(value.numpy())) + else: + digest.update(json.dumps(value, sort_keys=True).encode()) + return digest.hexdigest() + + +def _validate(args: argparse.Namespace) -> None: + for name in ( + "updates", + "batch_size", + "buffer_size", + "producer_envs", + "episode_budget", + "offline_episodes", + "prediction_horizon", + "max_episode_steps", + "encoder_threads", + ): + if type(getattr(args, name)) is not int or getattr(args, name) < 1: + raise ValueError(f"{name} must be a positive integer") + if ( + args.history_length != 2 + or not 0 < args.action_horizon <= args.prediction_horizon + ): + raise ValueError("Visual BC requires H=2 and a positive action horizon <=F") + if args.buffer_size % args.producer_envs or args.batch_size % args.buffer_size: + raise ValueError("C must divide B and be divisible by P") + if args.buffer_size * args.episode_budget < args.batch_size: + raise ValueError("A complete retained block must fund at least one batch") + if args.offline_episodes % args.buffer_size: + raise ValueError("Offline episode count must be a multiple of C") + if args.prefetch_policy not in ("budget", "eager"): + raise ValueError("prefetch_policy must be budget or eager") + if args.handoff_mode not in ("generation", "publication"): + raise ValueError("handoff_mode must be generation or publication") + if ( + args.arm == "ods_overlap" + and args.handoff_mode == "publication" + and args.producer_envs * args.episode_budget < args.batch_size + ): + raise ValueError( + "Publication handoff requires P*K >= B to fund a complete batch" + ) + if args.calibration_updates is not None and ( + type(args.calibration_updates) is not int or args.calibration_updates < 1 + ): + raise ValueError("calibration_updates must be a positive integer or None") + if args.arm == "ods_serial" and args.trace is None: + raise ValueError("Serial attribution requires the corresponding overlap trace") + if ( + not args.sim_device.startswith("cuda:") + or not args.train_device.startswith("cuda:") + or args.sim_device == args.train_device + ): + raise ValueError( + "The real experiment requires separate CUDA generation and training devices" + ) + + +def _prefetch_decision( + *, + remaining_updates: int, + batch_size: int, + remaining_credits: int, + handoff_seconds: float | None, + update_seconds: float | None, + policy: str = "budget", +) -> dict[str, Any]: + """Request necessary supply, or a refresh likely to reach a future update. + + K remains a maximum reuse budget. Sufficient current credits do not prevent + timely refreshes when production is faster than the remaining learner work. + Unknown timing is conservative; quota shortages always request supply. + """ + if ( + type(remaining_updates) is not int + or remaining_updates < 0 + or type(batch_size) is not int + or batch_size < 1 + or type(remaining_credits) is not int + or remaining_credits < 0 + or policy not in ("budget", "eager") + ): + raise ValueError( + "Prefetch requires nonnegative work/credits and a valid policy" + ) + for value in (handoff_seconds, update_seconds): + if value is not None and ( + type(value) not in (int, float) or not math.isfinite(value) or value <= 0 + ): + raise ValueError("Measured prefetch timings must be finite and positive") + remaining_seconds = ( + None + if update_seconds is None + else max(0, remaining_updates - 1) * update_seconds + ) + if remaining_updates == 0: + request, reason = False, "training_complete" + elif remaining_credits < remaining_updates * batch_size: + request, reason = True, "quota_required" + elif remaining_updates == 1: + request, reason = False, "last_update_funded" + elif policy == "eager": + request, reason = True, "eager_control" + elif handoff_seconds is None or update_seconds is None: + request, reason = False, "timing_not_observed" + elif handoff_seconds < remaining_seconds: + request, reason = True, "timely_refresh" + else: + request, reason = False, "arrival_after_remaining_work" + return { + "request": request, + "reason": reason, + "policy": policy, + "remaining_updates": remaining_updates, + "remaining_credits": remaining_credits, + "handoff_seconds": handoff_seconds, + "ready_update_seconds": update_seconds, + "remaining_refresh_window_seconds": remaining_seconds, + } + + +class _Experiment: + """Own one finite pipeline; retained data and recording staging are bounded.""" + + def __init__(self, args: argparse.Namespace) -> None: + import torch + + from .metrics import ResourceMonitor + from .training_metrics import TrainingMetrics + + self.args, self.output = args, args.output + self.output.mkdir(parents=True, exist_ok=True) + if (self.output / "config.json").exists(): + raise FileExistsError( + "Each timed experiment requires a fresh output directory" + ) + self.origin_ns = int( + os.environ.get("ODS_PROCESS_LAUNCH_MONOTONIC_NS", time.monotonic_ns()) + ) + self.deadline_ns = self.origin_ns + int(args.run_timeout * 1e9) + self.metrics = TrainingMetrics( + self.output, + origin_ns=self.origin_ns, + history=2, + future=args.prediction_horizon, + batch_size=args.batch_size, + reuse_cap=None if args.arm == "offline" else args.episode_budget, + ) + self.monitor = ResourceMonitor(self.output) + self.engine = self.flow = self.writer = self.reader = None + self.model = self.optimizer = self.block = None + self.metadata: dict[str, Any] = {} + self.statistics: dict[str, torch.Tensor] = {} + self.registry: list[dict[str, Any]] = [] + self.trace_updates: list[dict[str, Any]] = [] + self.retired_ledgers: list[dict[str, Any]] = [] + self.staging: list[Any] | None = None + self.staging_index = 0 + self.packet = None + self.training_completed_ns = self.origin_ns + self.generation = self.completed_updates = 0 + self.generation_starts: dict[int, int] = {} + self.writer_stats: dict[str, Any] = {} + self.recording_log_sha256: str | None = None + self.resources: dict[str, Any] = {} + self.cleanup: dict[str, Any] = {} + self.gpu_step_ms: list[float] = [] + self.ready_update_seconds: list[float] = [] + self.generation_seconds: list[float] = [] + self.snapshot_seconds: list[float] = [] + self.last_prefetch_decision: tuple[int, str, bool] | None = None + self.closed = False + self.retained_probe = None + self.retained_probe_hash: str | None = None + self.initial_episodes = None + self.source_run_id: int | None = None + self.cache_bytes = 0 + self.serial_trace = json.loads(args.trace.read_text()) if args.trace else None + self.retained_blocks: list[dict[str, Any]] = [] + self.peak_owned_episodes = 0 + self.recorded_generations: set[int] = set() + self.ready_episodes: dict[int, Any] = {} + self.update_completion_ns: list[int] = [] + + def now(self) -> int: + return time.monotonic_ns() + + def poll(self) -> None: + if self.now() >= self.deadline_ns: + raise TimeoutError("End-to-end pipeline exceeded its declared deadline") + if self.writer is not None: + self.writer.check_health() + self.pump_recording() + + def pump_recording(self) -> None: + from .payload import pack_episode + + if self.staging is None or self.writer is None: + return + if self.packet is None and self.staging_index < len(self.staging): + started = self.now() + self.packet = pack_episode(self.staging[self.staging_index], self.metadata) + self.staging[self.staging_index] = None + self.metrics.interval("recording_copy", started, self.now()) + if self.packet is not None and self.writer.try_submit(self.packet): + self.packet = None + self.staging_index += 1 + if self.staging_index == len(self.staging): + self.staging = None + + def flush_staging(self) -> None: + started = self.now() + while self.staging is not None: + self.poll() + if self.flow is not None: + self.flow.drain_events() + self.engine.check_health() + if self.staging is not None: + time.sleep(0.005) + self.metrics.interval("recording_staging_wait", started, self.now()) + + def stage(self, episodes: list[Any]) -> None: + if self.staging is None: + if self.packet is not None: + raise RuntimeError("A submission packet lost its staging owner") + self.staging, self.staging_index = list(episodes), 0 + else: + self.staging = self.staging[self.staging_index :] + list(episodes) + self.staging_index = 0 + self.check_retention() + self.pump_recording() + + def check_retention(self) -> int: + """Count the union of raw episode owners; archive packets are separate.""" + owned = {id(row) for row in self.staging or [] if row is not None} + owned.update(id(row) for row in self.ready_episodes.values()) + if self.block is not None: + owned.update(id(row["raw"]) for row in self.block.episodes.values()) + if len(owned) > self.args.buffer_size: + raise RuntimeError("Owned training/recording episodes exceeded C") + self.peak_owned_episodes = max(self.peak_owned_episodes, len(owned)) + return len(owned) + + def staging_slots(self) -> int: + return sum(row is not None for row in self.staging or []) + + def events(self, rows: list[dict[str, Any]]) -> None: + for row in rows: + details = dict(row) + stamp = int(details.pop("monotonic_ns")) + details["lifecycle_kind"] = details.pop("kind") + self.metrics.event("producer_lifecycle", stamp, **details) + + def requested(self, generation: int, started: float) -> None: + stamp = int(started * 1e9) + self.generation_starts[generation] = stamp + self.metrics.event("generation_requested", stamp, generation=generation) + + def take(self, generation: int, *, drain: bool = False) -> list[Any]: + from . import acquisition + + started = self.now() + self.flow.wait_for_generation(generation) + self.metrics.interval( + "generation_drain_wait" if drain else "data_wait", + started, + self.now(), + generation=generation, + ) + copied = self.now() + episodes = list(acquisition.iter_generation_episodes(self.engine, generation)) + return self.register_episodes(generation, episodes, copied) + + def register_episodes( + self, generation: int, episodes: list[Any], copied: int + ) -> list[Any]: + from .payload import _numeric, numeric_episode_digest + from .training_data import geometry_from_raw + + published = 0 + for episode in episodes: + ordinal = int(episode["ods_episode_id"][0]) + if ordinal != len(self.registry): + raise ValueError( + "Accepted episode ordinals differ from recording order" + ) + source = [int(episode["ods_run_id"][0]), ordinal] + timestamp = int(episode["ods_published_at_monotonic_ns"][0]) + geometry = geometry_from_raw(episode) + xy = geometry[0, 27:29].tolist() + digest = numeric_episode_digest( + { + "state": _numeric(episode, self.metadata), + "actions": episode["actions"], + "continuity_id": episode["continuity_id"], + } + ) + record = { + "ordinal": ordinal, + "source_id": source, + "generation": generation, + "length": len(episode["valid"]), + "initial_xy": xy, + "published_ns": timestamp, + "numeric_sha256": digest, + } + self.registry.append(record) + self.metrics.register_episode( + f"{source[0]}:{source[1]}", + ordinal=ordinal, + generation=generation, + initial_xy=xy, + published_ns=timestamp, + frame_count=record["length"], + metadata={"numeric_sha256": digest}, + ) + published = max(published, timestamp) + copied_end = self.now() + self.metrics.interval( + "snapshot_copy", copied, copied_end, generation=generation + ) + if len(self.registry) == generation * self.args.buffer_size: + if generation in self.recorded_generations: + raise RuntimeError("Generation service was recorded twice") + begin = self.generation_starts.get(generation, self.origin_ns) + self.metrics.interval("generation", begin, published, generation=generation) + self.generation_seconds.append((published - begin) / 1e9) + self.recorded_generations.add(generation) + self.snapshot_seconds.append((copied_end - copied) / 1e9) + self.generation = generation + return episodes + + def bootstrap(self, episodes: list[Any]) -> None: + import torch + + from .lerobot_io import StandardLeRobotWriter + from .metrics import write_json + from .training_data import make_statistics + from .visual_policy import VisualChunkBC + + started = self.now() + self.statistics = make_statistics(episodes, self.metadata) + torch.manual_seed(self.args.seed + 1234) + self.model = VisualChunkBC( + **self.statistics, + cameras=tuple(self.metadata["cameras"]), + prediction_horizon=self.args.prediction_horizon, + ).to(self.args.train_device) + self.optimizer = torch.optim.AdamW( + self.model.parameters(), lr=3e-4, weight_decay=1e-6, foreach=False + ) + torch.cuda.synchronize(self.args.train_device) + self.initial_model_hash = _digest(self.model.state_dict()) + self.normalization_hash = _digest(self.statistics) + self.save_checkpoint("initial.pt", 0) + self.metrics.interval("model_initialization", started, self.now()) + self.writer = StandardLeRobotWriter( + self.output / "lerobot_dataset", + self.metadata, + episode_count=( + self.args.offline_episodes if self.args.arm == "offline" else None + ), + queue_size=2, + events_path=self.output / "writer-sdk-events.jsonl", + encoder_threads=self.args.encoder_threads, + ) + self.stage(episodes) + write_json( + self.output / "initial-contract.json", + { + "model_sha256": self.initial_model_hash, + "normalization_sha256": self.normalization_hash, + "initial_episodes": self.registry.copy(), + "camera_order": list(self.metadata["cameras"]), + "training_budget": { + "updates": self.args.updates, + "chunks": self.args.updates * self.args.batch_size, + "targets": self.args.updates + * self.args.batch_size + * self.args.prediction_horizon, + }, + }, + ) + if self.serial_trace: + self.validate_trace() + + def record_writer_intervals(self) -> None: + """Import the joined SDK thread's real service spans, without callbacks.""" + if self.recording_log_sha256 is not None: + raise RuntimeError("Recording service intervals were already imported") + path = self.output / "writer-sdk-events.jsonl" + rows = [json.loads(line) for line in path.read_text().splitlines() if line] + phases = Counter(row.get("name") for row in rows) + if phases != { + "sdk_create": 1, + "sdk_episode_write_encode": len(self.registry), + "sdk_finalize": 1, + }: + raise ValueError("Recording service log is incomplete") + previous_end, ordinal, threads = self.origin_ns, 0, set() + for row in rows: + if ( + row.get("schema") != "ods-sdk-writer-service/v1" + or row.get("kind") != "interval" + or row.get("status") != "complete" + or row.get("worker_pid") != os.getpid() + or type(row.get("worker_thread_id")) is not int + or row["worker_thread_id"] <= 0 + or type(row.get("start_ns")) is not int + or type(row.get("end_ns")) is not int + or not previous_end <= row["start_ns"] <= row["end_ns"] <= self.now() + ): + raise ValueError("Recording service clock/owner/status is invalid") + if row["name"] == "sdk_episode_write_encode": + source = self.registry[ordinal] + if ( + row.get("source_id") != source["source_id"] + or row.get("episode_ordinal") != ordinal + ): + raise ValueError("Recording service source order differs") + ordinal += 1 + elif ( + row.get("source_id") is not None + or row.get("episode_ordinal") is not None + ): + raise ValueError("Dataset-wide recording service has an episode ID") + self.metrics.interval( + "sdk_recording", + row["start_ns"], + row["end_ns"], + sdk_phase=row["name"], + worker_pid=row["worker_pid"], + worker_thread_id=row["worker_thread_id"], + source_id=row.get("source_id"), + episode_ordinal=row.get("episode_ordinal"), + scope=row["scope"], + stage_intervals=row.get("stage_intervals"), + ) + previous_end = row["end_ns"] + threads.add(row["worker_thread_id"]) + if len(threads) != 1: + raise ValueError("Recording service did not have one SDK-owner thread") + self.recording_log_sha256 = hashlib.sha256(path.read_bytes()).hexdigest() + + def validate_trace(self) -> None: + trace = self.serial_trace + expected = { + key: getattr(self.args, key) + for key in ( + "vision", + "seed", + "updates", + "batch_size", + "buffer_size", + "producer_envs", + "episode_budget", + "prediction_horizon", + ) + } + if trace.get("format") == "ods-visual-training-trace-v2": + expected.update( + handoff_mode=self.args.handoff_mode, + encoder_threads=self.args.encoder_threads, + ) + blocks = trace.get("retained_blocks", []) + if not blocks or blocks[0].get("first_update") != 1: + raise ValueError("Serial v2 trace requires retained source boundaries") + previous_update, next_ordinal = 0, 0 + for block in blocks: + ordinals, first = block.get("ordinals", []), block.get("first_update") + if ( + not ordinals + or ordinals + != list(range(next_ordinal, next_ordinal + len(ordinals))) + or type(first) is not int + or not previous_update < first <= self.args.updates + or self.args.batch_size % len(ordinals) + or len(ordinals) * self.args.episode_budget < self.args.batch_size + or any( + o // self.args.buffer_size + 1 != block.get("generation") + for o in ordinals + ) + ): + raise ValueError("Serial retained-block contract is invalid") + previous_update, next_ordinal = first, next_ordinal + len(ordinals) + if ( + trace.get("format") + not in {"ods-visual-training-trace-v1", "ods-visual-training-trace-v2"} + or trace["recipe"] != expected + ): + raise ValueError("Serial replay recipe differs from overlap") + if ( + len(trace["updates"]) != self.args.updates + or trace["normalization_sha256"] != self.normalization_hash + or trace["initial_model_sha256"] != self.initial_model_hash + ): + raise ValueError("Serial replay normalization/initial model/budget differs") + + def save_checkpoint(self, name: str, step: int) -> None: + import numpy as np + import torch + + target = self.output / name + temporary = target.with_suffix(".tmp") + torch.save( + { + "model": self.model.state_dict(), + "optimizer": self.optimizer.state_dict(), + "normalization": self.statistics, + "config": _config(self.args), + "cameras": tuple(self.metadata["cameras"]), + "step": step, + "rng": { + "python": random.getstate(), + "numpy": np.random.get_state(), + "torch": torch.get_rng_state(), + "training_cuda": torch.cuda.get_rng_state(self.args.train_device), + }, + "source_run_id": self.source_run_id, + }, + temporary, + ) + temporary.replace(target) + + def update( + self, + batch: dict[str, Any], + plan: list[tuple[int, int]], + *, + requested_ns: int, + generation: int, + ) -> None: + import torch + + if self.retained_probe is None: + self.retained_probe = { + key: value + for key, value in batch.items() + if isinstance(value, torch.Tensor) + } + self.retained_probe.update( + {f"image/{key}": value for key, value in batch["images"].items()} + ) + self.retained_probe_hash = _digest(self.retained_probe) + transfer_start = self.now() + geometry = batch["geometry"].to(self.args.train_device) + actions = batch["actions"].to(self.args.train_device) + images = { + key: value.to(self.args.train_device) + for key, value in batch["images"].items() + } + valid = batch["valid"].to(self.args.train_device) + loss_mask = batch["loss_mask"].to(self.args.train_device) + torch.cuda.synchronize(self.args.train_device) + self.metrics.interval("device_transfer", transfer_start, self.now()) + training_start = self.now() + begin, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event( + enable_timing=True + ) + begin.record(torch.cuda.current_stream(self.args.train_device)) + self.optimizer.zero_grad(set_to_none=True) + arm, hand = self.model.loss( + geometry, actions, images, loss_mask=loss_mask, valid_mask=valid + ) + loss = (9 * arm + hand) / 10 + if not torch.isfinite(loss): + raise ValueError("Nonfinite real BC loss") + loss.backward() + torch.nn.utils.clip_grad_norm_( + self.model.parameters(), 1, error_if_nonfinite=True + ) + self.optimizer.step() + end.record(torch.cuda.current_stream(self.args.train_device)) + end.synchronize() + completed = self.now() + gpu_ms = begin.elapsed_time(end) + self.gpu_step_ms.append(gpu_ms) + self.metrics.interval( + "training", training_start, completed, gpu_milliseconds=gpu_ms + ) + self.completed_updates += 1 + self.training_completed_ns = completed + self.update_completion_ns.append(completed) + self.ready_update_seconds.append((completed - requested_ns) / 1e9) + if self.block is not None: + self.block.mark_trained(plan) + self.metrics.complete_update( + batch, + update=self.completed_updates, + losses={"arm": float(arm), "gripper": float(hand), "total": float(loss)}, + requested_ns=requested_ns, + train_started_ns=training_start, + completed_ns=completed, + ) + self.trace_updates.append( + { + "update": self.completed_updates, + "generation": generation, + "pairs": [list(pair) for pair in plan], + } + ) + if ( + self.completed_updates % 16 == 0 + or self.completed_updates == self.args.updates + ): + print( + json.dumps( + { + "update": self.completed_updates, + "budget": self.args.updates, + "generation": generation, + "generated_episodes": len(self.registry), + "loss": float(loss), + } + ), + flush=True, + ) + + def retire(self) -> None: + if self.block is None: + return + snapshot = self.block.snapshot() + if ( + snapshot["reserved_chunks"] + or snapshot["delivered_chunks"] != snapshot["trained_chunks"] + ): + raise ValueError("Retirement would discard an unfinished training batch") + self.retired_ledgers.append(snapshot) + for row in self.block.records: + self.metrics.retire_episode( + ":".join(map(str, row["source_id"])), retired_ns=self.now() + ) + self.block = None + + def install(self, episodes: list[Any], *, record: bool = True) -> None: + from .training_data import RetainedBlock + + self.block = RetainedBlock( + episodes, + self.metadata, + seed=self.args.seed + 10000 + self.generation, + episode_budget=self.args.episode_budget, + history_length=2, + prediction_horizon=self.args.prediction_horizon, + ) + if record: + self.stage(episodes) + self.check_retention() + self.retained_blocks.append( + { + "generation": self.generation, + "ordinals": [row["ordinal"] for row in self.block.records], + "first_update": self.completed_updates + 1, + } + ) + + def receive_publications(self, generation: int, count: int) -> list[Any]: + """Copy a complete producer subbatch only when owned slots are free.""" + from .acquisition import poll_generation_episodes + + if ( + self.block is not None + or self.check_retention() + count > self.args.buffer_size + ): + return [] + copied = self.now() + episodes = poll_generation_episodes( + self.engine, generation, len(self.registry), max_count=count + ) + if not episodes: + return [] + if len(episodes) != count: + raise RuntimeError( + "Producer publication was not an aligned complete subbatch" + ) + self.metrics.event( + "publication_handoff", + self.now(), + generation=generation, + count=count, + generation_complete=self.engine.get_flow_status()["completed_generation"] + >= generation, + owned_slots_before=self.staging_slots(), + ) + return self.register_episodes(generation, episodes, copied) + + def wait_publications(self, generation: int, count: int) -> list[Any]: + """Wait for free retention slots and publication, propagating failures.""" + started = self.now() + waiting_kind = "publication_wait" + deadline = min( + self.deadline_ns, started + int(self.args.generation_timeout * 1e9) + ) + while self.now() < deadline: + self.flow.poll() + kind = ( + "recording_capacity_wait" + if self.check_retention() + count > self.args.buffer_size + else "publication_wait" + ) + if kind != waiting_kind: + self.metrics.interval( + waiting_kind, started, self.now(), generation=generation + ) + started, waiting_kind = self.now(), kind + episodes = self.receive_publications(generation, count) + if episodes: + self.metrics.interval( + waiting_kind, started, self.now(), generation=generation + ) + return episodes + time.sleep(0.005) + raise TimeoutError( + "Published episode handoff or bounded recording staging timed out" + ) + + def drain_publications(self) -> None: + """Persist every accepted source of the final already-started generation.""" + target = self.engine.get_flow_status()["requested_generation"] + while len(self.registry) < target * self.args.buffer_size: + episodes = self.wait_publications(target, self.args.producer_envs) + self.stage(episodes) + del episodes + self.flow.wait_for_generation(target) + self.flush_staging() + + def collect_publications(self) -> None: + """Stream archive submission during the cold offline acquisition phase.""" + self.initial_episodes = None + while len(self.registry) < self.args.offline_episodes: + status = self.engine.get_flow_status() + target = status["requested_generation"] + if len(self.registry) == target * self.args.buffer_size: + self.flow.wait_for_generation(target) + target = self.flow.request_refill() + episodes = self.wait_publications(target, self.args.producer_envs) + self.stage(episodes) + del episodes + self.drain_publications() + + def offline(self) -> None: + from .lerobot_io import TrainingLeRobotReader, read_lerobot_manifest + from .training_data import build_offline_plan, unpack_geometry + + episodes, self.initial_episodes = self.initial_episodes, None + if self.args.handoff_mode == "publication": + del episodes + self.collect_publications() + while self.args.handoff_mode == "generation": + # Old staging owns its snapshot, so recording can overlap the next + # shared-cache fill even in the collect-then-train baseline. + target = ( + self.flow.request_refill() + if len(self.registry) < self.args.offline_episodes + else None + ) + self.flush_staging() + del episodes + if target is None: + break + episodes = self.take(target) + self.stage(episodes) + started = self.now() + self.writer_stats = self.writer.finalize( + expected_episode_count=len(self.registry), progress=self.flow.poll + ) + self.writer = None + self.metrics.interval("recording_finalize", started, self.now()) + self.record_writer_intervals() + self.stop_engine() + manifest = read_lerobot_manifest(self.output / "lerobot_dataset") + plans = build_offline_plan( + manifest, + self.args.updates, + self.args.batch_size, + self.args.seed + 10000, + history_length=2, + prediction_horizon=self.args.prediction_horizon, + ) + self.reader = TrainingLeRobotReader( + self.output / "lerobot_dataset", + plans, + history_length=2, + prediction_horizon=self.args.prediction_horizon, + num_workers=4, + prefetch_factor=1, + timeout=self.args.data_timeout, + seed=self.args.seed + 10000, + ) + for plan in plans: + self.poll() + started = self.now() + batch = self.reader.next_batch() + batch["geometry"] = unpack_geometry(batch["state"], self.metadata) + self.metrics.interval("batch_service", started, self.now()) + worker_spans = batch["_service"]["worker_fetch_intervals"] + if sum(row["item_count"] for row in worker_spans) != self.args.batch_size: + raise ValueError("SDK worker service does not cover the whole batch") + for row in worker_spans: + if ( + type(row["worker_pid"]) is not int + or row["worker_pid"] <= 0 + or type(row["item_count"]) is not int + or row["item_count"] <= 0 + or not self.origin_ns + <= row["start_ns"] + <= row["end_ns"] + <= self.now() + ): + raise ValueError("Invalid prefetched SDK worker service interval") + self.metrics.interval( + "sdk_decode", + row["start_ns"], + row["end_ns"], + worker_pid=row["worker_pid"], + item_count=row["item_count"], + batch_update=self.completed_updates + 1, + scope=row["scope"], + ) + self.update(batch, plan, requested_ns=started, generation=0) + del batch + delivered = self.reader.snapshot_budget() + trained = Counter( + pair[0] for row in self.trace_updates for pair in row["pairs"] + ) + self.retired_ledgers.append( + { + "episode_budget": None, + "episodes": [ + { + **row, + "ordinal": int(row["source_id"][1]), + "trained": trained[int(row["source_id"][1])], + } + for row in delivered["episodes"] + ], + } + ) + self.reader_stats = self.reader.stats + self.cleanup["reader"] = self.reader.close() + self.reader = None + + def prefetch(self) -> int | None: + """Decide at a batch boundary, before reserving the next training batch.""" + sampler = self.block.sampler + credits = sum( + self.args.episode_budget - sampler.delivered[key] - sampler.reserved[key] + for key in sampler.delivered + ) + recording_wait = 0.0 + if self.staging is not None: + stats = self.writer.stats + backlog = len(self.registry) - stats["written_episodes"] + must_finish = max(0, backlog - stats["queue_capacity"] - 1) + if must_finish: + recording_wait = ( + must_finish + * stats["sdk_write_encode_seconds"] + / stats["written_episodes"] + if stats["written_episodes"] + else None + ) + handoff = ( + max(max(self.generation_seconds[-3:]), recording_wait) + + max(self.snapshot_seconds[-3:]) + if self.generation_seconds + and self.snapshot_seconds + and recording_wait is not None + else None + ) + decision = _prefetch_decision( + remaining_updates=self.args.updates - self.completed_updates, + batch_size=self.args.batch_size, + remaining_credits=credits, + handoff_seconds=handoff, + update_seconds=( + min(self.ready_update_seconds[-8:]) + if len(self.ready_update_seconds) >= 3 + else None + ), + policy=self.args.prefetch_policy, + ) + signature = self.generation, decision["reason"], decision["request"] + if signature != self.last_prefetch_decision: + self.metrics.event( + "prefetch_decision", self.now(), generation=self.generation, **decision + ) + self.last_prefetch_decision = signature + return self.flow.request_refill() if decision["request"] else None + + def overlap_publications(self) -> None: + """Train as producer subbatches arrive, retaining at most C raw sources.""" + episodes, self.initial_episodes = self.initial_episodes, None + self.install(episodes, record=False) + del episodes + in_block = 0 + waiting_since: int | None = None + waiting_origin: int | None = None + waiting_kind = "publication_wait" + while self.completed_updates < self.args.updates: + self.flow.poll() + status = self.engine.get_flow_status() + target = status["requested_generation"] + unread = target * self.args.buffer_size - len(self.registry) + if not unread and status["completed_generation"] >= target: + requested = self.prefetch() + if requested is not None: + target = requested + unread = target * self.args.buffer_size - len(self.registry) + count = self.args.producer_envs + available = status["successful_episodes"] >= len(self.registry) + count + if ( + unread + and (self.block is None or in_block and available) + and self.staging_slots() + count <= self.args.buffer_size + ): + # Accepted count is only a readiness hint. The acquisition + # function validates actual publication and clones under lock. + self.retire() + episodes = self.receive_publications(target, count) + if episodes: + self.install(episodes) + del episodes + in_block = 0 + if self.block is not None: + requested_ns = self.now() + try: + plan = self.block.next_plan(self.args.batch_size) + except StopIteration: + plan = None + if plan is not None: + if waiting_since is not None: + self.metrics.interval( + waiting_kind, + waiting_since, + requested_ns, + generation=target, + ) + waiting_since = None + waiting_origin = None + batch = self.block.materialize(plan) + self.metrics.interval("batch_service", requested_ns, self.now()) + self.update( + batch, + plan, + requested_ns=requested_ns, + generation=self.generation, + ) + in_block += 1 + del batch + continue + kind = ( + "recording_capacity_wait" + if self.staging_slots() + count > self.args.buffer_size + else "publication_wait" + ) + if waiting_since is None: + waiting_since = waiting_origin = self.now() + waiting_kind = kind + elif kind != waiting_kind: + self.metrics.interval( + waiting_kind, waiting_since, self.now(), generation=target + ) + waiting_since, waiting_kind = self.now(), kind + if self.now() - waiting_origin > self.args.generation_timeout * 1e9: + raise TimeoutError( + "Publication or recorder backpressure exceeded its bound" + ) + time.sleep(0.005) + self.retire() + self.drain_publications() + + def overlap(self) -> None: + if self.args.handoff_mode == "publication": + return self.overlap_publications() + episodes, self.initial_episodes = self.initial_episodes, None + self.install(episodes, record=False) + del episodes + in_block = 0 + pending = self.prefetch() + while self.completed_updates < self.args.updates: + self.flow.poll() + if pending is None: + pending = self.prefetch() + ready = ( + pending is not None + and self.engine.get_flow_status()["completed_generation"] >= pending + ) + if in_block and ready: + self.flush_staging() + self.retire() + self.install(self.take(pending)) + in_block = 0 + pending = self.prefetch() + requested = self.now() + try: + plan = self.block.next_plan(self.args.batch_size) + except StopIteration: + if pending is None: + pending = self.flow.request_refill() + wait_started = self.now() + self.flow.wait_for_generation(pending) + self.metrics.interval( + "quota_wait", wait_started, self.now(), generation=pending + ) + continue + batch = self.block.materialize(plan) + self.metrics.interval("batch_service", requested, self.now()) + self.update(batch, plan, requested_ns=requested, generation=self.generation) + in_block += 1 + del batch + self.flush_staging() + self.retire() + if pending is not None: + # A whole already-requested generation is drained and recorded, + # even when none of its frames entered an optimizer update. + tail = self.take(pending, drain=True) + self.stage(tail) + self.flush_staging() + del tail + + def serial(self) -> None: + if self.serial_trace.get("retained_blocks"): + return self.serial_blocks() + episodes, self.initial_episodes = self.initial_episodes, None + self.install(episodes, record=False) + del episodes + for row in self.serial_trace["updates"]: + target = row["generation"] + if target != self.generation: + if target != self.generation + 1: + raise ValueError("Serial trace skips a production generation") + # The previous optimizer work has synchronized. Block only + # generation/training overlap, not generation/recording overlap. + self.flow.request_refill() + self.flush_staging() + self.retire() + self.install(self.take(target)) + self.flow.poll() + requested = self.now() + plan = [tuple(pair) for pair in row["pairs"]] + batch = self.block.materialize(plan, reserve=True) + self.metrics.interval("batch_service", requested, self.now()) + self.update(batch, plan, requested_ns=requested, generation=target) + del batch + while self.generation < self.serial_trace["generated_generations"]: + target = self.flow.request_refill() + self.flush_staging() + self.retire() + tail = self.take(target, drain=True) + self.stage(tail) + del tail + self.flush_staging() + self.retire() + expected = [ + (row["length"], row["numeric_sha256"]) + for row in self.serial_trace["episodes"] + ] + actual = [(row["length"], row["numeric_sha256"]) for row in self.registry] + if expected != actual or self.serial_trace["updates"] != self.trace_updates: + raise ValueError( + "Serial attribution did not reproduce the overlap data/work trace" + ) + + def serial_blocks(self) -> None: + """Replay exact retained-block boundaries with no generation during updates.""" + episodes, self.initial_episodes = self.initial_episodes, None + self.install(episodes, record=False) + del episodes + blocks = { + row["first_update"]: row for row in self.serial_trace["retained_blocks"] + } + for row in self.serial_trace["updates"]: + spec = blocks.get(row["update"]) + if spec is not None and row["update"] != 1: + self.retire() + target = spec["generation"] + current = self.engine.get_flow_status()["requested_generation"] + if target != current: + if ( + target != current + 1 + or len(self.registry) != current * self.args.buffer_size + ): + raise ValueError("Serial blocks skip unread source data") + if self.flow.request_refill() != target: + raise ValueError("Serial refill target differs from trace") + if self.args.handoff_mode == "publication": + # Archiving remains concurrent with generation. Hold only + # the bounded raw sources needed for later traced updates. + while len(self.registry) < target * self.args.buffer_size: + accepted = self.wait_publications( + target, self.args.producer_envs + ) + for raw in accepted: + ordinal = int(raw["ods_episode_id"][0]) + if ordinal in self.ready_episodes: + raise ValueError("Duplicate serial retained source") + self.ready_episodes[ordinal] = raw + del raw + self.stage(accepted) + del accepted + self.flow.wait_for_generation(target) + episodes = [ + self.ready_episodes.pop(ordinal) for ordinal in spec["ordinals"] + ] + self.install(episodes, record=False) + else: + self.flow.wait_for_generation(target) + self.flush_staging() + episodes = self.wait_publications(target, len(spec["ordinals"])) + self.install(episodes) + del episodes + if spec is not None and self.retained_blocks[-1] != spec: + raise ValueError("Serial retained source boundaries differ") + self.flow.poll() + requested = self.now() + plan = [tuple(pair) for pair in row["pairs"]] + batch = self.block.materialize(plan, reserve=True) + self.metrics.interval("batch_service", requested, self.now()) + self.update( + batch, plan, requested_ns=requested, generation=row["generation"] + ) + del batch + self.retire() + self.ready_episodes.clear() + self.drain_publications() + while self.generation < self.serial_trace["generated_generations"]: + self.flow.request_refill() + self.drain_publications() + expected = [ + (r["length"], r["numeric_sha256"]) for r in self.serial_trace["episodes"] + ] + actual = [(r["length"], r["numeric_sha256"]) for r in self.registry] + if expected != actual or self.serial_trace["updates"] != self.trace_updates: + raise ValueError("Serial data or update trace differs from overlap") + + def stop_engine(self) -> None: + if self.engine is None: + return + started = self.now() + self.engine_stats = self.engine.get_stats() + process = self.engine._sim_process + # Engine.stop() closes the multiprocessing.Process handle. Retain its + # joined Popen status so native aborts are still distinguishable from 0. + popen = None if process is None else process._popen + self.engine.stop() + self.flow.drain_events() + exitcode = None if popen is None else popen.returncode + self.cleanup["producer"] = { + "joined": self.engine._sim_process is None and exitcode is not None, + "exitcode": exitcode, + } + if exitcode != 0: + raise RuntimeError( + f"Producer shutdown was not clean: {self.cleanup['producer']}" + ) + self.engine = self.flow = None + self.metrics.interval("producer_shutdown", started, self.now()) + + def execute(self) -> dict[str, Any]: + import numpy as np + import torch + + from embodichain.data_pipeline.runtime import ( + ODSFlowController, + ODSFlowControllerCfg, + ) + from . import acquisition + from .metrics import write_json + from .payload import describe_raw_buffer + from .run import _environment, _schema + + args = self.args + torch.set_num_threads(1) + torch.set_default_dtype(torch.float32) + random.seed(args.seed) + np.random.seed(args.seed) + torch.manual_seed(args.seed) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = True + torch.cuda.set_device(args.train_device) + write_json(self.output / "config.json", _config(args)) + write_json(self.output / "environment.json", _environment()) + self.monitor.start() + self.engine = acquisition.make_engine( + args.vision, + args.sim_device, + producer_envs=args.producer_envs, + seed=args.seed, + buffer_size=args.buffer_size, + budget=None, + max_episode_steps=args.max_episode_steps, + initialization_timeout=args.initialization_timeout, + camera_profile=getattr(args, "camera_profile", "benchmark"), + ) + self.source_run_id = self.engine.run_id + self.cache_bytes = sum( + value.numel() * value.element_size() + for value in self.engine.shared_buffer.values( + include_nested=True, leaves_only=True + ) + ) + self.metadata = describe_raw_buffer(self.engine.shared_buffer) + self.metadata.update(action_schema=_schema()["action"], step_dt=1 / 30) + self.metadata["action_parameters"] = self.engine.cfg.gym_config["env"][ + "actions" + ] + write_json(self.output / "payload_schema.json", self.metadata) + self.flow = ODSFlowController( + self.engine, + ODSFlowControllerCfg(generation_timeout=args.generation_timeout), + progress=self.poll, + on_events=self.events, + on_refill=self.requested, + ) + self.generation_starts[1] = self.now() + self.engine.start() + write_json( + self.output / "roles.json", + {"learner_pid": os.getpid(), "producer_pid": self.engine._sim_process.pid}, + ) + self.flow.drain_events() + self.initial_episodes = self.take(1) + self.bootstrap(self.initial_episodes) + if args.arm == "offline": + self.offline() + elif args.arm == "ods_overlap": + self.overlap() + else: + self.serial() + if self.retained_probe_hash != _digest(self.retained_probe): + raise ValueError( + "A retained delivered batch changed after source replacement" + ) + self.retained_probe = None + if self.writer is not None: + started = self.now() + self.writer_stats = self.writer.finalize( + expected_episode_count=len(self.registry), progress=self.flow.poll + ) + self.writer = None + self.metrics.interval("recording_finalize", started, self.now()) + self.record_writer_intervals() + self.stop_engine() + if self.recording_log_sha256 is None: + raise ValueError("Recording SDK service intervals were not verified") + if self.completed_updates != args.updates: + raise ValueError("The actual optimizer budget is incomplete") + final_hash = _digest(self.model.state_dict()) + if final_hash == self.initial_model_hash: + raise ValueError("Real optimizer work did not update model parameters") + self.save_checkpoint("final.pt", self.completed_updates) + recipe = { + key: getattr(args, key) + for key in ( + "vision", + "seed", + "updates", + "batch_size", + "buffer_size", + "producer_envs", + "episode_budget", + "prediction_horizon", + ) + } + recipe.update( + handoff_mode=args.handoff_mode, encoder_threads=args.encoder_threads + ) + write_json( + self.output / "training-trace.json", + { + "format": "ods-visual-training-trace-v2", + "recipe": recipe, + "initial_model_sha256": self.initial_model_hash, + "normalization_sha256": self.normalization_hash, + "generated_generations": self.generation, + "episodes": self.registry, + "updates": self.trace_updates, + "retained_blocks": self.retained_blocks, + }, + ) + ledger_rows = [ + row for snapshot in self.retired_ledgers for row in snapshot["episodes"] + ] + observed = {tuple(row["source_id"]) for row in ledger_rows} + for source in self.registry: + if tuple(source["source_id"]) not in observed: + ledger_rows.append( + { + "source_id": source["source_id"], + "ordinal": source["ordinal"], + "reserved": 0, + "delivered": 0, + "trained": 0, + } + ) + if sum( + row["trained"] for row in ledger_rows + ) != args.updates * args.batch_size or any( + row["reserved"] + or row["delivered"] != row["trained"] + or (args.arm != "offline" and row["delivered"] > args.episode_budget) + for row in ledger_rows + ): + raise ValueError("Lifetime source ledger differs from completed training") + self.resources = self.monitor.close() + gpu_activity = self.resources.get("primary_gpu_utilization", {}).values() + sampled_roles = { + role + for device in gpu_activity + if device["sample_count"] > 0 + for role in device["roles"] + } + if ( + self.resources.get("observer_complete") is not True + or self.resources.get("cpu_utilization", {}).get("valid_interval_count", 0) + < 1 + or not {"simulation", "training"}.issubset(sampled_roles) + ): + raise RuntimeError("Timed CPU/GPU resource observations are incomplete") + dataset_bytes = sum( + path.stat().st_size + for path in (self.output / "lerobot_dataset").rglob("*") + if path.is_file() + ) + checkpoint_bytes = sum( + (self.output / name).stat().st_size for name in ("initial.pt", "final.pt") + ) + total_seconds = (self.now() - self.origin_ns) / 1e9 + result = self.metrics.summary( + total_seconds=total_seconds, + training_complete_seconds=(self.training_completed_ns - self.origin_ns) + / 1e9, + ledger={"episodes": ledger_rows}, + ) + publication_end = { + generation: max( + row["published_ns"] + for row in self.registry + if row["generation"] == generation + ) + for generation in self.recorded_generations + } + early_updates = sum( + row["generation"] > 1 + and completed < publication_end.get(row["generation"], 0) + for row, completed in zip(self.trace_updates, self.update_completion_ns) + ) + result.update( + timeline_contract="sdk-worker-service/v1", + writer_sdk_events_sha256=self.recording_log_sha256, + dataset_bytes=dataset_bytes, + checkpoint_bytes=checkpoint_bytes, + storage_scope="logical finalized dataset and initial/final checkpoint file sizes; excludes logs, temporary-file peaks and post-timing audit/evaluation artifacts", + generated_episodes=len(self.registry), + gpu_step_milliseconds=self.gpu_step_ms, + writer=self.writer_stats, + resources=self.resources, + initial_model_sha256=self.initial_model_hash, + final_model_sha256=final_hash, + normalization_sha256=self.normalization_hash, + cache_bytes=self.cache_bytes, + peak_owned_raw_episodes=self.peak_owned_episodes, + owned_raw_episode_limit=args.buffer_size, + updates_before_generation_fully_published=early_updates, + engine_stats=self.engine_stats, + ) + self.metrics.close() + self.closed = True + return { + "status": "complete", + "qualification_passed": True, + "config": _config(args), + "training_metrics": result, + "cleanup": self.cleanup, + } + + def failure(self, error: BaseException) -> dict[str, Any]: + errors = [] + for label, operation in ( + ("reader", lambda: self.reader.close() if self.reader else None), + ("writer", lambda: self.writer.abort() if self.writer else None), + ("producer", lambda: self.engine.stop() if self.engine else None), + ): + try: + operation() + except BaseException as cleanup_error: + errors.append(f"{label}: {cleanup_error!r}") + resources = self.monitor.close() + now = self.now() + partial = self.metrics.summary( + total_seconds=(now - self.origin_ns) / 1e9, + training_complete_seconds=(self.training_completed_ns - self.origin_ns) + / 1e9, + ) + partial["resources"] = resources + self.metrics.close() + self.closed = True + return { + "status": "failed", + "qualification_passed": False, + "config": _config(self.args), + "error": repr(error), + "traceback": traceback.format_exc(), + "cleanup_errors": errors, + "training_metrics": partial, + } + + +def run_one(args: argparse.Namespace) -> dict[str, Any]: + """Run one real finite learner pipeline and preserve failures as evidence.""" + from .metrics import write_json + + _validate(args) + experiment = _Experiment(args) + try: + result = experiment.execute() + except BaseException as error: + result = experiment.failure(error) + write_json(args.output / "summary.json", result) + return result + + +def main() -> None: + """Launch a qualified two-device suite or its private finite worker.""" + args = build_parser().parse_args() + if args.worker: + result = run_one(args) + else: + from .training_suite import run_suite + + result = run_suite(args) + if result.get("status") not in ("complete", "passed"): + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/examples/ods/data_stream_demo/training_audit.py b/examples/ods/data_stream_demo/training_audit.py new file mode 100644 index 000000000..31c30e272 --- /dev/null +++ b/examples/ods/data_stream_demo/training_audit.py @@ -0,0 +1,506 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Post-timing CPU qualification of saved visual training and standard videos.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import os +import time +import traceback +from collections import Counter +from pathlib import Path +from typing import Any + +import torch + +from .lerobot_io import ( + TrainingLeRobotReader, + audit_video_quality, + read_lerobot_manifest, +) +from .metrics import write_json +from .payload import numeric_episode_digest + +__all__ = ["audit_run"] + + +def _require(condition: bool, message: str) -> None: + if not condition: + raise ValueError(message) + + +def _digest(values: dict[str, Any]) -> str: + digest = hashlib.sha256() + for key, value in sorted(values.items()): + digest.update(key.encode()) + if isinstance(value, torch.Tensor): + value = value.detach().cpu().contiguous() + digest.update(str((tuple(value.shape), value.dtype)).encode()) + digest.update(memoryview(value.numpy())) + else: + digest.update(json.dumps(value, sort_keys=True).encode()) + return digest.hexdigest() + + +def _clean_shutdown(result: dict[str, Any]) -> None: + _require( + result.get("workers_joined") is True + and result.get("forced_termination") is False + and bool(result.get("worker_exitcodes")) + and all(code == 0 for code in result["worker_exitcodes"]), + f"Decoded SDK workers did not exit normally: {result}", + ) + + +def _audit(output: Path, result: dict[str, Any]) -> None: + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + summary = json.loads((output / "summary.json").read_text()) + config = json.loads((output / "config.json").read_text()) + trace = json.loads((output / "training-trace.json").read_text()) + contract = json.loads((output / "initial-contract.json").read_text()) + schema = json.loads((output / "payload_schema.json").read_text()) + _require( + summary.get("status") == "complete" + and summary.get("qualification_passed") is True, + "Training worker did not complete qualification", + ) + _require( + trace.get("format") + in {"ods-visual-training-trace-v1", "ods-visual-training-trace-v2"}, + "Unknown training trace format", + ) + _require( + all(config.get(key) == value for key, value in trace["recipe"].items()), + "Training trace recipe differs from the saved configuration", + ) + metrics, registry = summary["training_metrics"], trace["episodes"] + root = output / "lerobot_dataset" + manifest = read_lerobot_manifest(root) + episodes, payload = manifest["episodes"], manifest["payload"] + writer = metrics["writer"] + _require( + manifest.get("video", {}).get("encoder_threads") + == config.get("encoder_threads", 1), + "Recorded encoder thread setting differs from the run configuration", + ) + _require( + len(episodes) + == len(registry) + == writer["written_episodes"] + == writer["submitted_episodes"] + == metrics["generated_episodes"], + "Accepted, recorded and registered episode counts differ", + ) + _require( + writer.get("finalized") is True and writer.get("worker_joined") is True, + "Recording did not finalize and join", + ) + _require( + manifest.get("accepted_episode_count", len(episodes)) == len(episodes), + "Finalized manifest count differs", + ) + _require( + payload == schema and list(payload["cameras"]) == contract["camera_order"], + "Saved payload/camera order differs from the training contract", + ) + cameras = list(payload["cameras"]) + _require( + len(cameras) == (1 if config["vision"] == "single" else 2), + "Saved camera count differs from the configured visual condition", + ) + _require( + config["history_length"] == 2, "Visual audit requires two causal image frames" + ) + _require( + math.isclose(payload["step_dt"], 1 / manifest["fps"], rel_tol=0, abs_tol=1e-12), + "Saved step_dt and observation FPS differ", + ) + _require( + payload["action_schema"]["semantics"] == "policy_request" + and math.isclose( + payload["action_schema"]["step_dt"], + payload["step_dt"], + rel_tol=0, + abs_tol=1e-12, + ), + "Saved action semantics/control period differ", + ) + dataset = LeRobotDataset( + "local/ods-data-stream", root=root, video_backend="torchcodec", tolerance_s=1e-4 + ) + info = dataset.meta.info + _require( + info["total_episodes"] == len(episodes) + and info["total_frames"] == sum(row["length"] for row in episodes) + and info["total_frames"] == writer["valid_frames"] + and info["fps"] == manifest["fps"], + "Official SDK counts/FPS differ from the finalized sidecar", + ) + features = info["features"] + _require( + list(features["observation.state"]["shape"]) == [payload["state_width"]] + and list(features["action"]["shape"]) == [payload["action_width"]] + and features["observation.state"]["dtype"] == "float32" + and features["action"]["dtype"] == "float32", + "Official SDK numeric feature widths differ", + ) + names = [ + f"{'/'.join(feature['path'])}/{index}" + for feature in payload["state_features"] + for index in range(feature["width"]) + ] + _require( + features["observation.state"]["names"] == names + and features["action"]["names"] + == [f"action/{i}" for i in range(payload["action_width"])], + "Official SDK numeric feature order differs", + ) + saved_cameras = [ + key.removeprefix("observation.images.") + for key, value in features.items() + if value["dtype"] == "video" + ] + _require(saved_cameras == cameras, "Official SDK camera order differs") + for camera in cameras: + feature = features[f"observation.images.{camera}"] + _require( + list(feature["shape"]) == [*payload["cameras"][camera]["shape"][:2], 3] + and feature["info"]["video.codec"] == "av1" + and feature["info"]["video.fps"] == manifest["fps"], + "Recorded camera shape/codec/FPS differs", + ) + numeric = dataset.hf_dataset.select_columns( + ["observation.state", "action", "ods.provenance", "episode_index"] + ).with_format("numpy") + offset = 0 + source_by_ordinal = {} + for ordinal, (saved, registered) in enumerate(zip(episodes, registry)): + _require( + saved["index"] == registered["ordinal"] == ordinal + and saved["offset"] == offset, + "Saved episode order/offset differs from registration", + ) + for key in ("length", "source_id", "generation", "numeric_sha256"): + _require( + saved[key] == registered[key], + f"Saved {key} differs for episode {ordinal}", + ) + values = numeric[offset : offset + saved["length"]] + state = torch.as_tensor(values["observation.state"], dtype=torch.float32) + actions = torch.as_tensor(values["action"], dtype=torch.float32) + provenance = torch.as_tensor(values["ods.provenance"], dtype=torch.int64) + digest = numeric_episode_digest( + {"state": state, "actions": actions, "continuity_id": provenance[:, 4]} + ) + _require( + digest == saved["numeric_sha256"], + f"Actual Parquet numeric payload differs for episode {ordinal}", + ) + _require( + torch.as_tensor(values["episode_index"]).eq(ordinal).all().item() + and provenance[:, :2] + .eq(torch.tensor(saved["source_id"], dtype=torch.int64)) + .all() + .item() + and provenance[:, 2].eq(torch.arange(saved["length"])).all().item(), + "Actual saved episode/source/frame order differs", + ) + _require( + provenance[:, 3].eq(saved["generation"]).all().item() + and provenance[:, 5].eq(registered["published_ns"]).all().item(), + "Actual saved generation/publication differs", + ) + source_by_ordinal[ordinal] = ":".join(map(str, saved["source_id"])) + offset += saved["length"] + result["recording"] = { + "episodes": len(episodes), + "frames": offset, + "cameras": cameras, + "all_numeric_sha256_verified": True, + "sdk_metadata_verified": True, + } + updates, batch_size, future = ( + config["updates"], + config["batch_size"], + config["prediction_horizon"], + ) + _require( + len(trace["updates"]) == updates == metrics["completed_updates"], + "Optimizer update budget is incomplete", + ) + _require( + metrics["trained_chunks"] == updates * batch_size + and metrics["valid_target_occurrences"] == updates * batch_size * future, + "Actual training chunks/targets differ from the fixed budget", + ) + expected_counts, expected_anchors = Counter(), {} + for number, update in enumerate(trace["updates"], 1): + _require( + update["update"] == number and len(update["pairs"]) == batch_size, + "Training trace update/batch order differs", + ) + for ordinal, anchor in update["pairs"]: + _require( + type(ordinal) is int + and ordinal in source_by_ordinal + and type(anchor) is int + and 1 <= anchor <= episodes[ordinal]["length"] - future, + "Training trace has an illegal source/anchor", + ) + _require( + update["generation"] in (0, episodes[ordinal]["generation"]), + "Training trace generation differs", + ) + source = source_by_ordinal[ordinal] + expected_counts[source] += 1 + expected_anchors.setdefault(source, Counter())[anchor] += 1 + rows = {row["source_id"]: row for row in metrics["episodes"]} + _require( + set(rows) == set(source_by_ordinal.values()) + and len(rows) == len(metrics["episodes"]), + "Training metrics omitted or duplicated accepted sources", + ) + for ordinal, saved in enumerate(episodes): + source, length = source_by_ordinal[ordinal], saved["length"] + row = rows[source] + _require( + row["ordinal"] == ordinal + and row["frame_count"] == length + and row["trained_chunks"] == expected_counts[source], + "Per-source trained counts differ from the trace", + ) + anchors = expected_anchors.get(source, Counter()) + history, targets = Counter(), Counter() + for anchor, count in anchors.items(): + history.update({frame: count for frame in range(anchor - 1, anchor + 1)}) + targets.update({frame: count for frame in range(anchor, anchor + future)}) + for key, counts in ( + ("anchor_counts", anchors), + ("history_frame_counts", history), + ("target_frame_counts", targets), + ): + _require( + row[key] == [counts[frame] for frame in range(length)], + f"Actual {key} differs from completed training trace", + ) + ledger = metrics.get("ledger", {}).get("episodes", []) + seen = set() + for row in ledger: + source = ":".join(map(str, row["source_id"])) + _require( + source not in seen + and source in rows + and row["reserved"] == 0 + and row["delivered"] == row["trained"] == expected_counts[source], + "Retired delivery/training ledger differs from the trace", + ) + seen.add(source) + if config["arm"] != "offline": + _require( + all(source in seen for source, count in expected_counts.items() if count), + "Protected source ledger omitted trained sources", + ) + _require( + all(count <= config["episode_budget"] for count in expected_counts.values()) + and all(row["delivered"] <= config["episode_budget"] for row in ledger), + "Protected delivery/training cap exceeded", + ) + checkpoint = torch.load(output / "final.pt", map_location="cpu", weights_only=False) + _require( + checkpoint["step"] == updates and list(checkpoint["cameras"]) == cameras, + "Final checkpoint step/cameras differ", + ) + _require( + _digest(checkpoint["model"]) + == metrics["final_model_sha256"] + != metrics["initial_model_sha256"], + "Final model digest differs or parameters never updated", + ) + _require( + metrics["initial_model_sha256"] + == trace["initial_model_sha256"] + == contract["model_sha256"], + "Initial model contract differs from the training trace", + ) + _require( + _digest(checkpoint["normalization"]) + == metrics["normalization_sha256"] + == trace["normalization_sha256"] + == contract["normalization_sha256"], + "Frozen normalization differs", + ) + optimizer_steps = [ + int(state["step"]) for state in checkpoint["optimizer"]["state"].values() + ] + _require( + bool(optimizer_steps) and all(step == updates for step in optimizer_steps), + "Final optimizer state does not contain the completed update budget", + ) + del checkpoint + result["training"] = { + "updates": updates, + "chunks": updates * batch_size, + "valid_targets": updates * batch_size * future, + "trained_sources": len(expected_counts), + "untrained_saved_sources": len(episodes) - len(expected_counts), + "ledger_sources": len(ledger), + "checkpoint_verified": True, + } + plan = [[(ordinal, 1)] for ordinal in range(len(episodes))] + reader = TrainingLeRobotReader( + root, + plan, + history_length=2, + prediction_horizon=future, + num_workers=4, + prefetch_factor=1, + ) + decoded = 0 + try: + for ordinal in range(len(episodes)): + batch = reader.next_batch() + _require( + batch["valid"].all().item() + and batch["frame_ids"][0, 1, :2].tolist() + == episodes[ordinal]["source_id"], + "Representative decode returned a different source", + ) + for image in batch["images"].values(): + _require( + tuple(image.shape) == (1, 2, 3, 256, 256) + and image.dtype == torch.float32 + and torch.isfinite(image).all().item() + and image.min().item() >= 0 + and image.max().item() <= 1, + "Representative decoded RGB history differs from the training contract", + ) + decoded += 1 + try: + reader.next_batch() + except StopIteration: + pass + else: + raise ValueError("Finite audit decode plan did not exhaust") + result["decode"] = {"episodes": decoded, "stats": reader.stats} + finally: + result["decode_shutdown"] = reader.close() + _clean_shutdown(result["decode_shutdown"]) + quality = audit_video_quality(root) + result["video_quality"] = quality + _clean_shutdown(quality["worker_cleanup"]) + _require( + len(quality["samples"]) == 3 * len(cameras) + and all( + math.isfinite(row["mae_uint8"]) + and (row["psnr_db"] is None or math.isfinite(row["psnr_db"])) + for row in quality["samples"] + ), + "Raw-reference video quality audit is incomplete or nonfinite", + ) + + +def audit_run(output: Path, *, receipt: Path | None = None) -> dict[str, Any]: + """Audit a run after timing; preserve a failed JSON result without using GPUs. + + Args: + output: Completed run directory, containing its own standard corpus, + training trace, summary and checkpoint. + receipt: Fresh result path. Defaults to ``output/post-audit.json``; + explicitly choose another path to preserve an earlier audit. + + Returns: + Qualification, source-window SDK checks and full video-frame integrity. + + Raises: + FileExistsError: The receipt or its full-video artifacts already exist. + """ + from . import video_integrity + + output = Path(output).resolve() + receipt = output / "post-audit.json" if receipt is None else Path(receipt).resolve() + video_receipt = receipt.with_name(receipt.stem + "-video-integrity.json") + video_artifacts = video_receipt.parent / (video_receipt.stem + "-artifacts") + if any(path.exists() for path in (receipt, video_receipt, video_artifacts)): + raise FileExistsError("Use a fresh audit receipt and artifact path") + corpus = (output / "lerobot_dataset").resolve() + if receipt == corpus or corpus in receipt.parents: + raise ValueError("Audit receipts must stay outside the immutable corpus") + receipt.parent.mkdir(parents=True, exist_ok=True) + with receipt.open("x", encoding="utf-8") as stream: + json.dump({"passed": False, "status": "pending"}, stream) + stream.write("\n") + started = time.monotonic() + result: dict[str, Any] = { + "schema": "ods-visual-training-post-audit-v2", + "passed": False, + "qualification_passed": False, + "scope": "post-timing numeric/checkpoint, source-window SDK and all-frame video integrity audit; no duration contributes to pipeline time", + "audit_source_sha256": { + "training_audit": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), + "video_integrity": hashlib.sha256( + Path(video_integrity.__file__).read_bytes() + ).hexdigest(), + }, + } + try: + _audit(output, result) + result["full_video_integrity"] = video_integrity.verify_dataset( + output / "lerobot_dataset", video_receipt + ) + _require( + result["full_video_integrity"].get("passed") is True, + "Full recorded video-frame integrity failed", + ) + result.update(passed=True, qualification_passed=True) + except Exception as error: + result.update(error=repr(error), traceback=traceback.format_exc()) + result["seconds_outside_pipeline"] = time.monotonic() - started + write_json(receipt, result) + return result + + +def main() -> None: + """Run CPU qualification and exit unsuccessfully for incomplete evidence.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--run-dir", type=Path, required=True) + parser.add_argument("--output", type=Path) + args = parser.parse_args() + os.environ["CUDA_VISIBLE_DEVICES"] = "" + os.environ["HF_HUB_OFFLINE"] = "1" + os.environ["HF_DATASETS_OFFLINE"] = "1" + torch.set_num_threads(1) + result = audit_run(args.run_dir, receipt=args.output) + print( + json.dumps( + { + "passed": result["passed"], + "error": result.get("error"), + "seconds_outside_pipeline": result["seconds_outside_pipeline"], + } + ), + flush=True, + ) + raise SystemExit(0 if result["passed"] else 1) + + +if __name__ == "__main__": + main() diff --git a/examples/ods/data_stream_demo/training_data.py b/examples/ods/data_stream_demo/training_data.py new file mode 100644 index 000000000..5cdd60b5e --- /dev/null +++ b/examples/ods/data_stream_demo/training_data.py @@ -0,0 +1,478 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Owned visual training blocks, causal geometry and bounded source sampling.""" + +from __future__ import annotations + +import copy +from collections import Counter +from collections.abc import Mapping +from typing import Any + +import torch +from tensordict import TensorDict, TensorDictBase + +from examples.ods.simple_bc_demo.models import pose_targets +from examples.ods.simple_bc_demo.observations import ( + _geometry_statistics, + policy_observations, +) + +from .payload import _numeric, _provenance + +__all__ = [ + "unpack_geometry", + "geometry_from_raw", + "make_statistics", + "BalancedAnchorSampler", + "RetainedBlock", + "build_offline_plan", +] + +Plan = list[tuple[int, int]] + + +def unpack_geometry(state: torch.Tensor, metadata: Mapping[str, Any]) -> torch.Tensor: + """Reconstruct measured geometry from the canonical ordered numeric state.""" + if state.shape[-1] != metadata["state_width"] or state.dtype != torch.float32: + raise ValueError("Canonical float32 state width differs from payload metadata") + obs = TensorDict({}, batch_size=state.shape[:-1], device=state.device) + offset = 0 + for feature in metadata["state_features"]: + width = feature["width"] + obs[tuple(feature["path"])] = state[..., offset : offset + width].reshape( + *state.shape[:-1], *feature["shape"] + ) + offset += width + if offset != state.shape[-1]: + raise ValueError("State feature widths do not cover the canonical state") + return policy_observations(obs) + + +def geometry_from_raw(raw: TensorDictBase) -> torch.Tensor: + """Pack current measured features; optional invalid frames remain neutral.""" + return policy_observations(raw["obs"], valid_mask=raw.get("valid", None)) + + +def make_statistics( + episodes: list[TensorDictBase], metadata: Mapping[str, Any] +) -> dict[str, torch.Tensor]: + """Freeze statistics using only valid physical frames in the initial block.""" + if not episodes: + raise ValueError("Statistics require a nonempty initial block") + geometry, targets = [], [] + for episode in episodes: + values = geometry_from_raw(episode)[episode["valid"]] + if not len(values): + raise ValueError("Statistics cannot include an empty accepted episode") + if episode["actions"].shape[-1] != metadata["action_width"]: + raise ValueError("Recorded action width differs from payload metadata") + geometry.append(values) + targets.append(pose_targets(values, episode["actions"][episode["valid"]])) + values, target = torch.cat(geometry), torch.cat(targets) + mean, std = _geometry_statistics(values) + target_mean = target.mean(0) + target_std = target.std(0, correction=0).clamp_min(0.05) + target_mean[3:], target_std[3:] = 0, 1 + result = { + "proprio_mean": mean[:27].clone(), + "proprio_std": std[:27].clone(), + "target_mean": target_mean, + "target_std": target_std, + } + if not all(torch.isfinite(value).all() for value in result.values()): + raise ValueError("Initial physical frames produced nonfinite statistics") + return result + + +class BalancedAnchorSampler: + """One in-flight batch, balanced sources and per-source anchor permutations. + + Records contain ``ordinal``, ``length``, ``source_id`` and optionally exact + ``legal_anchors``. Permutations are exhausted before reshuffling. Reservations + protect K before delivery; training is acknowledged only by the caller. + """ + + def __init__( + self, + episodes: list[dict[str, Any]], + *, + history_length: int = 2, + prediction_horizon: int = 32, + seed: int = 29001, + episode_budget: int | None = None, + ) -> None: + if not episodes or any( + type(v) is not int or v < 1 for v in (history_length, prediction_horizon) + ): + raise ValueError("Sampling requires episodes and positive H/F") + if episode_budget is not None and ( + type(episode_budget) is not int or episode_budget < 1 + ): + raise ValueError("episode_budget must be a positive integer or None") + self.records = copy.deepcopy(episodes) + self.budget = episode_budget + self.generator = torch.Generator().manual_seed(seed) + self.order = torch.randperm(len(episodes), generator=self.generator).tolist() + self.cursor = 0 + self.anchors: dict[int, list[int]] = {} + self.permutations: dict[int, list[int]] = {} + self.positions: dict[int, int] = {} + for index, episode in enumerate(self.records): + ordinal = episode.setdefault("ordinal", episode.get("index", index)) + if type(ordinal) is not int or ordinal in self.anchors: + raise ValueError("Episode ordinals must be unique integers") + anchors = list( + episode.get( + "legal_anchors", + range( + history_length - 1, episode["length"] - prediction_horizon + 1 + ), + ) + ) + if ( + not anchors + or len(set(anchors)) != len(anchors) + or any( + type(a) is not int + or a < history_length - 1 + or a + prediction_horizon > episode["length"] + for a in anchors + ) + ): + raise ValueError( + "Each source requires distinct complete-window anchors" + ) + self.anchors[ordinal] = anchors + self.permutations[ordinal] = [] + self.positions[ordinal] = 0 + self.delivered = Counter({ordinal: 0 for ordinal in self.anchors}) + self.reserved = self.delivered.copy() + self.trained = self.delivered.copy() + self._pending: Plan | None = None + self._delivered: Plan | None = None + + def reserve_exact(self, plan: Plan) -> Plan: + """Reserve a caller-owned serial trace without changing its anchor order.""" + if self._pending is not None or self._delivered is not None: + raise RuntimeError("A previous batch is still reserved or not trained") + if not plan or any( + ordinal not in self.anchors or anchor not in self.anchors[ordinal] + for ordinal, anchor in plan + ): + raise ValueError("Exact plan contains an unknown source or illegal anchor") + counts = Counter(ordinal for ordinal, _ in plan) + if self.budget is not None and any( + self.delivered[ordinal] + count > self.budget + for ordinal, count in counts.items() + ): + raise StopIteration( + "Insufficient protected source credits for a complete batch" + ) + self.reserved.update(counts) + self._pending = list(plan) + return plan + + def reserve(self, batch_size: int) -> Plan: + """Reserve a balanced complete batch without spending delivery credits.""" + if type(batch_size) is not int or batch_size < 1: + raise ValueError("batch_size must be a positive integer") + if self._pending is not None or self._delivered is not None: + raise RuntimeError("A previous batch is still reserved or not trained") + if ( + self.budget is not None + and sum(self.budget - value for value in self.delivered.values()) + < batch_size + ): + raise StopIteration( + "Insufficient protected source credits for a complete batch" + ) + plan: Plan = [] + counts: Counter[int] = Counter() + while len(plan) < batch_size: + row = self.order[self.cursor % len(self.order)] + self.cursor += 1 + ordinal = self.records[row]["ordinal"] + if ( + self.budget is not None + and self.delivered[ordinal] + counts[ordinal] >= self.budget + ): + continue + if self.positions[ordinal] >= len(self.permutations[ordinal]): + indices = torch.randperm( + len(self.anchors[ordinal]), generator=self.generator + ).tolist() + self.permutations[ordinal] = [ + self.anchors[ordinal][index] for index in indices + ] + self.positions[ordinal] = 0 + anchor = self.permutations[ordinal][self.positions[ordinal]] + self.positions[ordinal] += 1 + plan.append((ordinal, anchor)) + counts[ordinal] += 1 + return self.reserve_exact(plan) + + def cancel(self, plan: Plan) -> None: + """Release a failed materialization without spending delivery credits.""" + if self._pending != plan: + raise RuntimeError("Cancellation differs from the reserved batch") + self.reserved.subtract(ordinal for ordinal, _ in plan) + self._pending = None + + def mark_delivered(self, plan: Plan) -> None: + """Commit delivery only after the entire batch has materialized.""" + if self._pending != plan: + raise RuntimeError("Delivery differs from the reserved batch") + self.reserved.subtract(ordinal for ordinal, _ in plan) + self.delivered.update(ordinal for ordinal, _ in plan) + self._pending, self._delivered = None, list(plan) + + def mark_trained(self, plan: Plan) -> None: + """Acknowledge exactly one successfully completed optimizer update.""" + if self._delivered != plan: + raise RuntimeError("Training acknowledgement differs from delivery") + self.trained.update(ordinal for ordinal, _ in plan) + self._delivered = None + + def snapshot(self) -> dict[str, Any]: + """Return immutable-source delivery/training counters for retirement.""" + return { + "episode_budget": self.budget, + "reserved_chunks": sum(self.reserved.values()), + "delivered_chunks": sum(self.delivered.values()), + "trained_chunks": sum(self.trained.values()), + "episodes": [ + { + "ordinal": row["ordinal"], + "source_id": list(row["source_id"]), + "reserved": self.reserved[row["ordinal"]], + "delivered": self.delivered[row["ordinal"]], + "trained": self.trained[row["ordinal"]], + } + for row in self.records + ], + } + + +class RetainedBlock: + """Own at most one producer block; gather only causal image histories. + + Ownership of the supplied raw episodes transfers to this object. Returned + batches have independent storage and remain valid after the block retires. + """ + + def __init__( + self, + episodes: list[TensorDictBase], + metadata: Mapping[str, Any], + *, + history_length: int = 2, + prediction_horizon: int = 32, + seed: int = 29001, + episode_budget: int | None = 128, + episode_ordinals: list[int] | None = None, + ) -> None: + if not episodes or ( + episode_ordinals is not None and len(episode_ordinals) != len(episodes) + ): + raise ValueError("A retained block requires aligned owned episodes") + self.metadata = copy.deepcopy(dict(metadata)) + self.history_length, self.prediction_horizon = ( + history_length, + prediction_horizon, + ) + self.length = history_length + prediction_horizon - 1 + self.episodes: dict[int, dict[str, Any]] = {} + records = [] + for index, raw in enumerate(episodes): + if raw.batch_dims != 1 or raw["actions"].device.type != "cpu": + raise ValueError("Retained episodes require one CPU time axis") + valid = raw["valid"].bool() + length = int(valid.sum()) + if not length or not torch.equal(valid, torch.arange(len(valid)) < length): + raise ValueError("Accepted episodes require a nonempty valid prefix") + provenance = _provenance(raw) + ids = provenance["frame_ids"][:length] + continuity = provenance["continuity_id"][:length] + if ( + (ids < 0).any() + or (ids[:, :2] != ids[0, :2]).any() + or (ids[1:, 2] != ids[:-1, 2] + 1).any() + or (provenance["generations"][:length] < 1).any() + or ( + provenance["generations"][:length] != provenance["generations"][0] + ).any() + ): + raise ValueError( + "Accepted source frames must be contiguous and immutable" + ) + windows = ( + continuity.unfold(0, self.length, 1) + if length >= self.length + else torch.empty((0, self.length), dtype=torch.int64) + ) + legal = ( + ((windows == windows[:, :1]).all(-1) & (windows >= 0).all(-1)) + .nonzero() + .flatten() + + history_length + - 1 + ) + ordinal = ( + int(ids[0, 1]) if episode_ordinals is None else episode_ordinals[index] + ) + if ordinal in self.episodes: + raise ValueError("Retained episode ordinals must be unique") + record = { + "ordinal": ordinal, + "length": length, + "source_id": ids[0, :2].tolist(), + "generation": int(provenance["generations"][0]), + "legal_anchors": legal.tolist(), + } + records.append(record) + self.episodes[ordinal] = { + "raw": raw, + "state": _numeric(raw, metadata), + "geometry": geometry_from_raw(raw), + **provenance, + } + self.sampler = BalancedAnchorSampler( + records, + history_length=history_length, + prediction_horizon=prediction_horizon, + seed=seed, + episode_budget=episode_budget, + ) + self.records = self.sampler.records + + def next_plan(self, batch_size: int) -> Plan: + """Reserve the next batch; the first batch must cover every source.""" + if not sum(self.sampler.delivered.values()) and ( + batch_size < len(self.records) or batch_size % len(self.records) + ): + raise ValueError("First retained batch must evenly cover every source") + return self.sampler.reserve(batch_size) + + def materialize(self, plan: Plan, *, reserve: bool = False) -> dict[str, Any]: + """Gather independent numeric windows and only H RGB image frames.""" + if reserve: + if not sum(self.sampler.delivered.values()) and ( + Counter(ordinal for ordinal, _ in plan) + != Counter( + { + record["ordinal"]: len(plan) // len(self.records) + for record in self.records + } + ) + ): + raise ValueError("First exact batch must evenly cover every source") + self.sampler.reserve_exact(plan) + if self.sampler._pending != plan: + raise RuntimeError("Materialization requires an exact reserved plan") + try: + numeric: dict[str, list[torch.Tensor]] = { + name: [] + for name in ( + "state", + "actions", + "geometry", + "frame_ids", + "generations", + "continuity_id", + "published_at_monotonic_ns", + ) + } + images = { + name: torch.empty( + (len(plan), self.history_length, 3, *feature["shape"][:2]), + dtype=torch.float32, + ) + for name, feature in self.metadata["cameras"].items() + } + divisor = torch.tensor(255.0, dtype=torch.float32, device="cpu") + for batch, (ordinal, anchor) in enumerate(plan): + episode = self.episodes[ordinal] + first = anchor - self.history_length + 1 + selection = slice(first, anchor + self.prediction_horizon) + for name in numeric: + value = ( + episode["raw"]["actions"] + if name == "actions" + else episode[name] + ) + numeric[name].append(value[selection]) + for name, feature in self.metadata["cameras"].items(): + rgb = episode["raw"][("obs", *feature["path"])][ + first : anchor + 1, ..., :3 + ].movedim(-1, -3) + if ( + rgb.dtype != torch.uint8 + or rgb.shape != images[name][batch].shape + ): + raise ValueError( + "Native camera layout differs from payload metadata" + ) + torch.div(rgb, divisor, out=images[name][batch]) + result = {name: torch.stack(values) for name, values in numeric.items()} + result["images"] = images + result["valid"] = torch.ones((len(plan), self.length), dtype=torch.bool) + result["loss_mask"] = ( + (torch.arange(self.length) >= self.history_length - 1) + .expand(len(plan), -1) + .clone() + ) + self.sampler.mark_delivered(plan) + return result + except BaseException: + self.sampler.cancel(plan) + raise + + def mark_trained(self, plan: Plan) -> None: + """Commit training only after the host observes a completed update.""" + self.sampler.mark_trained(plan) + + def snapshot(self) -> dict[str, Any]: + """Return lifetime quota counters, including retired-source identity.""" + return self.sampler.snapshot() + + +def build_offline_plan( + manifest: Mapping[str, Any], + updates: int, + batch_size: int, + seed: int, + *, + history_length: int = 2, + prediction_horizon: int = 32, +) -> list[Plan]: + """Build finite uncapped plans using the same balanced/permutation rule.""" + if type(updates) is not int or updates < 1: + raise ValueError("updates must be a positive integer") + sampler = BalancedAnchorSampler( + list(manifest["episodes"]), + history_length=history_length, + prediction_horizon=prediction_horizon, + seed=seed, + ) + result = [] + for _ in range(updates): + plan = sampler.reserve(batch_size) + sampler.mark_delivered(plan) + sampler.mark_trained(plan) + result.append(plan) + return result diff --git a/examples/ods/data_stream_demo/training_eval.py b/examples/ods/data_stream_demo/training_eval.py new file mode 100644 index 000000000..dacf4f987 --- /dev/null +++ b/examples/ods/data_stream_demo/training_eval.py @@ -0,0 +1,399 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Independent physical evaluation of a saved visual BC checkpoint. + +Run this entry point in a separate process. The comparison supervisor evaluates +after its timed pipeline; the learning driver also evaluates between updates. +The caller owns total-run timing and reserves/masks both devices explicitly. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +from pathlib import Path +import sys +import time +import traceback +from typing import Any + +import torch + +if __package__ in {None, ""}: + repository = Path(__file__).resolve().parents[3] + sys.path[:0] = [ + str(repository), + str(repository / "extensions"), + str(repository / "extensions/embodichain_tasks"), + ] + +from examples.ods.data_stream_demo.acquisition import ( + camera_profile_version, + camera_uids, + load_acquisition_config, +) +from examples.ods.data_stream_demo.visual_policy import VisualChunkBC +from examples.ods.simple_bc_demo.models import rotation_from_6d +from examples.ods.simple_bc_demo.observations import policy_observations +from embodichain.utils.math import axis_angle_from_quat, quat_from_matrix + +__all__ = ["evaluate_checkpoint"] + + +def _load_visual_policy( + checkpoint: Path, vision: str, device: str +) -> tuple[VisualChunkBC, dict[str, Any], int]: + if vision not in ("single", "dual"): + raise ValueError("Visual evaluation supports only single or dual cameras") + if not checkpoint.is_file(): + raise FileNotFoundError(checkpoint) + # Only load checkpoints written by this example: optimizer/RNG records may + # include trusted Python objects that are outside weights_only's allowlist. + state = torch.load(checkpoint, map_location="cpu", weights_only=False) + weights = state["model"] + config = state["config"] + cameras = tuple( + state.get("cameras", weights.get("_extra_state", {}).get("cameras", ())) + ) + if cameras != camera_uids(vision) or config.get("vision", vision) != vision: + raise ValueError( + "Checkpoint vision profile/camera order differs from evaluation" + ) + prediction_horizon = config.get("prediction_horizon", 32) + action_horizon = config.get("action_horizon", min(16, prediction_horizon)) + if ( + type(prediction_horizon) is not int + or prediction_horizon < 1 + or type(action_horizon) is not int + or not 1 <= action_horizon <= prediction_horizon + ): + raise ValueError( + "Checkpoint action horizon must fit its positive prediction horizon" + ) + normalization = state.get("normalization", {}) + statistics = {} + for key in ("proprio_mean", "proprio_std", "target_mean", "target_std"): + statistics[key] = torch.as_tensor( + normalization.get(key, weights[key]), dtype=torch.float32, device="cpu" + ) + if not torch.equal(statistics[key], weights[key].cpu()): + raise ValueError( + f"Checkpoint normalization differs from model buffer {key}" + ) + model = VisualChunkBC( + **statistics, cameras=cameras, prediction_horizon=prediction_horizon + ) + model.load_state_dict(weights) + model = model.to(device).eval() + return model, state, action_horizon + + +def _measured_images( + observation: Any, cameras: tuple[str, ...] +) -> dict[str, torch.Tensor]: + images = {} + for camera in cameras: + color = observation["sensor"][camera]["color"] + if ( + color.shape not in ((1, 256, 256, 3), (1, 256, 256, 4)) + or color.dtype != torch.uint8 + ): + raise ValueError( + f"Camera {camera} must provide one measured uint8 RGB(A) frame" + ) + # Own the historical image before the next simulator step mutates its + # camera buffers. CPU scaling matches the online/offline data adapters. + rgb = color[..., :3].movedim(-1, -3).to(device="cpu", dtype=torch.float32) + images[camera] = rgb.div_(255) + return images + + +def _relative_request(target: torch.Tensor, geometry: torch.Tensor) -> torch.Tensor: + rotation_error = rotation_from_6d(target[:, 3:9]) @ rotation_from_6d( + geometry[:, 21:27] + ).transpose(-1, -2) + return torch.cat( + ( + (target[:, :3] - geometry[:, 18:21]) / 0.02, + axis_angle_from_quat(quat_from_matrix(rotation_error)) / 0.1, + target[:, 9:], + ), + dim=-1, + ).clamp(-1, 1) + + +def evaluate_checkpoint( + checkpoint: Path, + vision: str, + sim_device: str = "cuda:0", + train_device: str = "cuda:1", + *, + scenes: int = 20, + seed: int = 197001, + output: Path | None = None, +) -> dict[str, Any]: + """Execute only the frozen visual student in a fresh physical environment. + + Args: + checkpoint: Trusted visual BC checkpoint produced by this demo. + vision: ``single`` or ``dual`` camera profile, matching the checkpoint. + sim_device: Explicit tensor/renderer device of the single environment. + train_device: Explicit device for frozen model inference. + scenes: Number of fixed, independently seeded initial cube states. + seed: CPU scene-generator seed, independent of training generation. + output: Optional JSON report path. All episode failures are retained. + + Returns: + Reloaded checkpoint identity, every physical outcome, aggregate metrics, + and successful environment shutdown evidence. + """ + if type(scenes) is not int or scenes < 1 or type(seed) is not int or seed < 0: + raise ValueError("Evaluation scenes must be positive and seed nonnegative") + checkpoint = Path(checkpoint) + model, state, action_horizon = _load_visual_policy(checkpoint, vision, train_device) + generator = torch.Generator(device="cpu").manual_seed(seed) + positions = ( + torch.tensor([-0.42, -0.08]) + + (torch.rand(scenes, 2, generator=generator) - 0.5) * 0.08 + ) + scene_digest = hashlib.sha256(positions.numpy().tobytes()).hexdigest() + profile = state["config"].get("camera_profile", "benchmark") + if state["config"].get( + "camera_profile_version", camera_profile_version(profile) + ) != camera_profile_version(profile): + raise ValueError( + "Checkpoint camera calibration version differs from this evaluator" + ) + config = load_acquisition_config( + vision, sim_device, producer_envs=1, seed=seed, camera_profile=profile + ) + config["max_episode_steps"] = 600 + os.environ["EMBODICHAIN_SIM_EXIT_PROCESS"] = "0" + # Import the runtime package before manager configuration, matching native + # worker bootstrap and ensuring its dynamic bindings are initialized. + from embodichain.lab.sim import SimulationManager + import gymnasium as gym + from embodichain.lab.gym.utils.gym_utils import config_to_cfg, get_manager_modules + from embodichain.lab.gym.utils.registration import discover_task_packages + + discover_task_packages() + cfg = config_to_cfg(config, manager_modules=get_manager_modules()) + cfg.init_rollout_buffer = False + cfg.filter_dataset_saving = True + cfg.max_episode_steps = 600 + env = None + records = [] + started = time.monotonic() + try: + env = gym.make(config["id"], cfg=cfg) + raw = env.unwrapped + raw._replay_no_auto_reset = True + for ordinal, xy in enumerate(positions): + scene_seed = seed + ordinal + observation, _ = env.reset( + seed=scene_seed, + options={"save_data": False, "cube_xy": [xy.tolist()]}, + ) + previous_geometry = None + previous_images = None + queued_targets = None + queue_index = 0 + replans = 0 + terminal = False + truncated = False + # The simulator updates task-state tensors that survive this loop + # and are reset outside it. no_grad suppresses policy autograd while + # keeping those persistent tensors writable at the next reset. + with torch.no_grad(): + for length in range(1, 601): + geometry = policy_observations(observation).to(train_device) + if geometry.shape != (1, 51): + raise ValueError("Evaluation geometry schema must be [1, 51]") + current_images = _measured_images(observation, model.cameras) + if queued_targets is None or queue_index >= action_horizon: + history = torch.stack( + ( + ( + geometry[:, :27] + if previous_geometry is None + else previous_geometry + ), + geometry[:, :27], + ), + dim=1, + ) + image_history = { + camera: torch.stack( + ( + ( + current_images[camera] + if previous_images is None + else previous_images[camera] + ), + current_images[camera], + ), + dim=1, + ).to(train_device) + for camera in model.cameras + } + queued_targets = ( + model(history, image_history) * model.target_std + + model.target_mean + ) + if not torch.isfinite(queued_targets).all(): + raise RuntimeError( + "Student produced nonfinite absolute targets" + ) + queue_index = 0 + replans += 1 + request = _relative_request( + queued_targets[:, queue_index], geometry + ) + queue_index += 1 + # Retain the immediately preceding measured observation at + # every step, including while executing a queued chunk. + previous_geometry = geometry[:, :27].clone() + previous_images = current_images + observation, _, terminated, timed_out, info = env.step( + request.to(raw.device) + ) + terminal = bool(torch.as_tensor(terminated).any()) + truncated = bool(torch.as_tensor(timed_out).any()) + if terminal or truncated: + break + success, failed, metrics = raw.compute_task_state() + record = { + "scene": ordinal, + "seed": scene_seed, + "cube_xy": xy.tolist(), + "success": bool(success[0]), + "failure": bool(failed[0]), + "lifted": bool(metrics["lifted"][0]), + "goal_error": float(metrics["goal_error"][0]), + "length": length, + "replans": replans, + "terminated": terminal, + "truncated": truncated, + } + records.append(record) + if output is not None: + output = Path(output) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text( + json.dumps( + { + "passed": False, + "status": "in_progress", + "checkpoint": str(checkpoint.resolve()), + "vision": vision, + "seed": seed, + "scene_digest": scene_digest, + "scenes": scenes, + "records": records, + }, + indent=2, + ) + + "\n" + ) + print( + f"visual evaluation {ordinal + 1}/{scenes}: success={record['success']} " + f"lifted={record['lifted']} error={record['goal_error']:.4f}", + flush=True, + ) + if getattr(raw, "_expert_engine", None) is not None: + raise RuntimeError( + "Independent student evaluation initialized an expert engine" + ) + finally: + try: + if env is not None: + env.unwrapped.close(exit_process=False) + finally: + SimulationManager.flush_cleanup_queue() + report = { + "passed": True, + "checkpoint": str(checkpoint.resolve()), + "checkpoint_sha256": hashlib.sha256(checkpoint.read_bytes()).hexdigest(), + "policy_step": int(state["step"]), + "vision": vision, + "cameras": list(model.cameras), + "camera_profile_version": camera_profile_version(profile), + "prediction_horizon": model.prediction_horizon, + "action_horizon": action_horizon, + "sim_device": sim_device, + "inference_device": train_device, + "seed": seed, + "scene_digest": scene_digest, + "scenes": scenes, + "records": records, + "success_rate": sum(record["success"] for record in records) / scenes, + "lift_rate": sum(record["lifted"] for record in records) / scenes, + "goal_error": sum(record["goal_error"] for record in records) / scenes, + "episode_length": sum(record["length"] for record in records) / scenes, + "wall_seconds": time.monotonic() - started, + "timing_scope": "this evaluation invocation; total-run inclusion is caller-owned", + "closed": True, + "expert_initialized": False, + } + if output is not None: + output = Path(output) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(report, indent=2) + "\n") + return report + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--vision", choices=("single", "dual"), required=True) + parser.add_argument("--sim-device", default="cuda:0") + parser.add_argument("--train-device", default="cuda:1") + parser.add_argument("--scenes", type=int, default=20) + parser.add_argument("--seed", type=int, default=197001) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--cpu-threads", type=int, default=1) + args = parser.parse_args() + if args.cpu_threads < 1: + parser.error("cpu threads must be positive") + torch.set_num_threads(args.cpu_threads) + try: + report = evaluate_checkpoint( + args.checkpoint, + args.vision, + args.sim_device, + args.train_device, + scenes=args.scenes, + seed=args.seed, + output=args.output, + ) + except BaseException: + args.output.parent.mkdir(parents=True, exist_ok=True) + partial = json.loads(args.output.read_text()) if args.output.is_file() else {} + args.output.write_text( + json.dumps( + {**partial, "passed": False, "error": traceback.format_exc()}, indent=2 + ) + + "\n" + ) + raise + print(json.dumps(report, indent=2), flush=True) + + +if __name__ == "__main__": + main() diff --git a/examples/ods/data_stream_demo/training_metrics.py b/examples/ods/data_stream_demo/training_metrics.py new file mode 100644 index 000000000..39bbba660 --- /dev/null +++ b/examples/ods/data_stream_demo/training_metrics.py @@ -0,0 +1,575 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Exact completed-training diversity and timing for the visual ODS experiment. + +Publication and delivery alone never increase the training coverage counters. +The owner calls :meth:`TrainingMetrics.complete_update` only after its device +has completed a real optimizer update. Images are neither inspected nor kept +by this recorder; compact per-source frame counters retain exact lifetimes. +""" + +from __future__ import annotations + +import json +import math +import time +from collections import Counter +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +from .metrics import write_json + +__all__ = ["TrainingMetrics"] + +_INITIAL_XY_LOW = np.array([-0.46, -0.12], dtype=np.float64) +_INITIAL_XY_HIGH = np.array([-0.38, -0.04], dtype=np.float64) +_INITIAL_BIN_METERS = 0.01 +_VOXEL_METERS = 0.02 + + +def _positive_int(value: int, name: str) -> int: + if type(value) is not int or value <= 0: + raise ValueError(f"{name} must be a positive integer") + return value + + +def _cpu_array(value: Any, name: str) -> np.ndarray: + if not isinstance(value, torch.Tensor) or value.device.type != "cpu": + raise ValueError(f"{name} must be a CPU tensor") + return value.detach().numpy() + + +def _finite(value: float, name: str) -> float: + result = float(value) + if not math.isfinite(result): + raise ValueError(f"{name} must be finite") + return result + + +def _quantiles(values: list[float]) -> dict[str, float] | None: + if not values: + return None + quantiles = np.percentile(values, [50, 95]) + return {"p50": float(quantiles[0]), "p95": float(quantiles[1])} + + +@dataclass +class _Episode: + source_id: str + ordinal: int + generation: int + initial_xy: tuple[float, float] + initial_bin: tuple[int, int] | None + published_ns: int + frame_count: int + accepted_ns: int | None + metadata: dict[str, Any] + anchors: np.ndarray = field(init=False) + history: np.ndarray = field(init=False) + targets: np.ndarray = field(init=False) + trained_chunks: int = 0 + first_completed_ns: int | None = None + retired_ns: int | None = None + + def __post_init__(self) -> None: + self.anchors = np.zeros(self.frame_count, dtype=np.int64) + self.history = np.zeros(self.frame_count, dtype=np.int64) + self.targets = np.zeros(self.frame_count, dtype=np.int64) + + +class TrainingMetrics: + """Record completed work, exact sample reuse and task-state coverage. + + Args: + output: Run directory for launch-relative JSONL and summary sidecars. + origin_ns: Supervisor's host-monotonic timestamp before launch. + history: Number of causal observations, including the anchor. + future: Number of supervised actions, including the anchor. + batch_size: Exact number of chunks in every optimizer update. + reuse_cap: Per-source lifetime training upper bound, or ``None`` for + the reusable offline baseline. Delivery is checked separately + against the sampler's ledger by the experiment owner. + + A source is identified by the first two columns of ``frame_ids`` as the + string ``"run_id:episode_id"``. Coverage grids use arena coordinates; + these privileged measurements never become policy inputs. + """ + + def __init__( + self, + output: Path, + *, + origin_ns: int, + history: int, + future: int, + batch_size: int, + reuse_cap: int | None, + ) -> None: + self.output = Path(output) + self.output.mkdir(parents=True, exist_ok=True) + self.origin_ns = _positive_int(origin_ns, "origin_ns") + self.history = _positive_int(history, "history") + self.future = _positive_int(future, "future") + self.batch_size = _positive_int(batch_size, "batch_size") + self.reuse_cap = ( + None if reuse_cap is None else _positive_int(reuse_cap, "reuse_cap") + ) + self.episodes: dict[str, _Episode] = {} + self.updates = 0 + self.chunks = 0 + self._initial_bins: Counter[tuple[int, int]] = Counter() + self._voxels: set[tuple[int, int, int]] = set() + self._curve: list[dict[str, Any]] = [] + self._intervals: list[dict[str, Any]] = [] + self._last_losses: dict[str, float] = {} + self._last_completed_ns = origin_ns + self._events = (self.output / "training-events.jsonl").open( + "w", encoding="utf-8", buffering=1 + ) + self._closed = False + + def _seconds(self, at_ns: int) -> float: + if type(at_ns) is not int or at_ns < self.origin_ns: + raise ValueError("Timestamp must use the supervisor monotonic clock") + return (at_ns - self.origin_ns) / 1e9 + + def event(self, kind: str, at_ns: int | None = None, **metadata: Any) -> None: + """Append a finite JSON event using the common host monotonic clock.""" + if self._closed: + raise RuntimeError("Training metrics are closed") + if not kind: + raise ValueError("An event needs a name") + timestamp = time.monotonic_ns() if at_ns is None else at_ns + row = { + **metadata, + "kind": kind, + "monotonic_ns": timestamp, + "time_seconds": self._seconds(timestamp), + } + self._events.write(json.dumps(row, allow_nan=False) + "\n") + + def register_episode( + self, + source_id: str, + *, + ordinal: int, + generation: int, + initial_xy: Sequence[float], + published_ns: int, + frame_count: int, + accepted_ns: int | None = None, + metadata: dict[str, Any] | None = None, + ) -> None: + """Register one accepted publication without counting it as training.""" + pieces = source_id.split(":") + if len(pieces) != 2 or any(int(part) < 0 for part in pieces): + raise ValueError("source_id must be a nonnegative run:episode identity") + if source_id in self.episodes: + raise ValueError(f"Duplicate immutable source identity: {source_id}") + if type(ordinal) is not int or ordinal < 0: + raise ValueError("ordinal must be a nonnegative integer") + if type(generation) is not int or generation < 1: + raise ValueError("generation must be a positive integer") + count = _positive_int(frame_count, "frame_count") + self._seconds(published_ns) + if accepted_ns is not None: + self._seconds(accepted_ns) + if accepted_ns > published_ns: + raise ValueError("Acceptance cannot follow source publication") + xy = np.asarray(initial_xy, dtype=np.float64) + if xy.shape != (2,) or not np.isfinite(xy).all(): + raise ValueError("initial_xy must contain two finite arena coordinates") + in_bounds = bool( + ((xy >= _INITIAL_XY_LOW - 1e-7) & (xy <= _INITIAL_XY_HIGH + 1e-7)).all() + ) + coordinates = (xy - _INITIAL_XY_LOW) / _INITIAL_BIN_METERS + nearest = np.rint(coordinates) + # Float32 task states can straddle an exact centimeter boundary by a + # few nanometers. Snap only that numerical error, never physical bins. + coordinates = np.where( + np.abs(coordinates - nearest) <= 1e-5, nearest, coordinates + ) + bins = np.minimum(7, np.maximum(0, np.floor(coordinates))).astype(np.int64) + initial_bin = tuple(int(v) for v in bins) if in_bounds else None + copied_metadata = json.loads(json.dumps(metadata or {}, allow_nan=False)) + episode = _Episode( + source_id=source_id, + ordinal=ordinal, + generation=generation, + initial_xy=(float(xy[0]), float(xy[1])), + initial_bin=initial_bin, + published_ns=published_ns, + frame_count=count, + accepted_ns=accepted_ns, + metadata=copied_metadata, + ) + self.event( + "episode_published", + published_ns, + source_id=source_id, + ordinal=ordinal, + generation=generation, + initial_xy=episode.initial_xy, + initial_xy_bin=initial_bin, + frame_count=count, + accepted_ns=accepted_ns, + metadata=copied_metadata, + ) + self.episodes[source_id] = episode + + def retire_episode(self, source_id: str, *, retired_ns: int) -> None: + """Record retirement while preserving a source's cumulative counters.""" + episode = self.episodes[source_id] + if retired_ns < episode.published_ns: + raise ValueError("Retirement precedes publication") + if episode.retired_ns is not None: + raise ValueError("Source was already retired") + self.event("episode_retired", retired_ns, source_id=source_id) + episode.retired_ns = retired_ns + + def interval(self, name: str, start_ns: int, end_ns: int, **metadata: Any) -> None: + """Record a stage interval without adding overlapping durations together.""" + start = self._seconds(start_ns) + end = self._seconds(end_ns) + if end < start: + raise ValueError("An interval cannot end before it starts") + row = { + **metadata, + "name": name, + "start_seconds": start, + "end_seconds": end, + "duration_seconds": end - start, + } + self.event("interval", end_ns, **row) + self._intervals.append(row) + + def _batch( + self, batch: Mapping[str, Any] + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + length = self.history + self.future - 1 + ids = _cpu_array(batch["frame_ids"], "frame_ids") + valid = _cpu_array(batch["valid"], "valid") + mask = _cpu_array(batch["loss_mask"], "loss_mask") + generations = _cpu_array(batch["generations"], "generations") + geometry = _cpu_array(batch["geometry"], "geometry") + if ids.shape != (self.batch_size, length, 3) or ids.dtype != np.int64: + raise ValueError("frame_ids shape/dtype differs from the training protocol") + if ( + valid.shape != (self.batch_size, length) + or valid.dtype != np.bool_ + or not valid.all() + ): + raise ValueError("Training requires complete, unpadded contexts") + expected_mask = np.broadcast_to( + np.arange(length) >= self.history - 1, valid.shape + ) + if mask.dtype != np.bool_ or not np.array_equal(mask, expected_mask): + raise ValueError("loss_mask must contain exactly the F future targets") + if generations.shape != valid.shape or generations.dtype != np.int64: + raise ValueError("Generation shape/dtype mismatch") + anchor = ids[:, self.history - 1] + offsets = np.arange(1 - self.history, self.future, dtype=np.int64) + if ( + (ids < 0).any() + or (ids[..., :2] != anchor[:, None, :2]).any() + or (ids[..., 2] != anchor[:, None, 2] + offsets).any() + or (generations != generations[:, self.history - 1, None]).any() + ): + raise ValueError("A training context crosses a source/frame boundary") + if "continuity_id" in batch: + continuity = _cpu_array(batch["continuity_id"], "continuity_id") + if ( + continuity.shape != valid.shape + or (continuity != continuity[:, self.history - 1, None]).any() + ): + raise ValueError("A training context crosses continuity boundaries") + if ( + geometry.shape != (self.batch_size, length, 51) + or not np.isfinite(geometry).all() + ): + raise ValueError("geometry must contain the 51 finite task-state features") + return ids, generations, geometry[:, self.history - 1, 27:30] + + def complete_update( + self, + batch: Mapping[str, Any], + *, + update: int, + losses: Mapping[str, float], + requested_ns: int, + train_started_ns: int, + completed_ns: int, + ) -> None: + """Count a synchronized, completed optimizer update exactly once. + + All validation happens before counters are changed. The input may keep + GPU image/model data, but the frame IDs and geometry must be owned CPU + tensors supplied by the data adapter. + """ + if update != self.updates + 1: + raise ValueError( + "Completed update numbers must start at one and be consecutive" + ) + for timestamp in (requested_ns, train_started_ns, completed_ns): + self._seconds(timestamp) + if not (requested_ns <= train_started_ns <= completed_ns): + raise ValueError("Training update timestamps are not causal") + if completed_ns < self._last_completed_ns: + raise ValueError( + "Completed optimizer updates cannot move backwards in time" + ) + finite_losses = {name: _finite(value, name) for name, value in losses.items()} + if not finite_losses: + raise ValueError("A completed update must record real finite losses") + ids, generations, xyz = self._batch(batch) + anchor_rows = ids[:, self.history - 1] + source_ids = [f"{row[0]}:{row[1]}" for row in anchor_rows] + counts = Counter(source_ids) + for source_id, count in counts.items(): + if source_id not in self.episodes: + raise ValueError(f"Training source was not registered: {source_id}") + episode = self.episodes[source_id] + rows = np.asarray( + [i for i, key in enumerate(source_ids) if key == source_id] + ) + if ( + train_started_ns < episode.published_ns + or (generations[rows] != episode.generation).any() + or (ids[rows, :, 2] >= episode.frame_count).any() + ): + raise ValueError( + "Training references unpublished or mismatched source data" + ) + if ( + self.reuse_cap is not None + and episode.trained_chunks + count > self.reuse_cap + ): + raise ValueError(f"Completed training exceeds source cap: {source_id}") + pairs = batch.get("_pairs") + if pairs is not None: + expected_pairs = [ + (self.episodes[source].ordinal, int(row[2])) + for source, row in zip(source_ids, anchor_rows) + ] + if [tuple(int(v) for v in pair) for pair in pairs] != expected_pairs: + raise ValueError( + "Batch source/anchor manifest does not match frame IDs" + ) + voxels = np.floor(np.asarray(xyz, dtype=np.float64) / _VOXEL_METERS).astype( + np.int64 + ) + chunks = [ + { + "source_id": source, + "ordinal": self.episodes[source].ordinal, + "generation": self.episodes[source].generation, + "anchor": int(row[2]), + "anchor_cube_xyz": [float(v) for v in position], + } + for source, row, position in zip(source_ids, anchor_rows, xyz) + ] + # The log is append-only. It records sufficient identities to recompute + # every count without retaining image tensors or per-frame Python sets. + for index, source in enumerate(source_ids): + episode = self.episodes[source] + anchor_frame = int(anchor_rows[index, 2]) + episode.trained_chunks += 1 + episode.anchors[anchor_frame] += 1 + episode.history[ids[index, : self.history, 2]] += 1 + episode.targets[ids[index, self.history - 1 :, 2]] += 1 + if episode.first_completed_ns is None: + episode.first_completed_ns = completed_ns + if episode.initial_bin is not None: + self._initial_bins[episode.initial_bin] += 1 + self._voxels.add(tuple(int(v) for v in voxels[index])) + self.updates = update + self.chunks += self.batch_size + self._last_completed_ns = completed_ns + self._last_losses = finite_losses + curve = { + "update": update, + "time_seconds": self._seconds(completed_ns), + **self._coverage(), + } + self.event( + "update_completed", + completed_ns, + update=update, + requested_seconds=self._seconds(requested_ns), + train_started_seconds=self._seconds(train_started_ns), + losses=finite_losses, + chunks=chunks, + coverage=curve, + ) + self._curve.append(curve) + + def _coverage(self) -> dict[str, Any]: + trained = [row for row in self.episodes.values() if row.trained_chunks] + proportions = [row.trained_chunks / self.chunks for row in trained] + unique_targets = sum(int(np.count_nonzero(row.targets)) for row in trained) + target_occurrences = self.chunks * self.future + return { + "completed_updates": self.updates, + "trained_chunks": self.chunks, + "valid_target_occurrences": target_occurrences, + "history_frame_occurrences": self.chunks * self.history, + "trained_sources": len(trained), + "effective_sources": ( + math.exp(-sum(p * math.log(p) for p in proportions)) if trained else 0.0 + ), + "initial_xy_occupied_cells": len(self._initial_bins), + "initial_xy_coverage_fraction": len(self._initial_bins) / 64, + "anchor_xyz_occupied_voxels": len(self._voxels), + "unique_anchors": sum( + int(np.count_nonzero(row.anchors)) for row in trained + ), + "unique_history_frames": sum( + int(np.count_nonzero(row.history)) for row in trained + ), + "unique_target_frames": unique_targets, + "target_duplicate_fraction": ( + 1 - unique_targets / target_occurrences if target_occurrences else 0.0 + ), + } + + def _source_rows(self) -> list[dict[str, Any]]: + rows = [] + for episode in sorted(self.episodes.values(), key=lambda row: row.ordinal): + first = episode.first_completed_ns + rows.append( + { + "source_id": episode.source_id, + "ordinal": episode.ordinal, + "generation": episode.generation, + "frame_count": episode.frame_count, + "initial_xy": episode.initial_xy, + "initial_xy_bin": episode.initial_bin, + "published_seconds": self._seconds(episode.published_ns), + "accepted_seconds": ( + self._seconds(episode.accepted_ns) + if episode.accepted_ns is not None + else None + ), + "first_completed_training_seconds": ( + self._seconds(first) if first is not None else None + ), + "publication_to_first_training_seconds": ( + (first - episode.published_ns) / 1e9 + if first is not None + else None + ), + "trained_chunks": episode.trained_chunks, + "unique_anchors": int(np.count_nonzero(episode.anchors)), + "unique_history_frames": int(np.count_nonzero(episode.history)), + "unique_target_frames": int(np.count_nonzero(episode.targets)), + "target_occurrences": int(episode.targets.sum()), + "anchor_counts": episode.anchors.tolist(), + "history_frame_counts": episode.history.tolist(), + "target_frame_counts": episode.targets.tolist(), + "retired_seconds": ( + self._seconds(episode.retired_ns) + if episode.retired_ns is not None + else None + ), + "metadata": episode.metadata, + } + ) + return rows + + def summary( + self, + *, + total_seconds: float, + training_complete_seconds: float, + ledger: Mapping[str, Any] | None = None, + ) -> dict[str, Any]: + """Return an auditable snapshot; failures may also snapshot partial work. + + ``ledger`` is copied verbatim as the sampler's separate reservation and + delivery evidence. The caller owns delivery-cap and exact-budget + qualification; this class enforces the completed-training cap itself. + """ + total = _finite(total_seconds, "total_seconds") + train_end = _finite(training_complete_seconds, "training_complete_seconds") + if total < train_end or train_end < 0: + raise ValueError("Completion times must satisfy 0 <= training <= total") + if self.updates and train_end + 1e-7 < self._seconds(self._last_completed_ns): + raise ValueError("Reported training completion precedes the final update") + sources = self._source_rows() + latencies = [ + row["publication_to_first_training_seconds"] + for row in sources + if row["publication_to_first_training_seconds"] is not None + ] + initial_weights = np.zeros((8, 8), dtype=np.int64) + for (x, y), count in self._initial_bins.items(): + initial_weights[x, y] = count + metrics = { + "schema": "ods-training-metrics/v1", + "history": self.history, + "future": self.future, + "batch_size": self.batch_size, + "reuse_cap": self.reuse_cap, + "origin_monotonic_ns": self.origin_ns, + "total_seconds": total, + "training_complete_seconds": train_end, + "post_training_drain_seconds": total - train_end, + "first_update_seconds": ( + self._curve[0]["time_seconds"] if self._curve else None + ), + **self._coverage(), + "published_sources": len(sources), + "untrained_sources": sum(row["trained_chunks"] == 0 for row in sources), + "initial_xy_grid": { + "low_meters": _INITIAL_XY_LOW.tolist(), + "high_meters": _INITIAL_XY_HIGH.tolist(), + "shape": [8, 8], + "bin_meters": _INITIAL_BIN_METERS, + "trained_chunk_weights_xy": initial_weights.tolist(), + "out_of_bounds_sources": [ + row["source_id"] for row in sources if row["initial_xy_bin"] is None + ], + }, + "anchor_xyz_voxel_meters": _VOXEL_METERS, + "anchor_xyz_voxels": [list(row) for row in sorted(self._voxels)], + "publication_to_first_training_seconds": _quantiles(latencies), + "max_source_trained_chunks": max( + (row["trained_chunks"] for row in sources), default=0 + ), + "training_cap_violations": sum( + self.reuse_cap is not None and row["trained_chunks"] > self.reuse_cap + for row in sources + ), + "episodes": sources, + "coverage_curve": self._curve.copy(), + "intervals": self._intervals.copy(), + "last_losses": self._last_losses.copy(), + "ledger": json.loads(json.dumps(ledger or {}, allow_nan=False)), + } + write_json(self.output / "training-sources.json", {"episodes": sources}) + return metrics + + def close(self) -> None: + """Close the append-only event stream exactly once.""" + if not self._closed: + self._events.close() + self._closed = True diff --git a/examples/ods/data_stream_demo/training_report.py b/examples/ods/data_stream_demo/training_report.py new file mode 100644 index 000000000..5057952cb --- /dev/null +++ b/examples/ods/data_stream_demo/training_report.py @@ -0,0 +1,1367 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Render fixed-budget visual training measurements, retaining failed runs. + +Run: ``python -m examples.ods.data_stream_demo.training_report OUTPUT_ROOT``. +Legacy delivery-only summaries are deliberately excluded from this report. +""" + +from __future__ import annotations + +import argparse +import json +import math +import re +from collections import Counter +from pathlib import Path +from statistics import median +from typing import Any + +__all__ = ["write_report"] + +_ARMS = ("offline8", "offline32", "offline128", "ods_overlap", "ods_serial") +_LABELS = { + "offline8": "Offline-8", + "offline32": "Offline-32", + "offline128": "Offline-128", + "ods_overlap": "ODS overlap", + "ods_serial": "ODS serial replay", +} +_COLORS = { + "offline8": "#4c78a8", + "offline32": "#f58518", + "offline128": "#b279a2", + "ods_overlap": "#159b80", + "ods_serial": "#c65656", +} +_MARKERS = { + "offline8": "o", + "offline32": "s", + "offline128": "^", + "ods_overlap": "D", + "ods_serial": "X", +} +_VISIONS = ("single", "dual") +_METRICS = ( + "total_seconds", + "training_complete_seconds", + "post_training_drain_seconds", + "first_update_seconds", + "published_sources", + "trained_sources", + "effective_sources", + "initial_xy_occupied_cells", + "anchor_xyz_occupied_voxels", + "unique_anchors", + "unique_target_frames", + "target_duplicate_fraction", + "max_source_trained_chunks", + "dataset_bytes", + "checkpoint_bytes", +) + + +def _json(path: Path) -> dict[str, Any]: + value = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, dict): + raise ValueError(f"Expected a JSON object in {path}") + return value + + +def _number(value: Any) -> float | None: + return ( + float(value) + if isinstance(value, (float, int)) + and not isinstance(value, bool) + and math.isfinite(value) + else None + ) + + +def _arm(config: dict[str, Any]) -> str: + arm = config.get("arm") + if arm == "offline": + return f"offline{config.get('offline_episodes')}" + return str(arm) + + +def _events(path: Path) -> list[dict[str, Any]]: + if not path.exists(): + return [] + records = [] + for line_number, line in enumerate( + path.read_text(encoding="utf-8").splitlines(), 1 + ): + if line.strip(): + row = json.loads(line) + if not isinstance(row, dict): + raise ValueError(f"Expected JSON object at {path}:{line_number}") + records.append(row) + return records + + +def _load(root: Path) -> list[dict[str, Any]]: + paths = sorted(root.rglob("summary.json")) if root.is_dir() else [root] + runs = [] + for path in paths: + summary = _json(path) + metrics = summary.get("training_metrics", summary.get("metrics", {})) + config = summary.get("config", {}) + if ( + not isinstance(metrics, dict) + or metrics.get("schema") != "ods-training-metrics/v1" + ): + # A child can fail before constructing its metric recorder. Keep + # these training failures; do not mix legacy delivery summaries. + if config.get("arm") not in {"offline", "ods_overlap", "ods_serial"}: + continue + metrics = {} + arm = _arm(config) + engineering_offline = ( + config.get("arm") == "offline" + and type(config.get("offline_episodes")) is int + and config["offline_episodes"] > 0 + and config.get("phase") in {"single", "smoke", "calibration"} + ) + if ( + arm not in _ARMS + and not engineering_offline + or config.get("vision") not in _VISIONS + ): + raise ValueError(f"Unknown visual training arm/vision in {path}") + status = summary.get("status") + if status not in {"complete", "failed", "unsupported"}: + raise ValueError(f"Unknown training status in {path}") + runs.append( + { + "path": str(path.resolve()), + "summary": summary, + "config": config, + "metrics": metrics, + "arm": arm, + "events": _events(path.parent / "training-events.jsonl"), + } + ) + if not runs: + raise ValueError(f"No fixed-budget visual training summaries found in {root}") + cells = Counter( + (run["arm"], run["config"]["vision"], str(run["config"].get("seed"))) + for run in runs + if not _qualification(run) + ) + for run in runs: + run["duplicate_cell"] = ( + cells[(run["arm"], run["config"]["vision"], str(run["config"].get("seed")))] + > 1 + ) + return sorted( + runs, + key=lambda run: ( + run["config"]["vision"], + run["arm"], + str(run["config"].get("seed")), + run["path"], + ), + ) + + +def _qualification(run: dict[str, Any]) -> list[str]: + summary, metrics, config = run["summary"], run["metrics"], run["config"] + reasons = [] + if config.get("phase") == "pilot" and config.get("prefetch_policy") != "budget": + reasons.append( + "pilot primary comparisons require explicit budget prefetch; eager is a separate diagnostic" + ) + if run.get("duplicate_cell"): + reasons.append("multiple qualified attempts for the same arm/vision/seed") + if summary["status"] != "complete": + reasons.append(f"status={summary['status']}") + if ( + summary.get("qualification_passed", metrics.get("qualification_passed")) + is not True + ): + reasons.append("no explicit qualification pass") + updates = config.get("updates") + batch_size = config.get("batch_size", metrics.get("batch_size")) + future = config.get("future", metrics.get("future")) + if any( + type(value) is not int or value < 1 for value in (updates, batch_size, future) + ): + reasons.append("missing positive training budget") + elif ( + metrics.get("completed_updates") != updates + or metrics.get("trained_chunks") != updates * batch_size + or metrics.get("valid_target_occurrences") != updates * batch_size * future + ): + reasons.append("incomplete or mismatched training work") + total = _number(metrics.get("total_seconds")) + train = _number(metrics.get("training_complete_seconds")) + if total is None or train is None or not (total >= train > 0): + reasons.append("invalid completion timing") + sources = metrics.get("episodes", []) + if not isinstance(sources, list) or sum( + row.get("trained_chunks", 0) for row in sources + ) != metrics.get("trained_chunks"): + reasons.append("source totals disagree with completed work") + if metrics.get("training_cap_violations") != 0: + reasons.append("training cap audit absent or failed") + actual = ( + config.get("worker") is True + or config.get("phase") in {"single", "smoke", "calibration", "report", "pilot"} + or "device_supervision" in summary + or "supervisor_total_seconds" in summary + or any( + key in metrics + for key in ( + "initial_model_sha256", + "normalization_sha256", + "gpu_step_milliseconds", + "engine_stats", + "timeline_contract", + ) + ) + ) + if actual and summary["status"] == "complete": + _, evidence_errors = _fairness(run) + reasons.extend(f"fairness evidence: {error}" for error in evidence_errors) + audit = summary.get("post_audit") + if ( + summary.get("post_audit_status") != "complete" + or not isinstance(audit, dict) + or audit.get("passed") is not True + ): + reasons.append("actual run lacks a completed passed full corpus audit") + expected = summary.get("requested_evaluation_scenes") + evaluation = summary.get("evaluation") + if ( + (type(expected) is not int or expected < 0) + and config.get("phase") not in {"report", "pilot"} + and isinstance(evaluation, dict) + ): + # Older smoke records predate the persisted request. Their actual + # completed evaluation is diagnostic evidence, never formal pairing. + expected = evaluation.get("scenes") + if type(expected) is not int or expected < 0: + reasons.append("missing actual requested evaluation scene count") + elif expected == 0: + if summary.get("evaluation_status") != "not_requested": + reasons.append( + "zero-scene evaluation lacks explicit not_requested status" + ) + elif not ( + summary.get("evaluation_status") == "complete" + and isinstance(evaluation, dict) + and evaluation.get("passed") is True + and evaluation.get("closed") is True + and evaluation.get("scenes") == expected + and len(evaluation.get("records", [])) == expected + ): + reasons.append("requested evaluation is pending, failed or incomplete") + if actual and summary["status"] == "complete": + reasons.extend(_new_formal_contract_errors(run)) + return reasons + + +def _qualified(run: dict[str, Any]) -> bool: + return not _qualification(run) + + +def _group(runs: list[dict[str, Any]], arm: str, vision: str) -> list[dict[str, Any]]: + return [ + run + for run in runs + if run["arm"] == arm and run["config"]["vision"] == vision and _qualified(run) + ] + + +def _stats(values: list[float]) -> dict[str, Any]: + return { + "n": len(values), + "median": median(values) if values else None, + "min": min(values) if values else None, + "max": max(values) if values else None, + } + + +def _aggregate(runs: list[dict[str, Any]]) -> dict[str, Any]: + return { + vision: { + arm: { + name: _stats( + [ + value + for run in _group(runs, arm, vision) + if (value := _number(run["metrics"].get(name))) is not None + ] + ) + for name in _METRICS + } + for arm in _ARMS + } + for vision in _VISIONS + } + + +def _allocation(value: Any) -> dict[str, Any]: + """Normalize assigned roles, preserving exact devices and CPU core IDs.""" + if not isinstance(value, dict): + raise ValueError("missing assigned primary GPU pair") + uuids = [value.get(role, {}).get("uuid") for role in ("simulation", "training")] + cpus = value.get("cpu_affinity") + if ( + any(not isinstance(uuid, str) or not uuid.startswith("GPU-") for uuid in uuids) + or uuids[0] == uuids[1] + or not isinstance(cpus, list) + or len(cpus) != 8 + or any(type(cpu) is not int or cpu < 0 for cpu in cpus) + or len(set(cpus)) != len(cpus) + ): + raise ValueError( + "expected two distinct role UUIDs and eight distinct CPU cores" + ) + return { + "simulation_uuid": uuids[0], + "training_uuid": uuids[1], + "cpu_affinity": sorted(cpus), + "cpu_cores": len(cpus), + } + + +def _launch_path(run: dict[str, Any]) -> Path: + directory = Path(run["path"]).parent + for root in (directory, *directory.parents): + relative = directory.relative_to(root) + candidate = ( + root + / "_supervision" + / ("single" if relative == Path(".") else relative) + / "training" + / "launch.json" + ) + if candidate.is_file(): + return candidate + raise ValueError("missing supervisor training launch.json") + + +def _fairness(run: dict[str, Any]) -> tuple[dict[str, Any], list[str]]: + """Check actual run evidence before allowing any comparative time ratio. + + Legacy CPU metric fixtures may still render and retain core qualification, + but missing GPU/hash/supervisor evidence never produces a speedup claim. + """ + metrics, config, summary = run["metrics"], run["config"], run["summary"] + evidence: dict[str, Any] = {} + if type(summary.get("requested_evaluation_scenes")) is int: + evidence["requested_evaluation_scenes"] = summary["requested_evaluation_scenes"] + reasons = [] + if "prefetch_policy" in config: + if config["prefetch_policy"] not in {"budget", "eager"}: + reasons.append("unknown prefetch policy") + evidence["prefetch_policy"] = config["prefetch_policy"] + for name, allowed in (("handoff_mode", {"publication", "generation"}),): + if name in config: + if config[name] not in allowed: + reasons.append(f"unknown {name}") + evidence[name] = config[name] + if "encoder_threads" in config: + if type(config["encoder_threads"]) is not int or config["encoder_threads"] < 1: + reasons.append("invalid encoder thread budget") + evidence["encoder_threads"] = config["encoder_threads"] + for key in ("initial_model_sha256", "normalization_sha256"): + value = metrics.get(key) + if not isinstance(value, str) or re.fullmatch(r"[0-9a-f]{64}", value) is None: + reasons.append(f"missing or malformed {key}") + else: + evidence[key] = value + budget = { + "updates": config.get("updates"), + "batch_size": config.get("batch_size"), + "history": config.get("history_length", config.get("history")), + "future": config.get("prediction_horizon", config.get("future")), + } + if any(type(value) is not int or value <= 0 for value in budget.values()): + reasons.append("missing positive declared U/B/H/F budget") + elif ( + metrics.get("completed_updates") != budget["updates"] + or metrics.get("batch_size") != budget["batch_size"] + or metrics.get("history") != budget["history"] + or metrics.get("future") != budget["future"] + or metrics.get("valid_target_occurrences") + != budget["updates"] * budget["batch_size"] * budget["future"] + ): + reasons.append("declared U/B/H/F differs from completed training") + else: + evidence["training_budget"] = budget + try: + initial = _json(Path(run["path"]).with_name("initial-contract.json")) + if any( + initial.get(contract_key) != evidence.get(metric_key) + for contract_key, metric_key in ( + ("model_sha256", "initial_model_sha256"), + ("normalization_sha256", "normalization_sha256"), + ) + ): + raise ValueError( + "initial contract disagrees with model/normalization hashes" + ) + rows = initial.get("initial_episodes", []) + prefix = [ + {key: row.get(key) for key in ("ordinal", "length", "numeric_sha256")} + for row in rows + ] + if ( + type(config.get("buffer_size")) is not int + or len(prefix) != config["buffer_size"] + or [row["ordinal"] for row in prefix] != list(range(len(prefix))) + or any( + type(row["length"]) is not int + or row["length"] <= 0 + or not isinstance(row["numeric_sha256"], str) + or re.fullmatch(r"[0-9a-f]{64}", row["numeric_sha256"]) is None + for row in prefix + ) + ): + raise ValueError("missing or malformed common initial C numeric prefix") + cameras = initial.get("camera_order") + if ( + not isinstance(cameras, list) + or len(cameras) != (1 if config["vision"] == "single" else 2) + or len(set(cameras)) != len(cameras) + ): + raise ValueError("missing or malformed initial camera order") + expected = { + "updates": metrics.get("completed_updates"), + "chunks": metrics.get("trained_chunks"), + "targets": metrics.get("valid_target_occurrences"), + } + if initial.get("training_budget") != expected: + raise ValueError("initial contract and completed work disagree") + evidence.update(initial_numeric_prefix=prefix, camera_order=cameras) + except (OSError, ValueError, TypeError, AttributeError) as error: + reasons.append(f"initial contract evidence: {error}") + try: + supervision = summary.get("device_supervision", {}) + assigned = _allocation(supervision.get("primary_pair")) + launch_path = _launch_path(run) + launch = _json(launch_path) + if ( + _allocation(launch.get("pair")) != assigned + or sorted(launch.get("cpu_affinity", [])) != assigned["cpu_affinity"] + ): + raise ValueError("summary and launch GPU/CPU assignments disagree") + arguments = launch.get("command", []) + expected_flags = { + "--updates": budget["updates"], + "--batch-size": budget["batch_size"], + "--history-length": budget["history"], + "--prediction-horizon": budget["future"], + "--sim-device": "cuda:0", + "--train-device": "cuda:1", + "--phase": config.get("phase"), + "--seed": config.get("seed"), + "--vision": config.get("vision"), + "--arm": config.get("arm"), + } + if "prefetch_policy" in config: + expected_flags["--prefetch-policy"] = config["prefetch_policy"] + for name in ("handoff_mode", "encoder_threads"): + if name in config: + expected_flags["--" + name.replace("_", "-")] = config[name] + for flag, value in expected_flags.items(): + if ( + not isinstance(arguments, list) + or arguments.count(flag) != 1 + or arguments.index(flag) + 1 >= len(arguments) + or arguments[arguments.index(flag) + 1] != str(value) + ): + raise ValueError(f"launch command disagrees with declared {flag}") + environment = launch.get("environment", {}) + if ( + environment.get("CUDA_VISIBLE_DEVICES") + != assigned["simulation_uuid"] + "," + assigned["training_uuid"] + or environment.get("ODS_GPU_UUID") != assigned["simulation_uuid"] + or any( + environment.get(key) != "1" + for key in ( + "OMP_NUM_THREADS", + "MKL_NUM_THREADS", + "OPENBLAS_NUM_THREADS", + ) + ) + ): + raise ValueError("launch UUID mask or CPU thread budget disagrees") + observed_cpu = (metrics.get("resources") or {}).get("cpu_utilization") or {} + if observed_cpu and ( + observed_cpu.get("assigned_cpu_count") != assigned["cpu_cores"] + or sorted(observed_cpu.get("assigned_cpu_affinity", [])) + != assigned["cpu_affinity"] + ): + raise ValueError("observed CPU affinity differs from supervisor assignment") + guard = supervision.get("guard", {}) + if ( + type(guard.get("samples")) is not int + or guard["samples"] <= 0 + or guard.get("unexpected_contexts") != [] + or not {assigned["simulation_uuid"], assigned["training_uuid"]}.issubset( + guard.get("seen_primary_uuids", []) + ) + ): + raise ValueError("missing or failed owned-device placement audit") + process = _json(Path(run["path"]).with_name("process.json")) + if ( + process.get("stage") != "training" + or process.get("returncode") != 0 + or any( + process.get(key) is not False + for key in ("timeout", "cancelled", "leaked_process_group") + ) + or process.get("device_guard") != guard + or launch.get("timed_pipeline") is not True + ): + raise ValueError("missing or inconsistent normal-exit process audit") + total = _number(metrics.get("total_seconds")) + if total is None or any( + (number := _number(value)) is None + or not math.isclose(number, total, rel_tol=1e-9, abs_tol=1e-6) + for value in ( + summary.get("supervisor_total_seconds"), + process.get("elapsed_seconds"), + ) + ): + raise ValueError("pipeline total differs from supervisor process timing") + origin = metrics.get("origin_monotonic_ns") + if ( + type(origin) is not int + or origin <= 0 + or origin != launch.get("process_launch_monotonic_ns") + or origin != process.get("process_launch_monotonic_ns") + or str(origin) != environment.get("ODS_PROCESS_LAUNCH_MONOTONIC_NS") + ): + raise ValueError("pipeline origin differs from supervisor launch clock") + auxiliary = supervision.get("aux_graphics_scope") + auxiliary_keys = ("uuid", "scope", "max_per_pid_bytes", "idle_required") + aux = ( + {key: auxiliary.get(key) for key in auxiliary_keys} + if isinstance(auxiliary, dict) + else None + ) + launch_auxiliary = launch.get("aux_graphics_scope") + launch_aux = ( + {key: launch_auxiliary.get(key) for key in auxiliary_keys} + if isinstance(launch_auxiliary, dict) + else None + ) + if aux != launch_aux: + raise ValueError("summary and launch auxiliary graphics scopes disagree") + evidence.update( + allocation=assigned, + auxiliary_graphics=aux, + launch_path=str(launch_path), + normal_exit_audited=True, + ) + except (OSError, ValueError, TypeError, AttributeError) as error: + reasons.append(f"resource/timing evidence: {error}") + return evidence, reasons + + +def _new_formal_contract_errors(run: dict[str, Any]) -> list[str]: + """Require chronological worker services and full video QA in new formal runs. + + Historical smoke/calibration/single engineering inputs retain their original + sparse-audit scope. Superseded formal inputs lack the new contract and may + be inspected as diagnostics, but cannot enter the refreshed formal claims. + """ + if run["config"].get("phase") not in {"report", "pilot"}: + return [] + metrics, summary = run["metrics"], run["summary"] + errors = [] + if metrics.get("timeline_contract") != "sdk-worker-service/v1": + errors.append("new formal run lacks sdk-worker-service/v1 timeline contract") + total = _number(metrics.get("total_seconds")) + intervals = metrics.get("intervals", []) + recording = [row for row in intervals if row.get("name") == "sdk_recording"] + decoding = [row for row in intervals if row.get("name") == "sdk_decode"] + for row in recording + decoding: + start, end = _number(row.get("start_seconds")), _number(row.get("end_seconds")) + if ( + start is None + or end is None + or total is None + or not 0 <= start < end <= total + 1e-6 + or type(row.get("worker_pid")) is not int + or row["worker_pid"] <= 0 + ): + errors.append("invalid chronological SDK worker service interval") + break + phases = Counter(row.get("sdk_phase") for row in recording) + sources = metrics.get("episodes", []) + published = metrics.get("published_sources") + if ( + type(published) is not int + or published < 1 + or not isinstance(sources, list) + or type(metrics.get("completed_updates")) is not int + or type(metrics.get("batch_size")) is not int + ): + errors.append( + "new formal run lacks valid source/update counts for worker evidence" + ) + return errors + if ( + phases + != Counter( + {"sdk_create": 1, "sdk_episode_write_encode": published, "sdk_finalize": 1} + ) + or len(recording) != (published or 0) + 2 + ): + errors.append( + "SDK recording service lacks create, every accepted write, or finalize" + ) + writes = [ + row for row in recording if row.get("sdk_phase") == "sdk_episode_write_encode" + ] + + def identity(value: Any) -> str: + return ( + value + if isinstance(value, str) + else ( + f"{value[0]}:{value[1]}" + if isinstance(value, (list, tuple)) and len(value) == 2 + else "invalid" + ) + ) + + expected_sources = {row.get("source_id"): row.get("ordinal") for row in sources} + actual_sources = { + identity(row.get("source_id")): row.get("episode_ordinal") for row in writes + } + if actual_sources != expected_sources or len(writes) != len(actual_sources): + errors.append( + "SDK write spans do not cover the exact immutable accepted sources" + ) + if any(row.get("status", "complete") != "complete" for row in recording): + errors.append("SDK recording worker reported failure") + if run["config"].get("arm") == "offline": + item_counts: Counter[int] = Counter() + for row in decoding: + update, count = row.get("batch_update"), row.get("item_count") + if type(update) is not int or type(count) is not int or count <= 0: + errors.append("invalid SDK temporal fetch update/item attribution") + break + item_counts[update] += count + expected_counts = Counter( + { + i: metrics.get("batch_size") + for i in range(1, metrics.get("completed_updates", 0) + 1) + } + ) + if not decoding or item_counts != expected_counts: + errors.append( + "SDK temporal fetch spans do not cover every delivered training batch" + ) + audit = summary.get("post_audit") or {} + receipt = audit.get("full_video_integrity") + sidecar = Path(run["path"]).with_name("video-integrity.json") + try: + saved = _json(sidecar) if sidecar.is_file() else None + if receipt is not None and saved is not None and receipt != saved: + errors.append("inline and sidecar full-video receipts disagree") + receipt = receipt if isinstance(receipt, dict) else saved + if not isinstance(receipt, dict): + raise ValueError("missing revisioned full-video-integrity receipt") + cameras = _json(Path(run["path"]).with_name("initial-contract.json"))[ + "camera_order" + ] + expected_frames = sum(row["frame_count"] for row in sources) + expected_cameras = {"observation.images." + camera for camera in cameras} + expected_decoded = {camera: expected_frames for camera in expected_cameras} + if not ( + receipt.get("revision") == "full-video-integrity/v1" + and receipt.get("passed") is True + and receipt.get("all_saved_video_frames_decoded") is True + and receipt.get("expected_episodes") == published + and receipt.get("expected_frames_per_camera") == expected_frames + and receipt.get("decoded_frames_per_camera") == expected_decoded + and receipt.get("total_decoded_camera_frames") + == expected_frames * len(cameras) + ): + raise ValueError( + "full-video receipt counts/status do not cover the saved corpus" + ) + verifier = receipt.get("verifier_sha256") + metadata = receipt.get("metadata_sha256") + valid_sha = ( + lambda value: isinstance(value, str) + and re.fullmatch(r"[0-9a-f]{64}", value) is not None + ) + if ( + not valid_sha(verifier) + or not isinstance(metadata, dict) + or not metadata + or not all(valid_sha(value) for value in metadata.values()) + ): + raise ValueError("full-video verifier/metadata fingerprints are missing") + files = receipt.get("files") + file_frames: Counter[str] = Counter() + episode_coverage: dict[str, set[int]] = { + camera: set() for camera in expected_cameras + } + if not isinstance(files, list) or not files: + raise ValueError("full-video receipt has no per-file checks") + for row in files: + if not ( + row.get("camera") in expected_cameras + and row.get("passed") is True + and row.get("probe_exitcode") == 0 + and row.get("decode_exitcode") == 0 + and type(row.get("expected_frames")) is int + and row["expected_frames"] > 0 + and row.get("decoded_frames") == row["expected_frames"] + and row.get("decoded_pixel_format") == "yuv420p" + and valid_sha(row.get("compressed_sha256")) + ): + raise ValueError( + "failed/incomplete per-file full-video decode evidence" + ) + file_frames[row["camera"]] += row["decoded_frames"] + episode_coverage[row["camera"]].update(row.get("episodes", [])) + if dict(file_frames) != expected_decoded or any( + coverage != set(range(published)) for coverage in episode_coverage.values() + ): + raise ValueError("per-file full-video evidence omits saved sources/frames") + except (OSError, ValueError, TypeError, KeyError) as error: + errors.append(f"new formal full-video integrity: {error}") + return errors + + +def _paired(runs: list[dict[str, Any]]) -> list[dict[str, Any]]: + rows = [] + pilot = all(run["config"].get("phase") == "pilot" for run in runs) + baselines = ( + ("offline32", "ods_serial") + if pilot + else ("offline8", "offline32", "offline128", "ods_serial") + ) + for vision in _VISIONS: + overlap = _group(runs, "ods_overlap", vision) + for arm in baselines: + baseline = _group(runs, arm, vision) + seeds = sorted( + {str(run["config"].get("seed")) for run in overlap + baseline} + ) + for seed in seeds: + left = [ + run for run in baseline if str(run["config"].get("seed")) == seed + ] + right = [ + run for run in overlap if str(run["config"].get("seed")) == seed + ] + reasons: list[str] = [] + evidence = {} + if ( + pilot + and len( + { + str(run["config"].get("seed")) + for run in runs + if run["config"]["vision"] == vision + } + ) + != 1 + ): + reasons.append("pilot scope requires exactly one seed per vision") + if len(left) != 1 or len(right) != 1: + reasons.append( + "requires exactly one qualified run per arm and seed" + ) + else: + for role, run in (("baseline", left[0]), ("overlap", right[0])): + facts, errors = _fairness(run) + evidence[role] = facts + reasons.extend(f"{role}: {error}" for error in errors) + if run["config"].get("phase") not in {"report", "pilot"}: + reasons.append( + f"{role}: comparative speedup requires phase=report or pilot, observed {run['config'].get('phase')!r}" + ) + audit = run["summary"].get("post_audit") + if ( + run["summary"].get("post_audit_status") != "complete" + or not isinstance(audit, dict) + or audit.get("passed") is not True + ): + reasons.append( + f"{role}: missing successful post-timing corpus audit" + ) + if left[0]["config"].get("phase") != right[0]["config"].get( + "phase" + ): + reasons.append("pilot and formal phases cannot be paired") + if evidence["baseline"].get("prefetch_policy") != evidence[ + "overlap" + ].get("prefetch_policy"): + reasons.append("paired prefetch_policy differs") + for key in ( + "handoff_mode", + "encoder_threads", + "training_budget", + "initial_model_sha256", + "normalization_sha256", + "initial_numeric_prefix", + "camera_order", + "allocation", + "auxiliary_graphics", + "requested_evaluation_scenes", + ): + if evidence["baseline"].get(key) != evidence["overlap"].get( + key + ): + reasons.append(f"paired {key} differs") + row = { + "vision": vision, + "baseline": arm, + "seed": seed, + "excluded_reason": "; ".join(reasons) if reasons else None, + "excluded_reasons": reasons, + "fairness_evidence": evidence, + } + if pilot: + row["scope"] = ( + "single-seed pilot; not a three-seed formal comparison" + ) + if not reasons: + offline = float(left[0]["metrics"]["total_seconds"]) + ods = float(right[0]["metrics"]["total_seconds"]) + row.update( + baseline_seconds=offline, + overlap_seconds=ods, + baseline_over_overlap=offline / ods, + overlap_time_saving_fraction=1 - ods / offline, + baseline_trained_sources=left[0]["metrics"]["trained_sources"], + overlap_trained_sources=right[0]["metrics"]["trained_sources"], + baseline_initial_xy_cells=left[0]["metrics"][ + "initial_xy_occupied_cells" + ], + overlap_initial_xy_cells=right[0]["metrics"][ + "initial_xy_occupied_cells" + ], + ) + rows.append(row) + return rows + + +def _save(figure: Any, output: Path, name: str) -> None: + for suffix in ("png", "svg", "pdf"): + figure.savefig(output / f"{name}.{suffix}", dpi=160, bbox_inches="tight") + + +def _style(axis: Any) -> None: + axis.spines[["top", "right"]].set_visible(False) + axis.grid(alpha=0.2) + axis.set_axisbelow(True) + + +def _tradeoff(runs: list[dict[str, Any]], plt: Any, output: Path) -> None: + fig, axes = plt.subplots(2, 2, figsize=(12, 8), squeeze=False) + for row, vision in enumerate(_VISIONS): + for col, (metric, label) in enumerate( + ( + ("initial_xy_occupied_cells", "Initial XY cells (of 64)"), + ("effective_sources", "Effective trained sources"), + ) + ): + axis = axes[row, col] + for arm in _ARMS: + group = _group(runs, arm, vision) + if group: + axis.scatter( + [run["metrics"]["total_seconds"] for run in group], + [run["metrics"][metric] for run in group], + color=_COLORS[arm], + marker=_MARKERS[arm], + s=55, + alpha=0.8, + label=_LABELS[arm], + ) + axis.set( + title=f"{vision}: time–coverage tradeoff", + xlabel="End-to-end time (s)", + ylabel=label, + ) + axis.set_ylim(bottom=0) + if col == 0: + axis.set_ylim(top=66) + _style(axis) + handles, labels = axes[0, 0].get_legend_handles_labels() + if handles: + fig.legend(handles, labels, loc="lower center", ncol=3) + fig.tight_layout(rect=(0, 0.08, 1, 1)) + _save(fig, output, "01_time_diversity") + plt.close(fig) + + +def _curves(runs: list[dict[str, Any]], plt: Any, output: Path) -> None: + fig, axes = plt.subplots(2, 2, figsize=(12, 8), squeeze=False) + for row, vision in enumerate(_VISIONS): + for col, (x_key, label) in enumerate( + ( + ("time_seconds", "Elapsed time (s)"), + ("update", "Completed optimizer updates"), + ) + ): + axis = axes[row, col] + for arm in _ARMS: + for index, run in enumerate(_group(runs, arm, vision)): + curve = run["metrics"].get("coverage_curve", []) + if curve: + axis.step( + [0] + [point[x_key] for point in curve], + [0] + + [point["initial_xy_occupied_cells"] for point in curve], + where="post", + color=_COLORS[arm], + alpha=0.55, + label=_LABELS[arm] if index == 0 else None, + ) + axis.set( + title=f"{vision}: coverage after completed training", + xlabel=label, + ylabel="Initial XY cells (of 64)", + ylim=(0, 66), + ) + _style(axis) + handles, labels = axes[0, 0].get_legend_handles_labels() + if handles: + fig.legend(handles, labels, loc="lower center", ncol=3) + fig.tight_layout(rect=(0, 0.08, 1, 1)) + _save(fig, output, "02_coverage_growth") + plt.close(fig) + + +def _stage(name: str) -> str: + lowered = name.lower() + if lowered == "sdk_recording": + return "SDK recording service" + if lowered == "sdk_decode": + return "SDK temporal fetch" + if "generation_drain_wait" in lowered: + return "finalize / cleanup" + if lowered.endswith("_wait") or "quota_wait" in lowered: + return "waiting" + if any( + word in lowered + for word in ("generation", "refill", "initial_fill", "acquisition") + ): + return "generation" + if any(word in lowered for word in ("finaliz", "drain", "close", "shutdown")): + return "finalize / cleanup" + if any( + word in lowered + for word in ("record", "decode", "copy", "sampl", "read", "batch", "transfer") + ): + return "data preparation" + if any(word in lowered for word in ("train", "optimizer", "backward", "forward")): + return "training" + return "other" + + +def _timeline(runs: list[dict[str, Any]], plt: Any, output: Path) -> None: + stages = ( + "generation", + "SDK recording service", + "SDK temporal fetch", + "data preparation", + "training", + "waiting", + "finalize / cleanup", + "other", + ) + colors = dict( + zip( + stages, + ( + "#4c78a8", + "#f2b134", + "#b279a2", + "#b3a9cf", + "#159b80", + "#b7a36a", + "#c65656", + "#999999", + ), + ) + ) + fig, axes = plt.subplots(2, 1, figsize=(14, 10), squeeze=False) + for row, vision in enumerate(_VISIONS): + axis = axes[row, 0] + eligible = [ + run for run in runs if run["config"]["vision"] == vision and _qualified(run) + ] + seeds = sorted({str(run["config"].get("seed")) for run in eligible}) + seed = seeds[0] if seeds else None + selected = [run for run in eligible if str(run["config"].get("seed")) == seed] + labels, locations = [], [] + y = 0 + for arm in _ARMS: + candidates = [run for run in selected if run["arm"] == arm] + if len(candidates) != 1: + continue + run = candidates[0] + intervals = run["metrics"].get("intervals", []) + for stage in stages: + spans = [ + (item["start_seconds"], item["end_seconds"] - item["start_seconds"]) + for item in intervals + if _stage(str(item.get("name", ""))) == stage + ] + if spans: + axis.broken_barh( + spans, (y, 0.7), facecolors=colors[stage], alpha=0.8 + ) + labels.append(f"{_LABELS[arm]} · {stage}") + locations.append(y + 0.35) + y += 1 + end = run["metrics"]["training_complete_seconds"] + total = run["metrics"]["total_seconds"] + axis.plot( + [end, end], + [y - len(stages), y - 0.3], + color="#222222", + linewidth=0.7, + linestyle=":", + ) + axis.plot( + [total, total], + [y - len(stages), y - 0.3], + color="#222222", + linewidth=0.7, + ) + y += 0.6 + axis.set_yticks(locations, labels, fontsize=6) + axis.invert_yaxis() + axis.set( + title=f"{vision}: actual stages, seed {seed}; dotted=last update, solid=normal exit", + xlabel="Elapsed time (s)", + ) + _style(axis) + fig.tight_layout() + _save(fig, output, "03_pipeline_timeline") + plt.close(fig) + + +def _reuse(runs: list[dict[str, Any]], plt: Any, output: Path) -> None: + fig, axes = plt.subplots(2, 1, figsize=(12, 8), squeeze=False) + for row, vision in enumerate(_VISIONS): + axis = axes[row, 0] + for arm in _ARMS: + counts = [ + int(episode["trained_chunks"]) + for run in _group(runs, arm, vision) + for episode in run["metrics"].get("episodes", []) + ] + if counts: + histogram = Counter(counts) + x = sorted(histogram) + axis.step( + x, + [histogram[count] for count in x], + where="mid", + color=_COLORS[arm], + label=_LABELS[arm], + ) + axis.scatter( + x, + [histogram[count] for count in x], + color=_COLORS[arm], + s=14, + alpha=0.65, + ) + caps = sorted( + { + run["metrics"].get("reuse_cap") + for run in runs + if run["config"]["vision"] == vision + and run["arm"] in {"ods_overlap", "ods_serial"} + and _qualified(run) + and run["metrics"].get("reuse_cap") is not None + } + ) + for cap in caps: + axis.axvline( + cap, color="#555555", linestyle="--", linewidth=1, label=f"ODS K={cap}" + ) + axis.set_xscale("symlog", linthresh=1) + axis.set( + title=f"{vision}: all published sources, including retired and never trained", + xlabel="Lifetime completed training chunks per source (symlog)", + ylabel="Sources across qualified seeds", + ) + axis.set_xlim(left=0) + axis.set_ylim(bottom=0) + _style(axis) + handles, labels = axis.get_legend_handles_labels() + if handles: + axis.legend(handles, labels, fontsize=8, ncol=3) + fig.tight_layout() + _save(fig, output, "04_lifetime_reuse") + plt.close(fig) + + +def _formatted(stats: dict[str, Any]) -> str: + if not stats["n"]: + return "—" + return f"{stats['median']:.3g} [{stats['min']:.3g}, {stats['max']:.3g}]" + + +def _resource_stats(runs: list[dict[str, Any]]) -> dict[str, Any]: + values: dict[str, list[float]] = { + "cpu_assigned_percent": [], + "simulation_gpu_board_percent": [], + "training_gpu_board_percent": [], + "sampled_pss_gib": [], + "sampled_owned_gpu_gib": [], + "dataset_mib": [], + "checkpoint_mib": [], + } + for run in runs: + resources = run["metrics"].get("resources") or {} + value = _number( + (resources.get("cpu_utilization") or {}).get( + "mean_assigned_capacity_percent" + ) + ) + if value is not None: + values["cpu_assigned_percent"].append(value) + for key, container, field, scale in ( + ("sampled_pss_gib", resources, "peak_pss_bytes", 1024**3), + ("sampled_owned_gpu_gib", resources, "peak_vram_bytes", 1024**3), + ("dataset_mib", run["metrics"], "dataset_bytes", 1024**2), + ("checkpoint_mib", run["metrics"], "checkpoint_bytes", 1024**2), + ): + value = _number(container.get(field)) + if value is not None: + values[key].append(value / scale) + pair = run["summary"].get("device_supervision", {}).get("primary_pair", {}) + boards = resources.get("primary_gpu_utilization") or {} + for role in ("simulation", "training"): + uuid = pair.get(role, {}).get("uuid") + value = _number(boards.get(uuid, {}).get("mean_sampled_gpu_percent")) + if value is not None: + values[role + "_gpu_board_percent"].append(value) + return {key: _stats(observed) for key, observed in values.items()} + + +def write_report(root: Path, output: Path | None = None) -> dict[str, Any]: + """Write four measured figures and qualified/unqualified JSON/Markdown. + + Args: + root: Directory recursively containing per-run ``summary.json`` files. + output: Report directory. Defaults to ``root/training-report``. + + Returns: + Finite JSON report, including every observed run and ratio exclusion. + """ + root = Path(root) + output = Path(output) if output is not None else root / "training-report" + if ( + (root / "main").is_dir() + and (root / "suite-summary.json").is_file() + and _json(root / "suite-summary.json").get("phase") == "pilot" + ): + root = root / "main" + output.mkdir(parents=True, exist_ok=True) + runs = _load(root) + aggregates = _aggregate(runs) + pairs = _paired(runs) + pilot = all(run["config"].get("phase") == "pilot" for run in runs) + display_arms = ("offline32", "ods_overlap", "ods_serial") if pilot else _ARMS + report = { + "schema": "ods-training-report/v1", + "input": str(root.resolve()), + "qualified_runs": sum(_qualified(run) for run in runs), + "observed_runs": len(runs), + "qualified_paired_comparisons": sum( + row["excluded_reason"] is None for row in pairs + ), + "qualification_scope": "Per-run completed-work qualification; formal paired ratios additionally require matching model/stat/initial-prefix/resource evidence and successful corpus audit.", + "aggregates": aggregates, + "resource_aggregates": { + vision: {arm: _resource_stats(_group(runs, arm, vision)) for arm in _ARMS} + for vision in _VISIONS + }, + "paired_comparisons": pairs, + "runs": [ + { + **run["summary"], + "path": run["path"], + "arm": run["arm"], + "qualification_reasons": _qualification(run), + } + for run in runs + ], + } + if pilot: + report["experiment_scope"] = ( + "single-seed systems pilot; eager diagnostics excluded; no policy-learning claim" + ) + (output / "report.json").write_text( + json.dumps(report, indent=2, allow_nan=False) + "\n", encoding="utf-8" + ) + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + for plot in (_tradeoff, _curves, _timeline, _reuse): + plot(runs, plt, output) + lines = [ + "# ODS visual training: fixed-work time and diversity", + "", + f"Qualified runs: **{report['qualified_runs']}/{len(runs)}**; qualified {'pilot' if pilot else 'formal'} pairs: **{report['qualified_paired_comparisons']}**. Failed, unsupported and unqualified runs remain in report.json and are excluded from aggregate values and paired time ratios.", + "", + "Each result completes the declared optimizer-update, batch and valid-target budget. Data amount and final coverage may differ. Ratios compare total pipeline time, including acquisition, recording, decode, actual training, checkpoint and normal teardown; they do not establish equal policy quality or a universal ODS speedup.", + "", + ] + if pilot: + lines += [ + "Single-seed systems pilot. Eager controls remain outside main data and primary ratios. Zero-scene evaluation is systems-only and provides no policy-quality evidence.", + "", + ] + for vision in _VISIONS: + lines += [ + f"## {vision}", + "", + "Values are median [minimum, maximum] across qualified seeds.", + "", + "| Path | n | Total s | Trained sources | Effective sources | XY cells / 64 | Unique targets |", + "| --- | ---: | ---: | ---: | ---: | ---: | ---: |", + ] + for arm in display_arms: + group = aggregates[vision][arm] + lines.append( + f"| {_LABELS[arm]} | {group['total_seconds']['n']} | {_formatted(group['total_seconds'])} | {_formatted(group['trained_sources'])} | {_formatted(group['effective_sources'])} | {_formatted(group['initial_xy_occupied_cells'])} | {_formatted(group['unique_target_frames'])} |" + ) + lines += [ + "", + "| Baseline / ODS overlap | Paired seeds | Total-time ratio median [min, max] |", + "| --- | ---: | ---: |", + ] + for arm in ( + ("offline32", "ods_serial") + if pilot + else ("offline8", "offline32", "offline128", "ods_serial") + ): + values = [ + row["baseline_over_overlap"] + for row in pairs + if row["vision"] == vision + and row["baseline"] == arm + and row["excluded_reason"] is None + ] + lines.append( + f"| {_LABELS[arm]} | {len(values)} | {_formatted(_stats(values))} |" + ) + lines += [ + "", + "| Path | CPU mean / assigned capacity % | Simulation GPU board mean % | Training GPU board mean % |", + "| --- | ---: | ---: | ---: |", + ] + for arm in display_arms: + resource = report["resource_aggregates"][vision][arm] + lines.append( + f"| {_LABELS[arm]} | {_formatted(resource['cpu_assigned_percent'])} | {_formatted(resource['simulation_gpu_board_percent'])} | {_formatted(resource['training_gpu_board_percent'])} |" + ) + lines += [ + "", + "CPU values use continuing-PID sampled interval deltas normalized by the assigned cores; new or disappeared processes may be undercounted. GPU values are whole-board NVML samples and include foreign work; they do not attribute kernel utilization to this run. Missing observations remain unknown (—), not zero. See raw resources for sample counts and process-coverage gaps.", + "", + "| Path | Sampled PSS max GiB | Sampled owned GPU max GiB | Dataset MiB | Checkpoints MiB |", + "| --- | ---: | ---: | ---: | ---: |", + ] + for arm in display_arms: + resource = report["resource_aggregates"][vision][arm] + lines.append( + f"| {_LABELS[arm]} | {_formatted(resource['sampled_pss_gib'])} | {_formatted(resource['sampled_owned_gpu_gib'])} | {_formatted(resource['dataset_mib'])} | {_formatted(resource['checkpoint_mib'])} |" + ) + lines += [ + "", + "Memory maxima are sampled observations, not guaranteed allocator/process peaks. Disk sizes are logical finalized dataset and initial/final checkpoint bytes; they exclude temporary disk peaks, logs and post-timing audit/evaluation artifacts.", + "", + ] + lines += [ + "## Figures and interpretation", + "", + "![End-to-end time and actual trained-state coverage](01_time_diversity.png)", + "", + "![Coverage after completed optimizer updates](02_coverage_growth.png)", + "", + "![Measured pipeline stages](03_pipeline_timeline.png)", + "", + "![Lifetime source reuse, including retired sources](04_lifetime_reuse.png)", + "", + "Coverage uses completed training only. Initial XY bins are fixed 1 cm cells over the task's 8 cm × 8 cm reset area; anchor-state voxels are 2 cm. Source identity growth alone is not physical diversity. Overlapping prediction windows share target frames and are deduplicated by immutable frame identity. Published but untrained data does not increase these metrics.", + "", + "AV1 decoding is lossy. Compare source quality audits and held-out evaluation alongside timing. Differences between online raw pixels and decoded pixels prevent a claim of identical optimization trajectories. Serial trace replay isolates overlap only when its source/anchor manifest audit passes. A faster small offline run or broader large offline corpus is a valid result.", + "", + "SDK recording spans cover real create/write/encode/bookkeeping/finalize service; SDK temporal fetch spans include worker initialization, dataset reads, decode, conversion and validation. Neither lane claims pure encoder/decoder CPU work or GPU kernel activity. Owner packing/copy and waiting remain separate. Intervals may overlap and are not summed into total time. Timeline rows use the first observed qualified seed per vision, without selecting the fastest run. PNG, SVG and PDF versions share measured inputs.", + "", + "Formal ratios require matching initial-model/normalization SHA256, initial-C numeric prefix, camera order, declared/completed training budget, ordered primary GPU UUIDs and the same eight CPU cores. Retained supervisor launch masks/thread limits, placement guards, process exit and clocks must agree. New formal inputs additionally require chronological sdk-worker-service/v1 intervals and a revisioned full-video-integrity/v1 receipt for every saved frame; the original numeric/H2 audit remains required. Full-video YUV sequential integrity is distinct from lossless RGB fidelity or every SDK temporal seek. Missing evidence blocks claims; smoke/calibration or legacy CPU rendering fixtures are not formal comparisons. Post-timing corpus audit and evaluation durations are excluded from pipeline time.", + "", + "## Excluded paired claims", + "", + ] + excluded_pairs = [row for row in pairs if row["excluded_reason"] is not None] + if excluded_pairs: + for row in excluded_pairs: + lines.append( + f"- {row['vision']} / {row['baseline']} / seed {row['seed']}: {row['excluded_reason']}." + ) + else: + lines.append("Every observed paired comparison passed its evidence gates.") + lines += [ + "", + "## Retained unqualified runs", + "", + ] + excluded = [run for run in runs if not _qualified(run)] + if excluded: + for run in excluded: + lines.append(f"- `{run['path']}`: {', '.join(_qualification(run))}.") + else: + lines.append("No unqualified runs were observed in this input directory.") + (output / "report.md").write_text("\n".join(lines) + "\n", encoding="utf-8") + return report + + +def _main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("root", type=Path) + parser.add_argument("--output", type=Path) + args = parser.parse_args() + result = write_report(args.root, args.output) + print( + json.dumps( + { + "qualified_runs": result["qualified_runs"], + "observed_runs": result["observed_runs"], + } + ) + ) + + +if __name__ == "__main__": + _main() diff --git a/examples/ods/data_stream_demo/training_suite.py b/examples/ods/data_stream_demo/training_suite.py new file mode 100644 index 000000000..11d1c0f54 --- /dev/null +++ b/examples/ods/data_stream_demo/training_suite.py @@ -0,0 +1,952 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Supervise visual BC pipelines on disjoint generation/training GPU pairs.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import os +import signal +import subprocess +import sys +import time +from pathlib import Path +from typing import Any + +import psutil + +from .suite import ( + _gpu_contexts, + _gpu_reading, + _hashes, + _memory_bytes, + _terminate, + _write, +) + +__all__ = ["run_suite"] + +_ROOT = Path(__file__).resolve().parents[3] +_SETTINGS = ( + "updates", + "prefetch_policy", + "handoff_mode", + "encoder_threads", + "batch_size", + "buffer_size", + "producer_envs", + "episode_budget", + "history_length", + "prediction_horizon", + "action_horizon", + "max_episode_steps", + "initialization_timeout", + "generation_timeout", + "data_timeout", + "run_timeout", +) +_AUX_LIMIT = 8 * 1024**2 +_LIMITS = [ + "Each primary GPU passes three zero-utilization readings and has at least 16 GiB free; admission is not an exclusive reservation.", + "One generation/training UUID pair and eight disjoint CPUs are held for each seed/vision group; at most two groups run concurrently.", + "Outside a pipeline's primary pair, an explicit auxiliary UUID permits only owned pure-G contexts with known per-PID memory <=8 MiB. It may also be an explicitly selected idle primary of another pipeline; per-process-group limits still apply.", + "Periodic process-tree/process-group GPU audits can miss short-lived contexts; matching native UUID placement still needs qualification.", + "Total pipeline time ends after normal training-worker exit; independent evaluation is outside that time.", +] + + +def _devices(args: argparse.Namespace) -> tuple[list[dict[str, Any]], dict[str, Any]]: + readings = [] + for index in range(3): + readings.append({"monotonic_ns": time.monotonic_ns(), "gpus": _gpu_reading()}) + if index < 2: + time.sleep(2) + _write( + args.output / "gpu-preflight.json", + {"readings": readings, "limitations": _LIMITS}, + ) + by_id = [{gpu["index"]: gpu for gpu in row["gpus"]} for row in readings] + requested, aux_id = args.gpu_ids, args.aux_graphics_gpu_id + ids = list(requested) if requested is not None else sorted(by_id[-1]) + if len(set(ids)) != len(ids) or any(type(i) is not int or i < 0 for i in ids): + raise ValueError( + "Primary physical GPU IDs must be distinct nonnegative integers" + ) + if aux_id is not None and aux_id < 0: + raise ValueError("Auxiliary GPU must be a nonnegative physical ID") + selected, rejected = [], [] + for gpu_id in ids: + if gpu_id == aux_id and requested is None: + continue + if all( + gpu_id in row + and row[gpu_id]["utilization_percent"] == 0 + and row[gpu_id]["free_mib"] >= 16384 + for row in by_id + ): + selected.append(by_id[-1][gpu_id]) + else: + rejected.append( + { + "index": gpu_id, + "reason": "not present, not idle in all three readings, or below 16 GiB free", + } + ) + auxiliary = None + if aux_id is not None: + if not all(aux_id in row and row[aux_id]["free_mib"] >= 16384 for row in by_id): + raise RuntimeError( + "Passive graphics auxiliary GPU requires 16 GiB free in all three readings" + ) + auxiliary = { + **by_id[-1][aux_id], + "scope": "passive pure-G driver footprint only", + "max_per_pid_bytes": _AUX_LIMIT, + "idle_required": False, + } + cpus = sorted(os.sched_getaffinity(0)) + count = min(2, len(selected) // 2, len(cpus) // 8) + pairs = [ + { + "simulation": selected[2 * i], + "training": selected[2 * i + 1], + "cpu_affinity": cpus[8 * i : 8 * i + 8], + } + for i in range(count) + ] + audit = { + "readings": readings, + "pairs": pairs, + "rejected": rejected, + "aux_graphics_gpu": auxiliary, + "limitations": _LIMITS, + } + _write(args.output / "gpu-preflight.json", audit) + if not pairs: + raise RuntimeError( + "Need two idle primary GPUs with 16 GiB free and eight available CPUs" + ) + _gpu_contexts() + return pairs, audit + + +def _groups(args: argparse.Namespace) -> list[list[dict[str, Any]]]: + phase = "calibration" if args.suite == "calibrate" else args.suite + settings = { + key: getattr( + args, + key, + { + "history_length": 2, + "action_horizon": 16, + "prefetch_policy": "budget", + "handoff_mode": "publication", + "encoder_threads": 1, + }.get(key), + ) + for key in _SETTINGS + } + if args.suite == "smoke": + settings.update( + buffer_size=2, + producer_envs=1, + updates=4, + batch_size=4, + episode_budget=4, + prediction_horizon=32, + ) + elif args.suite == "calibrate": + for key in ("buffer_size", "producer_envs", "batch_size", "episode_budget"): + if type(settings[key]) is not int or settings[key] < 1: + raise ValueError(f"Calibration {key} must be a positive integer") + updates = getattr(args, "calibration_updates", None) + if updates is not None and (type(updates) is not int or updates < 1): + raise ValueError("calibration_updates must be a positive integer") + settings["updates"] = ( + updates + if updates is not None + else max( + 16, + math.ceil( + settings["buffer_size"] + * settings["episode_budget"] + / settings["batch_size"] + ) + + 8, + ) + ) + seeds = (19001, 19002, 19003) if args.suite == "report" else (args.seed,) + groups = [] + for seed in seeds: + for vision in ( + (args.vision,) if args.suite == "single" else ("single", "dual") + ): + if args.suite == "single": + paths = [(args.arm, args.offline_episodes)] + elif args.suite == "calibrate": + paths = [("ods_overlap", args.offline_episodes)] + else: + counts = ( + (8, 32, 128) + if args.suite == "report" + else ((32,) if args.suite == "pilot" else (4,)) + ) + paths = [("offline", n) for n in counts] + if (seed + (vision == "dual")) % 2: + paths.insert(0, ("ods_overlap", args.offline_episodes)) + else: + paths.append(("ods_overlap", args.offline_episodes)) + paths.append(("ods_serial", args.offline_episodes)) + if args.suite == "pilot": + eager = ("ods_overlap", args.offline_episodes) + # Controls remain separate from primary data and alternate their + # position; Serial still follows its primary Overlap trace. + paths.insert( + len(paths) if (seed + (vision == "dual")) % 2 else 0, eager + ) + jobs = [] + for index, (arm, episodes) in enumerate(paths): + control = args.suite == "pilot" and index == ( + len(paths) - 1 if (seed + (vision == "dual")) % 2 else 0 + ) + label = f"offline{episodes}" if arm == "offline" else arm + output = ( + args.output + if args.suite == "single" + else ( + args.output / "controls" / "eager" / vision / f"seed{seed}" + if control + else args.output + / ("main" if args.suite == "pilot" else "") + / label + / vision + / f"seed{seed}" + ) + ) + trace = ( + args.trace + if args.suite == "single" + else ( + args.output + / ("main" if args.suite == "pilot" else "") + / "ods_overlap" + / vision + / f"seed{seed}" + / "training-trace.json" + if arm == "ods_serial" + else None + ) + ) + scenes = ( + 0 + if args.suite == "calibrate" + else (2 if args.suite == "smoke" else args.evaluation_scenes) + ) + job_settings = settings.copy() + if args.suite == "pilot": + job_settings["prefetch_policy"] = "eager" if control else "budget" + jobs.append( + { + "arm": arm, + "vision": vision, + "seed": seed, + "phase": phase, + "output": output, + "offline_episodes": episodes, + "trace": trace, + "evaluation_scenes": scenes, + "settings": job_settings, + } + ) + groups.append(jobs) + return groups + + +def _config(job: dict[str, Any]) -> dict[str, Any]: + return { + **job["settings"], + **{ + key: job[key] + for key in ( + "arm", + "vision", + "seed", + "phase", + "offline_episodes", + "evaluation_scenes", + ) + }, + } + + +def _mark(job: dict[str, Any], status: str, reason: str) -> dict[str, Any]: + path = job["output"] / "summary.json" + if path.exists(): + try: + summary = json.loads(path.read_text()) + except ValueError: + path.rename(path.with_name("summary.partial.json")) + summary = {} + else: + summary = {} + summary.setdefault("config", _config(job)) + summary.update(status=status, qualification_passed=False, supervisor_error=reason) + summary.setdefault("error", {"phase": "supervisor", "message": reason}) + _write(path, summary) + return summary + + +def _command(job: dict[str, Any], stage: str = "training") -> list[str]: + options: dict[str, Any] = {} + if stage == "audit": + module = "examples.ods.data_stream_demo.training_audit" + options.update(run_dir=job["output"], output=job["output"] / "post-audit.json") + else: + options.update(vision=job["vision"], sim_device="cuda:0", train_device="cuda:1") + if stage == "evaluation": + module = "examples.ods.data_stream_demo.training_eval" + options.update( + checkpoint=job["output"] / "final.pt", + output=job["output"] / "evaluation.json", + scenes=job["evaluation_scenes"], + ) + else: + module = "examples.ods.data_stream_demo.train" + options.update( + suite="single", + phase=job["phase"], + arm=job["arm"], + seed=job["seed"], + output=job["output"], + offline_episodes=job["offline_episodes"], + evaluation_scenes=0, + **job["settings"], + ) + if job["trace"] is not None: + options["trace"] = job["trace"] + command = [sys.executable, "-m", module] + ( + ["--worker"] if stage == "training" else [] + ) + for name, value in options.items(): + command += ["--" + name.replace("_", "-"), str(value)] + return command + + +def _trace_ready(job: dict[str, Any]) -> bool: + if job["trace"] is None or not Path(job["trace"]).is_file(): + return False + if job["phase"] == "single": + return True + try: + original = json.loads(Path(job["trace"]).with_name("summary.json").read_text()) + return ( + original.get("status") == "complete" + and original.get("qualification_passed") is True + ) + except (OSError, ValueError): + return False + + +def _launch( + slot: dict[str, Any], + job: dict[str, Any], + root: Path, + auxiliary: dict[str, Any] | None, + *, + stage: str = "training", +) -> dict[str, Any]: + pair = slot["pair"] + relative = ( + Path("single") if job["output"] == root else job["output"].relative_to(root) + ) + supervision = root / "_supervision" / relative / stage + supervision.mkdir(parents=True, exist_ok=False) + environment = os.environ.copy() + environment.update( + CUDA_VISIBLE_DEVICES=pair["simulation"]["uuid"] + + "," + + pair["training"]["uuid"], + ODS_GPU_UUID=pair["simulation"]["uuid"], + SDL_VIDEODRIVER="dummy", + OMP_NUM_THREADS="1", + MKL_NUM_THREADS="1", + OPENBLAS_NUM_THREADS="1", + ) + for name in ("DISPLAY", "WAYLAND_DISPLAY", "ODS_AUX_GRAPHICS_GPU_UUID"): + environment.pop(name, None) + if auxiliary is not None: + environment["ODS_AUX_GRAPHICS_GPU_UUID"] = auxiliary["uuid"] + command = _command(job, stage) + stream = (supervision / "process.log").open("w", encoding="utf-8") + started_ns = time.monotonic_ns() + environment["ODS_PROCESS_LAUNCH_MONOTONIC_NS"] = str(started_ns) + try: + process = subprocess.Popen( + command, + cwd=_ROOT, + env=environment, + stdout=stream, + stderr=subprocess.STDOUT, + start_new_session=True, + preexec_fn=lambda: os.sched_setaffinity(0, pair["cpu_affinity"]), + ) + except BaseException: + stream.close() + raise + launch = { + "command": command, + "pair": pair, + "aux_graphics_scope": auxiliary, + "process_launch_monotonic_ns": started_ns, + "pgid": process.pid, + "cpu_affinity": pair["cpu_affinity"], + "timed_pipeline": stage == "training", + "environment": { + key: environment.get(key) + for key in ( + "CUDA_VISIBLE_DEVICES", + "ODS_GPU_UUID", + "ODS_AUX_GRAPHICS_GPU_UUID", + "OMP_NUM_THREADS", + "MKL_NUM_THREADS", + "OPENBLAS_NUM_THREADS", + "SDL_VIDEODRIVER", + "ODS_PROCESS_LAUNCH_MONOTONIC_NS", + ) + }, + } + _write(supervision / "launch.json", launch) + print( + f"Started {job['phase']} {job['arm']}/{job['vision']}/seed{job['seed']} {stage} on GPUs {pair['simulation']['index']},{pair['training']['index']}", + flush=True, + ) + return { + "process": process, + "stream": stream, + "started_ns": started_ns, + "deadline": started_ns / 1e9 + float(job["settings"]["run_timeout"]), + "supervision": supervision, + "evaluation": stage == "evaluation", + "stage": stage, + "job": job, + "allocation": pair, + "auxiliary_scope": auxiliary, + "guard": { + "samples": 0, + "unexpected_contexts": [], + "allowed_aux_contexts": [], + "seen_primary_uuids": [], + }, + "error": None, + } + + +def _audit(slots: list[dict[str, Any]], auxiliary: dict[str, Any] | None) -> None: + contexts = _gpu_contexts() + for slot in slots: + item = slot.get("active") + if item is None: + continue + process = item["process"] + owners = {process.pid} + try: + owners.update( + child.pid + for child in psutil.Process(process.pid).children(recursive=True) + ) + except (psutil.NoSuchProcess, psutil.AccessDenied): + pass + owned = [] + for row in contexts: + try: + matches = row["pid"] in owners or os.getpgid(row["pid"]) == process.pid + except ProcessLookupError: + matches = False + if matches: + owned.append(row) + primary = {slot["pair"][role]["uuid"] for role in ("simulation", "training")} + aux_uuid = auxiliary["uuid"] if auxiliary else None + invalid_aux = { + row["pid"] + for row in owned + if row["gpu_uuid"] == aux_uuid + and ( + row.get("type") != "G" + or _memory_bytes(row.get("used_memory")) is None + or _memory_bytes(row.get("used_memory")) > _AUX_LIMIT + ) + } + allowed, unexpected = [], [] + for row in owned: + if item["stage"] == "audit": + unexpected.append(row) + elif row["gpu_uuid"] in primary: + if row["gpu_uuid"] not in item["guard"]["seen_primary_uuids"]: + item["guard"]["seen_primary_uuids"].append(row["gpu_uuid"]) + elif ( + aux_uuid + and row["gpu_uuid"] == aux_uuid + and row["pid"] not in invalid_aux + ): + allowed.append(row) + else: + unexpected.append(row) + item["guard"]["samples"] += 1 + for row in allowed: + if row not in item["guard"]["allowed_aux_contexts"]: + item["guard"]["allowed_aux_contexts"].append(row) + if unexpected: + item["guard"]["unexpected_contexts"].extend(unexpected) + item["error"] = ( + f"Owned GPU placement violates pair/passive scope: {unexpected}" + ) + with (item["supervision"] / "device-contexts.jsonl").open( + "a", encoding="utf-8" + ) as stream: + stream.write( + json.dumps( + { + "monotonic_ns": time.monotonic_ns(), + "owned_contexts": owned, + "allowed_aux_contexts": allowed, + "unexpected_contexts": unexpected, + }, + allow_nan=False, + ) + + "\n" + ) + + +def _group_snapshot(pgid: int) -> dict[str, Any]: + """Capture identities only for members of this owned process group.""" + result: dict[str, Any] = { + "monotonic_ns": time.monotonic_ns(), + "pgid": pgid, + "exists": False, + "members": [], + "inspection_errors": [], + } + try: + os.killpg(pgid, 0) + except ProcessLookupError: + return result + except PermissionError as error: + result["inspection_errors"].append(repr(error)) + result["exists"] = True + for member in psutil.process_iter(): + try: + if os.getpgid(member.pid) != pgid: + continue + except ProcessLookupError: + continue + except PermissionError as error: + result["inspection_errors"].append( + {"pid": member.pid, "error": repr(error)} + ) + continue + row: dict[str, Any] = {"pid": member.pid} + try: + with member.oneshot(): + row.update( + ppid=member.ppid(), + create_time=member.create_time(), + name=member.name(), + state=member.status(), + ) + try: + row["cmdline"] = member.cmdline() + except (psutil.NoSuchProcess, psutil.AccessDenied) as error: + row["cmdline_error"] = repr(error) + except psutil.NoSuchProcess: + continue + except psutil.AccessDenied as error: + row["inspection_error"] = repr(error) + result["members"].append(row) + result["members"].sort(key=lambda row: row["pid"]) + states = [row.get("state") for row in result["members"]] + result["member_kind"] = ( + "unresolved" + if result["inspection_errors"] + or any("inspection_error" in row for row in result["members"]) + else ( + "zombie_only" + if states and all(state == psutil.STATUS_ZOMBIE for state in states) + else "live_or_unknown" if states else "unresolved" + ) + ) + return result + + +def _await_group_exit(pgid: int, *, timeout: float) -> dict[str, Any]: + """Allow brief normal helper cleanup; any group at the deadline fails.""" + started = time.monotonic() + initial = _group_snapshot(pgid) + latest = initial + polls = 1 + while latest["exists"] and time.monotonic() < started + timeout: + time.sleep(min(0.05, max(0.0, started + timeout - time.monotonic()))) + latest = _group_snapshot(pgid) + polls += 1 + return { + "grace_seconds": timeout, + "wait_seconds": time.monotonic() - started, + "poll_count": polls, + "group_disappeared": not latest["exists"], + "initial": initial, + "final": latest, + } + + +def _finish( + item: dict[str, Any], *, cancelled: bool +) -> tuple[dict[str, Any], dict[str, Any]] | None: + process, job = item["process"], item["job"] + timeout = time.monotonic() >= item["deadline"] + reason = item["error"] or ( + "suite cancellation" + if cancelled + else "process deadline exceeded" if timeout else None + ) + if process.poll() is None and reason is None: + return None + if reason is not None: + _terminate(process) + process.wait() + shutdown = _await_group_exit(process.pid, timeout=2.0) + leaked = not shutdown["group_disappeared"] + if leaked: + kind = shutdown["final"].get("member_kind", "unresolved") + reason = reason or ( + "Owned process group remains after normal-exit grace " + f"({kind}; see process_group_shutdown identities)" + ) + # A zombie cannot execute work or respond to signals. Its presence still + # fails the group-empty boundary; classification grants no exemption. + if kind != "zombie_only": + _terminate(process) + shutdown["termination_attempted"] = True + shutdown["after_termination"] = _await_group_exit(process.pid, timeout=2.0) + else: + shutdown["termination_attempted"] = False + ended_ns = time.monotonic_ns() + item["stream"].close() + elapsed = (ended_ns - item["started_ns"]) / 1e9 + if process.returncode != 0: + reason = reason or f"Child exit={process.returncode}" + stage = item["stage"] + artifact = { + "training": "summary.json", + "audit": "post-audit.json", + "evaluation": "evaluation.json", + }[stage] + path = job["output"] / artifact + if not path.exists(): + reason = reason or f"Child omitted {path.name}" + if stage != "training": + summary = json.loads((job["output"] / "summary.json").read_text()) + if reason is None: + payload = json.loads(path.read_text()) + valid = payload.get("passed") is True + if stage == "evaluation": + valid = ( + valid + and payload.get("closed") is True + and payload.get("scenes") == job["evaluation_scenes"] + and len(payload.get("records", [])) == job["evaluation_scenes"] + ) + if not valid: + reason = f"Post-timing {stage} did not complete its contract" + else: + summary["post_audit" if stage == "audit" else "evaluation"] = payload + key = "post_audit" if stage == "audit" else "evaluation" + summary[key + "_status"] = "failed" if reason else "complete" + summary[key + "_seconds_outside_pipeline"] = elapsed + if reason: + summary.update( + quality_error=reason, qualification_passed=False, status="failed" + ) + else: + summary = ( + _mark(job, "failed", reason) if reason else json.loads(path.read_text()) + ) + if ( + summary.get("status") == "complete" + and summary.get("qualification_passed") is not True + ): + summary = _mark(job, "failed", "Worker did not pass training qualification") + metrics = summary.get("training_metrics") + if isinstance(metrics, dict): + train_end = float(metrics.get("training_complete_seconds", 0)) + if train_end > elapsed: + summary = _mark( + job, "failed", "Child training timestamps exceed supervisor total" + ) + metrics = summary.get("training_metrics", {}) + metrics.update( + total_seconds=elapsed, post_training_drain_seconds=elapsed - train_end + ) + summary["supervisor_total_seconds"] = elapsed + summary["requested_evaluation_scenes"] = job["evaluation_scenes"] + summary["post_audit_status"] = "pending" + summary["evaluation_status"] = ( + "pending" if job["evaluation_scenes"] else "not_requested" + ) + summary["device_supervision"] = { + "primary_pair": item["allocation"], + "aux_graphics_scope": item["auxiliary_scope"], + "guard": item["guard"], + } + result = { + "output": str(job["output"]), + "status": summary["status"], + "returncode": process.returncode, + "elapsed_seconds": elapsed, + "stage": stage, + "evaluation": item["evaluation"], + "timeout": timeout, + "cancelled": cancelled, + "leaked_process_group": leaked, + "process_group_shutdown": shutdown, + "pgid": process.pid, + "process_launch_monotonic_ns": item["started_ns"], + "device_guard": item["guard"], + } + process_artifact = ( + "process.json" if stage == "training" else stage + "-process.json" + ) + _write(path.parent / process_artifact, result) + _write(job["output"] / "summary.json", summary) + return summary, result + + +def _calibration( + root: Path, results: list[dict[str, Any]], source_hashes: dict[str, str] +) -> dict[str, Any]: + from .training_throughput import write_calibration + + return write_calibration(root, results, source_hashes) + + +def run_suite(args: argparse.Namespace) -> dict[str, Any]: + """Run a finite matrix with normal-exit timing and separately timed evaluation.""" + args.output = Path(args.output).resolve() + if args.output.exists() and any(args.output.iterdir()): + raise FileExistsError( + "Use a fresh output directory for every training experiment" + ) + args.output.mkdir(parents=True, exist_ok=True) + groups = _groups(args) + source_hashes = _hashes() + for path in _ROOT.joinpath("embodichain/learning/models").glob("*.py"): + source_hashes[str(path.relative_to(_ROOT))] = hashlib.sha256( + path.read_bytes() + ).hexdigest() + results, cancelled = [], False + try: + pairs, audit = _devices(args) + except Exception as error: + for group in groups: + for job in group: + _mark(job, "unsupported", f"Device admission failed: {error!r}") + results.append({"output": str(job["output"]), "status": "unsupported"}) + summary = { + "status": "unsupported", + "error": repr(error), + "results": results, + "source_sha256": source_hashes, + } + _write(args.output / "suite-summary.json", summary) + return summary + slots = [{"pair": pair, "jobs": [], "active": None} for pair in pairs] + + def cancel(signum: int, frame: Any) -> None: + nonlocal cancelled + cancelled = True + + previous = { + sig: signal.signal(sig, cancel) for sig in (signal.SIGINT, signal.SIGTERM) + } + last_guard = last_progress = 0.0 + try: + while groups or any(slot["jobs"] or slot["active"] for slot in slots): + for slot in slots: + if slot["active"] is None and not cancelled: + if not slot["jobs"] and groups: + slot["jobs"] = groups.pop(0) + if slot["jobs"]: + job = slot["jobs"].pop(0) + if job["arm"] == "ods_serial" and not _trace_ready(job): + _mark( + job, + "unsupported", + "Matching overlap training trace is unavailable", + ) + results.append( + {"output": str(job["output"]), "status": "unsupported"} + ) + continue + try: + slot["active"] = _launch( + slot, job, args.output, audit["aux_graphics_gpu"] + ) + except Exception as error: + _mark(job, "failed", f"Launch failed: {error!r}") + results.append( + {"output": str(job["output"]), "status": "failed"} + ) + now = time.monotonic() + if now - last_guard >= 2 and any(slot["active"] for slot in slots): + try: + _audit(slots, audit["aux_graphics_gpu"]) + except Exception as error: + for slot in slots: + if slot["active"]: + slot["active"][ + "error" + ] = f"Cannot audit GPU placement: {error!r}" + last_guard = now + for slot in slots: + item = slot["active"] + if item is None: + continue + try: + finished = _finish(item, cancelled=cancelled) + except Exception as error: + _terminate(item["process"]) + item["stream"].close() + failed = _mark( + item["job"], "failed", f"Invalid child artifacts: {error!r}" + ) + finished = failed, { + "output": str(item["job"]["output"]), + "status": "failed", + "evaluation": item["evaluation"], + "error": repr(error), + } + if finished is None: + continue + summary, result = finished + job = item["job"] + if summary.get("status") == "complete" and not cancelled: + next_stage = ( + "audit" + if item["stage"] == "training" + else ( + "evaluation" + if item["stage"] == "audit" and job["evaluation_scenes"] > 0 + else None + ) + ) + if next_stage is not None: + try: + slot["active"] = _launch( + slot, + job, + args.output, + audit["aux_graphics_gpu"], + stage=next_stage, + ) + slot["active"]["pipeline_result"] = ( + result + if item["stage"] == "training" + else item["pipeline_result"] + ) + if item["stage"] == "audit": + slot["active"]["audit_result"] = result + continue + except Exception as error: + summary.update( + status="failed", + qualification_passed=False, + quality_error=f"{next_stage} launch failed: {error!r}", + ) + _write(job["output"] / "summary.json", summary) + result["status"] = "failed" + if item["stage"] != "training": + result = { + **item["pipeline_result"], + item["stage"] + "_process": result, + "status": summary["status"], + } + if "audit_result" in item: + result["audit_process"] = item["audit_result"] + results.append({**result, "status": summary["status"]}) + slot["active"] = None + print( + f"Finished {job['arm']}/{job['vision']}/seed{job['seed']}: {summary['status']}", + flush=True, + ) + if cancelled: + for group in groups + [slot["jobs"] for slot in slots]: + for job in group: + _mark(job, "failed", "Suite cancelled before launch") + results.append( + {"output": str(job["output"]), "status": "failed"} + ) + groups.clear() + for slot in slots: + slot["jobs"].clear() + if now - last_progress >= 30: + print( + f"Visual suite progress: {len(results)} finished; {sum(slot['active'] is not None for slot in slots)} active; {len(groups)} pending seed/vision groups", + flush=True, + ) + last_progress = now + if groups or any(slot["jobs"] or slot["active"] for slot in slots): + time.sleep(0.25) + finally: + for slot in slots: + if slot["active"]: + _terminate(slot["active"]["process"]) + slot["active"]["stream"].close() + for sig, handler in previous.items(): + signal.signal(sig, handler) + summary = { + "status": ( + "complete" + if all(result["status"] == "complete" for result in results) + else "failed" + ), + "phase": args.suite, + "results": results, + "gpu_preflight": audit, + "source_sha256": source_hashes, + "limitations": _LIMITS, + } + _write(args.output / "suite-summary.json", summary) + if args.suite == "calibrate": + summary["calibration"] = _calibration(args.output, results, source_hashes) + elif args.suite in {"report", "pilot"}: + from .training_report import write_report + + report_root = args.output / "main" if args.suite == "pilot" else args.output + summary["report"] = write_report(report_root, args.output / "training-report") + if args.suite == "pilot": + summary["scope"] = ( + "Single-seed systems pilot: six main outcomes, two eager diagnostics; no three-seed or policy-learning claim" + ) + summary["control_outputs"] = [ + str(job["output"]) + for group in _groups(args) + for job in group + if job["settings"]["prefetch_policy"] == "eager" + ] + _write(args.output / "suite-summary.json", summary) + return summary diff --git a/examples/ods/data_stream_demo/training_throughput.py b/examples/ods/data_stream_demo/training_throughput.py new file mode 100644 index 000000000..1f70bc471 --- /dev/null +++ b/examples/ods/data_stream_demo/training_throughput.py @@ -0,0 +1,399 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Measure generator, recorder and hot-learner service without changing reuse caps.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +from pathlib import Path +from typing import Any + +__all__ = ["analyze_run", "write_calibration", "compare_calibrations"] + + +def _read(path: Path) -> dict[str, Any]: + return json.loads(path.read_text(encoding="utf-8")) + + +def _write(path: Path, value: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(value, indent=2, allow_nan=False) + "\n") + + +def _duration(rows: list[dict[str, Any]]) -> float: + """Return interval union length, never a sum of overlapping services.""" + spans = sorted((float(r["start_seconds"]), float(r["end_seconds"])) for r in rows) + merged: list[list[float]] = [] + for start, end in spans: + if not (math.isfinite(start) and math.isfinite(end) and 0 <= start <= end): + raise ValueError("Service intervals must have finite causal clocks") + if merged and start <= merged[-1][1]: + merged[-1][1] = max(end, merged[-1][1]) + else: + merged.append([start, end]) + return sum(end - start for start, end in merged) + + +def _rate(count: int, seconds: float) -> float | None: + return count / seconds if count > 0 and seconds > 0 else None + + +def analyze_run(directory: Path, source_hashes: dict[str, str]) -> dict[str, Any]: + """Describe measured rates and whether a later source generation was trained. + + Generator and recording rates use accepted episodes per respective busy + interval union. The hot learner includes parent batch preparation, transfer + and synchronized updates, excluding quota/data waits. These independent + service times are not additive wall-clock components. + """ + directory = Path(directory) + summary = _read(directory / "summary.json") + config, metrics = summary["config"], summary.get("training_metrics", {}) + event_path = directory / "training-events.jsonl" + events = ( + [json.loads(s) for s in event_path.read_text().splitlines() if s] + if event_path.is_file() + else [] + ) + updates = [row for row in events if row.get("kind") == "update_completed"] + intervals = metrics.get("intervals", []) + generation = [row for row in intervals if row.get("name") == "generation"] + generation_ids = {row.get("generation") for row in generation} + accepted = [ + row + for row in metrics.get("episodes", []) + if row.get("generation") in generation_ids + ] + recording = [ + row + for row in intervals + if row.get("name") == "sdk_recording" + and row.get("sdk_phase") == "sdk_episode_write_encode" + ] + hot = [ + row + for row in intervals + if row.get("name") in {"batch_service", "device_transfer", "training"} + ] + generation_seconds, recording_seconds, hot_seconds = ( + _duration(rows) for rows in (generation, recording, hot) + ) + train_seconds = sum( + float(row["time_seconds"]) - float(row["train_started_seconds"]) + for row in updates + ) + all_generation_rate = _rate(len(accepted), generation_seconds) + later_generation = [ + row + for row in generation + if isinstance(row.get("generation"), int) and row["generation"] > 1 + ] + later_ids = {row["generation"] for row in later_generation} + later_accepted = [row for row in accepted if row["generation"] in later_ids] + later_seconds = _duration(later_generation) + later_rate = _rate(len(later_accepted), later_seconds) + gamma = later_rate if later_rate is not None else all_generation_rate + recording_rate, hot_rate = _rate(len(recording), recording_seconds), _rate( + len(updates), hot_seconds + ) + effective = min(gamma, recording_rate) if gamma and recording_rate else None + batch, cap = config.get("batch_size"), config.get("episode_budget") + supply = ( + effective * cap / batch + if effective + and type(batch) is int + and batch > 0 + and type(cap) is int + and cap > 0 + else None + ) + later_updates = sum( + any(c.get("generation", 0) > 1 for c in row.get("chunks", [])) + for row in updates + ) + trained_generations = sorted( + {c["generation"] for row in updates for c in row.get("chunks", [])} + ) + reasons = [] + if ( + summary.get("status") != "complete" + or summary.get("qualification_passed") is not True + ): + reasons.append("run did not complete its engineering qualification") + if len(updates) != metrics.get("completed_updates"): + reasons.append("completed-update event count disagrees with summary") + if not (1 in trained_generations and any(g > 1 for g in trained_generations)): + reasons.append("training did not consume both initial and later generations") + if not later_generation: + reasons.append("no measured post-initial generation service") + if later_updates < 8: + reasons.append( + "fewer than eight updates consumed later-generation data; steady-state evidence is short" + ) + if any(value is None for value in (gamma, recording_rate, hot_rate)): + reasons.append( + "generator, SDK recording or hot learner service evidence is missing" + ) + if len(recording) != metrics.get("published_sources") or len( + accepted + ) != metrics.get("published_sources"): + reasons.append("service spans do not cover every accepted episode") + for name in ("batch_service", "device_transfer", "training"): + if sum(row.get("name") == name for row in hot) != len(updates): + reasons.append(f"{name} service count differs from completed updates") + pair = summary.get("device_supervision", {}).get("primary_pair", {}) + hardware = { + "simulation_uuid": pair.get("simulation", {}).get("uuid"), + "training_uuid": pair.get("training", {}).get("uuid"), + "cpu_affinity": pair.get("cpu_affinity"), + "auxiliary_graphics": ( + { + key: summary.get("device_supervision", {}) + .get("aux_graphics_scope", {}) + .get(key) + for key in ("uuid", "scope", "max_per_pid_bytes", "idle_required") + } + if isinstance( + summary.get("device_supervision", {}).get("aux_graphics_scope"), dict + ) + else None + ), + } + if not all( + hardware.get(k) for k in ("simulation_uuid", "training_uuid", "cpu_affinity") + ): + reasons.append("actual GPU pair or CPU affinity is missing") + recipe = { + key: config.get(key) + for key in ( + "vision", + "updates", + "batch_size", + "buffer_size", + "episode_budget", + "history_length", + "prediction_horizon", + "action_horizon", + "max_episode_steps", + "seed", + ) + } + recipe["prefetch_policy"] = config.get("prefetch_policy", "eager") + recipe["handoff_mode"] = config.get("handoff_mode", "generation") + recipe["encoder_threads"] = config.get("encoder_threads", 1) + contract = {"hardware": hardware, "recipe": recipe, "source_sha256": source_hashes} + implied = ( + math.ceil(hot_rate * batch / effective) + if hot_rate and effective and type(batch) is int and batch > 0 + else None + ) + return { + "vision": config["vision"], + "status": summary["status"], + "output": str(directory), + "producer_envs": config.get("producer_envs"), + "updates": len(updates), + "accepted_episode_rate": all_generation_rate, + "generation_episodes_per_busy_second": gamma, + "generation_rate_scope": ( + "post-initial generation service" + if later_rate is not None + else "initial generation only; not eligible for P selection" + ), + "initial_fill_service_seconds": _duration( + [row for row in generation if row.get("generation") == 1] + ), + "later_generation_observed_seconds": later_seconds, + "later_generation_service_count": len(later_ids), + "later_generation_accepted_episodes": len(later_accepted), + "recording_episode_service_count": len(recording), + "learner_updates_per_second": _rate(len(updates), train_seconds), + "recording_episodes_per_busy_second": recording_rate, + "hot_learner_updates_per_second": hot_rate, + "effective_data_episodes_per_second": effective, + "supply_updates_per_second": supply, + "hot_demand_over_supply": hot_rate / supply if hot_rate and supply else None, + "implied_episode_budget_for_full_hot_utilization": implied, + "implied_budget_within_declared_cap": ( + implied <= cap if implied is not None and type(cap) is int else None + ), + "generation_observed_seconds": generation_seconds, + "recording_observed_seconds": recording_seconds, + "training_observed_seconds": train_seconds, + "hot_learner_observed_seconds": hot_seconds, + "trained_generations": trained_generations, + "later_generation_updates": later_updates, + "later_generation_consumed": any(g > 1 for g in trained_generations), + "selection_eligible": not reasons, + "later_generation_evidence_sufficient": 1 in trained_generations + and any(g > 1 for g in trained_generations) + and later_updates >= 8 + and bool(later_generation), + "long_run_stationarity_established": False, + "selection_exclusion_reasons": reasons, + "comparison_contract": contract, + "comparison_contract_sha256": hashlib.sha256( + json.dumps(contract, sort_keys=True).encode() + ).hexdigest(), + "limits": "Busy-service estimates, not steady-state end-to-end guarantees. Generation selection uses post-initial services; a default 40-update run may measure only one such generation. Implied K is diagnostic only: never increase a protected reuse cap automatically.", + } + + +def _method() -> str: + return """# Three-stage calibration method + +Measure the accepted generator rate γ and the SDK episode write/encode busy +rate ρ independently. Preserve the all-generation rate for compatibility, but +use post-initial generations for effective supply and P selection: the initial +fill includes simulation startup. Report the later-generation service count; +the default 40 updates may observe only one later generation. + SDK create/finalize and queue waits remain separate +costs. Effective data service is min(γ, ρ). Pure model-update rate is retained; +the hot learner rate ν instead includes batch preparation, transfer and real +synchronized updates, excluding data/quota waits. Merge overlapping intervals +within each service; never sum independent stage times into total wall time. + +At fixed batch B and protected lifetime episode budget K, the supply ceiling +is min(γ,ρ) K/B updates/s. Demand/supply is ν divided by that ceiling. The +implied ceil(ν B/min(γ,ρ)) is a diagnostic, not permission or a recommendation +to raise K. A demand/supply ratio above one indicates expected data waiting. + +Consume the initial cache and at least eight later-generation updates before +using a candidate for P selection; this is a short service-capacity estimate, +not a steady-state end-to-end throughput guarantee. Explicit shorter calibrations remain visible and are +excluded from selection. With the same hardware, source and recipe, select +the smallest measured P reaching 90% of the best effective data rate. Do not +compare different camera modes, hardware placements or training recipes. +""" + + +def write_calibration( + root: Path, results: list[dict[str, Any]], source_hashes: dict[str, str] +) -> dict[str, Any]: + """Write per-run three-stage observations and their measurement definitions.""" + packet = { + "schema": "ods-visual-training-calibration/v2", + "observations": [ + analyze_run(Path(r["output"]), source_hashes) for r in results + ], + "source_sha256": source_hashes, + "limits": "Calibration only; service spans overlap and must not be added as wall time. A short or unqualified calibration never recommends a producer count or a larger K.", + } + _write(root / "calibration.json", packet) + (root / "calibration.md").write_text(_method(), encoding="utf-8") + return packet + + +def compare_calibrations(roots: list[Path], output: Path) -> dict[str, Any]: + """Select the smallest measured P near the best rate only for matching inputs.""" + observations = [ + row + for root in roots + for row in _read(Path(root) / "calibration.json")["observations"] + ] + selections = [] + for vision in sorted({row["vision"] for row in observations}): + candidates = [row for row in observations if row["vision"] == vision] + eligible = [ + row + for row in candidates + if row.get("selection_eligible") is True + and type(row.get("producer_envs")) is int + and row["producer_envs"] > 0 + and isinstance(row.get("comparison_contract"), dict) + and row.get("effective_data_episodes_per_second", 0) > 0 + ] + fingerprints = { + json.dumps(row["comparison_contract"], sort_keys=True) for row in eligible + } + reasons = [] + if not eligible: + reasons.append("no completed, later-generation, three-stage calibration") + elif len(fingerprints) != 1: + reasons.append( + "hardware, source or training recipe differs; producer counts are not comparable" + ) + chosen = None + best = max( + (row["effective_data_episodes_per_second"] for row in eligible), + default=None, + ) + if not reasons: + qualifying = [ + row + for row in eligible + if row["effective_data_episodes_per_second"] >= 0.9 * best + or math.isclose( + row["effective_data_episodes_per_second"], 0.9 * best, rel_tol=1e-12 + ) + ] + chosen = min( + qualifying, + key=lambda row: ( + row["producer_envs"], + -row["effective_data_episodes_per_second"], + ), + ) + selections.append( + { + "vision": vision, + "comparable": not reasons, + "reasons": reasons, + "best_effective_episode_rate": best if not reasons else None, + "selected_producer_envs": chosen["producer_envs"] if chosen else None, + "selected_calibration": chosen["output"] if chosen else None, + "threshold_fraction": 0.9, + "candidates": candidates, + } + ) + packet = { + "schema": "ods-three-stage-producer-selection/v1", + "roots": [str(root) for root in roots], + "selections": selections, + "limits": "Selects among measured P only; never changes K or hardware. Repeat the chosen recipe in a complete pilot.", + } + output.mkdir(parents=True, exist_ok=True) + _write(output / "producer-selection.json", packet) + lines = [ + _method(), + "", + "| Vision | Comparable | Selected P | Reason |", + "| --- | --- | ---: | --- |", + ] + lines += [ + f"| {row['vision']} | {row['comparable']} | {row['selected_producer_envs'] if row['selected_producer_envs'] is not None else '—'} | {'; '.join(row['reasons']) or 'smallest measured P at ≥90% of best effective supply'} |" + for row in selections + ] + (output / "producer-selection.md").write_text( + "\n".join(lines) + "\n", encoding="utf-8" + ) + return packet + + +def _main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("roots", nargs="+", type=Path) + parser.add_argument("--output", required=True, type=Path) + args = parser.parse_args() + print(json.dumps(compare_calibrations(args.roots, args.output), allow_nan=False)) + + +if __name__ == "__main__": + _main() diff --git a/examples/ods/data_stream_demo/video_integrity.py b/examples/ods/data_stream_demo/video_integrity.py new file mode 100644 index 000000000..098daa384 --- /dev/null +++ b/examples/ods/data_stream_demo/video_integrity.py @@ -0,0 +1,543 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Strict, post-timing integrity checks for every saved LeRobot video frame. + +The verifier software-decodes AV1 videos sequentially, hashes raw YUV420p +frames, and reconciles actual counts with official LeRobot episode metadata. +It does not establish lossless RGB fidelity or every SDK random-seek position. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import os +import shutil +import subprocess +import time +import traceback +from collections import defaultdict +from fractions import Fraction +from pathlib import Path +from typing import Any + +__all__ = ["verify_dataset"] + +_AUDIT_REVISION = "full-video-integrity/v1" + + +def _digest(path: Path) -> str: + result = hashlib.sha256() + with path.open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + result.update(block) + return result.hexdigest() + + +def _require(condition: bool, message: str) -> None: + if not condition: + raise ValueError(message) + + +def _inventory( + root: Path, +) -> tuple[dict[str, Any], dict[str, dict[str, Any]], dict[str, str]]: + import pyarrow.parquet as pq + + _require(not (root / "INCOMPLETE").exists(), "Dataset is marked INCOMPLETE") + info_file = root / "meta/info.json" + info = json.loads(info_file.read_text()) + _require( + info["total_episodes"] > 0 and info["total_frames"] > 0, + "Official corpus must contain accepted episodes and frames", + ) + cameras = { + key: value + for key, value in info["features"].items() + if value["dtype"] == "video" + } + _require(bool(cameras), "Official metadata contains no recorded video cameras") + fps = float(info["fps"]) + _require(fps > 0 and math.isfinite(fps), "Invalid official FPS") + columns = ["episode_index", "length", "dataset_from_index", "dataset_to_index"] + for camera in cameras: + columns.extend( + f"videos/{camera}/{field}" + for field in ("chunk_index", "file_index", "from_timestamp", "to_timestamp") + ) + metadata_files = [info_file, *sorted((root / "meta/episodes").rglob("*.parquet"))] + _require(len(metadata_files) > 1, "Official episode metadata is missing") + hashes = {str(path.relative_to(root)): _digest(path) for path in metadata_files} + files: dict[str, dict[str, Any]] = {} + count = frames = 0 + for meta in metadata_files[1:]: + for batch in pq.ParquetFile(meta).iter_batches( + batch_size=64, columns=columns, use_threads=False + ): + for episode in batch.to_pylist(): + length = int(episode["length"]) + _require( + episode["episode_index"] == count + and episode["dataset_from_index"] == frames + and episode["dataset_to_index"] == frames + length, + "Official episode order/length offsets differ", + ) + _require(length > 0, "Official episode has no frames") + for camera, feature in cameras.items(): + relative = info["video_path"].format( + video_key=camera, + chunk_index=int(episode[f"videos/{camera}/chunk_index"]), + file_index=int(episode[f"videos/{camera}/file_index"]), + ) + path = (root / relative).resolve() + _require( + root.resolve() in path.parents, + "Video metadata escapes the corpus", + ) + _require(path.is_file(), f"Missing registered video: {relative}") + start = float(episode[f"videos/{camera}/from_timestamp"]) + stop = float(episode[f"videos/{camera}/to_timestamp"]) + _require( + math.isclose( + stop - start, length / fps, rel_tol=0, abs_tol=1e-4 + ), + "Episode video duration differs from its official frame count", + ) + group = files.setdefault( + relative, + { + "camera": camera, + "expected_frames": 0, + "episodes": [], + "spans": [], + "shape": list(feature["shape"]), + "fps": fps, + }, + ) + _require( + group["camera"] == camera, + "Multiple cameras reference one video file", + ) + group["expected_frames"] += length + group["episodes"].append(count) + group["spans"].append([start, stop]) + count += 1 + frames += length + _require( + count == info["total_episodes"] and frames == info["total_frames"], + "Official episode and info frame/episode totals disagree", + ) + actual = {str(path.relative_to(root)) for path in (root / "videos").rglob("*.mp4")} + _require( + actual == set(files), + f"Unregistered or missing MP4 files: actual={sorted(actual)}, registered={sorted(files)}", + ) + for relative, group in files.items(): + end = 0.0 + for start, stop in sorted(group["spans"]): + _require( + math.isclose(start, end, rel_tol=0, abs_tol=1e-4), + f"Video spans overlap or contain gaps: {relative}", + ) + end = stop + _require( + math.isclose(end, group["expected_frames"] / fps, rel_tol=0, abs_tol=1e-4), + f"Video span total differs: {relative}", + ) + _require( + all( + sum( + group["expected_frames"] + for group in files.values() + if group["camera"] == camera + ) + == frames + for camera in cameras + ), + "A camera is missing official episode frames", + ) + return info, files, hashes + + +def _decode_file( + root: Path, + relative: str, + expected: dict[str, Any], + artifacts: Path, + index: int, + *, + ffmpeg: str, + ffprobe: str, + threads: int, + timeout: float, + sink: list[dict[str, Any]], +) -> dict[str, Any]: + path = root / relative + before = _digest(path) + record: dict[str, Any] = { + "file": relative, + "camera": expected["camera"], + "expected_frames": expected["expected_frames"], + "episodes": expected["episodes"], + "compressed_sha256": before, + "decoded_pixel_format": "yuv420p", + "frame_hash_scope": "SHA256 of FFmpeg decoded rawvideo YUV420p planes; not RGB or BC-normalized tensors", + } + sink.append(record) + probe_command = [ + ffprobe, + "-v", + "error", + "-select_streams", + "v", + "-show_entries", + "stream=index,codec_name,codec_type,width,height,pix_fmt,r_frame_rate,nb_frames,time_base,duration", + "-of", + "json", + str(path), + ] + probe = subprocess.run( + probe_command, capture_output=True, text=True, timeout=timeout, check=False + ) + record.update(probe_command=probe_command, probe_exitcode=probe.returncode) + (artifacts / f"{index:04d}.probe.stderr").write_text(probe.stderr) + _require( + probe.returncode == 0, f"ffprobe failed for {relative}: {probe.stderr[-2000:]}" + ) + streams = json.loads(probe.stdout)["streams"] + _require(len(streams) == 1, f"Expected exactly one video stream in {relative}") + stream = streams[0] + _require(stream["codec_name"] == "av1", f"Unexpected recorded codec in {relative}") + _require( + [stream["height"], stream["width"], 3] == expected["shape"], + f"Camera dimensions differ from official metadata in {relative}", + ) + _require( + math.isclose( + float(Fraction(stream["r_frame_rate"])), + expected["fps"], + rel_tol=0, + abs_tol=1e-8, + ), + f"Actual video FPS differs from official metadata: {relative}", + ) + _require( + stream.get("pix_fmt") == "yuv420p", + f"Unexpected AV1 stored pixel format in {relative}", + ) + record["probe"] = stream + framehash = artifacts / f"{index:04d}.decoded-frame-sha256.txt" + error_log = artifacts / f"{index:04d}.decode.stderr" + command = [ + ffmpeg, + "-nostdin", + "-hide_banner", + "-loglevel", + "error", + "-xerror", + "-err_detect", + "explode", + "-hwaccel", + "none", + "-threads:v", + str(threads), + "-i", + str(path), + "-map", + "0:v:0", + "-an", + "-sn", + "-dn", + "-fps_mode", + "passthrough", + "-c:v", + "rawvideo", + "-pix_fmt", + "yuv420p", + "-threads:v", + str(threads), + "-filter_threads", + "1", + "-filter_complex_threads", + "1", + "-f", + "framehash", + "-hash", + "sha256", + "-y", + str(framehash), + ] + started = time.monotonic() + with error_log.open("wb") as stderr: + process = subprocess.run( + command, + stdout=subprocess.DEVNULL, + stderr=stderr, + timeout=timeout, + check=False, + ) + record.update( + decode_command=command, + decode_exitcode=process.returncode, + decode_seconds=time.monotonic() - started, + decoded_frame_hashes=str(framehash), + decode_error_log=str(error_log), + ) + errors = error_log.read_text(errors="replace") + _require( + process.returncode == 0 and not errors.strip(), + f"Strict full decode failed for {relative} (exit={process.returncode}): {errors[-2000:]}", + ) + count = 0 + previous_pts = None + with framehash.open() as values: + for line in values: + if line.startswith("#") or not line.strip(): + continue + fields = [value.strip() for value in line.split(",")] + _require( + len(fields) == 6 and fields[0] == "0" and len(fields[-1]) == 64, + "Malformed decoded frame-hash record", + ) + pts = int(fields[2]) + _require( + previous_pts is None or pts > previous_pts, + "Decoded frames have duplicate or reversed presentation timestamps", + ) + previous_pts = pts + count += 1 + record["decoded_frames"] = count + _require( + count == expected["expected_frames"], + f'Actual decoded frame count differs for {relative}: decoded={count}, official={expected["expected_frames"]}', + ) + if str(stream.get("nb_frames", "")).isdigit(): + _require( + int(stream["nb_frames"]) == count, + f"Container declared frame count differs from full decode: {relative}", + ) + _require( + _digest(path) == before, + f"Recorded video changed while being audited: {relative}", + ) + record["decoded_frame_hash_file_sha256"] = _digest(framehash) + record["passed"] = True + return record + + +def verify_dataset( + root: Path | str, + output: Path | str, + *, + ffmpeg: str | None = None, + ffprobe: str | None = None, + threads: int = 1, + timeout: float = 600.0, +) -> dict[str, Any]: + """Software-decode all saved AV1 frames and retain an independent receipt. + + Args: + root: Finalized LeRobot corpus containing official ``meta/info.json``, + episode Parquet metadata and AV1/YUV420p camera files. + output: New receipt JSON path outside the immutable corpus. The sibling + ``-artifacts`` directory must also be new. Existing + successful or failed evidence is never overwritten. + ffmpeg: Software decoder executable, resolved from PATH when omitted. + ffprobe: Stream metadata executable, resolved from PATH when omitted. + threads: FFmpeg input/output codec threads, from one to eight. Files + are decoded sequentially. The caller controls process CPU affinity; + the CLI caps its inherited CPU budget at eight cores. + timeout: Positive finite timeout in seconds for each probe/decode call. + + Returns: + A persisted receipt with ``passed`` and ``all_saved_video_frames_decoded`` + flags, official/actual counts, decoder commands/versions, source hashes + and artifact paths. Decode or metadata failures return a failed receipt. + Frame hashes cover decoded raw YUV420p planes, rather than RGB tensors; + success proves software decode/count integrity, not lossless camera + fidelity, learned policy quality or all SDK random-seek positions. + + Raises: + FileExistsError: Receipt or artifact directory already exists. + ValueError: Output would modify the immutable input corpus. + """ + root, output = Path(root).resolve(), Path(output).resolve() + _require( + root not in output.parents and output != root, + "Audit output must be outside the immutable corpus", + ) + artifacts = output.parent / (output.stem + "-artifacts") + if output.exists() or artifacts.exists(): + raise FileExistsError(f"Audit evidence already exists: {output} or {artifacts}") + output.parent.mkdir(parents=True, exist_ok=True) + artifacts.mkdir(exist_ok=False) + result: dict[str, Any] = { + "revision": _AUDIT_REVISION, + "passed": False, + "all_saved_video_frames_decoded": False, + "dataset": str(root), + "scope": "post-timing sequential software decode/count integrity; no lossless RGB, policy-quality or all-SDK-seek claim", + "decoded_pixel_format": "yuv420p", + "frame_hash_scope": "SHA256 of decoded rawvideo YUV420p planes; not RGB or BC-normalized tensors", + "files": [], + "cpu_affinity": ( + sorted(os.sched_getaffinity(0)) + if hasattr(os, "sched_getaffinity") + else None + ), + "threads_per_file": threads, + "files_in_parallel": 1, + "verifier_sha256": _digest(Path(__file__)), + } + started = time.monotonic() + try: + _require( + type(threads) is int and 1 <= threads <= 8, + "Software decoder threads must be in [1,8]", + ) + _require( + math.isfinite(timeout) and timeout > 0, + "Probe/decode timeout must be finite and positive", + ) + ffmpeg = ffmpeg or shutil.which("ffmpeg") + ffprobe = ffprobe or shutil.which("ffprobe") + _require( + ffmpeg is not None and ffprobe is not None, + "FFmpeg/ffprobe executables are required", + ) + for name, executable in (("ffmpeg", ffmpeg), ("ffprobe", ffprobe)): + version = subprocess.run( + [executable, "-version"], + capture_output=True, + text=True, + check=True, + timeout=min(timeout, 30.0), + ) + result[name + "_version"] = version.stdout.splitlines()[0] + info, files, metadata_hashes = _inventory(root) + result.update( + expected_episodes=info["total_episodes"], + expected_frames_per_camera=info["total_frames"], + metadata_sha256=metadata_hashes, + ) + for index, (relative, expected) in enumerate(sorted(files.items())): + _decode_file( + root, + relative, + expected, + artifacts, + index, + ffmpeg=ffmpeg, + ffprobe=ffprobe, + threads=threads, + timeout=timeout, + sink=result["files"], + ) + counts: dict[str, int] = defaultdict(int) + for record in result["files"]: + counts[record["camera"]] += record["decoded_frames"] + _require( + all(value == info["total_frames"] for value in counts.values()), + "Actual per-camera decoded total differs from official SDK total", + ) + for relative, value in metadata_hashes.items(): + _require( + _digest(root / relative) == value, + f"Official metadata changed during audit: {relative}", + ) + result.update( + passed=True, + decoded_frames_per_camera=dict(counts), + total_decoded_camera_frames=sum(counts.values()), + all_saved_video_frames_decoded=True, + ) + except Exception as error: + result.update(error=repr(error), traceback=traceback.format_exc()) + result["seconds_outside_training_pipeline"] = time.monotonic() - started + with output.open("x", encoding="utf-8") as stream: + json.dump(result, stream, indent=2, allow_nan=False) + stream.write("\n") + return result + + +def _main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dataset", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--cpu-count", type=int, default=8) + parser.add_argument("--threads", type=int, default=1) + parser.add_argument("--timeout", type=float, default=600) + parser.add_argument("--ffmpeg", default=shutil.which("ffmpeg")) + parser.add_argument("--ffprobe", default=shutil.which("ffprobe")) + args = parser.parse_args() + os.environ.update( + CUDA_VISIBLE_DEVICES="", + OMP_NUM_THREADS="1", + MKL_NUM_THREADS="1", + OPENBLAS_NUM_THREADS="1", + ARROW_NUM_THREADS="1", + ) + _require( + args.ffmpeg is not None and args.ffprobe is not None, + "FFmpeg/ffprobe executables are required", + ) + _require( + 1 <= args.cpu_count <= 8 and 1 <= args.threads <= args.cpu_count, + "CPU budget must be 1..8 and cover decoder threads", + ) + if hasattr(os, "sched_getaffinity"): + available = sorted(os.sched_getaffinity(0)) + os.sched_setaffinity(0, available[-args.cpu_count :]) + try: + result = verify_dataset( + args.dataset, + args.output, + ffmpeg=args.ffmpeg, + ffprobe=args.ffprobe, + threads=args.threads, + timeout=args.timeout, + ) + except (FileExistsError, ValueError) as error: + print( + json.dumps( + {"passed": False, "error": repr(error), "receipt_written": False} + ), + flush=True, + ) + raise SystemExit(2) from error + print( + json.dumps( + { + key: result.get(key) + for key in ( + "passed", + "error", + "expected_frames_per_camera", + "decoded_frames_per_camera", + "seconds_outside_training_pipeline", + ) + } + ), + flush=True, + ) + raise SystemExit(0 if result["passed"] else 1) + + +if __name__ == "__main__": + _main() diff --git a/examples/ods/data_stream_demo/visual_policy.py b/examples/ods/data_stream_demo/visual_policy.py new file mode 100644 index 000000000..93106a3b4 --- /dev/null +++ b/examples/ods/data_stream_demo/visual_policy.py @@ -0,0 +1,172 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Example-local visual action-chunk BC with a shared trainable CNN.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import torch +from torch import nn + +from embodichain.learning.models import ResNet10Encoder +from examples.ods.simple_bc_demo.models import pose_targets + +__all__ = ["VisualChunkBC"] + + +class VisualChunkBC(nn.Module): + """Predict TCP/gripper chunks using two frames of vision and proprioception. + + Vision is flattened in history-major, configured-camera-minor order. One + encoder is shared by all histories and views. Future geometry is accepted + only by :meth:`loss` for constructing targets, never by :meth:`forward`. + """ + + history_length = 2 + arm_command_mode = "pose" + observation_mode = "visual" + + def __init__( + self, + proprio_mean: torch.Tensor, + proprio_std: torch.Tensor, + target_mean: torch.Tensor, + target_std: torch.Tensor, + cameras: tuple[str, ...], + prediction_horizon: int = 32, + ) -> None: + super().__init__() + if ( + not isinstance(cameras, tuple) + or len(cameras) not in (1, 2) + or any(not isinstance(camera, str) or not camera for camera in cameras) + or len(set(cameras)) != len(cameras) + ): + raise ValueError( + "Visual BC requires one or two distinct ordered camera names" + ) + if type(prediction_horizon) is not int or prediction_horizon < 1: + raise ValueError("Prediction horizon must be a positive integer") + for name, value, width, scale in ( + ("proprio_mean", proprio_mean, 27, False), + ("proprio_std", proprio_std, 27, True), + ("target_mean", target_mean, 10, False), + ("target_std", target_std, 10, True), + ): + if ( + value.shape != (width,) + or value.dtype != torch.float32 + or not torch.isfinite(value).all() + or (scale and (value <= 0).any()) + ): + raise ValueError( + f"Invalid finite float32 normalization statistic: {name}" + ) + self.register_buffer(name, value.detach().clone()) + self.cameras = cameras + self.prediction_horizon = prediction_horizon + self.encoder = ResNet10Encoder() + input_dim = 2 * (27 + len(cameras) * self.encoder.output_dim) + self.head = nn.Sequential( + nn.Linear(input_dim, 512), + nn.ReLU(), + nn.Linear(512, 512), + nn.ReLU(), + nn.Linear(512, prediction_horizon * 10), + ) + + def get_extra_state(self) -> dict[str, Any]: + """Save the ordered input contract alongside tensor parameters.""" + return { + "version": 1, + "cameras": self.cameras, + "history_length": self.history_length, + "prediction_horizon": self.prediction_horizon, + } + + def set_extra_state(self, state: dict[str, Any]) -> None: + """Reject checkpoint reuse under a different camera or temporal contract.""" + if state != self.get_extra_state(): + raise ValueError( + "Visual BC checkpoint camera order or horizon does not match" + ) + + def forward( + self, + proprio: torch.Tensor, + images: Mapping[str, torch.Tensor], + ) -> torch.Tensor: + """Return normalized ``[B, F, 10]`` targets from two causal frames.""" + if ( + proprio.ndim != 3 + or proprio.shape[0] < 1 + or proprio.shape[1:] != (2, 27) + or proprio.dtype != torch.float32 + or proprio.device != self.proprio_mean.device + ): + raise ValueError("Visual BC proprioception must be float32 [B, 2, 27]") + if set(images) != set(self.cameras): + raise ValueError( + "Visual BC camera keys must exactly match configured cameras" + ) + batch_size = proprio.shape[0] + expected = (batch_size, 2, 3, 256, 256) + for camera in self.cameras: + image = images[camera] + if ( + image.shape != expected + or image.dtype != torch.float32 + or image.device != proprio.device + ): + raise ValueError(f"Camera {camera!r} requires float32 {expected}") + # [B, H, V, C, W, W] -> [B*H*V, C, W, W]; one shared backbone call. + stacked = torch.stack([images[camera] for camera in self.cameras], dim=2) + visual = self.encoder(stacked.flatten(0, 2)).reshape(batch_size, -1) + normalized = ((proprio - self.proprio_mean) / self.proprio_std).clamp(-10, 10) + combined = torch.cat((normalized.flatten(1), visual), dim=-1) + return self.head(combined).reshape(batch_size, self.prediction_horizon, 10) + + def loss( + self, + geometry: torch.Tensor, + actions: torch.Tensor, + images: Mapping[str, torch.Tensor], + *, + loss_mask: torch.Tensor, + valid_mask: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Reuse pose-target MSE with geometry kept outside policy inputs.""" + expected = (geometry.shape[0], self.prediction_horizon + 1) + if ( + geometry.shape != (*expected, 51) + or actions.shape != (*expected, 7) + or loss_mask.shape != expected + or valid_mask.shape != expected + or loss_mask.dtype != torch.bool + or valid_mask.dtype != torch.bool + ): + raise ValueError("Visual BC geometry/actions/masks must match H=2 and F") + if not valid_mask.all() or loss_mask[:, 0].any() or not loss_mask[:, 1:].all(): + raise ValueError("Visual BC trains only complete valid future windows") + predictions = self(geometry[:, :2, :27], images) + targets = ( + pose_targets(geometry[:, 1:], actions[:, 1:]) - self.target_mean + ) / self.target_std + errors = (predictions - targets).square() + return errors[..., :9].mean(), errors[..., 9].mean() diff --git a/pyproject.toml b/pyproject.toml index 9812e2926..850fb1193 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,10 @@ dependencies = [ ] [project.optional-dependencies] +vision = [ + # Visual ODS reference stack: torchvision 0.22 pairs with PyTorch 2.7. + "torchvision>=0.22,<0.23", +] policy-deploy = [ # ORT 1.27+ defaults to CUDA 13; EmbodiChain's supported stack uses CUDA 12. "onnxruntime-gpu>=1.20,<1.27", diff --git a/tests/learning/models/test_resnet_encoder.py b/tests/learning/models/test_resnet_encoder.py new file mode 100644 index 000000000..b3cf4e936 --- /dev/null +++ b/tests/learning/models/test_resnet_encoder.py @@ -0,0 +1,196 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""CPU contracts for the trainable RGB encoder shared by learning clients.""" + +from __future__ import annotations + +from collections import OrderedDict +from copy import deepcopy +from functools import partial +import subprocess +import sys + +import pytest +import torch +from torch import nn +from torchvision.models.resnet import BasicBlock, ResNet + +from embodichain.learning.models import ResNet10Encoder + + +@pytest.fixture(scope="module", autouse=True) +def bounded_torch_threads(): + previous = torch.get_num_threads() + torch.set_num_threads(min(previous, 4)) + yield + torch.set_num_threads(previous) + + +def test_nonvisual_model_import_does_not_load_torchvision(): + subprocess.run( + [ + sys.executable, + "-c", + "import sys; from embodichain.learning.models import MLP; " + "assert 'torchvision' not in sys.modules", + ], + check=True, + capture_output=True, + text=True, + ) + + +def test_reference_structure_and_all_stages_receive_gradients(): + torch.manual_seed(12) + encoder = ResNet10Encoder() + assert [ + len(getattr(encoder.resnet_backbone, f"layer{i}")) for i in range(1, 5) + ] == [1] * 4 + assert not any( + isinstance(module, (nn.BatchNorm2d, nn.Dropout)) for module in encoder.modules() + ) + assert not any( + name.startswith("resnet_backbone.fc") for name in encoder.state_dict() + ) + for module in encoder.modules(): + if isinstance(module, nn.GroupNorm): + assert module.num_groups == 4 and module.eps == 1e-5 + images = torch.rand(1, 3, 256, 256, requires_grad=True) + output = encoder(images) + assert output.shape == (1, 256) and output.abs().max() <= 1 + output.square().mean().backward() + assert images.grad is not None and images.grad.norm() > 0 + for name, parameter in encoder.named_parameters(): + assert parameter.requires_grad, name + assert parameter.grad is not None and torch.isfinite(parameter.grad).all(), name + assert parameter.grad.norm() > 0, name + + +def test_reference_forward_intermediates_and_gradient_parity(): + """Independent torchvision backbone and pinned RLInf's explicit pooling.""" + torch.manual_seed(91) + encoder = ResNet10Encoder() + torchvision_backbone = ResNet( + BasicBlock, [1, 1, 1, 1], norm_layer=partial(nn.GroupNorm, 4, eps=1e-5) + ) + reference_backbone = nn.Sequential( + OrderedDict(list(torchvision_backbone.named_children())[:-2]) + ) + reference_backbone.load_state_dict(encoder.resnet_backbone.state_dict()) + reference_pool = encoder.pooling_layer.kernel.detach().clone().requires_grad_() + reference_head = deepcopy(encoder.mlp) + inputs = torch.rand(1, 3, 256, 256, requires_grad=True) + reference_inputs = inputs.detach().clone().requires_grad_() + actual = encoder(inputs) + features = reference_backbone( + (reference_inputs - encoder.image_mean) / encoder.image_std + ) + assert features.shape == (1, 512, 8, 8) + torch.testing.assert_close( + features, + encoder.resnet_backbone((inputs - encoder.image_mean) / encoder.image_std), + atol=0, + rtol=0, + ) + reference_embedding = ( + (features.unsqueeze(-1) * reference_pool.unsqueeze(0)) + .sum((2, 3)) + .reshape(1, -1) + ) + torch.testing.assert_close( + encoder.pooling_layer(features), reference_embedding, atol=5e-5, rtol=2e-5 + ) + expected = reference_head(reference_embedding) + torch.testing.assert_close(actual, expected, atol=5e-6, rtol=5e-5) + weights = torch.linspace(-1, 1, 256).reshape(1, -1) + (actual * weights).sum().backward() + (expected * weights).sum().backward() + torch.testing.assert_close(inputs.grad, reference_inputs.grad, atol=2e-5, rtol=5e-4) + torch.testing.assert_close( + encoder.pooling_layer.kernel.grad, reference_pool.grad, atol=2e-5, rtol=5e-4 + ) + for (_, actual_parameter), (_, expected_parameter) in zip( + encoder.resnet_backbone.named_parameters(), + reference_backbone.named_parameters(), + strict=True, + ): + torch.testing.assert_close( + actual_parameter.grad, expected_parameter.grad, atol=3e-4, rtol=5e-4 + ) + + +def test_rgb_normalization_once_and_noncontiguous_input_ownership(): + encoder = ResNet10Encoder().eval() + rgb = torch.tensor([0.0, 0.5, 1.0]).view(1, 3, 1, 1).expand(1, 3, 256, 256) + assert not rgb.is_contiguous() + before = rgb.clone() + seen = [] + handle = encoder.resnet_backbone.conv1.register_forward_pre_hook( + lambda _module, arguments: seen.append(arguments[0].detach().clone()) + ) + with torch.no_grad(): + original_output = encoder(rgb) + preserved = original_output.clone() + encoder(torch.zeros_like(rgb)) + handle.remove() + torch.testing.assert_close(seen[0], (rgb - encoder.image_mean) / encoder.image_std) + torch.testing.assert_close(rgb, before, atol=0, rtol=0) + torch.testing.assert_close(original_output, preserved, atol=0, rtol=0) + + +def test_batch32_and_individual_image_encoding_agree(): + torch.manual_seed(4) + encoder = ResNet10Encoder().eval() + image = torch.rand(1, 3, 256, 256) + noncontiguous = image.expand(32, -1, -1, -1) + with torch.no_grad(): + batch_features = encoder(noncontiguous) + single = encoder(image) + assert batch_features.shape == (32, 256) + torch.testing.assert_close( + batch_features, single.expand_as(batch_features), atol=5e-6, rtol=2e-5 + ) + + +def test_checkpoint_roundtrip_preserves_outputs_and_trainability(): + torch.manual_seed(3) + encoder = ResNet10Encoder().eval() + image = torch.rand(1, 3, 256, 256) + with torch.no_grad(): + expected = encoder(image) + restored = ResNet10Encoder().eval() + restored.load_state_dict(deepcopy(encoder.state_dict())) + with torch.no_grad(): + torch.testing.assert_close(restored(image), expected, atol=0, rtol=0) + assert all(parameter.requires_grad for parameter in restored.parameters()) + + +@pytest.mark.parametrize( + "shape,dtype", + [ + ((0, 3, 256, 256), torch.float32), + ((1, 256, 256, 3), torch.float32), + ((1, 4, 256, 256), torch.float32), + ((1, 3, 128, 128), torch.float32), + ((1, 3, 256, 256), torch.uint8), + ((1, 3, 256, 256), torch.float64), + ], +) +def test_wrong_image_metadata_is_rejected(shape, dtype): + encoder = ResNet10Encoder() + with pytest.raises(ValueError, match="float32"): + encoder(torch.zeros(shape, dtype=dtype))