diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index d7994406b0..c901b41039 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -569,6 +569,7 @@ jobs: "embodichain_tasks/configs/tasks/manipulation/tableware/pour_water/env.yaml", "embodichain_robots/__init__.py", "embodichain_robots/_catalog.py", + "embodichain_models/__init__.py", "embodichain_robots/_asset_catalog.py", "embodichain/data/_asset_registry.py", "embodichain_robots/franka/assets.py", @@ -607,6 +608,7 @@ jobs: "embodichain_tasks/", "embodichain_robots/", "embodichain_benchmarks/", + "embodichain_models/", f"embodichain-{version}.data/data/share/embodichain/examples/", dist_info, ) diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index affe772885..d01eba53f1 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -268,6 +268,25 @@ topics: tests/gym/envs/task_program/test_simulation_environment.py] related_topics: [scene-services, simulation-system, task-programs, robot-workspace, motion-planning, gen-sim, data-pipeline] status: active +- id: neural-dynamics + title: Neural Dynamics Models and Training + aliases: [neural dynamics, neural robot dynamics, NeRD, 神经动力学, 动力学模型训练, NeRD 示例] + keywords: [NeRDStateCodec, NeRDModel, DynamicsGenerationCfg, + DynamicsShardWriter, DynamicsWindowDataset, DynamicsTrainer, embodichain_models, + learning/nerd, nerd reference assets, + 官方训练精度对比] + paths: [topics/neural-dynamics/neural-dynamics.md] + source_of_truth: + - extensions/embodichain_models/embodichain_models/dynamics/nerd/ + - examples/learning/nerd/__main__.py + - examples/learning/nerd/_resources.py + - examples/learning/nerd/_support/ + - examples/learning/nerd/configs/ + watch_paths: [extensions/embodichain_models/, examples/learning/nerd/, + docs/source/features/neural_dynamics.md, + docs/source/api_reference/model_catalogs.rst] + related_topics: [supervised-learning, learning-layers, data-pipeline, data-assets] + status: active - id: learning-layers title: Differentiable Learning Layers aliases: [learning layers, differentiable IK layer, 可微 IK, 训练层, 学习模块分层] @@ -485,7 +504,7 @@ topics: 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/, 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] + related_topics: [manager-functor, env-framework, task-programs, sensor-system, neural-dynamics, supervised-learning] status: active - id: robot-workspace title: Robot Workspace @@ -525,3 +544,15 @@ topics: watch_paths: [embodichain/data/asset_catalog/, tests/data/asset_catalog/, embodichain/lab/task_program/semantics/asset_requirements.py, tests/lab/task_program/semantics/test_asset_requirements.py, embodichain/lab/task_program/semantics/asset_matching.py, tests/lab/task_program/semantics/test_asset_matching.py, embodichain/lab/task_program/integrations/asset_materialization.py, tests/gym/envs/task_program/test_asset_materialization.py, embodichain/lab/gym/envs/asset_generalization.py, tests/gym/envs/test_asset_generalization.py, examples/expansion/asset_generalization/, examples/expansion/asset_catalog/, examples/expansion/README.md] related_topics: [data-assets, task-programs, gen-sim, scene-services, manager-functor, env-framework] status: active + +- id: supervised-learning + title: Supervised Learning + aliases: [supervised learning, SL learning, 监督学习, 统一监督学习] + keywords: [train-sl, SupervisedBatch, LossResult, SupervisedObjective, SupervisedLearner, SupervisedTrainer, TensorBatchSource] + paths: [topics/supervised-learning/supervised-learning.md] + source_of_truth: + - embodichain/learning/sl/ + - embodichain/cli/main.py + watch_paths: [embodichain/learning/sl/, tests/learning/sl/] + related_topics: [learning-layers, data-pipeline, rl-learning, neural-dynamics] + status: active diff --git a/agent_context/topics/data-pipeline/data-pipeline.md b/agent_context/topics/data-pipeline/data-pipeline.md index f795d24b6f..612ba903f2 100644 --- a/agent_context/topics/data-pipeline/data-pipeline.md +++ b/agent_context/topics/data-pipeline/data-pipeline.md @@ -23,6 +23,7 @@ belongs to [data-assets](../data-assets/data-assets.md). | Offline collection retries and final partial batch | `embodichain/lab/scripts/run_task.py` | | File-only validation and grouped quality splits | `embodichain/data_pipeline/datasets/inspection.py`; `embodichain dataset validate\|split` | | Durable commit diagnostics and metadata repair | `embodichain/data_pipeline/recording/journal.py`; `embodichain dataset recover` | +| NeRD-specific HDF5 windows and teacher collection | `extensions/embodichain_models/embodichain_models/dynamics/nerd/data.py` and `examples/learning/nerd/_support/_generation.py`; contracts in [neural dynamics](../neural-dynamics/neural-dynamics.md) | | Isolated LeRobot 0.6.1 language/depth conversion | `tools/lerobot_export/` (Python >=3.12, independent dependencies) | Read [online sampling](online-sampling.md) for worker states, refill, shared diff --git a/agent_context/topics/learning-layers/learning-layers.md b/agent_context/topics/learning-layers/learning-layers.md index fd46bc6486..e1d5ff3c23 100644 --- a/agent_context/topics/learning-layers/learning-layers.md +++ b/agent_context/topics/learning-layers/learning-layers.md @@ -9,8 +9,13 @@ 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, -TensorDict or simulation. RL policy interfaces, actor-critic composition and +and configuration-driven MLP construction. These shared components do not import +RL, TensorDict or simulation. Concrete NeRD models, codecs and the shared-SL +recipe live in the bundled `embodichain_models.dynamics.nerd` extension; +teacher collection, historical training and research belong to +`examples/learning/nerd/`. See +[neural dynamics](../neural-dynamics/neural-dynamics.md) for model, SL integration, +dataset and checkpoint ownership. 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) for those contracts. @@ -21,6 +26,9 @@ There is no registered Robot solver adapter: callers provide root-frame targets and explicitly configure the layer's TCP and limits. Existing Robot solver binding and arena conversion remain in [IK solvers](../ik-solvers/ik-solvers.md). +Supervised optimizer and offline-run ownership belongs to +[supervised learning](../supervised-learning/supervised-learning.md). + ## Gradient and state boundaries Targets are root-frame TCP matrices; seed batch axes match the target or contain diff --git a/agent_context/topics/neural-dynamics/neural-dynamics.md b/agent_context/topics/neural-dynamics/neural-dynamics.md new file mode 100644 index 0000000000..747cb0b625 --- /dev/null +++ b/agent_context/topics/neural-dynamics/neural-dynamics.md @@ -0,0 +1,95 @@ +# Neural dynamics models and training + +## Entry points and ownership + +The main wheel bundles `embodichain_models` with the same version and dependencies. +It is an import package under `extensions/`, not a separate distribution or a +runtime registry. Shared MLPs and feature normalization stay in core +`embodichain.learning.models`; generic supervised orchestration stays in +[SL](../supervised-learning/supervised-learning.md). + +The [model catalog contract](../../../extensions/embodichain_models/README.md) +owns domain and public-import naming. Policy models predict actions; dynamics +models predict future state or geometry. RWM and RigidFormer follow-ups use +`dynamics.rwm` and `dynamics.rigidformer`; only `dynamics.nerd` is implemented here. +World-model naming does not introduce a parallel `world_models` domain. + +| Request | Owner | +|---|---| +| NeRD network, portable checkpoint, physical codec | `extensions/embodichain_models/embodichain_models/dynamics/nerd/{model,checkpoint,codec}.py` | +| Committed HDF5 shards and bounded windows | Model package `data.py`, private `_resident.py` | +| Encoding, frozen moments and offline errors | Model package `integration.py` | +| Objective and common trainer factory | Model package `supervised.py`, `sl_recipe.py` | +| Prepare/train/export/eval/generate commands | `examples/learning/nerd/__main__.py`; `embodichain run-example learning/nerd` | +| Pinned datasets, exports and reference assets | Example `_resources.py`; explicit `prepare` command | +| Analytical teacher and physical profile | Example `_support/{_generation,_newton,_reference_assets}.py` | +| Historical trainer and checkpoint workflow | Example `_support/{_trainer,_workflow}.py` | +| Original-repository comparison and acceptance | Example `_support/` runners; protocol in `research/` | + +The showcase defaults to `sl.train.yaml` and the installable +`embodichain_models.dynamics.nerd.sl_recipe:build_supervised_trainer` factory. +Direct `embodichain train-sl --config FILE` works independently of example +bootstrap. An explicitly selected `dynamics.train.*.yaml` uses the historical +trainer for controlled comparisons and its legacy checkpoint contract. + +The extension imports neither examples nor simulation. Generation uses Newton's +built-in Featherstone/MuJoCo solvers directly without a DexSim custom solver or +Scene API. Native learned-solver deployment and functional Gym integration are +separate consumers. There are no NeRD-specific core `compute/dynamics`, +`data_pipeline/dynamics`, `learning/dynamics` modules or core `dynamics` command. + +## Resource and physical identity + +`prepare --robot cartpole|ant --resource-root DIR` downloads the pinned dataset, +portable model and raw reference assets. Revisions and SHA256 checks are owned by +`_resources.py`; `--assets-only` skips data/model downloads and +`--local-files-only` forbids network access. The default resource root is +`EMBODICHAIN_NERD_RESOURCE_ROOT`, otherwise the showcase cache under +`XDG_CACHE_HOME` or `~/.cache`. Downloads do not live in Git or the wheel; model +third-party licensing remains in the model package. + +A trajectory has T+1 states and T effort/contact inputs. Splits belong to whole +trajectories; windows never cross episodes or splits. Only manifest-committed +shards are visible. Generation commits its cursor and rejection evidence with +shards; rejection-only progress preserves data fingerprints. Legacy progress +sidecars remain readable, while collector source and physical identity still +must match to continue collection. A path move does not authorize cross-version +resume or reconstruct missing historical rejection evidence. + +The direct teacher records importer profile, asset/physical-array hashes, solver +settings and actual Newton/Warp versions. Its `nerd-direct-newton-builder-v1` +signature differs from published Scene-generated data, which remains readable +without simulation. `warp_reference` reconstructs historical asset parameters +and Ant friction mapping; matching that name does not prove identical physics +across model-building paths, versions or solvers. + +The codec owns COM-to-spatial-twist transport and relative targets in the frame +at t. Ant uses 25 ordered plane-contact slots; Cartpole is contact-free. Controls +are generalized efforts with zero free-base force. Frozen train-only moments +are reused for evaluation, and supplied statistics must bind the training data +fingerprint. Do not relabel physical records or replace their normalizers. + +## Training and validation + +SL checkpoints use the common trainer format and resume through the same +recipe/signature. The historical trainer keeps `embodichain.nerd.training.v1`: +`last.pt` retains optimizer, scheduler, RNG, committed sampling state and one +owned CPU snapshot of the selected best model; automatic `best.pt` is export-only. +Both trainers use shared run-directory/history admission from SL. Common `train` +does not turn its final `last.pt` into a validation-selected model. + +The showcase `export` accepts common SL and historical NeRD checkpoints and +writes the same portable tensor/JSON model format. Precision, ordered feature +layout, normalization and attention policy are part of recovery/export identity. +Historical fine-tuning and shared sample-order options require a historical +configuration. Only explicitly trusted legacy conversion unpickles model objects. + +Frozen source revisions, budgets and gates live in the +[acceptance protocol](../../../examples/learning/nerd/research/acceptance.md). +Historical evidence remains bound to its recorded source, data and runtime; +relocating code does not rerun it. NeRD validation uses the runnable showcase +and its reference/acceptance workflows; `tests/learning/sl/` covers the common +training contracts. The [feature guide](../../../docs/source/features/neural_dynamics.md) +owns user commands; the [model API](../../../docs/source/api_reference/model_catalogs.rst) +owns exported definitions. Transition error, accumulated rollout error and +paper-scale accuracy remain separate claims. diff --git a/agent_context/topics/supervised-learning/supervised-learning.md b/agent_context/topics/supervised-learning/supervised-learning.md new file mode 100644 index 0000000000..982c189bd7 --- /dev/null +++ b/agent_context/topics/supervised-learning/supervised-learning.md @@ -0,0 +1,45 @@ +# Supervised learning (SL) + +`embodichain.learning.sl` owns common supervised optimizer updates and offline +training orchestration. Shared MLPs and feature normalization remain in +`learning.models`; concrete model families and their model-specific adapters live +in the bundled `embodichain_models` extension. Generic numerical state and geometry +operations remain in `compute`; source storage belongs to data adapters. + +## Entry points + +- `SupervisedBatch` and `LossResult` in `learning/sl/types.py` preserve structured + inputs, targets, masks and objective-specific reductions. +- `SupervisedLearner.update()` in `learning/sl/learner.py` owns a finite optimizer + update. Its optional numerical context covers forward, backward and clipping. +- `TensorBatchSource` snapshots fixed CPU tensor trees and returns owned samples + under the caller's explicit CPU generator. +- `SupervisedTrainer` in `learning/sl/trainer.py` composes a source and learner, + saving optimizer, RNG, sampling progress, source identity and recipe semantics. +- `embodichain train-sl --config FILE [--resume FILE]` lazily resolves a + `factory: module:function`. The factory receives the full model/data/objective/ + optimizer/trainer/evaluation configuration and returns a `SupervisedTrainer`. +- The [NeRD recipe](../neural-dynamics/neural-dynamics.md) is an installable factory + at `embodichain_models.dynamics.nerd.sl_recipe:build_supervised_trainer`. + The core does not import the model extension or example helpers. The NeRD + showcase's default `train` command uses this same common recipe. + +Objectives own loss masks, normalization, integration and multi-step graph +construction. The optimizer does not infer action semantics or detach histories. +Schema, action/frame/timestep and numerical-policy identifiers belong in each +recipe's checkpoint signature. Tensor sources additionally bind content identity. +`learning/sl/_artifacts.py` owns run identities and output-history admission for +both SL and historical NeRD trainers. A fresh run claims `run.json`; resuming an +occupied directory requires the checkpoint's matching run identity. Checkpoints +without an identity require an artifact-free destination. Newer or incomplete +metric history is retained and requires a fresh resume destination. +Validation callbacks preserve RNG, sampling state and model buffers, including +checkpoints saved inside a callback. Only development data may select checkpoints. +Acquisition and producer lifecycles remain application-owned. + +## Validation + +`tests/learning/sl/` checks independent masked updates, graph propagation, +nonfinite update boundaries, owned samples, simulator-free imports and exact +CPU recovery with dropout, optimizer and sampling state. Model-specific +qualification belongs to each model's acceptance recipe. diff --git a/docs/architecture/curated.json b/docs/architecture/curated.json index 00aa7240a4..4eeee8a4b8 100644 --- a/docs/architecture/curated.json +++ b/docs/architecture/curated.json @@ -2619,7 +2619,7 @@ "symbol": "", "start_line": 24, "end_line": 35, - "excerpt": "\"\"\"Depth quantization/dequantization for compressed depth sidecar videos.\n\nDepth maps are packed into 12-bit integer codes (``uint16``, values\n``0\u2026DEPTH_QMAX``) so they fit the ``gray12le`` pixel format and can be encoded\nlosslessly by HEVC (``libx265``). Logarithmic quantization is the default\nbecause it allocates more quanta to near-range depth, matching the\n``1/depth`` error profile of typical depth sensors.\n\nThe math is ported from lerobot 0.6.0 (itself adapted from BEHAVIOR-1K's\n``obs_utils.py``) and depends only on :mod:`av`, :mod:`numpy` and\n:mod:`torch` -- all of which run on Python 3.10--3.12. See issue #424.\n\"\"\"" + "excerpt": "\"\"\"Depth quantization/dequantization for compressed depth sidecar videos.\n\nDepth maps are packed into 12-bit integer codes (``uint16``, values\n``0…DEPTH_QMAX``) so they fit the ``gray12le`` pixel format and can be encoded\nlosslessly by HEVC (``libx265``). Logarithmic quantization is the default\nbecause it allocates more quanta to near-range depth, matching the\n``1/depth`` error profile of typical depth sensors.\n\nThe math is ported from lerobot 0.6.0 (itself adapted from BEHAVIOR-1K's\n``obs_utils.py``) and depends only on :mod:`av`, :mod:`numpy` and\n:mod:`torch` -- all of which run on Python 3.10--3.12. See issue #424.\n\"\"\"" } ], "documentation": [ @@ -3339,6 +3339,97 @@ "label": "Shared numerical geometry and layout" } ] + }, + { + "id": "neural-dynamics", + "label": "NeRD model extension", + "kind": "package", + "topic_ids": [ + "neural-dynamics" + ], + "summary": "Bundles the causal NeRD network, physical codec, trajectory windows and supervised recipe in the main wheel.", + "details": "embodichain_models.dynamics.nerd owns the concrete model, portable checkpoints, physical representations, HDF5 windows and installable SL recipe. Generic MLP/normalization and supervised orchestration stay in core. The model extension imports neither the showcase nor simulation; collection, historical training and reference research belong to the example.", + "boundaries": [ + "The core does not import concrete model implementations.", + "Data/model/reference-asset payloads are external resources; numerical and physical identities remain explicit." + ], + "evidence": [ + { + "path": "extensions/embodichain_models/embodichain_models/dynamics/nerd/model.py", + "symbol": "NeRDModel", + "start_line": 201, + "end_line": 201, + "excerpt": "class NeRDModel(nn.Module):" + } + ], + "documentation": [ + { + "docname": "api_reference/model_catalogs", + "label": "Official model library" + }, + { + "docname": "features/neural_dynamics", + "label": "NeRD data and training workflow" + } + ] + }, + { + "id": "nerd-showcase", + "label": "NeRD showcase", + "kind": "module", + "topic_ids": [ + "neural-dynamics" + ], + "summary": "Prepares pinned external resources and routes default NeRD training through shared SL.", + "details": "run-example learning/nerd owns preparation, train/export/eval/generate commands and private teacher, historical-trainer and reference-research workflows. The default train preset resolves the installable model SL factory; explicit historical configurations retain their checkpoint protocol. Downloads are verified outside Git and the wheel.", + "boundaries": [ + "Historical reports remain bound to the source and runtime that produced them.", + "Reference assets and teacher/domain metadata are not production model registrations." + ], + "evidence": [ + { + "path": "examples/learning/nerd/__main__.py", + "symbol": "main", + "start_line": 81, + "end_line": 81, + "excerpt": "def main(argv: Sequence[str] | None = None) -> None:" + } + ], + "documentation": [ + { + "docname": "features/neural_dynamics", + "label": "NeRD showcase commands" + } + ] + }, + { + "id": "supervised-learning", + "label": "Supervised learning (SL)", + "kind": "package", + "topic_ids": [ + "supervised-learning" + ], + "summary": "Shares structured batches, objective-owned losses, optimizer updates and offline recovery across policy and dynamics fitting.", + "details": "The SL learner commits finite single-device optimizer updates; the offline trainer owns sampling, validation, logging and checkpoint recovery. Concrete model packages supply codecs, normalization, objectives and installable recipe factories. The core does not import those implementations. Training does not create a simulation World or data producer.", + "boundaries": [ + "Source identity and recipe signatures bind feature/action semantics before recovery.", + "Policy execution, physics collection and APG rollout scheduling remain separate consumers." + ], + "evidence": [ + { + "path": "embodichain/learning/sl/learner.py", + "symbol": "SupervisedLearner", + "start_line": 34, + "end_line": 34, + "excerpt": "class SupervisedLearner:" + } + ], + "documentation": [ + { + "docname": "api_reference/embodichain/embodichain.learning.sl", + "label": "SL batches, updates and recovery" + } + ] } ], "edges": [ @@ -5681,6 +5772,42 @@ "excerpt": " sim = SimulationManager(\n SimulationManagerCfg(\n headless=self.headless,\n sim_device=str(self.device),\n num_envs=batch_size,\n arena_space=2.0,\n )\n )" } ] + }, + { + "id": "nerd-recipe-constructs-sl-learner", + "source": "neural-dynamics", + "target": "supervised-learning", + "relation": "constructs", + "description": "The NeRD recipe constructs the shared learner with its physical objective and numerical policy.", + "provenance": "source-reviewed", + "scope": "The installable NeRD supervised recipe only; the core does not resolve or import model families eagerly.", + "evidence": [ + { + "path": "extensions/embodichain_models/embodichain_models/dynamics/nerd/sl_recipe.py", + "symbol": "build_supervised_trainer", + "start_line": 167, + "end_line": 167, + "excerpt": " learner = SupervisedLearner(" + } + ] + }, + { + "id": "nerd-showcase-calls-sl-workflow", + "source": "nerd-showcase", + "target": "supervised-learning", + "relation": "calls", + "description": "The showcase delegates a factory-based training configuration to the common offline workflow.", + "provenance": "source-reviewed", + "scope": "The train command when its selected configuration declares a factory; explicit historical configurations take the historical trainer path.", + "evidence": [ + { + "path": "examples/learning/nerd/__main__.py", + "symbol": "main", + "start_line": 273, + "end_line": 273, + "excerpt": " result = train_from_config(" + } + ] } ], "views": [ @@ -5774,7 +5901,10 @@ "grasp-tools", "urdf-assembly", "convex-decomposition", - "scene-services" + "scene-services", + "neural-dynamics", + "supervised-learning", + "nerd-showcase" ], "edge_ids": [ "adapter-calls-compiler", @@ -5860,7 +5990,9 @@ "scene-services-references-manager", "scene-recipes-import-expansion-contracts", "cli-calls-motion-benchmarks", - "motion-benchmarks-constructs-simulation" + "motion-benchmarks-constructs-simulation", + "nerd-recipe-constructs-sl-learner", + "nerd-showcase-calls-sl-workflow" ], "groups": [ { @@ -5989,7 +6121,10 @@ "differentiable-trainer", "differentiable-collector", "observation-normalizer", - "policy-evaluation" + "policy-evaluation", + "neural-dynamics", + "supervised-learning", + "nerd-showcase" ] }, { @@ -6299,7 +6434,7 @@ { "id": "data-learning", "label": "Data & Learning", - "description": "Demonstration recording, online sampling, and separate standard and differentiable training routes.", + "description": "Demonstration recording, online sampling, shared supervised training and concrete model recipes, alongside standard and differentiable RL routes.", "node_ids": [ "embodied-env", "managers", @@ -6329,7 +6464,10 @@ "rollout-buffer", "differentiable-trainer", "differentiable-collector", - "observation-normalizer" + "observation-normalizer", + "supervised-learning", + "neural-dynamics", + "nerd-showcase" ], "edge_ids": [ "cli-calls-training-assembly", @@ -6356,7 +6494,9 @@ "differentiable-trainer-constructs-observation-normalizer", "rl-algorithms-holds-ppo", "rl-algorithms-holds-grpo", - "rl-algorithms-holds-apg" + "rl-algorithms-holds-apg", + "nerd-recipe-constructs-sl-learner", + "nerd-showcase-calls-sl-workflow" ], "groups": [ { @@ -6396,7 +6536,10 @@ "cli", "rl-routing", "learning-env", - "policy-evaluation" + "policy-evaluation", + "supervised-learning", + "neural-dynamics", + "nerd-showcase" ] }, { diff --git a/docs/scripts/build_architecture.py b/docs/scripts/build_architecture.py index 00deed57fe..7af23ac160 100644 --- a/docs/scripts/build_architecture.py +++ b/docs/scripts/build_architecture.py @@ -71,6 +71,8 @@ def _module_name(path: str) -> str: path = path.removeprefix("extensions/") if path.startswith("embodichain_tasks/embodichain_tasks/"): path = path.removeprefix("embodichain_tasks/") + if path.startswith("embodichain_models/embodichain_models/"): + path = path.removeprefix("embodichain_models/") return path.removesuffix(".py").replace("/", ".").removesuffix(".__init__") diff --git a/docs/scripts/check_api_docs.py b/docs/scripts/check_api_docs.py index 7ddabfbdd1..3077622e42 100644 --- a/docs/scripts/check_api_docs.py +++ b/docs/scripts/check_api_docs.py @@ -126,6 +126,10 @@ def is_aligned(self) -> bool: DEFAULT_PACKAGE_ROOTS = ( PackageRoot("embodichain", REPO_ROOT / "embodichain"), + PackageRoot( + "embodichain_models", + REPO_ROOT / "extensions" / "embodichain_models" / "embodichain_models", + ), PackageRoot("embodichain_robots", REPO_ROOT / "extensions" / "embodichain_robots"), PackageRoot( "embodichain_benchmarks", REPO_ROOT / "extensions" / "embodichain_benchmarks" diff --git a/docs/source/api_reference/embodichain/embodichain.learning.rst b/docs/source/api_reference/embodichain/embodichain.learning.rst index 10f9c879a4..662af80727 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.rst +++ b/docs/source/api_reference/embodichain/embodichain.learning.rst @@ -1,6 +1,5 @@ embodichain.learning ==================== - .. automodule:: embodichain.learning .. rubric:: Submodules @@ -10,6 +9,7 @@ embodichain.learning layers models rl + sl Reinforcement Learning ---------------------- @@ -59,8 +59,21 @@ Shared trainable models builders. RL-specific ``Policy``, ``ActorCritic``, ``ActorOnly`` and policy registration/builders remain in ``learning.rl.policies``. Shared components and standalone IK layers do not import RL, simulation or TensorDict. +Concrete model families and their adapters are documented in the +:doc:`official model library `. .. toctree:: :maxdepth: 1 embodichain.learning.models + +Supervised learning +------------------- + +Policies and learned dynamics share tensor-tree batches, objective-owned losses, +optimizer updates and committed-update recovery through ``learning.sl``. + +.. toctree:: + :maxdepth: 1 + + embodichain.learning.sl diff --git a/docs/source/api_reference/embodichain/embodichain.learning.sl.rst b/docs/source/api_reference/embodichain/embodichain.learning.sl.rst new file mode 100644 index 0000000000..47aa97a49d --- /dev/null +++ b/docs/source/api_reference/embodichain/embodichain.learning.sl.rst @@ -0,0 +1,133 @@ +Supervised learning (SL) +======================== + +The SL layer composes objective-owned losses with a common optimizer update and +offline run lifecycle. Batches retain their tensor-tree structure and causal +masks; source adapters, objectives and downstream execution define their semantics. +Importing this package does not initialize simulation or reinforcement learning. + +embodichain.learning.sl +----------------------- + +.. currentmodule:: embodichain.learning.sl + +.. autosummary:: + + LossResult + SupervisedBatch + SupervisedObjective + SupervisedLearner + SupervisedSource + TensorBatchSource + SupervisedTrainer + SupervisedTrainerCfg + +.. autoclass:: LossResult + :members: + +.. autoclass:: SupervisedBatch + :members: + +.. autoclass:: SupervisedObjective + :members: + +.. autoclass:: SupervisedLearner + :members: + +.. autoclass:: SupervisedSource + :members: + +.. autoclass:: TensorBatchSource + :members: + +.. autoclass:: SupervisedTrainer + :members: + +.. autoclass:: SupervisedTrainerCfg + :members: + +embodichain.learning.sl.learner +------------------------------- + +.. currentmodule:: embodichain.learning.sl.learner + +.. autosummary:: + + SupervisedLearner + +.. autoclass:: SupervisedLearner + :members: + +embodichain.learning.sl.sources +------------------------------- + +.. currentmodule:: embodichain.learning.sl.sources + +.. autosummary:: + + SupervisedSource + TensorBatchSource + +.. autoclass:: SupervisedSource + :members: + +.. autoclass:: TensorBatchSource + :members: + +embodichain.learning.sl.trainer +------------------------------- + +.. currentmodule:: embodichain.learning.sl.trainer + +.. autosummary:: + + SupervisedTrainerCfg + SupervisedTrainer + +.. autoclass:: SupervisedTrainerCfg + :members: + +.. autoclass:: SupervisedTrainer + :members: + +embodichain.learning.sl.types +----------------------------- + +.. currentmodule:: embodichain.learning.sl.types + +.. autosummary:: + + SupervisedBatch + LossResult + SupervisedObjective + +.. autoclass:: SupervisedBatch + :members: + +.. autoclass:: LossResult + :members: + +.. autoclass:: SupervisedObjective + :members: + +embodichain.learning.sl.workflow +-------------------------------- + +.. currentmodule:: embodichain.learning.sl.workflow + +.. autosummary:: + + train_from_config + cli + +.. autofunction:: train_from_config + +.. autofunction:: cli + +Model-specific recipes +---------------------- + +Concrete models own their physical codecs, objectives and source adapters in +the :doc:`official model library `. The NeRD +recipe connects committed trajectory windows to this common trainer without +loading simulation or requiring the example runner. diff --git a/docs/source/api_reference/index.rst b/docs/source/api_reference/index.rst index 43cd5cd8a6..b43ee9e034 100644 --- a/docs/source/api_reference/index.rst +++ b/docs/source/api_reference/index.rst @@ -41,9 +41,10 @@ The core ``embodichain`` framework includes these top-level packages: independently of the simulation loop. ``learning`` - Learning systems, currently the ``rl`` subpackage: on-policy RL algorithms - (PPO/GRPO), rollout buffers, collectors, policy/model builders, and the - training entry point. + Shared trainable components, supervised updates and offline orchestration, + plus on-policy RL algorithms (PPO/GRPO), rollout buffers, collectors and + RL policy builders. Concrete model families live in the bundled + ``embodichain_models`` library. ``utils`` Shared utilities: the ``@configclass`` decorator, logging, math/tensor @@ -92,3 +93,4 @@ documentation. CI runs this same checker after style checks and before tests. :maxdepth: 2 robot_catalogs + model_catalogs diff --git a/docs/source/api_reference/model_catalogs.rst b/docs/source/api_reference/model_catalogs.rst new file mode 100644 index 0000000000..ba6c17cf3c --- /dev/null +++ b/docs/source/api_reference/model_catalogs.rst @@ -0,0 +1,254 @@ +Official model library +====================== + +The main ``embodichain`` distribution bundles ``embodichain_models`` with the +same release version and dependencies. This package owns concrete model +families, their physical representations, model-specific data adapters and +supervised recipes. General MLPs and feature normalization remain in +:mod:`embodichain.learning.models`; shared optimizer updates, recovery and +training orchestration remain in :doc:`embodichain/embodichain.learning.sl`. + +Import a selected family explicitly. Importing the model-library root does not +load every model or initialize simulation. NeRD's raw teacher assets, data +generation and original-repository comparison runners belong to the +``learning/nerd`` showcase, outside the runtime model package. + +NeRD public entry points +------------------------- + +``embodichain_models.dynamics.nerd`` lazily exposes the causal network, its +physical codec, portable checkpoints and shared-SL integration. Inputs are named +history tensors; predictions are physical state-increment targets. Model +evaluation retains input gradients. Frozen training statistics travel with the +network as buffers. + +.. currentmodule:: embodichain_models.dynamics.nerd + +.. autosummary:: + + NeRDModelCfg + NeRDModel + FeatureNormalizer + export_model + load_model + convert_legacy_checkpoint + NeRDStateCodec + NeRDObjective + make_nerd_batch + NeRDWindowSource + build_supervised_trainer + +.. autoclass:: NeRDModelCfg + +.. autoclass:: NeRDModel + +.. autoclass:: FeatureNormalizer + +.. autofunction:: export_model + +.. autofunction:: load_model + +.. autofunction:: convert_legacy_checkpoint + +.. autoclass:: NeRDStateCodec + +.. autoclass:: NeRDObjective + +.. autofunction:: make_nerd_batch + +.. autoclass:: NeRDWindowSource + +.. autofunction:: build_supervised_trainer + +Network and normalization +-------------------------- + +Ordered input widths and Transformer settings define the network schema. +``forward`` predicts every history position; ``evaluate`` selects the last valid +position for each row. Causal attention preserves valid prefix outputs under +right padding; the objective owns the padded loss mask. Feature conversion +follows the network's floating dtype without detaching input gradients. + +.. currentmodule:: embodichain_models.dynamics.nerd.model + +.. autosummary:: + + NeRDModelCfg + NeRDModel + FeatureNormalizer + +.. autoclass:: NeRDModelCfg + +.. autoclass:: NeRDModel + :members: forward, evaluate, set_normalization + +.. autoclass:: FeatureNormalizer + :members: forward + +Portable checkpoints +--------------------- + +Exports contain tensor weights and JSON deployment metadata, including a weight +checksum, numerical profile and attention backend. Loading restores the saved +floating precision and returns an evaluation-mode model. Legacy Python-object +conversion requires explicit trust and accepts the supported official +Transformer architecture. + +.. currentmodule:: embodichain_models.dynamics.nerd.checkpoint + +.. autosummary:: + + export_model + load_model + convert_legacy_checkpoint + +.. autofunction:: export_model + +.. autofunction:: load_model + +.. autofunction:: convert_legacy_checkpoint + +Physical state, contact and history encoding +-------------------------------------------- + +The Cartpole and Ant codec converts native Newton COM velocities to NeRD's +angular-first spatial twist and expresses relative targets in the current body +frame. Quaternion helpers use xyzw ordering. Contact features use fixed plane +slots; history append and reset return new tensors while retaining gradients +through unaffected values. + +.. currentmodule:: embodichain_models.dynamics.nerd.codec + +.. autosummary:: + + NeRDStateCodec + quat_multiply + quat_rotate + quat_inverse + plane_contact_features + append_history + reset_history + +.. autoclass:: NeRDStateCodec + :members: native_to_legacy, legacy_to_native, encode, targets, decode + +.. autofunction:: quat_multiply + +.. autofunction:: quat_rotate + +.. autofunction:: quat_inverse + +.. autofunction:: plane_contact_features + +.. autofunction:: append_history + +.. autofunction:: reset_history + +Committed trajectories and windows +----------------------------------- + +The single-writer HDF5 format stores T+1 states and T transitions. Only +manifest-committed shards are visible to readers. Windows respect trajectory +and split boundaries; optional startup windows carry right-padding masks. +Readers bound open shard handles and offer an explicitly budgeted resident +cache. Operational generation progress is excluded from the data fingerprint. + +.. currentmodule:: embodichain_models.dynamics.nerd.data + +.. autosummary:: + + DynamicsShardWriter + DynamicsWindowDataset + trajectory_split + load_manifest + dataset_fingerprint + +.. autoclass:: DynamicsShardWriter + :members: committed_ids, add_trajectory, commit, close + +.. autoclass:: DynamicsWindowDataset + :members: get_batch, enable_device_cache, close + +.. autofunction:: trajectory_split + +.. autofunction:: load_manifest + +.. autofunction:: dataset_fingerprint + +Encoding, statistics and evaluation +----------------------------------- + +Raw windows become named feature histories, physical increment targets and a +valid-position mask. Statistics use only the training split and weight each +valid overlapping-window occurrence; artifacts bind the dataset fingerprint. +The normalized MSE uses frozen target variance. Offline evaluation reports +valid-position prediction and decoded state errors without selecting on test. + +.. currentmodule:: embodichain_models.dynamics.nerd.integration + +.. autosummary:: + + encode_batch + normalized_mse + compute_statistics + save_statistics + load_statistics + evaluate_model + +.. autofunction:: encode_batch + +.. autofunction:: normalized_mse + +.. autofunction:: compute_statistics + +.. autofunction:: save_statistics + +.. autofunction:: load_statistics + +.. autofunction:: evaluate_model + +Shared supervised objective +--------------------------- + +The adapter preserves the raw reader's sampling contract when creating a +``SupervisedBatch``. Its objective returns a graph-connected normalized +state-increment loss and the number of valid positions used for aggregation. +Computing the loss does not update normalization statistics. + +.. currentmodule:: embodichain_models.dynamics.nerd.supervised + +.. autosummary:: + + NeRDObjective + make_nerd_batch + +.. autoclass:: NeRDObjective + :members: compute_loss + +.. autofunction:: make_nerd_batch + +Offline SL recipe +------------------ + +``NeRDWindowSource`` uses the trainer's explicit CPU sampling generator and +encodes copied windows on the requested device. The factory combines this +source, the physical objective, Adam scheduling and held-out evaluation with +the common SL trainer. Its checkpoints use the SL format; historical NeRD +training checkpoints remain owned by the showcase's historical trainer. + +Use the installable factory path +``embodichain_models.dynamics.nerd.sl_recipe:build_supervised_trainer`` in an +``embodichain train-sl`` configuration. It is available independently of the +working directory and example-runner bootstrap. + +.. currentmodule:: embodichain_models.dynamics.nerd.sl_recipe + +.. autosummary:: + + NeRDWindowSource + build_supervised_trainer + +.. autoclass:: NeRDWindowSource + :members: sample + +.. autofunction:: build_supervised_trainer diff --git a/docs/source/conf.py b/docs/source/conf.py index c37ec8abea..323191504b 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -36,6 +36,9 @@ os.environ.setdefault("AZURE_OPENAI_API_KEY", "mock-api-key-for-docs-build") sys.path.insert(0, os.path.abspath("../..")) +sys.path.insert( + 0, str(Path(__file__).resolve().parents[2] / "extensions" / "embodichain_models") +) sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "scripts")) diff --git a/docs/source/features/neural_dynamics.md b/docs/source/features/neural_dynamics.md new file mode 100644 index 0000000000..2300b89876 --- /dev/null +++ b/docs/source/features/neural_dynamics.md @@ -0,0 +1,174 @@ +# Neural Dynamics Models + +A neural dynamics model predicts physical state changes under joint efforts and +contact. EmbodiChain bundles the simulation-independent NeRD network, codec, +trajectory reader and SL recipe in `embodichain_models.dynamics.nerd`. Shared +supervised updates and recovery remain in `embodichain.learning.sl`; general +MLPs remain in `embodichain.learning.models`. + +[NeRD (Neural Robot Dynamics)](https://neural-robot-dynamics.github.io/) uses a +causal Transformer over state, effort and contact histories. The current +reference systems are contact-free Cartpole and free-base Ant on a static plane. +The `learning/nerd` showcase owns preparation, collection, training commands and +reference research. Its default training command uses the common SL recipe. + +## Comparison videos + +These recordings compare the learned NeRD transition with its matching Newton +Featherstone solver. Each video shows passive motion, shared random efforts, and +a PD controller. Passive and random-effort comparisons use the same initial +state and applied effort sequence. In the PD segment, each simulator computes +feedback from its own state, so the comparison shows closed-loop behavior. + +```{raw} html +
+ +
Cartpole: analytical and learned dynamics.
+
+
+ +
Ant: analytical and learned contact dynamics.
+
+``` + +Download the [Cartpole video](https://huggingface.co/DexForceAI/embodichain_model/resolve/main/neural_dynamics/nerd/media/cartpole-comparison.mp4) +or [Ant video](https://huggingface.co/DexForceAI/embodichain_model/resolve/main/neural_dynamics/nerd/media/ant-comparison.mp4). +The recordings visualize the released reference models; they are not a claim of +paper-level accuracy or equivalent behavior in arbitrary robot scenes. + +## Pinned resources + +The showcase prepares these resources at fixed revisions: + +| Resource | Repository | Revision | +| --- | --- | --- | +| Dataset | [DexForceAI/embodichain_datasets](https://huggingface.co/datasets/DexForceAI/embodichain_datasets) | `b9ec58f36da99cc99a93f97adfe74ebc9fd66586` | +| Portable model | [DexForceAI/embodichain_model](https://huggingface.co/DexForceAI/embodichain_model) | `78e581526570069410bf269b38c0c52c800b24a7` | +| Raw teacher assets | [DexForceAI/embodichain_data](https://huggingface.co/DexForceAI/embodichain_data/tree/97e202e41fe7fb504656a6b163291305e400780d/demo/nerd/reference_assets/v1) | `97e202e41fe7fb504656a6b163291305e400780d` | + +Datasets and models use `neural_dynamics/nerd//v1`, where `` is +`cartpole` or `ant`. Raw assets use `demo/nerd/reference_assets/v1`; the verified +archive includes the URDF/XML, attribution and license. These downloaded +payloads are stored outside Git and the wheel. The model implementation and its +third-party license ship in the main EmbodiChain wheel with the same version. + +Dataset manifests publish trajectory-level train/validation/test splits and +physical provenance. Trajectories have T+1 states and T effort/contact inputs; +windows cannot cross trajectories or splits. Portable models contain +`weights.pt` and `manifest.json`, with the network, normalization, codec, +timestep and physical identity. They are inference exports, not resumable +training checkpoints. Dataset/model cards describe dependency and evaluation +limits; long-horizon Ant predictions can drift. + +## Prepare and evaluate + +These commands also work from an installed wheel outside the checkout: + +```bash +embodichain run-example learning/nerd prepare --robot cartpole --resource-root ./outputs/nerd-resources +embodichain run-example learning/nerd eval --robot cartpole --resource-root ./outputs/nerd-resources --device cpu +``` + +Use `--robot ant` for Ant. `prepare` verifies pinned manifests, shards, weights +and reference asset hashes. Add `--assets-only` to prepare only teacher assets, +or `--local-files-only` to require an existing verified cache without network +access. Without `--resource-root`, the showcase uses +`EMBODICHAIN_NERD_RESOURCE_ROOT`, otherwise the cache under `XDG_CACHE_HOME` or +`~/.cache/embodichain/examples/nerd`. Evaluation writes +`outputs/nerd//offline.json` unless overridden with `--output`. + +Offline evaluation measures held-out transitions and does not start a simulator. +It does not establish accumulated rollout accuracy. Published Scene-generated +trajectories retain their original physical identity; the independent direct +Newton collector uses `nerd-direct-newton-builder-v1`. A shared schema or profile +name does not make different solvers, versions or construction paths identical. + +## Train, recover and export + +After preparation, train through the default common SL recipe: + +```bash +embodichain run-example learning/nerd train --robot cartpole --resource-root ./outputs/nerd-resources --device cuda:0 --output-dir ./outputs/nerd/cartpole/train +embodichain run-example learning/nerd export --checkpoint ./outputs/nerd/cartpole/train/last.pt --output ./outputs/nerd/cartpole/export +embodichain run-example learning/nerd eval --robot cartpole --resource-root ./outputs/nerd-resources --device cpu --model ./outputs/nerd/cartpole/export --output ./outputs/nerd/cartpole/trained-eval.json +``` + +`--device cpu` is also supported for training. The shipped `sl.train.yaml` +configures 1,000 updates; this bounded example does not claim full convergence. +The default output is `outputs/nerd//train`. Resume with the same recipe +and resources using `train --resume PATH/last.pt`; choose an artifact-free output +for a new run. Recovery checks data, recipe, feature order and run identity. +Different or newer metric history is preserved and requires a fresh destination. +The final SL checkpoint is not a claim of validation-selected accuracy. + +The generic entry point accepts the same installable factory without example +bootstrap. Copy `examples/learning/nerd/configs//sl.train.yaml`, set +`data.path` to the prepared dataset and `trainer.output_dir` to a fresh output, +then run: + +```bash +embodichain train-sl --config /path/to/recipe.yaml +embodichain train-sl --config /path/to/recipe.yaml --resume /path/to/train/last.pt +``` + +Relative paths in a common SL recipe resolve from the working directory. Its +factory is `embodichain_models.dynamics.nerd.sl_recipe:build_supervised_trainer`. +Export accepts both SL checkpoints and historical NeRD checkpoints and preserves +floating dtype, normalization, ordered features and attention policy. Export +destinations must be new directories. + +## Teacher collection and historical experiments + +Prepare reference assets before invoking the independent collector: + +```bash +embodichain run-example learning/nerd prepare --robot cartpole --resource-root ./outputs/nerd-resources --assets-only +embodichain run-example learning/nerd generate --robot cartpole --resource-root ./outputs/nerd-resources --device cpu --output-dir ./outputs/nerd/cartpole/generated +``` + +The default collection preset is a bounded smoke workflow. Generation uses +Newton's built-in analytical solvers without DexSim's custom solver/Scene +extension. Its manifest commits the candidate cursor and rejection evidence +with shards. Rejection-only progress preserves data identity; legacy sidecars +remain readable, but resuming requires matching collector source and physical +identity. + +Explicit `dynamics.train.*.yaml` configurations in the showcase select the +historical trainer. It retains its NeRD checkpoint format and comparison +sampling/scheduling, including historical fine-tuning and shared sample-order +options. Its `best.pt` is a model-only selected snapshot for export; `last.pt` +is resumable and retains the selected best weights. These options are separate +from default SL training. See the showcase's [research protocol](https://github.com/DexForce/EmbodiChain/blob/main/examples/learning/nerd/research/acceptance.md) +for frozen source revisions, paired-seed gates and historical evidence limits. +Moving code or rerunning a smoke does not rerun those frozen experiments. + +## Use the model in PyTorch + +```python +from embodichain_models.dynamics.nerd import load_model + +model, manifest = load_model( + "outputs/nerd-resources/models/neural_dynamics/nerd/cartpole/v1", + device="cpu", +) +``` + +The codec and `integration.encode_batch` construct named feature windows from +state, effort and contact records. `model(inputs)` predicts every token; +`model.evaluate(inputs, valid_lengths)` selects each row's last valid history +position. Both preserve input gradients. The [official model API](../api_reference/model_catalogs.rst) +describes the runtime model, data and objective components. + +Live Scene stepping and differentiable Gym integration are separate consumers, +tracked in the [neural dynamics solver PR](https://github.com/DexForce/EmbodiChain/pull/822). +The comparison videos above remain reference recordings from that integration; +they do not claim that this offline showcase provides a native simulation solver +or reproduces paper-scale training. diff --git a/docs/source/guides/neural_dynamics.md b/docs/source/guides/neural_dynamics.md new file mode 100644 index 0000000000..c0d657ff21 --- /dev/null +++ b/docs/source/guides/neural_dynamics.md @@ -0,0 +1,11 @@ +--- +orphan: true +--- + +# Neural dynamics + +See {doc}`/features/neural_dynamics` for the bundled `embodichain_models` runtime, +pinned resources and `embodichain run-example learning/nerd` workflow. Training +defaults to the common SL recipe. The {doc}`/api_reference/model_catalogs` page +owns public model, codec, data and objective APIs; historical comparison tools +remain in the NeRD showcase. diff --git a/docs/source/index.rst b/docs/source/index.rst index 993bb22d07..813d0c8a44 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -46,6 +46,7 @@ Table of Contents :glob: features/online_data.md + features/neural_dynamics.md features/generative_sim/index* features/workspace_analyzer/index* features/interaction/index* diff --git a/embodichain/cli/main.py b/embodichain/cli/main.py index 8a72100137..0c90341580 100644 --- a/embodichain/cli/main.py +++ b/embodichain/cli/main.py @@ -37,6 +37,11 @@ class Command: COMMANDS = ( + Command( + name="train-sl", + target="embodichain.learning.sl.workflow:cli", + help="Train a supervised policy or dynamics recipe from YAML or JSON.", + ), Command( name="data", target="embodichain.data.download:main", diff --git a/embodichain/learning/__init__.py b/embodichain/learning/__init__.py index 373314c615..47b846ace2 100644 --- a/embodichain/learning/__init__.py +++ b/embodichain/learning/__init__.py @@ -24,7 +24,7 @@ from importlib import import_module from types import ModuleType -__all__ = ["layers", "models", "rl"] +__all__ = ["layers", "models", "rl", "sl"] def __getattr__(name: str) -> ModuleType: diff --git a/embodichain/learning/sl/__init__.py b/embodichain/learning/sl/__init__.py new file mode 100644 index 0000000000..d63ae13c3a --- /dev/null +++ b/embodichain/learning/sl/__init__.py @@ -0,0 +1,35 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Supervised learning (SL) for policies and learned state transitions.""" + +from __future__ import annotations + +from .types import LossResult, SupervisedBatch, SupervisedObjective +from .learner import SupervisedLearner +from .sources import SupervisedSource, TensorBatchSource +from .trainer import SupervisedTrainer, SupervisedTrainerCfg + +__all__ = [ + "LossResult", + "SupervisedBatch", + "SupervisedObjective", + "SupervisedLearner", + "SupervisedSource", + "TensorBatchSource", + "SupervisedTrainer", + "SupervisedTrainerCfg", +] diff --git a/embodichain/learning/sl/_artifacts.py b/embodichain/learning/sl/_artifacts.py new file mode 100644 index 0000000000..7b1f465078 --- /dev/null +++ b/embodichain/learning/sl/_artifacts.py @@ -0,0 +1,147 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Directory ownership shared by supervised training entry points.""" + +from __future__ import annotations + +import json +from pathlib import Path +import uuid + +__all__: list[str] = [] + + +class _RunArtifacts: + """Bind training artifacts to a durable run identity, independent of paths. + + This is an admission check, not a multi-writer lock. A caller owns one + training writer; concurrent distributed writers must coordinate externally. + Metadata validation never deserializes an existing destination checkpoint. + """ + + def __init__(self, output: Path) -> None: + self.output = output + self.run_id = uuid.uuid4().hex + self._marker = output / "run.json" + self._claimed = False + + def _occupied(self) -> bool: + return ( + any( + path.exists() + for path in ( + self.output / "metrics.jsonl", + self.output / "last.pt", + self.output / "best.pt", + ) + ) + or next(self.output.glob("step_*.pt"), None) is not None + ) + + def _read_identity(self) -> str: + try: + record = json.loads(self._marker.read_text()) + identity = record["run_id"] + if ( + record.get("version") != 1 + or not isinstance(identity, str) + or not identity + ): + raise ValueError("Invalid training run identity") + return identity + except (OSError, ValueError, KeyError, TypeError) as error: + raise ValueError( + "Invalid training run identity; use a fresh output directory" + ) from error + + def claim(self) -> None: + """Admit one run without replacing another run's artifacts.""" + if self._claimed: + return + self.output.mkdir(parents=True, exist_ok=True) + if self._marker.exists(): + if self._read_identity() != self.run_id: + raise FileExistsError( + "Training output belongs to another run; use a fresh output directory" + ) + self._claimed = True + return + if self._occupied(): + raise FileExistsError( + "Training output contains another run; use a fresh output directory" + ) + try: + with self._marker.open("x") as stream: + json.dump({"version": 1, "run_id": self.run_id}, stream) + stream.write("\n") + except FileExistsError: + if self._read_identity() != self.run_id: + raise FileExistsError( + "Training output belongs to another run; use a fresh output directory" + ) from None + self._claimed = True + + def validate_resume(self, saved_run_id: str | None) -> str: + """Check destination identity without mutating files or runtime state. + + Checkpoints without an identity can be recovered into an artifact-free + destination and acquire an identity with their next save. Existing + directories require an explicit matching run marker and checkpoint ID. + """ + if saved_run_id is not None and ( + not isinstance(saved_run_id, str) or not saved_run_id + ): + raise ValueError("Invalid checkpoint run identity") + occupied = self._marker.exists() or self._occupied() + if occupied and ( + saved_run_id is None + or not self._marker.exists() + or self._read_identity() != saved_run_id + ): + raise FileExistsError( + "Resume destination belongs to another run; use a fresh output directory" + ) + return saved_run_id or self.run_id + + def adopt_resume(self, run_id: str) -> None: + """Adopt an identity only after all checkpoint checks have succeeded.""" + self.run_id = run_id + self._claimed = False + + def validate_history(self, committed_step: int, *, key: str) -> None: + """Keep existing history intact when admitting a checkpoint boundary.""" + if type(committed_step) is not int or committed_step < 0: + raise ValueError("Invalid checkpoint update cursor") + path = self.output / "metrics.jsonl" + if not path.exists(): + return + with path.open() as stream: + for line in stream: + if not line.strip(): + continue + try: + step = json.loads(line)[key] + if type(step) is not int or step < 0: + raise ValueError("Invalid history update cursor") + except (ValueError, KeyError, TypeError) as error: + raise ValueError( + "Training output history is incomplete or invalid; use a fresh output directory" + ) from error + if step > committed_step: + raise ValueError( + "Training output history is ahead of this checkpoint; use a fresh output directory" + ) diff --git a/embodichain/learning/sl/learner.py b/embodichain/learning/sl/learner.py new file mode 100644 index 0000000000..33af0e32a9 --- /dev/null +++ b/embodichain/learning/sl/learner.py @@ -0,0 +1,204 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Single-device optimizer updates shared by policy and dynamics fitting.""" + +from __future__ import annotations + +import math +from collections.abc import Callable, Iterable, Mapping +from contextlib import AbstractContextManager, nullcontext +from typing import Any + +import torch +from torch import nn + +from .types import SupervisedBatch, SupervisedObjective + +__all__ = ["SupervisedLearner"] + + +class SupervisedLearner: + """Own one optimizer while objectives own mathematical and graph semantics. + + Args: + model: Caller-constructed model on its intended device. + objective: Loss builder, including any sequence unrolling or integration. + optimizer: Optimizer over the model's trainable parameters. + max_grad_norm: Positive clipping bound, or zero to measure without clipping. + amp_dtype: Optional bfloat16 autocast for the forward objective. Numerical + objectives may explicitly leave autocast for FP32 state integration. + clip_grad: Optional numerical-profile-specific clipping implementation. + numerical_context: Context factory spanning forward, backward and clipping. + + Sampling, normalizer updates, learning-rate schedules and model publication + belong to callers. An update is committed only after optimizer.step succeeds. + """ + + def __init__( + self, + model: nn.Module, + objective: SupervisedObjective, + optimizer: torch.optim.Optimizer, + *, + max_grad_norm: float = 0.0, + amp_dtype: torch.dtype | None = None, + clip_grad: ( + Callable[[Iterable[nn.Parameter], float], torch.Tensor] | None + ) = None, + numerical_context: Callable[[], AbstractContextManager[Any]] | None = None, + ) -> None: + if not math.isfinite(max_grad_norm) or max_grad_norm < 0: + raise ValueError("max_grad_norm must be finite and nonnegative") + if amp_dtype not in (None, torch.bfloat16): + raise ValueError("Only optional bfloat16 autocast is supported") + self.model, self.objective, self.optimizer = model, objective, optimizer + self.max_grad_norm, self.amp_dtype = float(max_grad_norm), amp_dtype + self.clip_grad = clip_grad + self.numerical_context = numerical_context or nullcontext + self.num_updates = 0 + + def update(self, batch: SupervisedBatch) -> dict[str, float | int]: + """Backpropagate one objective and commit one finite parameter update. + + Args: + batch: Owned training inputs and supervision on the model's device. + + Returns: + Loss, pre-clipping gradient norm, effective weight, update count and + objective metrics. Returned statistics do not retain the graph. + + Raises: + FloatingPointError: Loss or gradient norm is not finite; optimizer + parameters and the committed update counter are not advanced. + """ + self.model.train() + self.optimizer.zero_grad(set_to_none=True) + parameter = next(self.model.parameters()) + autocast = ( + torch.autocast(parameter.device.type, dtype=self.amp_dtype) + if self.amp_dtype is not None + else nullcontext() + ) + with self.numerical_context(): + with autocast: + result = self.objective.compute_loss(self.model, batch) + if result.loss.ndim != 0: + raise ValueError("A supervised objective must return a scalar loss") + if not math.isfinite(float(result.weight)) or result.weight <= 0: + raise ValueError("Objective weight must be finite and positive") + result.loss.backward() + parameters = list(self.model.parameters()) + if self.clip_grad is not None: + norm = self.clip_grad(parameters, self.max_grad_norm) + else: + norm = torch.nn.utils.clip_grad_norm_( + parameters, + self.max_grad_norm or float("inf"), + error_if_nonfinite=True, + ) + observed = ( + torch.stack( + ( + result.loss.detach(), + torch.as_tensor(norm, device=result.loss.device).detach(), + ) + ) + .cpu() + .tolist() + ) + if not all(math.isfinite(value) for value in observed): + raise FloatingPointError("Non-finite supervised loss or gradient norm") + self.optimizer.step() + self.num_updates += 1 + metrics = {name: float(value) for name, value in result.metrics.items()} + metrics.update( + loss=observed[0], + grad_norm=observed[1], + weight=float(result.weight), + num_updates=self.num_updates, + ) + return metrics + + def _optimizer_layout(self) -> list[list[str]]: + names = { + id(parameter): name for name, parameter in self.model.named_parameters() + } + try: + return [ + [names[id(parameter)] for parameter in group["params"]] + for group in self.optimizer.param_groups + ] + except KeyError as error: + raise ValueError( + "Optimizer parameters must belong to the learner model" + ) from error + + def state_dict(self) -> dict[str, Any]: + """Return model, optimizer and committed update state. + + Returns: + A PyTorch-style state mapping; callers serialize it before mutation. + """ + return { + "model": self.model.state_dict(), + "optimizer": self.optimizer.state_dict(), + "optimizer_type": f"{type(self.optimizer).__module__}.{type(self.optimizer).__qualname__}", + "optimizer_layout": self._optimizer_layout(), + "num_updates": self.num_updates, + } + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + """Restore an update boundary after checking model tensor compatibility. + + Args: + state: State produced by :meth:`state_dict`. + """ + current = self.model.state_dict() + saved = state["model"] + if current.keys() != saved.keys(): + raise ValueError("Supervised checkpoint model keys differ") + for key, value in current.items(): + if isinstance(value, torch.Tensor) and ( + not isinstance(saved[key], torch.Tensor) + or value.shape != saved[key].shape + or value.dtype != saved[key].dtype + ): + raise ValueError(f"Supervised checkpoint model tensor differs: {key}") + elif not isinstance(value, torch.Tensor) and value != saved[key]: + raise ValueError(f"Supervised checkpoint model metadata differs: {key}") + count = state["num_updates"] + if type(count) is not int or count < 0: + raise ValueError("Invalid committed update count") + optimizer_type = ( + f"{type(self.optimizer).__module__}.{type(self.optimizer).__qualname__}" + ) + if ( + state.get("optimizer_type") != optimizer_type + or state.get("optimizer_layout") != self._optimizer_layout() + ): + raise ValueError( + "Supervised checkpoint optimizer type or parameter layout differs" + ) + groups = state["optimizer"].get("param_groups", []) + if len(groups) != len(self.optimizer.param_groups) or any( + len(group.get("params", [])) != len(actual["params"]) + for group, actual in zip(groups, self.optimizer.param_groups) + ): + raise ValueError("Supervised checkpoint optimizer parameter groups differ") + self.model.load_state_dict(saved) + self.optimizer.load_state_dict(state["optimizer"]) + self.num_updates = count diff --git a/embodichain/learning/sl/sources.py b/embodichain/learning/sl/sources.py new file mode 100644 index 0000000000..1e2c3259d3 --- /dev/null +++ b/embodichain/learning/sl/sources.py @@ -0,0 +1,169 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Adapters for fixed CPU tensor corpora without simulator dependencies.""" + +from __future__ import annotations + +from collections.abc import Iterator, Mapping +import hashlib +import json +import copy +from typing import Any, Protocol +import torch +from .types import SupervisedBatch, _map_tensors + +__all__ = ["SupervisedSource", "TensorBatchSource"] + + +class SupervisedSource(Protocol): + """Owned batches drawn with an explicit caller-owned CPU sampling RNG.""" + + source_id: str + + def sample(self, batch_size: int, *, generator: torch.Generator) -> SupervisedBatch: + """Draw a batch without hidden waits, production or lifecycle changes. + + Args: + batch_size: Number of training examples. + generator: Caller-owned sampling generator. + + Returns: + An independently owned batch with source-specific tensor axes. + """ + ... + + +class TensorBatchSource: + """Fixed tensor-tree examples with explicit content identity. + + Args: + batch: Resident CPU examples; every tensor has the same first batch axis. + source_id: Optional durable corpus identity. A content hash is always + retained so reusing an ID cannot silently replace training data. + + Batch metadata must be finite JSON data so units and feature declarations + enter content identity. Mapping keys in tensor trees retain their types. + Construction takes a detached owned snapshot. This source has no simulator, + action-schema, sequence-length or image-layout assumption. + """ + + def __init__(self, batch: SupervisedBatch, *, source_id: str = "tensor") -> None: + sizes = [] + digest = hashlib.sha256() + + def snapshot(tensor: torch.Tensor) -> torch.Tensor: + if tensor.device.type != "cpu" or tensor.ndim == 0: + raise ValueError("Tensor corpora require batched CPU tensors") + sizes.append(tensor.shape[0]) + value = tensor.detach().clone().contiguous() + digest.update(f"{tuple(value.shape)}:{value.dtype}".encode()) + digest.update(value.view(torch.uint8).numpy().tobytes()) + return value + + self.batch = SupervisedBatch( + _map_tensors(batch.inputs, snapshot), + _map_tensors(batch.targets, snapshot), + _map_tensors(batch.masks, snapshot), + copy.deepcopy(dict(batch.metadata)), + ) + if not sizes or min(sizes) < 1 or len(set(sizes)) != 1: + raise ValueError("All corpus tensors require the same nonempty batch axis") + + def key_signature(key: Any) -> Any: + if isinstance(key, tuple): + return {"kind": "tuple", "items": [key_signature(item) for item in key]} + if key is None or type(key) in (str, int, float, bool): + return {"kind": type(key).__name__, "value": key} + raise ValueError( + "Tensor-tree mapping keys must be scalar primitives or tuples" + ) + + def structure(value: Any) -> Any: + if isinstance(value, torch.Tensor): + return {"shape": list(value.shape), "dtype": str(value.dtype)} + if isinstance(value, Mapping): + return { + "kind": "mapping", + "items": [ + [key_signature(key), structure(item)] + for key, item in value.items() + ], + } + if isinstance(value, (tuple, list)): + return { + "kind": "tuple" if isinstance(value, tuple) else "list", + "items": [structure(item) for item in value], + } + return value + + digest.update( + json.dumps( + { + "inputs": structure(batch.inputs), + "targets": structure(batch.targets), + "masks": structure(batch.masks), + "metadata": batch.metadata, + }, + sort_keys=True, + allow_nan=False, + ).encode() + ) + self.size = sizes[0] + self.source_id = f"{source_id}:{digest.hexdigest()}" + + def _gather(self, indices: torch.Tensor) -> SupervisedBatch: + gather = lambda tensor: tensor[indices] + return SupervisedBatch( + _map_tensors(self.batch.inputs, gather), + _map_tensors(self.batch.targets, gather), + _map_tensors(self.batch.masks, gather), + { + **copy.deepcopy(self.batch.metadata), + "source_id": self.source_id, + "indices": indices.clone(), + }, + ) + + def sample(self, batch_size: int, *, generator: torch.Generator) -> SupervisedBatch: + """Sample examples with replacement using an explicit RNG. + + Args: + batch_size: Positive number of examples. + generator: CPU sampling RNG. + + Returns: + Owned tensor trees and sampled source indices. + """ + if type(batch_size) is not int or batch_size < 1: + raise ValueError("batch_size must be positive") + return self._gather( + torch.randint(self.size, (batch_size,), generator=generator) + ) + + def iter_batches(self, batch_size: int) -> Iterator[SupervisedBatch]: + """Visit all examples once in stable source order, retaining the tail. + + Args: + batch_size: Positive maximum batch size. + + Yields: + Owned consecutive batches, including the final partial batch. + """ + if type(batch_size) is not int or batch_size < 1: + raise ValueError("batch_size must be positive") + for start in range(0, self.size, batch_size): + yield self._gather(torch.arange(start, min(start + batch_size, self.size))) diff --git a/embodichain/learning/sl/trainer.py b/embodichain/learning/sl/trainer.py new file mode 100644 index 0000000000..d3d46cc816 --- /dev/null +++ b/embodichain/learning/sl/trainer.py @@ -0,0 +1,277 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Offline training orchestration and reproducible committed-update recovery.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +import json +import math +from pathlib import Path +import random +from typing import Any +import numpy as np +import torch +from embodichain.utils import configclass +from ._artifacts import _RunArtifacts +from .learner import SupervisedLearner +from .sources import SupervisedSource + +__all__ = ["SupervisedTrainerCfg", "SupervisedTrainer"] + + +@configclass +class SupervisedTrainerCfg: + """Configure an offline, single-device supervised run. + + Args: + total_updates: Absolute optimizer-update budget, including resumed work. + batch_size: Source examples per update. + device: Tensor destination; runtime adapters remain caller-owned. + seed: Independent source-sampling seed. Initialize models before building + the trainer using the recipe's explicitly chosen model seed. + output_dir: Local metrics and checkpoint directory. + checkpoint_interval: Periodic save frequency; zero saves only at the end. + validation_interval: Optional validation callback frequency. + """ + + total_updates: int = 1000 + batch_size: int = 64 + device: str = "cpu" + seed: int = 0 + output_dir: str = "outputs/sl" + checkpoint_interval: int = 1000 + validation_interval: int = 0 + + +class SupervisedTrainer: + """Compose sampling and a learner without owning acquisition or simulation. + + Args: + learner: Shared single-update learner. + source: Fixed data source with explicit identity and RNG-based sampling. + cfg: Update budget, device and observation settings. + signature: Recipe semantics, including feature/action units, frames, + model configuration and numerical policy where applicable. + scheduler: Optional scheduler stepped after each successful update. + evaluate: Optional callback receiving the committed learner. It owns + validation metrics and must not select a model on a test split. + """ + + def __init__( + self, + learner: SupervisedLearner, + source: SupervisedSource, + cfg: SupervisedTrainerCfg | None = None, + *, + signature: Mapping[str, Any] | None = None, + scheduler: Any = None, + evaluate: Callable[[SupervisedLearner], Mapping[str, float]] | None = None, + ) -> None: + self.cfg = cfg or SupervisedTrainerCfg() + self.cfg.validate() + for name in ( + "total_updates", + "batch_size", + "checkpoint_interval", + "validation_interval", + "seed", + ): + value = getattr(self.cfg, name) + if type(value) is not int or value < (1 if name == "batch_size" else 0): + raise ValueError(f"Invalid supervised configuration: {name}") + self.learner, self.source = learner, source + self.signature = json.loads( + json.dumps(dict(signature or {}), sort_keys=True, allow_nan=False) + ) + self.scheduler, self.evaluate = scheduler, evaluate + self.sampler = torch.Generator(device="cpu").manual_seed(self.cfg.seed) + self.output = Path(self.cfg.output_dir) + self.output.mkdir(parents=True, exist_ok=True) + self._artifacts = _RunArtifacts(self.output) + self.learner.model.to(self.cfg.device) + self._validation_boundary: dict[str, Any] | None = None + + def fit(self, *, stop_after: int | None = None) -> dict[str, Any]: + """Train up to the configured or an earlier committed update boundary. + + Args: + stop_after: Optional absolute early-stop update for resumable runs. + + Returns: + Completed update count, final metrics and last-checkpoint path. + """ + end = ( + self.cfg.total_updates + if stop_after is None + else min(stop_after, self.cfg.total_updates) + ) + if end < self.learner.num_updates: + raise ValueError("stop_after precedes the restored update cursor") + self._artifacts.claim() + self._artifacts.validate_history(self.learner.num_updates, key="num_updates") + metrics = {} + while self.learner.num_updates < end: + candidate = torch.Generator(device="cpu") + candidate.set_state(self.sampler.get_state()) + batch = self.source.sample(self.cfg.batch_size, generator=candidate).to( + self.cfg.device + ) + metrics = self.learner.update(batch) + self.sampler.set_state(candidate.get_state()) + if self.scheduler is not None: + self.scheduler.step() + step = self.learner.num_updates + if ( + self.evaluate is not None + and self.cfg.validation_interval + and step % self.cfg.validation_interval == 0 + ): + metrics.update(self._evaluate_step()) + with (self.output / "metrics.jsonl").open("a") as handle: + handle.write(json.dumps(metrics, allow_nan=False) + "\n") + if ( + self.cfg.checkpoint_interval + and step % self.cfg.checkpoint_interval == 0 + ): + self.save(self.output / f"step_{step}.pt") + path = self.output / "last.pt" + self.save(path) + return { + "num_updates": self.learner.num_updates, + "metrics": metrics, + "checkpoint": str(path), + } + + def _evaluate_step(self) -> dict[str, float]: + boundary = { + "rng": self._rng_state(), + "sampler": self.sampler.get_state(), + "buffers": { + name: value.detach().clone() + for name, value in self.learner.model.named_buffers() + }, + } + was_training = self.learner.model.training + self._validation_boundary = boundary + try: + self.learner.model.eval() + assert self.evaluate is not None + with torch.no_grad(): + return { + f"validation/{key}": float(value) + for key, value in self.evaluate(self.learner).items() + } + finally: + with torch.no_grad(): + for name, value in self.learner.model.named_buffers(): + value.copy_(boundary["buffers"][name]) + self.sampler.set_state(boundary["sampler"]) + self._restore_rng(boundary["rng"]) + self.learner.model.train(was_training) + self._validation_boundary = None + + def _rng_state(self) -> dict[str, Any]: + device = torch.device(self.cfg.device) + return { + "python": random.getstate(), + "numpy": np.random.get_state(), + "torch": torch.get_rng_state(), + "cuda": torch.cuda.get_rng_state(device) if device.type == "cuda" else None, + } + + def _restore_rng(self, state: Mapping[str, Any]) -> None: + random.setstate(state["python"]) + np.random.set_state(state["numpy"]) + torch.set_rng_state(state["torch"]) + if state["cuda"] is not None: + torch.cuda.set_rng_state(state["cuda"], self.cfg.device) + + def save(self, path: str | Path) -> None: + """Atomically save a complete committed learning and sampling boundary. + + Args: + path: Local checkpoint destination, replaced atomically. + """ + path = Path(path) + self._artifacts.claim() + path.parent.mkdir(parents=True, exist_ok=True) + state = { + "version": 1, + "run_id": self._artifacts.run_id, + "learner": self.learner.state_dict(), + "sampler": self.sampler.get_state(), + "rng": self._rng_state(), + "signature": self.signature, + "source_id": self.source.source_id, + "config": self.cfg.to_dict(), + "scheduler_type": ( + None + if self.scheduler is None + else f"{type(self.scheduler).__module__}.{type(self.scheduler).__qualname__}" + ), + "scheduler": ( + None if self.scheduler is None else self.scheduler.state_dict() + ), + } + if self._validation_boundary is not None: + boundary = self._validation_boundary + state["sampler"] = boundary["sampler"] + state["rng"] = boundary["rng"] + for name, value in boundary["buffers"].items(): + if name in state["learner"]["model"]: + state["learner"]["model"][name] = value + temporary = path.with_suffix(path.suffix + ".tmp") + torch.save(state, temporary) + temporary.replace(path) + + def resume(self, path: str | Path) -> None: + """Restore the same run after validating data, recipe and configuration. + + Args: + path: Trusted training checkpoint from this trainer. + """ + state = torch.load(path, map_location="cpu", weights_only=False) + if ( + state.get("version") != 1 + or state["signature"] != self.signature + or state["source_id"] != self.source.source_id + ): + raise ValueError("Supervised checkpoint recipe or dataset identity differs") + current, saved = self.cfg.to_dict(), dict(state["config"]) + for key in ("output_dir",): + current.pop(key, None) + saved.pop(key, None) + if current != saved or (state["scheduler"] is None) != (self.scheduler is None): + raise ValueError("Supervised checkpoint training configuration differs") + scheduler_type = ( + None + if self.scheduler is None + else f"{type(self.scheduler).__module__}.{type(self.scheduler).__qualname__}" + ) + if state.get("scheduler_type") != scheduler_type: + raise ValueError("Supervised checkpoint scheduler type differs") + run_id = self._artifacts.validate_resume(state.get("run_id")) + self._artifacts.validate_history( + state["learner"]["num_updates"], key="num_updates" + ) + self.learner.load_state_dict(state["learner"]) + if self.scheduler is not None: + self.scheduler.load_state_dict(state["scheduler"]) + self.sampler.set_state(state["sampler"]) + self._restore_rng(state["rng"]) + self._artifacts.adopt_resume(run_id) diff --git a/embodichain/learning/sl/types.py b/embodichain/learning/sl/types.py new file mode 100644 index 0000000000..4d3a85ec74 --- /dev/null +++ b/embodichain/learning/sl/types.py @@ -0,0 +1,109 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Tensor-tree batches and objective results for supervised learning.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any, Protocol + +import torch +from torch import nn + +__all__ = ["SupervisedBatch", "LossResult", "SupervisedObjective"] + + +def _map_tensors(value: Any, function: Any) -> Any: + if isinstance(value, torch.Tensor): + return function(value) + if isinstance(value, Mapping): + return {key: _map_tensors(item, function) for key, item in value.items()} + if isinstance(value, tuple): + return tuple(_map_tensors(item, function) for item in value) + if isinstance(value, list): + return [_map_tensors(item, function) for item in value] + return value + + +@dataclass(frozen=True) +class SupervisedBatch: + """Separate measured inputs, supervision and masks without imposing axes. + + Args: + inputs: Tensor or nested tensor tree supplied to the objective. + targets: Supervision tree. Its history may differ from the input history. + masks: Named masks interpreted by the objective, not the optimizer. + metadata: Source identities, units, timing and other host-side context. + + The producer owns storage and causal alignment. This container does not + detach tensors: differentiable multi-step objectives retain their graph. + """ + + inputs: Any + targets: Any + masks: Mapping[str, torch.Tensor] = field(default_factory=dict) + metadata: Mapping[str, Any] = field(default_factory=dict) + + def to(self, device: str | torch.device) -> SupervisedBatch: + """Move tensor trees, preserving dtypes and host metadata. + + Args: + device: Destination for inputs, targets and masks. + + Returns: + A batch whose differentiable tensor transfers preserve gradients. + """ + move = lambda tensor: tensor.to(device) + return SupervisedBatch( + _map_tensors(self.inputs, move), + _map_tensors(self.targets, move), + _map_tensors(self.masks, move), + self.metadata, + ) + + +@dataclass(frozen=True) +class LossResult: + """An objective-owned scalar loss and sample-weighted observations. + + Args: + loss: Scalar connected to the parameters to optimize. + metrics: Scalar diagnostics; they do not enter backward automatically. + weight: Effective observation count for evaluation aggregation. Loss + reduction and sequence weighting remain the objective's decision. + """ + + loss: torch.Tensor + metrics: Mapping[str, torch.Tensor | float] = field(default_factory=dict) + weight: float = 1.0 + + +class SupervisedObjective(Protocol): + """Compute task-specific losses without owning optimizer or data lifecycle.""" + + def compute_loss(self, model: nn.Module, batch: SupervisedBatch) -> LossResult: + """Build a scalar objective, optionally unrolling a differentiable model. + + Args: + model: Trainable module, already on the batch's device. + batch: Explicit inputs, targets, masks and sample context. + + Returns: + Loss and statistics using this objective's numerical conventions. + """ + ... diff --git a/embodichain/learning/sl/workflow.py b/embodichain/learning/sl/workflow.py new file mode 100644 index 0000000000..9d6c68275f --- /dev/null +++ b/embodichain/learning/sl/workflow.py @@ -0,0 +1,78 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Lazy recipe-based CLI for standalone supervised learning.""" + +from __future__ import annotations + +import argparse +from collections.abc import Sequence +import importlib +import json +from pathlib import Path +from typing import Any + +from .trainer import SupervisedTrainer + +__all__ = ["train_from_config", "cli"] + + +def train_from_config( + config: dict[str, Any], *, resume: str | None = None +) -> dict[str, Any]: + """Build a recipe and run the common offline trainer. + + Args: + config: JSON/YAML recipe with ``factory: module:function`` and model, + data, objective, optimizer, trainer and evaluation recipe settings. + The factory receives this mapping and returns a SupervisedTrainer. + resume: Optional trusted training checkpoint. + + Returns: + The trainer's committed update and artifact report. + """ + factory = config.get("factory", "") + module, separator, name = factory.partition(":") + if not separator or not module or not name: + raise ValueError("SL configuration requires factory: module:function") + trainer = getattr(importlib.import_module(module), name)(config) + if not isinstance(trainer, SupervisedTrainer): + raise TypeError("SL recipe factory must return a SupervisedTrainer") + if resume is not None: + trainer.resume(resume) + return trainer.fit() + + +def cli(argv: Sequence[str] | None = None) -> None: + """Train a supervised recipe from a YAML/JSON file and optional checkpoint. + + Args: + argv: Explicit CLI arguments, or the current process arguments. + """ + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", required=True, type=Path) + parser.add_argument("--resume") + args = parser.parse_args(argv) + import yaml + + config = yaml.safe_load(args.config.read_text()) + if not isinstance(config, dict): + raise ValueError("SL configuration must be a mapping") + print(json.dumps(train_from_config(config, resume=args.resume), allow_nan=False)) + + +if __name__ == "__main__": + cli() diff --git a/examples/README.md b/examples/README.md index f6d897a562..f252b0577d 100644 --- a/examples/README.md +++ b/examples/README.md @@ -6,6 +6,7 @@ Use the path of any Python script under this directory, with or without `.py`: embodichain run-example --list embodichain run-example sim/motion/planners/curobo_planner --help embodichain run-example ods/simple_bc_demo/train --help +embodichain run-example learning/nerd --help embodichain run-example sim/atomic_action --skill place --variants 8 embodichain run-example sim/atomic_action/coordinated_pickment --help ``` @@ -30,6 +31,13 @@ distributions preserve this same tree and its resources; wheels install it under distribution metadata and works outside the repository. Example code is not installed inside the `embodichain` Python package. +The [NeRD showcase](learning/nerd/README.md) prepares pinned external resources, +trains through the shared SL recipe, exports models and evaluates saved +transitions. Its private support code owns teacher collection, historical +training and reference experiments. Reusable models and codecs are installed +as `embodichain_models`; downloaded datasets, weights and raw reference assets +are stored outside the repository and wheel. + The [Atomic Skill example catalog](sim/atomic_action/README.md) owns all per-skill, coordinated-manipulation and recovery demonstrations. The two source-checkout walkthroughs under `scripts/tutorials/atomic_action/` teach single-arm and dual-arm diff --git a/examples/learning/nerd/README.md b/examples/learning/nerd/README.md new file mode 100644 index 0000000000..10a7e2ccdb --- /dev/null +++ b/examples/learning/nerd/README.md @@ -0,0 +1,91 @@ +# NeRD showcase + +This example prepares pinned resources, fits causal robot dynamics through the +common SL trainer, exports portable models and evaluates held-out transitions. +Cartpole and Ant are the supported reference systems. Teacher collection and +historical trainer comparisons remain private support workflows. + +Reusable networks, physical codecs, committed trajectory readers and the NeRD +SL recipe live in `embodichain_models.dynamics.nerd`, bundled in the main wheel. +Core `embodichain.learning.sl` owns the optimizer and training lifecycle; shared +MLPs remain in `embodichain.learning.models`. The runtime model package does not +import this example or simulation. + +## Default workflow + +The launcher works from a source checkout or installed wheel and keeps the +working directory. Use a caller-owned resource directory: + +```bash +embodichain run-example learning/nerd prepare --robot cartpole --resource-root ./outputs/nerd-resources +embodichain run-example learning/nerd eval --robot cartpole --resource-root ./outputs/nerd-resources --device cpu +embodichain run-example learning/nerd train --robot cartpole --resource-root ./outputs/nerd-resources --device cuda:0 --output-dir ./outputs/nerd/cartpole/train +embodichain run-example learning/nerd export --checkpoint ./outputs/nerd/cartpole/train/last.pt --output ./outputs/nerd/cartpole/export +embodichain run-example learning/nerd eval --robot cartpole --resource-root ./outputs/nerd-resources --device cpu --model ./outputs/nerd/cartpole/export --output ./outputs/nerd/cartpole/trained-eval.json +``` + +Use `--robot ant` for Ant or `--device cpu` for CPU training. Commands expose +`--help`. `train` defaults to the robot's `sl.train.yaml`, with 1,000 updates; +this bounded recipe is not a convergence qualification. Exported `last.pt` +weights are the final committed weights, not an implicit validation selection. +Resume with the same recipe using `train --resume PATH/last.pt`; each new run +and portable export requires a fresh destination. + +`prepare` pins dataset revision `b9ec58f36da99cc99a93f97adfe74ebc9fd66586`, +model revision `78e581526570069410bf269b38c0c52c800b24a7` and reference-asset +revision `97e202e41fe7fb504656a6b163291305e400780d`. `_resources.py` owns their +repository IDs, paths and manifest/member checksums. Raw reference files and +license/attribution are downloaded from +`DexForceAI/embodichain_data/demo/nerd/reference_assets/v1`; they are not in Git +or the wheel. Use `prepare --assets-only` for collection-only setup and +`--local-files-only` for offline cache checks. Without `--resource-root`, the +example uses `EMBODICHAIN_NERD_RESOURCE_ROOT` or the showcase cache under +`XDG_CACHE_HOME`/`~/.cache`. + +## Custom configurations + +The [Cartpole](configs/cartpole/README.md) and [Ant](configs/ant/README.md) +configuration directories contain common SL recipes and historical presets. +For the generic entry point, copy `sl.train.yaml`, set `data.path` to the +prepared dataset and `trainer.output_dir` to a fresh output, then run: + +```bash +embodichain train-sl --config /path/to/recipe.yaml +``` + +The factory is the installable +`embodichain_models.dynamics.nerd.sl_recipe:build_supervised_trainer`; it does +not require example bootstrap. Paths in a common SL recipe resolve from the +working directory. `--resource-root` replaces paths only for the showcase's +implicit default presets, not an explicitly supplied configuration. + +Explicit `dynamics.train.*.yaml` files select the historical trainer, retaining +its checkpoint format, sample ordering and numerical schedule for comparisons. +Only that configuration accepts historical `--fine-tune` and `--sample-order`. +Its model-only `best.pt` supports export; optimizer recovery uses `last.pt`. +Historical configuration paths resolve relative to the file, and explicit CLI +path overrides resolve from the working directory. + +## Teacher and research workflows + +```bash +embodichain run-example learning/nerd prepare --robot cartpole --resource-root ./outputs/nerd-resources --assets-only +embodichain run-example learning/nerd generate --robot cartpole --resource-root ./outputs/nerd-resources --device cpu --output-dir ./outputs/nerd/cartpole/generated +``` + +Collection uses Newton's built-in analytical solver without a custom DexSim +solver or Scene initialization. Its direct-builder physical identity differs +from the published Scene-generated dataset. Preserve recorded asset, solver, +version and normalizer identities when comparing models. Continuing collection +requires a matching collector source; readable legacy manifests do not imply +cross-version collection resume. + +[Reference experiments](research/README.md) and the +[frozen acceptance protocol](research/acceptance.md) own research restrictions, +source revisions and quality gates. Relocating an implementation or passing a +new smoke does not rerun earlier frozen evidence. Offline prediction, live +rollout and paper-scale accuracy are separate claims. + +See the [feature guide](../../../docs/source/features/neural_dynamics.md) and +[model API](../../../docs/source/api_reference/model_catalogs.rst) for the +resource format and runtime interfaces. diff --git a/examples/learning/nerd/__init__.py b/examples/learning/nerd/__init__.py new file mode 100644 index 0000000000..051f680fe4 --- /dev/null +++ b/examples/learning/nerd/__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. +# ---------------------------------------------------------------------------- + +"""NeRD model and showcase support.""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/examples/learning/nerd/__main__.py b/examples/learning/nerd/__main__.py new file mode 100644 index 0000000000..7bf671174f --- /dev/null +++ b/examples/learning/nerd/__main__.py @@ -0,0 +1,334 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Lazy commands for neural dynamics collection, learning and evaluation.""" + +from __future__ import annotations + +import argparse +from collections.abc import Sequence +import json +import os +from pathlib import Path +from typing import Any + +__all__: list[str] = [] + + +def _read_config(path: Path) -> dict[str, Any]: + if not path.is_file(): + raise ValueError(f"Configuration file does not exist: {path}") + if path.suffix.lower() == ".json": + value = json.loads(path.read_text()) + elif path.suffix.lower() in {".yaml", ".yml"}: + import yaml + + value = yaml.safe_load(path.read_text()) + else: + raise ValueError("Dynamics configuration must be .json, .yaml or .yml") + if not isinstance(value, dict): + raise ValueError("Dynamics configuration must contain a mapping") + return value + + +def _path(value: str, base: Path) -> Path: + if not isinstance(value, str) or not value.strip(): + raise ValueError("Configured paths must be nonempty strings") + path = Path(value).expanduser() + return (base / path).resolve() if not path.is_absolute() else path.resolve() + + +def _resolve_paths( + config: dict[str, Any], base: Path, operation: str +) -> dict[str, Any]: + config = dict(config) + inputs = { + "dataset": "directory", + "statistics": "file", + "sample_order": "file", + "resume": "file", + "fine_tune": "directory", + } + if operation == "eval": + inputs["model"] = "directory" + if operation == "generate": + inputs.pop("resume") # generation resume is a Boolean, not a checkpoint. + for name, kind in inputs.items(): + if name in config and config[name] is not None: + path = _path(config[name], base) + if not (path.is_file() if kind == "file" else path.is_dir()): + raise ValueError(f"{name} {kind} does not exist: {path}") + config[name] = str(path) + for name in ("output_dir", "output"): + if name in config: + config[name] = str(_path(config[name], base)) + return config + + +def main(argv: Sequence[str] | None = None) -> None: + """Prepare, train and evaluate the NeRD showcase without core special cases.""" + parser = argparse.ArgumentParser( + prog="embodichain run-example learning/nerd", description=__doc__ + ) + subparsers = parser.add_subparsers(dest="operation", required=True) + prepare = subparsers.add_parser( + "prepare", help="Download pinned showcase resources." + ) + prepare.add_argument("--robot", choices=("cartpole", "ant"), default="cartpole") + prepare.add_argument("--resource-root", type=Path) + prepare.add_argument("--assets-only", action="store_true") + prepare.add_argument("--local-files-only", action="store_true") + for name, help_text in ( + ("generate", "Generate committed teacher trajectory shards."), + ("train", "Train, fine-tune or resume a NeRD model."), + ("eval", "Evaluate an exported model on held-out trajectories."), + ): + child = subparsers.add_parser(name, help=help_text) + child.add_argument("--config", type=Path) + child.add_argument("--robot", choices=("cartpole", "ant"), default="cartpole") + child.add_argument("--resource-root", type=Path) + child.add_argument("--device", help="Override the configured CPU/CUDA device.") + if name in {"generate", "train"}: + child.add_argument( + "--seed", type=int, help="Override the collection/training seed." + ) + child.add_argument( + "--output-dir", + type=Path, + help="Override output directory (relative to current directory).", + ) + if name == "generate": + child.add_argument( + "--resume", + action="store_true", + default=None, + help="Continue from committed trajectory shards.", + ) + elif name == "train": + restart = child.add_mutually_exclusive_group() + restart.add_argument( + "--resume", + type=Path, + help="Restore an identical committed training run.", + ) + restart.add_argument( + "--fine-tune", + type=Path, + help="Initialize weights/statistics from a standard model export.", + ) + child.add_argument( + "--statistics", + type=Path, + help="Use a shared frozen normalization artifact.", + ) + child.add_argument( + "--sample-order", + type=Path, + help="Shared batch-index tensor (.pt) or memory-mapped integer matrix (.npy).", + ) + else: + child.add_argument( + "--model", type=Path, help="Override the exported model directory." + ) + child.add_argument( + "--output", + type=Path, + help="Override JSON report path (relative to current directory).", + ) + export = subparsers.add_parser( + "export", help="Export portable model weights and metadata." + ) + export.add_argument("--checkpoint", type=Path, required=True) + export.add_argument("--output", type=Path, required=True) + convert = subparsers.add_parser( + "convert-checkpoint", + help="Convert an explicitly trusted official Python-object checkpoint.", + ) + convert.add_argument("--checkpoint", type=Path, required=True) + convert.add_argument("--output", type=Path, required=True) + convert.add_argument( + "--manifest", + type=Path, + required=True, + help="JSON deployment manifest with codec and dt.", + ) + convert.add_argument( + "--trust-pickle", + action="store_true", + required=True, + help="Explicitly trust the source pickle, which can execute Python code.", + ) + args = parser.parse_args(argv) + try: + from ._resources import _resource_root + + resource_root = ( + args.resource_root.expanduser().resolve() + if getattr(args, "resource_root", None) is not None + else _resource_root() + ) + os.environ["EMBODICHAIN_NERD_RESOURCE_ROOT"] = str(resource_root) + if args.operation == "prepare": + from ._resources import prepare_resources + + result = prepare_resources( + resource_root, + args.robot, + assets_only=args.assets_only, + local_files_only=args.local_files_only, + ) + print( + json.dumps({key: str(value) for key, value in result.items()}, indent=2) + ) + return + if args.operation in {"export", "convert-checkpoint"}: + checkpoint = args.checkpoint.expanduser().resolve() + if not checkpoint.is_file(): + raise ValueError(f"Checkpoint does not exist: {checkpoint}") + output = args.output.expanduser().resolve() + if args.operation == "export": + from examples.learning.nerd._support._workflow import run_export + + result = run_export(checkpoint, output) + else: + manifest = json.loads(args.manifest.expanduser().read_text()) + if ( + not isinstance(manifest, dict) + or not isinstance(manifest.get("codec"), dict) + or "dt" not in manifest + ): + raise ValueError( + "Conversion manifest requires codec mapping and dt" + ) + from embodichain_models.dynamics.nerd import convert_legacy_checkpoint + + convert_legacy_checkpoint( + checkpoint, output, manifest, trust_pickle=args.trust_pickle + ) + result = {"output": str(output), "qualification": "unqualified"} + else: + supplied_config = args.config is not None + if supplied_config: + path = args.config.expanduser().resolve() + else: + preset = { + "generate": "dynamics.generate.smoke.yaml", + "train": "sl.train.yaml", + "eval": "dynamics.eval.pretrained.yaml", + }[args.operation] + path = Path(__file__).resolve().parent / "configs" / args.robot / preset + config = _read_config(path) + if not supplied_config: + prefix = Path("neural_dynamics/nerd") / args.robot / "v1" + if "factory" in config: + config["data"] = { + **config["data"], + "path": str(resource_root / "datasets" / prefix), + } + config["trainer"] = { + **config["trainer"], + "output_dir": str(Path("outputs/nerd") / args.robot / "train"), + } + elif args.operation == "eval": + config.update( + dataset=str(resource_root / "datasets" / prefix), + model=str(resource_root / "models" / prefix), + output=str( + ( + Path("outputs/nerd") / args.robot / "offline.json" + ).resolve() + ), + ) + else: + config["output_dir"] = str( + (Path("outputs/nerd") / args.robot / "generated").resolve() + ) + if args.operation == "train" and "factory" in config: + from embodichain.learning.sl.workflow import train_from_config + + if args.fine_tune is not None or args.sample_order is not None: + raise ValueError( + "Historical fine-tune/sample-order options require a historical training configuration" + ) + trainer = dict(config.get("trainer", {})) + for name in ("device", "seed"): + value = getattr(args, name) + if value is not None: + trainer[name] = value + if args.output_dir is not None: + trainer["output_dir"] = str(args.output_dir.expanduser().resolve()) + config["trainer"] = trainer + if args.statistics is not None: + config["data"] = { + **config["data"], + "statistics": str(args.statistics.expanduser().resolve()), + } + result = train_from_config( + config, + resume=( + str(args.resume.expanduser().resolve()) if args.resume else None + ), + ) + print(json.dumps(result, indent=2, allow_nan=False)) + return + if args.device is not None: + if args.operation == "train": + config["train"] = {**config.get("train", {}), "device": args.device} + else: + config["device"] = args.device + if getattr(args, "seed", None) is not None: + if args.operation == "train": + config["train"] = {**config.get("train", {}), "seed": args.seed} + else: + config["seed"] = args.seed + if getattr(args, "output_dir", None) is not None: + config["output_dir"] = str(args.output_dir.expanduser().resolve()) + if getattr(args, "output", None) is not None: + config["output"] = str(args.output.expanduser().resolve()) + for field in ("resume", "fine_tune", "statistics", "sample_order", "model"): + value = getattr(args, field, None) + if value is not None: + config[field] = ( + str(value.expanduser().resolve()) + if isinstance(value, Path) + else value + ) + config = _resolve_paths(config, path.parent, args.operation) + if args.operation == "generate": + from examples.learning.nerd._support._generation import ( + generate_dataset, + ) + + result = generate_dataset(config) + elif args.operation == "train": + from examples.learning.nerd._support._workflow import run_train + + result = run_train(config) + else: + mode = config.pop("mode", "offline") + if mode == "offline": + from examples.learning.nerd._support._workflow import run_eval + + result = run_eval(config) + else: + raise ValueError("This workflow supports offline evaluation only") + except (OSError, ValueError, KeyError, TypeError) as error: + parser.exit(2, f"{error}\n") + print(json.dumps(result, indent=2, ensure_ascii=False, allow_nan=False)) + + +if __name__ == "__main__": + main() diff --git a/examples/learning/nerd/_resources.py b/examples/learning/nerd/_resources.py new file mode 100644 index 0000000000..2cd739b2e3 --- /dev/null +++ b/examples/learning/nerd/_resources.py @@ -0,0 +1,339 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Explicit, pinned resources for the NeRD showcase, without SDK registration.""" + +from __future__ import annotations + +import hashlib +import json +import os +from pathlib import Path, PurePosixPath +import shutil +import stat +import tempfile +from typing import Any +import zipfile + +__all__ = ["prepare_reference_assets", "prepare_resources"] + +_ASSET_REPO = "DexForceAI/embodichain_data" +_ASSET_REVISION = "97e202e41fe7fb504656a6b163291305e400780d" +_ASSET_PREFIX = "demo/nerd/reference_assets/v1" +_ASSET_ROOT = "nerd_reference_assets" +_UPSTREAM_REVISION = "476dc67fccc6a99be293bf611f897b1c06f13e30" +_MANIFEST_SHA256 = "10dd356b1c818111c31d4db8b5011fa9c0df1bfc55b01632ed992ecd9806711e" +_ARCHIVE_SHA256 = "008d90d395422be81f8bbd9707fc05fafd358b44794b29ff0f6158a75cd1da22" +_ASSET_FILES = { + "ATTRIBUTION.md": "4d86dc2651a4e24903162e8bac2a96122952d49d91c7ffb52a7fa28119438ca5", + "LICENSE.md": "59899c6091b540582ed617e8eeaac4919dc985ccfc35459ee9752b699be5205b", + "ant.xml": "23e9a720465a71fde3fd9fd2d5a18f8c11ec642f3dca8921c66cb4e169421fa7", + "cartpole_single.urdf": "f31276531c03050a3ac3a3e24443899f38a4424586db17140978383433dec237", +} +_DATASET_REPO = "DexForceAI/embodichain_datasets" +_DATASET_REVISION = "b9ec58f36da99cc99a93f97adfe74ebc9fd66586" +_MODEL_REPO = "DexForceAI/embodichain_model" +_MODEL_REVISION = "78e581526570069410bf269b38c0c52c800b24a7" +_PAYLOAD_MANIFEST_SHA256 = { + "dataset": { + "cartpole": "1cba6f65535a4e84c23cab4e633623ed0d5ef4f8b85e67781c5a4316c4f14775", + "ant": "5fe070411f3b61c87915ec61923d976b399059f0d93e071f119e58ed0e68127a", + }, + "model": { + "cartpole": "6e507149c9aacdc3f47665f91c3eb141cbcc0bd3bdcd697700a081d08077cf6a", + "ant": "52f0975aed5cd5946647f8ae60c52be307d51956dda4d7e9b9b3e2414cd23e9f", + }, +} + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _require_hash(path: Path, expected: str) -> None: + if not path.is_file() or _sha256(path) != expected: + raise ValueError(f"Resource SHA256 mismatch: {path}") + + +def _valid_assets(path: Path) -> bool: + if not path.is_dir() or path.is_symlink(): + return False + try: + if {entry.name for entry in path.iterdir()} != set(_ASSET_FILES): + return False + return all( + not (path / name).is_symlink() + and (path / name).is_file() + and _sha256(path / name) == digest + for name, digest in _ASSET_FILES.items() + ) + except OSError: + return False + + +def _asset_payloads(archive: Path, manifest_path: Path) -> dict[str, bytes]: + _require_hash(manifest_path, _MANIFEST_SHA256) + _require_hash(archive, _ARCHIVE_SHA256) + manifest = json.loads(manifest_path.read_text()) + if ( + manifest.get("schema_version") != 1 + or manifest.get("license") != "Apache-2.0" + or manifest["upstream"]["revision"] != _UPSTREAM_REVISION + or manifest["archive"]["root"] != _ASSET_ROOT + or manifest["archive"]["sha256"] != _ARCHIVE_SHA256 + or manifest["archive"]["size"] != archive.stat().st_size + or set(manifest["files"]) != set(_ASSET_FILES) + ): + raise ValueError("Unexpected NeRD reference asset manifest") + directory = f"{_ASSET_ROOT}/" + expected = {directory, *(f"{directory}{name}" for name in _ASSET_FILES)} + payloads = {} + with zipfile.ZipFile(archive) as bundle: + names = bundle.namelist() + if ( + not names + or names[0] != directory + or len(names) != len(set(names)) + or set(names) != expected + ): + raise ValueError("Unexpected NeRD archive member layout") + for entry in bundle.infolist(): + name = entry.filename + parts = name.rstrip("/").split("/") + mode = entry.external_attr >> 16 + if ( + PurePosixPath(name).is_absolute() + or "\\" in name + or ":" in name + or any(part in ("", ".", "..") for part in parts) + or stat.S_ISLNK(mode) + ): + raise ValueError("Unsafe NeRD archive member") + if entry.is_dir(): + if name != directory or entry.file_size: + raise ValueError("Unexpected NeRD archive directory") + continue + if stat.S_IFMT(mode) not in (0, stat.S_IFREG): + raise ValueError("Unsupported NeRD archive member type") + filename = PurePosixPath(name).name + record = manifest["files"][filename] + if ( + record["path"] != name + or record["sha256"] != _ASSET_FILES[filename] + or entry.file_size != record["size"] + or entry.file_size > 1024 * 1024 + ): + raise ValueError("Unexpected NeRD asset member identity") + value = bundle.read(entry) + if hashlib.sha256(value).hexdigest() != _ASSET_FILES[filename]: + raise ValueError(f"NeRD member SHA256 mismatch: {filename}") + payloads[filename] = value + return payloads + + +def _resource_root() -> Path: + """Resolve the showcase resource root without preparing any resources.""" + default = ( + Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")) + / "embodichain" + / "examples" + / "nerd" + ) + return ( + Path(os.environ.get("EMBODICHAIN_NERD_RESOURCE_ROOT", default)) + .expanduser() + .resolve() + ) + + +def prepare_reference_assets( + cache_dir: str | Path | None = None, *, local_files_only: bool = False +) -> Path: + """Download, verify and repair the four pinned reference asset files. + + Args: + cache_dir: Explicit asset cache root, taking precedence over the example's + ``EMBODICHAIN_NERD_RESOURCE_ROOT/assets``. Without either setting, use + the showcase cache under ``XDG_CACHE_HOME`` or ``~/.cache``. + local_files_only: Use verified extracted files or the existing Hub cache; + missing downloads fail without network access. + + Returns: + Directory containing the two robot files, license and attribution. + """ + from filelock import FileLock + + if cache_dir is not None: + base = Path(cache_dir).expanduser() + else: + base = _resource_root() / "assets" + base = base.resolve() + locks = base / ".locks" + locks.mkdir(parents=True, exist_ok=True) + with FileLock(str(locks / f"reference-assets-{_ASSET_REVISION}.lock")): + revision = base / "reference_assets" / _ASSET_REVISION + if revision.is_symlink() or not revision.resolve().is_relative_to(base): + raise ValueError("Reference asset revision must remain inside its cache") + target = revision / _ASSET_ROOT + if _valid_assets(target): + return target + from huggingface_hub import hf_hub_download + + options = { + "repo_id": _ASSET_REPO, + "repo_type": "dataset", + "revision": _ASSET_REVISION, + # Mirrors may lag this small bundle or omit required Hub metadata. + "endpoint": "https://huggingface.co", + "cache_dir": base / "hub", + "local_files_only": local_files_only, + } + manifest = Path( + hf_hub_download(filename=f"{_ASSET_PREFIX}/manifest.json", **options) + ) + archive = Path( + hf_hub_download(filename=f"{_ASSET_PREFIX}/reference_assets.zip", **options) + ) + payloads = _asset_payloads(archive, manifest) + revision.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(prefix=".staging-", dir=revision) as temporary: + staged = Path(temporary) / _ASSET_ROOT + staged.mkdir() + for name, value in payloads.items(): + (staged / name).write_bytes(value) + if not _valid_assets(staged): + raise ValueError("Extracted NeRD reference assets failed verification") + if target.is_symlink() or (target.exists() and not target.is_dir()): + target.unlink() + elif target.exists(): + shutil.rmtree(target) + staged.replace(target) + return target + + +def _resource_file(root: Path, relative: str) -> Path: + if ( + not isinstance(relative, str) + or "\\" in relative + or ":" in relative + or PurePosixPath(relative).is_absolute() + or any(part in ("", ".", "..") for part in relative.split("/")) + ): + raise ValueError("Unsafe published resource filename") + target = root.joinpath(*PurePosixPath(relative).parts) + if not target.resolve().is_relative_to(root.resolve()): + raise ValueError("Published resource resolves outside its directory") + return target + + +def _verify_payload(path: Path, robot: str, kind: str) -> dict[str, Any]: + _require_hash(path / "manifest.json", _PAYLOAD_MANIFEST_SHA256[kind][robot]) + manifest = json.loads((path / "manifest.json").read_text()) + if kind == "dataset": + if ( + manifest.get("schema_version") != 1 + or manifest["metadata"]["codec"]["robot"] != robot + ): + raise ValueError("Unexpected NeRD dataset manifest") + shards = manifest["shards"] + if not shards or len({row["file"] for row in shards}) != len(shards): + raise ValueError("NeRD dataset shards must be nonempty and unique") + for row in shards: + _require_hash(_resource_file(path, row["file"]), row["sha256"]) + else: + if manifest.get("format_version") != 1 or manifest["codec"]["robot"] != robot: + raise ValueError("Unexpected NeRD model manifest") + _require_hash(path / "weights.pt", manifest["weights_sha256"]) + return manifest + + +def _prepare_payload(root: Path, robot: str, kind: str) -> tuple[Path, dict[str, Any]]: + prefix = f"neural_dynamics/nerd/{robot}/v1" + path = root / prefix + return path, _verify_payload(path, robot, kind) + + +def prepare_resources( + root: Path, + robot: str, + *, + assets_only: bool = False, + local_files_only: bool = False, +) -> dict[str, Path]: + """Prepare pinned showcase assets and optional published data/model bundles. + + Args: + root: Caller-owned resource root for assets, datasets and models. + robot: ``cartpole`` or ``ant``. + assets_only: Prepare assets without starting model or dataset downloads. + local_files_only: Forbid remote access when a required cache is missing. + + Returns: + Verified ``assets`` directory; also ``dataset`` and ``model`` directories + when ``assets_only`` is false. + """ + robot = robot.lower() + if robot not in ("cartpole", "ant"): + raise ValueError("NeRD resources support cartpole or ant") + root = Path(root).expanduser().resolve() + resources = { + "assets": prepare_reference_assets( + root / "assets", local_files_only=local_files_only + ) + } + if assets_only: + return resources + manifests = {} + for kind, directory, repo, revision in ( + ("dataset", "datasets", _DATASET_REPO, _DATASET_REVISION), + ("model", "models", _MODEL_REPO, _MODEL_REVISION), + ): + base = root / directory + try: + path, manifest = _prepare_payload(base, robot, kind) + except (OSError, ValueError, KeyError, TypeError): + from huggingface_hub import snapshot_download + + snapshot_download( + repo_id=repo, + repo_type="dataset" if kind == "dataset" else "model", + revision=revision, + allow_patterns=[f"neural_dynamics/nerd/{robot}/v1/*"], + local_dir=base, + local_files_only=local_files_only, + max_workers=4, + ) + path, manifest = _prepare_payload(base, robot, kind) + resources[kind], manifests[kind] = path, manifest + metadata = manifests["dataset"]["metadata"] + for field in ( + "codec", + "dt", + "asset_sha256", + "physics_signature", + "gravity", + "joint_names", + ): + if manifests["model"].get(field) != metadata.get(field): + raise ValueError(f"Published NeRD model/dataset {field} mismatch") + filename = "cartpole_single.urdf" if robot == "cartpole" else "ant.xml" + if metadata.get("asset_sha256") != _ASSET_FILES[filename]: + raise ValueError("Published NeRD resource/reference asset mismatch") + return resources diff --git a/examples/learning/nerd/_support/__init__.py b/examples/learning/nerd/_support/__init__.py new file mode 100644 index 0000000000..b1f8e18a59 --- /dev/null +++ b/examples/learning/nerd/_support/__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. +# ---------------------------------------------------------------------------- + +"""Simulation-independent neural dynamics workflows.""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/examples/learning/nerd/_support/_acceptance.py b/examples/learning/nerd/_support/_acceptance.py new file mode 100644 index 0000000000..db2cae354c --- /dev/null +++ b/examples/learning/nerd/_support/_acceptance.py @@ -0,0 +1,585 @@ +# ---------------------------------------------------------------------------- +# 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 independent NeRD trainers on immutable published data. + +Run: + python -m embodichain run-example learning/nerd/_support/_acceptance \ + --reference-source /path/to/pr855 --datasets /path/to/nerd-datasets \ + --models /path/to/nerd-models --device cuda:0 --output /path/to/results + +Dataset and model roots contain ``cartpole/v1`` and ``ant/v1``. No data is +generated, no ODS engine is started, and no physics rollout result is inferred +from these offline prediction and paired-training checks. +""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib +import json +import os +from pathlib import Path +import resource +import subprocess +import sys +import time +from typing import Any + + +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") + + +def _sha(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _source_modules(source: Path) -> dict[str, str]: + """Select only the requested checkout's original or extension module layout.""" + if ( + source + / "extensions/embodichain_models/embodichain_models/dynamics/nerd/model.py" + ).is_file(): + return { + "codec": "embodichain_models.dynamics.nerd.codec", + "data": "embodichain_models.dynamics.nerd.data", + "model": "embodichain_models.dynamics.nerd", + "trainer": "examples.learning.nerd._support._trainer", + "qualification": "examples.learning.nerd._support._qualification", + } + return { + "codec": "embodichain.compute.dynamics", + "data": "embodichain.data_pipeline.dynamics.shards", + "model": "embodichain.learning.models.nerd", + "trainer": "embodichain.learning.dynamics.trainer", + "qualification": "embodichain.learning.dynamics.qualification", + } + + +def _source_paths(source: Path) -> list[str]: + """Keep each independent worker's imports bound to its requested source.""" + paths = [str(source)] + if _source_modules(source)["model"] == "embodichain_models.dynamics.nerd": + paths.insert(0, str(source / "extensions/embodichain_models")) + return paths + + +def _source_identity(source: Path) -> dict[str, Any]: + def git(*args: str) -> bytes: + return subprocess.check_output(["git", "-C", str(source), *args]) + + if _source_modules(source)["model"] == "embodichain_models.dynamics.nerd": + modules = ( + "examples/learning/nerd/_support/_trainer.py", + *( + "extensions/embodichain_models/embodichain_models/dynamics/nerd/" + name + for name in ( + "model.py", + "checkpoint.py", + "codec.py", + "data.py", + "_resident.py", + "_numerics.py", + "integration.py", + "supervised.py", + "sl_recipe.py", + ) + ), + "embodichain/learning/sl/learner.py", + "embodichain/learning/sl/types.py", + ) + else: + modules = ( + "embodichain/learning/dynamics/trainer.py", + "embodichain/learning/dynamics/supervised.py", + "embodichain/learning/sl/learner.py", + "embodichain/learning/sl/types.py", + ) + return { + "commit": git("rev-parse", "HEAD").decode().strip(), + "tracked_diff_sha256": hashlib.sha256(git("diff", "HEAD")).hexdigest(), + "modules": { + name: _sha(source / name) for name in modules if (source / name).is_file() + }, + } + + +def _worker(request_path: Path) -> None: + request = json.loads(request_path.read_text()) + if "runner_sha256" in request and _sha(Path(__file__)) != request["runner_sha256"]: + raise RuntimeError("Acceptance runner changed after protocol freeze") + if ( + "source_identity" in request + and _source_identity(Path(request["source"])) != request["source_identity"] + ): + raise RuntimeError("Trainer source changed after protocol freeze") + source = Path(request["source"]) + sys.path[:0] = _source_paths(source) + import torch + + modules = _source_modules(source) + codec_module = importlib.import_module(modules["codec"]) + data_module = importlib.import_module(modules["data"]) + trainer_module = importlib.import_module(modules["trainer"]) + model_module = importlib.import_module(modules["model"]) + NeRDStateCodec = codec_module.NeRDStateCodec + DynamicsWindowDataset = data_module.DynamicsWindowDataset + dataset_fingerprint = data_module.dataset_fingerprint + load_manifest = data_module.load_manifest + DynamicsTrainCfg = trainer_module.DynamicsTrainCfg + DynamicsTrainer = trainer_module.DynamicsTrainer + evaluate_model = trainer_module.evaluate_model + NeRDModel, load_model = model_module.NeRDModel, model_module.load_model + + torch.set_num_threads(request["threads"]) + started = time.perf_counter() + device = request["device"] + cuda = torch.device(device).type == "cuda" + if cuda: + torch.cuda.set_device(device) + torch.cuda.reset_peak_memory_stats(device) + output = Path(request["output"]) + output.mkdir(parents=True, exist_ok=True) + pretrained, _ = load_model(request["model"]) + manifest = load_manifest(request["dataset"]) + codec = NeRDStateCodec(**manifest["metadata"]["codec"]) + result: dict[str, Any] = {"mode": request["mode"]} + train = validation = test = None + try: + if request["mode"] == "offline": + test = DynamicsWindowDataset( + request["dataset"], "test", pretrained.cfg.history_length + ) + pretrained.to(device) + result["metrics"] = evaluate_model( + pretrained, codec, test, request["batch_size"], device + ) + else: + seed = request["seed"] + torch.manual_seed(seed) + model = NeRDModel(pretrained.cfg) + statistics = { + "schema_version": 1, + "weighting": "shared_published_training_statistics", + "dataset_fingerprint": dataset_fingerprint(manifest), + "inputs": { + name: {"mean": rms.mean.clone(), "var": rms.var.clone()} + for name, rms in pretrained.input_rms.items() + }, + "output": { + "mean": pretrained.output_rms.mean.clone(), + "var": pretrained.output_rms.var.clone(), + }, + } + train = DynamicsWindowDataset( + request["dataset"], "train", model.cfg.history_length + ) + validation = DynamicsWindowDataset( + request["dataset"], "validation", model.cfg.history_length + ) + cfg = DynamicsTrainCfg( + total_steps=4_995_000, + batch_size=request["batch_size"], + learning_rate=1e-3, + final_learning_rate=1e-4, + max_grad_norm=1.0, + seed=seed, + device=device, + validation_interval=5000, + validation_batches=100, + checkpoint_interval=5000, + lr_update_interval=5000, + lr_offset=1, + log_interval=100, + resident_cache_bytes=request["resident_cache_bytes"], + numerical_profile="reference", + ) + trainer = DynamicsTrainer( + model, + codec, + train, + validation, + cfg, + output, + manifest=manifest["metadata"], + statistics=statistics, + initialize_normalization=not request.get("resume"), + ) + if request.get("resume"): + trainer.resume(output / "last.pt") + last_update: dict[str, Any] = {} + if trainer.global_step == 0: + captured: dict[str, Any] = {"raw_gradients": {}} + handles = [ + parameter.register_hook( + lambda gradient, name=name: captured["raw_gradients"].update( + {name: gradient.detach().cpu().clone()} + ) + ) + for name, parameter in model.named_parameters() + ] + handles.append( + model.register_forward_hook( + lambda module, inputs, prediction: captured.update( + prediction=prediction.detach().cpu().clone() + ) + ) + ) + first = trainer.step() + last_update = first + for handle in handles: + handle.remove() + captured.update( + loss=torch.tensor(first["loss"]), + grad_norm=torch.tensor(first["grad_norm"]), + clipped_gradients={ + name: parameter.grad.detach().cpu().clone() + for name, parameter in model.named_parameters() + if parameter.grad is not None + }, + updated_parameters={ + name: parameter.detach().cpu().clone() + for name, parameter in model.named_parameters() + }, + ) + torch.save(captured, output / "first_update.pt") + result["training"] = trainer.train(stop_after=request["stop_after"]) + last_update = result["training"]["last"] or last_update + # Diagnostic only: read already clipped gradients after the final + # update, without another forward, backward, hook or optimizer call. + torch.save( + { + "num_updates": trainer.global_step, + "pre_clip_grad_norm": last_update.get("grad_norm"), + "clipped_gradients": { + name: parameter.grad.detach().cpu().clone() + for name, parameter in model.named_parameters() + if parameter.grad is not None + }, + }, + output / "final_gradients.pt", + ) + result["metrics"] = evaluate_model( + model, codec, validation, request["batch_size"], device + ) + if cuda: + torch.cuda.synchronize(device) + result.update( + elapsed_seconds=time.perf_counter() - started, + process_peak_rss_mib=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + / 1024, + torch_peak_gpu_mib=( + torch.cuda.max_memory_allocated(device) / 1024**2 if cuda else None + ), + source_identity=_source_identity(Path(request["source"])), + ) + _write(output / "result.json", result) + finally: + for dataset in (train, validation, test): + if dataset is not None: + dataset.close() + + +def _run(request: dict[str, Any]) -> dict[str, Any]: + output = Path(request["output"]) + output.mkdir(parents=True, exist_ok=True) + suffix = f"{request.get('stop_after', 'offline')}-{bool(request.get('resume'))}" + request_path = output / f"request-{suffix}.json" + _write(request_path, request) + environment = dict(os.environ) + environment.update( + PYTHONPATH=os.pathsep.join(_source_paths(Path(request["source"]))), + OMP_NUM_THREADS=str(request["threads"]), + MKL_NUM_THREADS=str(request["threads"]), + OPENBLAS_NUM_THREADS=str(request["threads"]), + ) + with (output / f"worker-{suffix}.log").open("w") as log: + subprocess.run( + [ + sys.executable, + str(Path(__file__).resolve()), + "--worker", + str(request_path), + ], + cwd=request["source"], + env=environment, + stdout=log, + stderr=subprocess.STDOUT, + check=True, + ) + return json.loads((output / "result.json").read_text()) + + +def _state_equal(first: Any, second: Any) -> bool: + import torch + + if isinstance(first, torch.Tensor): + return isinstance(second, torch.Tensor) and torch.equal(first, second) + if isinstance(first, dict): + return first.keys() == second.keys() and all( + _state_equal(value, second[key]) for key, value in first.items() + ) + if isinstance(first, (list, tuple)): + return len(first) == len(second) and all( + _state_equal(a, b) for a, b in zip(first, second) + ) + return first == second + + +def _parse_seeds(value: str) -> list[int]: + try: + seeds = [int(item) for item in value.split(",")] + except ValueError as error: + raise argparse.ArgumentTypeError( + "seeds must be comma-separated integers" + ) from error + if len(set(seeds)) != len(seeds): + raise argparse.ArgumentTypeError("seeds must be distinct") + return seeds + + +def main() -> None: + """Run the full offline and five-seed matched-update acceptance matrix.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--worker", type=Path) + parser.add_argument("--reference-source", type=Path) + parser.add_argument( + "--candidate-source", type=Path, default=Path(__file__).resolve().parents[4] + ) + parser.add_argument("--datasets", type=Path) + parser.add_argument("--models", type=Path) + parser.add_argument("--output", type=Path) + parser.add_argument("--device", default="cpu") + parser.add_argument("--batch-size", type=int, default=512) + parser.add_argument("--steps", type=int, default=1000) + parser.add_argument("--seeds", type=_parse_seeds, default="0,1,2,3,4") + parser.add_argument("--robots", nargs="+", default=["cartpole", "ant"]) + parser.add_argument("--threads", type=int, default=4) + parser.add_argument("--resident-cache-bytes", type=int, default=4 * 1024**3) + args = parser.parse_args() + if args.worker: + _worker(args.worker) + return + if any( + value is None + for value in (args.reference_source, args.datasets, args.models, args.output) + ): + parser.error("reference-source, datasets, models and output are required") + if args.steps < 2 or args.batch_size < 1: + parser.error("steps must be >=2 and batch-size positive") + if args.output.exists() and any(args.output.iterdir()): + parser.error( + "output must be a fresh directory; completed results are immutable" + ) + import torch + + candidate = args.candidate_source.resolve() + sys.path[:0] = _source_paths(candidate) + modules = _source_modules(candidate) + load_manifest = importlib.import_module(modules["data"]).load_manifest + qualification = importlib.import_module(modules["qualification"]) + compare_training_steps = qualification.compare_training_steps + paired_noninferiority = qualification.paired_noninferiority + + seeds = args.seeds + sources = { + "reference": str(args.reference_source.resolve()), + "candidate": str(args.candidate_source.resolve()), + } + source_identities = { + name: _source_identity(Path(source)) for name, source in sources.items() + } + rows, checks = [], {} + for robot in args.robots: + dataset, model = args.datasets / robot / "v1", args.models / robot / "v1" + load_manifest(dataset, verify=True) + base = { + "dataset": str(dataset.resolve()), + "model": str(model.resolve()), + "device": args.device, + "batch_size": args.batch_size, + "threads": args.threads, + "resident_cache_bytes": args.resident_cache_bytes, + "runner_sha256": _sha(Path(__file__)), + } + _write( + args.output / robot / "artifacts.json", + { + "dataset_manifest_sha256": _sha(dataset / "manifest.json"), + "model_manifest_sha256": _sha(model / "manifest.json"), + "model_weights_sha256": _sha(model / "weights.pt"), + }, + ) + offline = {} + for variant, source in sources.items(): + request = dict( + base, + source=source, + source_identity=source_identities[variant], + mode="offline", + output=str((args.output / robot / variant / "offline").resolve()), + ) + offline[variant] = _run(request) + rows.append( + dict(robot=robot, variant=variant, seed="published", **offline[variant]) + ) + metric_names = [ + key + for key, value in offline["reference"]["metrics"].items() + if isinstance(value, (int, float)) + ] + checks[f"{robot}/published_offline"] = all( + abs( + offline["reference"]["metrics"][key] + - offline["candidate"]["metrics"][key] + ) + <= 1e-5 + 1e-4 * abs(offline["reference"]["metrics"][key]) + for key in metric_names + ) + errors = {variant: [] for variant in sources} + for seed in seeds: + paths = {} + for variant, source in sources.items(): + output = (args.output / robot / variant / f"seed-{seed}").resolve() + paths[variant] = output + report = _run( + dict( + base, + source=source, + source_identity=source_identities[variant], + mode="train", + seed=seed, + stop_after=args.steps, + output=str(output), + ) + ) + rows.append(dict(robot=robot, variant=variant, seed=seed, **report)) + errors[variant].append(report["metrics"]["normalized_mse"]) + captures = [ + torch.load(paths[variant] / "first_update.pt", weights_only=True) + for variant in sources + ] + parity = compare_training_steps(*captures) + _write(args.output / robot / f"seed-{seed}-first-update.json", parity) + checks[f"{robot}/seed-{seed}/first_update"] = parity["passed"] + states = [ + torch.load( + paths[variant] / "last.pt", weights_only=True, map_location="cpu" + ) + for variant in sources + ] + checks[f"{robot}/seed-{seed}/final_parameters_exact"] = _state_equal( + states[0]["model"], states[1]["model"] + ) + if seed == seeds[0]: + resumed = ( + args.output / robot / "candidate" / f"seed-{seed}-resumed" + ).resolve() + request = dict( + base, + source=sources["candidate"], + source_identity=source_identities["candidate"], + mode="train", + seed=seed, + stop_after=args.steps // 2, + output=str(resumed), + ) + _run(request) + _run(dict(request, stop_after=args.steps, resume=True)) + state = torch.load( + resumed / "last.pt", weights_only=True, map_location="cpu" + ) + checks[f"{robot}/resume_exact"] = all( + _state_equal(state[key], states[1][key]) + for key in ( + "model", + "optimizer", + "global_step", + "sampler_state", + "rng", + ) + ) + gate = paired_noninferiority(errors["reference"], errors["candidate"]) + _write(args.output / robot / "paired-validation.json", gate) + checks[f"{robot}/paired_noninferiority"] = gate["status"] == "passed" + passed = all(checks.values()) + _write( + args.output / "summary.json", {"passed": passed, "checks": checks, "rows": rows} + ) + lines = [ + "# NeRD SL acceptance", + "", + f"Passed: {passed}", + "", + "## Time & Memory", + "", + "| Robot | Trainer | Seed | Seconds | Process peak RSS MiB | Torch peak GPU MiB |", + "| --- | --- | --- | --- | --- | --- |", + ] + lines += [ + f"| {row['robot']} | {row['variant']} | {row['seed']} | {row['elapsed_seconds']:.3f} | {row['process_peak_rss_mib']:.1f} | {row['torch_peak_gpu_mib']} |" + for row in rows + ] + lines += [ + "", + "## Success & Other Metrics", + "", + "| Check | Passed |", + "| --- | --- |", + ] + lines += [f"| {name} | {value} |" for name, value in checks.items()] + lines += [ + "", + "## Leaderboard", + "", + "| Robot | Trainer | Mean validation normalized MSE |", + "| --- | --- | --- |", + ] + for robot in args.robots: + ranked = [] + for variant in sources: + values = [ + row["metrics"]["normalized_mse"] + for row in rows + if row["robot"] == robot + and row["variant"] == variant + and row["seed"] != "published" + ] + ranked.append((sum(values) / len(values), variant)) + lines += [ + f"| {robot} | {variant} | {value:.9g} |" + for value, variant in sorted(ranked) + ] + lines += [ + "", + "Peak RSS is the worker-process high water mark; GPU memory covers Torch only. Setup, training and validation are included in elapsed time. Published normalization is shared; these fits do not claim full-budget convergence or physics rollout reproduction.", + "", + ] + (args.output / "report.md").write_text("\n".join(lines)) + print(args.output / "summary.json", flush=True) + if not passed: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/examples/learning/nerd/_support/_benchmark_dynamics_input.py b/examples/learning/nerd/_support/_benchmark_dynamics_input.py new file mode 100644 index 0000000000..3bfc12c962 --- /dev/null +++ b/examples/learning/nerd/_support/_benchmark_dynamics_input.py @@ -0,0 +1,320 @@ +# ---------------------------------------------------------------------------- +# 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 individual and grouped dynamics window I/O on real and stress data. + +Run: python -m embodichain run-example learning/nerd/_support/_benchmark_dynamics_input --dataset DATASET + --stress-dir /tmp/nerd-input-stress --output /tmp/input-benchmark.json + +CPU read, collation and encoding are timed together; index generation, source +verification and correctness checks are excluded. Stress data repeats an actual +trajectory only to exercise many-shard I/O; it is not a physics accuracy dataset. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +import platform +import statistics +import time +import tracemalloc +from typing import Any + +__all__ = ["benchmark_inputs", "create_stress_dataset"] + + +def _assert_equal(left: Any, right: Any) -> None: + import torch + + if isinstance(left, dict): + assert left.keys() == right.keys() + for key in left: + _assert_equal(left[key], right[key]) + elif isinstance(left, torch.Tensor): + assert left.dtype == right.dtype and torch.equal(left, right) + else: + assert left == right + + +def _tensor_bytes(value: Any) -> int: + import torch + + if isinstance(value, dict): + return sum(_tensor_bytes(item) for item in value.values()) + return ( + value.numel() * value.element_size() if isinstance(value, torch.Tensor) else 0 + ) + + +def create_stress_dataset( + source: str | Path, + output: str | Path, + *, + trajectories: int = 512, + trajectories_per_shard: int = 4, +) -> Path: + """Repeat one real trajectory into independent shards for bounded-I/O stress.""" + import h5py + from embodichain_models.dynamics.nerd.data import ( + load_manifest, + DynamicsShardWriter, + dataset_fingerprint, + ) + + if min(trajectories, trajectories_per_shard) < 1: + raise ValueError("Stress sizes must be positive") + source = Path(source) + output = Path(output) + manifest = load_manifest(source, verify=True) + first_shard = manifest["shards"][0] + first = first_shard["trajectories"][0] + with h5py.File(source / first_shard["file"], "r") as file: + root = file[first["group"]] + arrays = { + kind: {name: value[()] for name, value in root[kind].items()} + for kind in ("states", "transitions") + } + with DynamicsShardWriter( + output, + { + **manifest["metadata"], + "purpose": "repeated_real_trajectory_io_stress_only", + "source_fingerprint": dataset_fingerprint(manifest), + }, + trajectories_per_shard=trajectories_per_shard, + ) as writer: + for index in range(trajectories): + writer.add_trajectory( + f"stress-{index}", + arrays["states"], + arrays["transitions"], + split="train", + ) + return output + + +def benchmark_inputs( + root: str | Path, + *, + batch_sizes: tuple[int, ...] = (1, 16, 512), + history: int = 10, + warmup: int = 2, + repeats: int = 10, + max_open_shards: int = 4, + read_page_steps: int = 256, + unique_trajectories: bool = False, +) -> dict[str, Any]: + """Measure identical batches using two readers with identical cache limits.""" + import h5py + import numpy as np + import torch + from torch.utils.data import default_collate + from embodichain_models.dynamics.nerd.codec import NeRDStateCodec + from embodichain_models.dynamics.nerd.data import DynamicsWindowDataset + from examples.learning.nerd._support._trainer import encode_batch + + if ( + min(*batch_sizes, history, warmup, repeats, max_open_shards, read_page_steps) + < 1 + ): + raise ValueError("Benchmark workload sizes must be positive") + torch.set_num_threads(1) + readers = { + name: DynamicsWindowDataset( + root, + "train", + history, + max_open_shards=max_open_shards, + read_page_steps=read_page_steps, + verify=True, + ) + for name in ("individual", "grouped") + } + codec = NeRDStateCodec(**readers["individual"].manifest["metadata"]["codec"]) + reference = readers["individual"] + records = [] + try: + for batch_size in batch_sizes: + generator = torch.Generator().manual_seed(42) + selections = [] + for _ in range(warmup + repeats): + if unique_trajectories: + if batch_size > len(reference._rows): + raise ValueError( + "Stress needs at least one distinct trajectory per requested sample" + ) + rows = torch.randperm(len(reference._rows), generator=generator)[ + :batch_size + ].tolist() + indices = [] + for row in rows: + start = reference._ends[row - 1] if row else 0 + width = reference._ends[row] - start + indices.append( + start + int(torch.randint(width, (), generator=generator)) + ) + else: + indices = torch.randint( + len(reference), (batch_size,), generator=generator + ).tolist() + selections.append(indices) + expected = default_collate([reference[index] for index in selections[0]]) + actual = readers["grouped"].get_batch(selections[0]) + _assert_equal(actual, expected) + # Compare actual encoder results explicitly; no tolerance changes. + encoded_old = encode_batch(expected, codec, "cpu") + encoded_new = encode_batch(actual, codec, "cpu") + for old, new in zip(encoded_old, encoded_new): + _assert_equal(old, new) + del actual, expected, encoded_old, encoded_new + for name, dataset in readers.items(): + + def load(indices: list[int]) -> dict[str, Any]: + return ( + dataset.get_batch(indices) + if name == "grouped" + else default_collate([dataset[index] for index in indices]) + ) + + samples = [] + for iteration, indices in enumerate(selections): + start = time.perf_counter() + batch = load(indices) + encoded = encode_batch(batch, codec, "cpu") + elapsed = time.perf_counter() - start + if iteration >= warmup: + samples.append(elapsed) + del batch, encoded + # Separate memory/I/O instrumentation from latency samples. + dataset.close() + original_handle = dataset._handle + opens = 0 + + def observed_handle(filename: str) -> Any: + nonlocal opens + if filename not in dataset._handles: + opens += 1 + return original_handle(filename) + + dataset._handle = observed_handle + tracemalloc.start() + batch = load(selections[-1]) + peak = tracemalloc.get_traced_memory()[1] + tracemalloc.stop() + output_bytes = _tensor_bytes(batch) + del batch + dataset._handle = original_handle + records.append( + { + "implementation": name, + "batch_size": batch_size, + "history": history, + "median_seconds": statistics.median(samples), + "min_seconds": min(samples), + "max_seconds": max(samples), + "samples_seconds": samples, + "python_numpy_peak_bytes_cold_handles": peak, + "returned_tensor_bytes": output_bytes, + "cold_batch_file_opens": opens, + "retained_file_handles": len(dataset._handles), + "exact_output_and_encoding_equal": True, + } + ) + return { + "dataset": str(Path(root).resolve()), + "fingerprint": reference.fingerprint, + "robot": codec.robot, + "scope": "CPU read+collate+encode; excludes indices,verification,training", + "unique_trajectories": unique_trajectories, + "trajectory_count": len(reference._rows), + "shard_count": len(reference.manifest["shards"]), + "max_open_shards": max_open_shards, + "read_page_steps": read_page_steps, + "max_source_frames_per_field_read": read_page_steps + history, + "cache_policy": "bounded file LRU only; no retained trajectory/data blocks", + "memory_scope": "tracemalloc Python/NumPy allocations, excludes native HDF5 caches/Torch allocator", + "dtype": "preserved source dtype; codec float32", + "cpu_threads": 1, + "python": platform.python_version(), + "torch": str(torch.__version__), + "h5py": h5py.__version__, + "numpy": np.__version__, + "warmup": warmup, + "repeats": repeats, + "runs": records, + } + finally: + for dataset in readers.values(): + dataset.close() + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dataset", type=Path, action="append", required=True) + parser.add_argument("--batch-sizes", type=int, nargs="+", default=[1, 16, 512]) + parser.add_argument("--history", type=int, default=10) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--repeats", type=int, default=10) + parser.add_argument("--stress-dir", type=Path) + parser.add_argument("--stress-trajectories", type=int, default=512) + parser.add_argument("--stress-shard-size", type=int, default=4) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + results = [ + benchmark_inputs( + path, + batch_sizes=tuple(args.batch_sizes), + history=args.history, + warmup=args.warmup, + repeats=args.repeats, + ) + for path in args.dataset + ] + if args.stress_dir: + stress = create_stress_dataset( + args.dataset[-1], + args.stress_dir, + trajectories=args.stress_trajectories, + trajectories_per_shard=args.stress_shard_size, + ) + results.append( + benchmark_inputs( + stress, + batch_sizes=tuple(args.batch_sizes), + history=args.history, + warmup=args.warmup, + repeats=args.repeats, + unique_trajectories=True, + ) + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(results, indent=2, allow_nan=False)) + for result in results: + print(result["dataset"], "unique", result["unique_trajectories"]) + for record in result["runs"]: + print( + { + key: value + for key, value in record.items() + if key != "samples_seconds" + } + ) + + +if __name__ == "__main__": + main() diff --git a/examples/learning/nerd/_support/_benchmark_dynamics_sampling.py b/examples/learning/nerd/_support/_benchmark_dynamics_sampling.py new file mode 100644 index 0000000000..aa3c8837c8 --- /dev/null +++ b/examples/learning/nerd/_support/_benchmark_dynamics_sampling.py @@ -0,0 +1,193 @@ +# ---------------------------------------------------------------------------- +# 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 exact minibatch-order storage and replay with bounded RNG checkpoints. + +Run: python -m embodichain run-example learning/nerd/_support/_benchmark_dynamics_sampling --output-dir PATH +""" + +from __future__ import annotations + +import argparse +import gc +import json +from pathlib import Path +import platform +import statistics +import time + +import numpy as np +import psutil +import torch + +from examples.learning.nerd._support._sample_order import ( + _open_order, + _write_checkpointed_order, +) + +__all__: list[str] = [] + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--steps", type=int, default=4097) + parser.add_argument("--full-budget", action="store_true") + args = parser.parse_args() + root = args.output_dir + root.mkdir(parents=True, exist_ok=False) + torch.set_num_threads(1) + rows = [] + for batch in [1, 16, 512]: + generator = torch.Generator().manual_seed(7) + matrix = np.lib.format.open_memmap( + root / f"matrix-{batch}.npy", + mode="w+", + dtype=np.int64, + shape=(args.steps, batch), + ) + for start in range(0, args.steps, 1024): + matrix[start : start + 1024] = torch.randint( + 1000000, (min(1024, args.steps - start), batch), generator=generator + ).numpy() + matrix.flush() + del matrix + start = time.perf_counter() + descriptor, state = _write_checkpointed_order( + root / f"order-{batch}.json", + dataset_size=1000000, + steps=args.steps, + batch_size=batch, + seed=7, + ) + prepare_ms = (time.perf_counter() - start) * 1000 + expected = np.load(root / f"matrix-{batch}.npy", mmap_mode="r") + stream = _open_order(descriptor) + for index in range(args.steps): + np.testing.assert_array_equal(stream[index], expected[index]) + del stream, expected + gc.collect() + for kind, path in [ + ("matrix", root / f"matrix-{batch}.npy"), + ("checkpoint", descriptor), + ]: + samples = [] + deltas = [] + # Completed CPU replay windows; setup/verification outside timing. + for repeat in range(4): + order = _open_order(path) + rss = psutil.Process().memory_info().rss + start = time.perf_counter() + for index in range(args.steps): + value = order[index] + elapsed = (time.perf_counter() - start) * 1000 / args.steps + if repeat: + samples.append(elapsed) + deltas.append((psutil.Process().memory_info().rss - rss) / 1024**2) + del order, value + gc.collect() + size = path.stat().st_size + ( + state.stat().st_size if kind == "checkpoint" else 0 + ) + rows.append( + { + "implementation": kind, + "batch_size": batch, + "cost_time_ms": statistics.median(samples), + "cpu_delta_mb": statistics.median(deltas), + "gpu_delta_mb": 0, + "peak_gpu_mb": 0, + "storage_bytes": size, + "success_rate": 1.0, + "prepare_ms": prepare_ms if kind == "checkpoint" else None, + } + ) + full = None + if args.full_budget: + start = time.perf_counter() + descriptor, state = _write_checkpointed_order( + root / "full-order.json", + dataset_size=1000000, + steps=4995000, + batch_size=512, + seed=7, + ) + full = { + "steps": 4995000, + "batch": 512, + "actual_storage_bytes": descriptor.stat().st_size + state.stat().st_size, + "dense_index_bytes": 4995000 * 512 * 8, + "prepare_seconds": time.perf_counter() - start, + } + result = { + "torch": str(torch.__version__), + "python": platform.python_version(), + "cpu": platform.processor(), + "threads": 1, + "steps": args.steps, + "warmup_windows": 1, + "repeats": 3, + "scope": "CPU exact order replay; RSS deltas are not peak memory; GPU unused.", + "rows": rows, + "full_budget": full, + } + (root / "results.json").write_text(json.dumps(result, indent=2)) + lines = [ + "# Sampling order benchmark", + "", + result["scope"], + "", + "## Time & Memory", + "", + "| Algorithm | Batch | cost_time_ms | cpu_delta_mb | gpu_delta_mb | peak_gpu_mb |", + "|---|---:|---:|---:|---:|---:|", + ] + for row in rows: + lines.append( + f"| {row['implementation']} | {row['batch_size']} | {row['cost_time_ms']:.6f} | {row['cpu_delta_mb']:.3f} | 0 | 0 |" + ) + lines += [ + "", + "## Success & Other Metrics", + "", + "| Algorithm | Batch | success_rate | storage_bytes |", + "|---|---:|---:|---:|", + ] + for row in rows: + lines.append( + f"| {row['implementation']} | {row['batch_size']} | 1.0 | {row['storage_bytes']} |" + ) + lines += [ + "", + "## Leaderboard", + "", + "| Algorithm | Batch | success_rate |", + "|---|---:|---:|", + ] + for row in sorted(rows, key=lambda x: -x["success_rate"]): + lines.append(f"| {row['implementation']} | {row['batch_size']} | 1.0 |") + if full: + lines += [ + "", + f"Measured full-budget checkpoint storage: {full['actual_storage_bytes']} bytes versus {full['dense_index_bytes']} dense index bytes; setup {full['prepare_seconds']:.3f}s.", + ] + report = root / "REPORT.md" + report.write_text("\n".join(lines) + "\n") + print(report) + + +if __name__ == "__main__": + main() diff --git a/examples/learning/nerd/_support/_benchmark_dynamics_training.py b/examples/learning/nerd/_support/_benchmark_dynamics_training.py new file mode 100644 index 0000000000..197ac5dd82 --- /dev/null +++ b/examples/learning/nerd/_support/_benchmark_dynamics_training.py @@ -0,0 +1,309 @@ +# ---------------------------------------------------------------------------- +# 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 official-width resident NeRD optimizer updates and budget lower bounds. + +Run: python -m embodichain run-example learning/nerd/_support/_benchmark_dynamics_training --source /path/to/main + --trust-external-code --output /path/to/result.json + +Synthetic, seeded device-resident features measure computation only. Data loading, +state/contact encoding, validation, exports and simulation are excluded; projected +full training times are optimistic lower bounds, never accuracy qualification. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +import platform +import statistics +import subprocess +import sys +import time +from typing import Any + +__all__ = ["benchmark_training"] + + +def benchmark_training( + source: str | Path, + *, + device: str = "cuda:0", + batch_size: int = 512, + history: int = 10, + warmup: int = 20, + repeats: int = 100, + trust_external_code: bool = False, +) -> dict[str, Any]: + """Time synchronized completed FP32 updates using the real two model classes.""" + if not trust_external_code: + raise ValueError("Reference model code requires explicit trust") + if min(batch_size, history, warmup, repeats) < 1: + raise ValueError("Workload sizes must be positive") + import torch + from examples.learning.nerd._support._reference import reference_identity + from examples.learning.nerd._support._trainer import normalized_mse + from embodichain_models.dynamics.nerd import NeRDModel, NeRDModelCfg + + identity = reference_identity(source) + if device.startswith("cuda") and not torch.cuda.is_available(): + return {"status": "skipped", "reason": "CUDA unavailable"} + sys.path.insert(0, str(Path(source).resolve())) + from models.models import ModelMixedInput + from utils.running_mean_std import RunningMeanStd + + torch.set_num_threads(1) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + sync = ( + (lambda: torch.cuda.synchronize(device)) + if device.startswith("cuda") + else (lambda: None) + ) + records = [] + for robot in ("cartpole", "ant"): + torch.manual_seed(42) + dims = ( + {"states_embedding": 4, "joint_f": 2} + if robot == "cartpole" + else { + "states_embedding": 29, + "contact_normals": 75, + "contact_points_1": 75, + "contact_depths": 25, + "joint_f": 8, + "gravity_dir": 3, + } + ) + width = 192 if robot == "cartpole" else 384 + output = 4 if robot == "cartpole" else 29 + inputs = { + key: torch.randn(batch_size, history, size, device=device) + for key, size in dims.items() + } + target = torch.randn(batch_size, history, output, device=device) + migrated = NeRDModel( + NeRDModelCfg( + input_dims=dims, + output_dim=output, + n_embd=width, + n_layer=6, + n_head=12, + history_length=history, + block_size=max(32, history), + ) + ).to(device) + rename = lambda key: "joint_acts" if key == "joint_f" else key + official = ModelMixedInput( + input_sample={rename(key): value for key, value in inputs.items()}, + output_dim=output, + input_cfg={"low_dim": [rename(key) for key in dims]}, + network_cfg={ + "encoder": { + "low_dim": { + "activation": "relu", + "layer_sizes": [], + "layernorm": False, + } + }, + "model": { + "mlp": { + "activation": "relu", + "layer_sizes": [64], + "layernorm": False, + } + }, + "normalize_input": True, + "normalize_output": True, + "transformer": { + "n_embd": width, + "n_layer": 6, + "n_head": 12, + "block_size": max(32, history), + "bias": False, + "dropout": 0.0, + }, + }, + device=device, + ) + official.load_state_dict( + { + key: value + for key, value in migrated.state_dict().items() + if not key.startswith(("input_rms.", "output_rms.")) + } + ) + stats = {} + for key, size in dims.items(): + stats[rename(key)] = RunningMeanStd(shape=(size,), device=device) + official.set_input_rms(stats) + official.set_output_rms(RunningMeanStd(shape=(output,), device=device)) + with torch.no_grad(): + torch.testing.assert_close( + official({rename(k): v for k, v in inputs.items()}), + migrated(inputs), + atol=1e-5, + rtol=1e-4, + ) + for name, model in (("official", official), ("embodichain", migrated)): + optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) + model.train() + if device.startswith("cuda"): + torch.cuda.reset_peak_memory_stats(device) + + def update() -> torch.Tensor: + optimizer.zero_grad(set_to_none=True) + prediction = model( + {rename(k): v for k, v in inputs.items()} + if name == "official" + else dict(inputs) + ) + if name == "official": + weight = torch.rsqrt(model.output_rms.var + 1e-5) + loss = torch.nn.functional.mse_loss( + prediction * weight, target * weight + ) + else: + loss = normalized_mse(prediction, target, model.output_rms.var) + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + optimizer.step() + return loss.detach() + + start = time.perf_counter() + last = update() + sync() + first = time.perf_counter() - start + for _ in range(warmup): + last = update() + sync() + samples = [] + for _ in range(repeats): + sync() + start = time.perf_counter() + last = update() + sync() + samples.append(time.perf_counter() - start) + if not torch.isfinite(last): + raise FloatingPointError("Non-finite timing workload") + median = statistics.median(samples) + records.append( + { + "robot": robot, + "implementation": name, + "median_seconds_per_update": median, + "mean_seconds_per_update": statistics.mean(samples), + "min_seconds_per_update": min(samples), + "max_seconds_per_update": max(samples), + "first_update_seconds": first, + "samples_seconds": samples, + "peak_cuda_bytes": ( + torch.cuda.max_memory_allocated(device) + if device.startswith("cuda") + else None + ), + "last_loss": float(last), + "one_4995000_update_run_hours": median * 4995000 / 3600, + } + ) + del optimizer + delta = max( + (dict(official.named_parameters())[key] - value).abs().max().item() + for key, value in migrated.named_parameters() + ) + for record in records[-2:]: + record["matched_final_parameter_max_abs"] = delta + del official, migrated, inputs, target, last + if device.startswith("cuda"): + torch.cuda.empty_cache() + # 2 dataset rows x 5 seeds for each implementation and robot. + total = sum( + record["median_seconds_per_update"] * 4995000 * 10 for record in records + ) + try: + driver = subprocess.check_output( + ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"], + text=True, + ).strip() + except (OSError, subprocess.CalledProcessError): + driver = None + return { + "status": "completed", + "scope": "device_resident_forward_normalized_mse_backward_clip_adam", + "excluded": [ + "data I/O", + "feature encoding", + "validation", + "checkpoint/export", + "teacher generation", + "rollout evaluation", + ], + "precision": "float32, TF32 disabled", + "device": device, + "gpu": ( + torch.cuda.get_device_name(device) if device.startswith("cuda") else None + ), + "torch": str(torch.__version__), + "cuda_runtime": torch.version.cuda, + "driver": driver, + "python": platform.python_version(), + "cpu_threads": 1, + "batch_size": batch_size, + "history": history, + "warmup_updates": warmup, + "measured_updates": repeats, + "source": identity, + "runs": records, + "full_40_run_matrix_optimistic_hours": total / 3600, + "full_40_run_matrix_optimistic_days": total / 86400, + "accuracy_qualification": "not_measured_synthetic_performance_inputs", + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source", type=Path, required=True) + parser.add_argument("--trust-external-code", action="store_true", required=True) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--repeats", type=int, default=100) + parser.add_argument("--batch-size", type=int, default=512) + parser.add_argument("--history", type=int, default=10) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + result = benchmark_training( + args.source, + device=args.device, + batch_size=args.batch_size, + history=args.history, + warmup=args.warmup, + repeats=args.repeats, + trust_external_code=args.trust_external_code, + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(result, indent=2, allow_nan=False)) + print( + json.dumps( + {key: value for key, value in result.items() if key != "runs"}, indent=2 + ) + ) + for record in result.get("runs", []): + print({key: value for key, value in record.items() if key != "samples_seconds"}) + + +if __name__ == "__main__": + main() diff --git a/examples/learning/nerd/_support/_compare_official_model.py b/examples/learning/nerd/_support/_compare_official_model.py new file mode 100644 index 0000000000..e8af2aa9c1 --- /dev/null +++ b/examples/learning/nerd/_support/_compare_official_model.py @@ -0,0 +1,278 @@ +# ---------------------------------------------------------------------------- +# 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 the migrated model with frozen outputs from the tested official env. + +No Newton or DexSim teacher is used. Invoke with the robot's isolated old Warp +on PYTHONPATH. Reference data/outputs come from official_model_reference.py. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +from pathlib import Path +import platform +import sys + +__all__: list[str] = [] + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source", type=Path, required=True) + parser.add_argument("--reference", type=Path, required=True) + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--canonical-oracle", type=Path) + parser.add_argument("--trust-external-code", action="store_true", required=True) + args = parser.parse_args() + if args.output.exists(): + raise FileExistsError(args.output) + identity = json.loads((args.reference / "reference-identity.json").read_text()) + for name, digest in identity["files"].items(): + if hashlib.sha256((args.reference / name).read_bytes()).hexdigest() != digest: + raise ValueError("Reference artifact checksum mismatch: " + name) + import h5py + import numpy as np + import torch + import warp as wp + from embodichain_models.dynamics.nerd.codec import NeRDStateCodec + from embodichain_models.dynamics.nerd import load_model + from examples.learning.nerd._support._reference import ( + reference_identity, + import_legacy_datasets, + _official_batch, + ) + from examples.learning.nerd._support._reference_rollout import ( + _cpu_stream_compatibility, + ) + from examples.learning.nerd._support._trainer import encode_batch + from embodichain_models.dynamics.nerd.data import DynamicsWindowDataset + + source = reference_identity(args.source) + if ( + source["commit"] != identity["source_commit"] + or wp.__version__ != identity["warp"] + ): + raise ValueError("Original source or Warp identity mismatch") + common, manifest = load_model(args.model) + if manifest["legacy_checkpoint_sha256"] != identity["checkpoint_sha256"]: + raise ValueError("Port must retain the same original checkpoint identity") + torch.set_num_threads(1) + codec = NeRDStateCodec(**manifest["codec"]) + robot = identity["robot"] + canonical = args.output.parent / (args.output.stem + "-data") + rotation = ( + [0.5, 0.5, 0.5, 0.5] + if robot == "cartpole" + else [math.sqrt(0.5), 0, 0, math.sqrt(0.5)] + ) + import_legacy_datasets( + {"test": args.reference / "test.raw.h5"}, + canonical, + { + "robot": robot, + "dt": 1 / 60, + "codec": manifest["codec"], + "reference_identity": identity, + }, + world_rotation_xyzw=rotation, + gravity_direction=tuple(identity["gravity_transport"]["direction"]), + ) + results = {} + oracle_identity = None + oracle_windows = None + if args.canonical_oracle is not None: + oracle_identity = json.loads( + (args.canonical_oracle / "reference-identity.json").read_text() + ) + oracle_file = args.canonical_oracle / "official-windows.npz" + if ( + oracle_identity["checkpoint_sha256"] != identity["checkpoint_sha256"] + or hashlib.sha256(oracle_file.read_bytes()).hexdigest() + != oracle_identity["predictions_sha256"] + ): + raise ValueError("Common-frame oracle identity mismatch") + with np.load(oracle_file, allow_pickle=False) as values: + oracle_windows = { + key: torch.from_numpy(values[key].copy()) for key in values.files + } + with np.load(args.reference / "official-windows.npz") as fixtures: + for length in [1, 9, 10]: + prefix = str(length) + "__" + expected = torch.from_numpy(fixtures[prefix + "prediction"]) + direct_features = { + ("joint_f" if name == "joint_acts" else name): torch.from_numpy( + fixtures[prefix + "features__" + name] + ) + for name in identity["model_inputs"] + } + with torch.no_grad(): + direct = common(direct_features) + ds = DynamicsWindowDataset(canonical, "test", length, verify=True) + offsets = {} + offset = 0 + for shard in ds.manifest["shards"]: + for row in shard["trajectories"]: + offsets[row["metadata"]["legacy_index"]] = offset + offset += row["length"] - length + 1 + ids = [ + offsets[int(episode)] + int(start) + for episode, start in fixtures[prefix + "indices"] + ] + batch = ds.get_batch(ids) + features, target, mask = encode_batch(batch, codec, "cpu") + with torch.no_grad(): + actual = common(features) + expected_target = torch.from_numpy(fixtures[prefix + "target"]) + source_frame_prediction = expected + if oracle_windows is not None: + for key, value in _official_batch(batch, codec).items(): + actual_hash = hashlib.sha256(value.numpy().tobytes()).hexdigest() + if ( + actual_hash + != oracle_identity["window_field_sha256"][ + str(length) + "__raw__" + key + ] + ): + raise ValueError( + "Reference and candidate did not receive identical floating-point inputs" + ) + expected = oracle_windows[prefix + "prediction"] + expected_target = oracle_windows[prefix + "target"] + direct_features = { + ("joint_f" if name == "joint_acts" else name): oracle_windows[ + prefix + "features__" + name + ] + for name in identity["model_inputs"] + } + variance = common.output_rms.var + official_loss = ( + (expected - expected_target).square() / (variance + 1e-5) + ).mean() + port_loss = ((actual - target).square() / (variance + 1e-5)).mean() + results[str(length)] = { + "windows": len(ids), + "direct_model_max_abs": float( + (direct - source_frame_prediction).abs().max() + ), + "reference_frame_rotation_effect": float( + (expected - source_frame_prediction).abs().max() + ), + "codec_and_model_max_abs": float((actual - expected).abs().max()), + "target_max_abs": float((target - expected_target).abs().max()), + "feature_max_abs": { + key: float((features[key] - direct_features[key]).abs().max()) + for key in features + }, + "official_normalized_mse": float(official_loss), + "migrated_normalized_mse": float(port_loss), + "relative_mse_change": float(port_loss / official_loss - 1), + "forward_tolerance_passed": bool( + torch.allclose(actual, expected, atol=1e-5, rtol=1e-4) + ), + } + ds.close() + sys.path.insert(0, source["root"]) + from envs.neural_environment import NeuralEnvironment + + class PortableAdapter(torch.nn.Module): + def __init__(self, model): + super().__init__() + self.common = model + + def evaluate(self, values): + return self.common.evaluate( + { + ("joint_f" if key == "joint_acts" else key): v + for key, v in values.items() + } + ) + + def init_rnn(self, batch_size): + del batch_size + + with np.load(args.reference / "official-rollout.npz") as trace: + initial = trace["states"][0].copy() + controls = trace["joint_acts"].copy() + expected_rollout = trace["states"].copy() + cfg = { + "name": "TransformerNeuralIntegrator", + "num_states_history": 10, + "states_frame": "body" if robot == "ant" else "world", + "anchor_frame_step": "every", + "states_embedding_type": "identical", + "prediction_type": "relative", + "orientation_prediction_parameterization": "quaternion", + } + env = NeuralEnvironment( + env_name=robot.capitalize(), + num_envs=identity["count"], + neural_model=PortableAdapter(common), + neural_integrator_cfg=cfg, + default_env_mode="neural", + warp_env_cfg={"seed": 500, "random_reset": False}, + device="cpu", + render=False, + ) + outputs = [] + with _cpu_stream_compatibility(wp, "cpu"): + env.reset(initial_states=torch.from_numpy(initial)) + outputs.append(env.states.numpy().copy()) + for force in controls: + state = env.step_with_joint_act(torch.from_numpy(force), env_mode="neural") + outputs.append(state.numpy().copy()) + outputs = np.stack(outputs) + np.savez(args.output.with_suffix(".npz"), states=outputs) + rollout = {} + for h in [1, 10, 100, 1000]: + if h >= len(outputs): + continue + delta = np.abs(outputs[1 : h + 1] - expected_rollout[1 : h + 1]) + rollout[str(h)] = { + "max_abs_state_difference": float(delta.max()), + "mean_abs_state_difference": float(delta.mean()), + "all_finite": bool(np.isfinite(outputs[: h + 1]).all()), + } + report = { + "scope": "Original pretrained model versus migrated copy of the same weights; official supported environment generates comparison data. No Newton teacher.", + "reference": identity, + "common_frame_oracle": oracle_identity, + "migrated_environment": { + "python": platform.python_version(), + "torch": str(torch.__version__), + "warp": wp.__version__, + }, + "model_sha256": hashlib.sha256( + (args.model / "weights.pt").read_bytes() + ).hexdigest(), + "same_original_checkpoint": True, + "histories": results, + "rollouts": rollout, + "current_simulator_imports": [ + key for key in sys.modules if key.split(".")[0] in {"dexsim", "newton"} + ], + "qualification": "bounded_real_data_model_port_comparison_not_full_training_quality", + } + args.output.write_text(json.dumps(report, indent=2, allow_nan=False)) + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/examples/learning/nerd/_support/_comparison.py b/examples/learning/nerd/_support/_comparison.py new file mode 100644 index 0000000000..dfafe1a70e --- /dev/null +++ b/examples/learning/nerd/_support/_comparison.py @@ -0,0 +1,532 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Reproducible, immutable two-teacher/two-trainer NeRD experiments.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +from pathlib import Path +import subprocess +import sys +from typing import Any + +from ._reference import OFFICIAL_COMMIT, _hash, reference_identity + +__all__ = ["prepare_comparison", "run_comparison", "run_candidate"] +_TEACHERS = ("official_warp", "current_newton") +_TRAINERS = ("official", "embodichain") + + +def _json(path: Path, value: Any, *, exclusive: bool = False) -> None: + text = json.dumps(value, sort_keys=True, indent=2, allow_nan=False) + "\n" + if exclusive: + with path.open("x") as stream: + stream.write(text) + else: + temporary = path.with_suffix(path.suffix + ".pending") + temporary.write_text(text) + os.replace(temporary, path) + + +def prepare_comparison(config: dict[str, Any]) -> Path: + """Freeze common initialization, statistics, batch orders and run manifests. + + ``datasets`` maps both teacher names to canonical dataset paths or ``None``. + Missing inputs create explicitly pending cells. ``scope=smoke`` runs cannot + satisfy full-budget qualification, even if every training process succeeds. + No official source is executed during preparation. + """ + import numpy as np + import torch + from embodichain_models.dynamics.nerd.codec import NeRDStateCodec + from embodichain_models.dynamics.nerd.data import DynamicsWindowDataset + from embodichain_models.dynamics.nerd import NeRDModel, NeRDModelCfg, export_model + from ._trainer import ( + DynamicsTrainCfg, + compute_statistics, + save_statistics, + load_statistics, + ) + from ._sample_order import _write_checkpointed_order + + allowed = { + "output_dir", + "datasets", + "source", + "reference_dependencies", + "python", + "seeds", + "model", + "train", + "statistics", + "scope", + "cpu_threads", + "dataset_origins", + } + if config.keys() - allowed: + raise ValueError(f"Unknown comparison options: {config.keys()-allowed}") + if set(config["datasets"]) != set(_TEACHERS): + raise ValueError( + "datasets must name official_warp and current_newton, using null for unavailable data" + ) + scope = config.get("scope", "smoke") + if scope not in {"smoke", "full"}: + raise ValueError("scope must be smoke or full") + seeds = config.get("seeds", [0, 1, 2, 3, 4]) + if ( + not seeds + or len(set(seeds)) != len(seeds) + or any(not isinstance(seed, int) or seed < 0 for seed in seeds) + ): + raise ValueError("seeds must be distinct nonnegative integers") + root = Path(config["output_dir"]).resolve() + root.mkdir(parents=True, exist_ok=True) + if any(root.iterdir()): + raise FileExistsError("Comparison output must be a new empty directory") + source = reference_identity(config["source"]) + # Resolving the executable symlink discards pyvenv.cfg and silently selects + # the base interpreter's packages instead of the requested virtualenv. + python = str(Path(config.get("python", sys.executable)).expanduser().absolute()) + dependencies = ( + str(Path(config["reference_dependencies"]).resolve()) + if config.get("reference_dependencies") + else None + ) + architecture = NeRDModelCfg(**config["model"]) + cfg = DynamicsTrainCfg(**config["train"]) + architecture.validate() + cfg.validate() + if cfg.total_steps < 1 or cfg.batch_size < 1: + raise ValueError("Training sizes must be positive") + cells = [] + datasets = {} + artifacts = {} + for teacher in _TEACHERS: + path = config["datasets"][teacher] + if path is None: + datasets[teacher] = {"status": "pending", "reason": "dataset unavailable"} + for seed in seeds: + for trainer in _TRAINERS: + cells.append( + { + "id": f"{teacher}-{trainer}-{seed}", + "teacher": teacher, + "trainer": trainer, + "seed": seed, + "status": "pending", + "reason": "dataset unavailable", + } + ) + continue + dataset = DynamicsWindowDataset( + path, "train", architecture.history_length, verify=True + ) + if not len(dataset): + raise ValueError(f"No full training windows: {teacher}") + codec = NeRDStateCodec(**dataset.manifest["metadata"]["codec"]) + teacher_dir = root / teacher + teacher_dir.mkdir() + stats_source = config.get("statistics", {}).get(teacher) + stats = ( + load_statistics(stats_source) + if stats_source + else compute_statistics(dataset, codec, cfg.batch_size) + ) + stats_path = teacher_dir / "statistics.pt" + save_statistics(stats_path, stats) + artifacts[str(stats_path.relative_to(root))] = _hash(stats_path) + datasets[teacher] = { + "status": "available", + "path": str(Path(path).resolve()), + "fingerprint": dataset.fingerprint, + "origin": config.get("dataset_origins", {}).get(teacher), + "normalizer": ( + "shared_artifact" + if stats_source + else "shared_window_population_moments" + ), + } + for seed in seeds: + seed_dir = teacher_dir / f"seed-{seed}" + seed_dir.mkdir() + order_path, rng_path = _write_checkpointed_order( + seed_dir / "order.json", + dataset_size=len(dataset), + steps=cfg.total_steps, + batch_size=cfg.batch_size, + seed=seed, + ) + artifacts[str(rng_path.relative_to(root))] = _hash(rng_path) + order_hash = _hash(order_path) + artifacts[str(order_path.relative_to(root))] = order_hash + torch.manual_seed(seed) + model = NeRDModel(architecture) + model.set_normalization( + { + key: (item["mean"], item["var"]) + for key, item in stats["inputs"].items() + }, + stats["output"]["mean"], + stats["output"]["var"], + ) + initial = seed_dir / "initial-model" + export_model(initial, model, dataset.manifest["metadata"]) + for filename in ("manifest.json", "weights.pt"): + artifacts[str((initial / filename).relative_to(root))] = _hash( + initial / filename + ) + for trainer in _TRAINERS: + run_id = f"{teacher}-{trainer}-{seed}" + directory = seed_dir / trainer + spec = { + "source": source["root"], + "dataset": str(Path(path).resolve()), + "dataset_fingerprint": dataset.fingerprint, + "initial_model": str(initial), + "initial_model_sha256": _hash(initial / "weights.pt"), + "sample_order": str(order_path), + "sample_order_sha256": order_hash, + "train": {**cfg.to_dict(), "seed": seed}, + "output_dir": str(directory), + "cpu_threads": config.get("cpu_threads", 1), + } + spec_path = seed_dir / f"{trainer}.json" + _json(spec_path, spec, exclusive=True) + artifacts[str(spec_path.relative_to(root))] = _hash(spec_path) + worker = "_reference" if trainer == "official" else "_comparison" + command = [ + python, + str(Path(__file__).with_name("_worker.py")), + worker, + "--spec", + str(spec_path), + ] + if trainer == "official": + command.append("--trust-external-code") + else: + command.append("--candidate") + cells.append( + { + "id": run_id, + "teacher": teacher, + "trainer": trainer, + "seed": seed, + "status": "prepared", + "spec": str(spec_path), + "output_dir": str(directory), + "command": command, + } + ) + dataset.close() + plan = { + "format": "embodichain.nerd.comparison.v1", + "scope": scope, + "source": source, + "reference_dependencies": dependencies, + "python": python, + "seeds": seeds, + "model": architecture.to_dict(), + "feature_order": list(architecture.input_dims), + "train": cfg.to_dict(), + "comparison_protocol": { + "kind": "controlled_matched_data_noninferiority", + "initialization": "shared exported weights and normalizers", + "training_order": "shared frozen torch.randint-with-replacement indices", + "statistics": "shared provided artifact or train-window population moments", + "checkpoint_selection": "one validation split; ordered first configured batches; partial final batch included; no epoch-zero best selection", + "preprocessing_execution": "shared current frame; scripted coordinate math; three real-layout CUDA preprocessing warmups without optimizer/RNG advancement", + "paper_reproduction": "not established by this controlled protocol", + }, + "datasets": datasets, + "cells": cells, + "artifacts": artifacts, + "adapter_source_sha256": _hash(Path(__file__)), + "worker_bootstrap_sha256": _hash(Path(__file__).with_name("_worker.py")), + "acceptance": { + "relative_margin": 0.05, + "absolute_tolerance": 1e-5, + "resamples": 10000, + "bootstrap_seed": 0, + "required_seeds": 5, + "official_updates": 4995000, + "primary_metrics": ["q_error_norm", "qd_error_norm"], + "rollout_primary_metrics": ["q_mae", "qd_mae"], + "policy_return_relative_margin": 0.05, + "failure_rate_increase": 0.0, + "nonfinite_rate_increase": 0.0, + }, + } + _json(root / "experiment.json", plan, exclusive=True) + _json( + root / "experiment.identity.json", + {"sha256": _hash(root / "experiment.json")}, + exclusive=True, + ) + return root / "experiment.json" + + +def _read_plan(path: str | Path) -> tuple[Path, dict[str, Any]]: + path = Path(path).resolve() + root = path.parent + identity = json.loads((root / "experiment.identity.json").read_text()) + if _hash(path) != identity["sha256"]: + raise ValueError("Experiment manifest changed after preparation") + plan = json.loads(path.read_text()) + if plan.get("format") != "embodichain.nerd.comparison.v1": + raise ValueError("Unsupported comparison manifest") + for name, digest in plan["artifacts"].items(): + file = (root / name).resolve() + if not file.is_relative_to(root) or not file.is_file() or _hash(file) != digest: + raise ValueError(f"Immutable comparison artifact changed: {name}") + return root, plan + + +def run_candidate( + spec: dict[str, Any], + *, + resume: str | Path | None = None, + stop_after: int | None = None, +) -> dict[str, Any]: + """Execute the migrated trainer on the same frozen run specification.""" + import torch + from embodichain_models.dynamics.nerd.data import ( + load_manifest, + dataset_fingerprint, + ) + from ._workflow import run_train, run_export, run_eval + + torch.set_num_threads(int(spec.get("cpu_threads", 1))) + torch.set_float32_matmul_precision("highest") + if ( + dataset_fingerprint(load_manifest(spec["dataset"], verify=True)) + != spec["dataset_fingerprint"] + ): + raise ValueError("Candidate dataset changed") + if _hash(Path(spec["sample_order"])) != spec["sample_order_sha256"]: + raise ValueError("Candidate sample order changed") + if ( + _hash(Path(spec["initial_model"]) / "weights.pt") + != spec["initial_model_sha256"] + ): + raise ValueError("Common initialization changed") + directory = Path(spec["output_dir"]) + initial_manifest = json.loads( + (Path(spec["initial_model"]) / "manifest.json").read_text() + ) + execution = ( + {"resume": str(resume)} + if resume is not None + else {"fine_tune": spec["initial_model"]} + ) + result = run_train( + { + "dataset": spec["dataset"], + "output_dir": str(directory), + "model": initial_manifest["model_config"], + "sample_order": spec["sample_order"], + "train": spec["train"], + "stop_after": stop_after, + **execution, + } + ) + if result["status"] != "completed": + result["completed_updates"] = result["global_step"] + _json(directory / "result.json", result) + return result + for name in ("best", "last"): + run_export(directory / f"{name}.pt", directory / f"{name}-export") + result.update( + completed_updates=result["global_step"], + test=run_eval( + { + "dataset": spec["dataset"], + "model": str(directory / "best-export"), + "device": spec["train"]["device"], + "batch_size": spec["train"]["batch_size"], + } + ), + dataset_fingerprint=spec["dataset_fingerprint"], + sample_order_sha256=spec["sample_order_sha256"], + seed=spec["train"]["seed"], + qualification="matched_training_run_only", + best_model_sha256=_hash(directory / "best-export" / "weights.pt"), + last_model_sha256=_hash(directory / "last-export" / "weights.pt"), + ) + _json(directory / "result.json", result) + return result + + +def run_comparison( + path: str | Path, *, trust_external_code: bool = False, max_runs: int | None = None +) -> dict[str, Any]: + """Execute prepared cells sequentially, retaining real command/stdout evidence. + + Completed verified outputs are reused. Incomplete directories are not silently + restarted; create a new experiment or recover its run explicitly. This function + does not claim rollout or paper qualification from offline training results. + """ + root, plan = _read_plan(path) + if not trust_external_code: + raise ValueError("Comparison requires explicit trust_external_code=True") + if max_runs is not None and max_runs < 1: + raise ValueError("max_runs must be positive") + statuses = [] + executed = 0 + for cell in plan["cells"]: + if cell["status"] != "prepared": + statuses.append( + {"id": cell["id"], "status": "pending", "reason": cell["reason"]} + ) + continue + output = Path(cell["output_dir"]) + receipt = root / (cell["id"] + ".receipt.json") + if receipt.exists(): + previous = json.loads(receipt.read_text()) + if ( + previous["status"] == "completed" + and _hash(output / "result.json") == previous["result_sha256"] + ): + statuses.append(previous) + continue + statuses.append( + { + "id": cell["id"], + "status": "pending", + "reason": "previous run needs explicit recovery", + } + ) + continue + if max_runs is not None and executed >= max_runs: + statuses.append( + {"id": cell["id"], "status": "pending", "reason": "execution limit"} + ) + continue + if output.exists() and any(output.iterdir()): + statuses.append( + { + "id": cell["id"], + "status": "pending", + "reason": "unreceipted output needs explicit recovery", + } + ) + continue + output.mkdir(parents=True, exist_ok=True) + spec = json.loads(Path(cell["spec"]).read_text()) + environment = { + **os.environ, + "OMP_NUM_THREADS": str(spec["cpu_threads"]), + "MKL_NUM_THREADS": str(spec["cpu_threads"]), + } + pythonpaths = [] + if cell["trainer"] == "official": + if plan["reference_dependencies"]: + pythonpaths.append(plan["reference_dependencies"]) + pythonpaths.append(plan["source"]["root"]) + source_root = Path(__file__).resolve().parents[4] + if (source_root / "pyproject.toml").is_file(): + pythonpaths.extend( + (str(source_root / "extensions/embodichain_models"), str(source_root)) + ) + if environment.get("PYTHONPATH"): + pythonpaths.append(environment["PYTHONPATH"]) + environment["PYTHONPATH"] = os.pathsep.join(pythonpaths) + if spec["train"]["device"] == "cpu": + environment["CUDA_VISIBLE_DEVICES"] = "" + _json( + output / "command.json", + { + "argv": cell["command"], + "PYTHONPATH": environment["PYTHONPATH"], + "source_commit": OFFICIAL_COMMIT, + }, + exclusive=True, + ) + with ( + (output / "stdout.log").open("w") as stdout, + (output / "stderr.log").open("w") as stderr, + ): + process = subprocess.run( + cell["command"], + env=environment, + stdout=stdout, + stderr=stderr, + check=False, + ) + executed += 1 + result_file = output / "result.json" + status = "failed" + if process.returncode == 0 and result_file.is_file(): + result = json.loads(result_file.read_text()) + if ( + result.get("status") == "completed" + and result.get("completed_updates") == spec["train"]["total_steps"] + ): + status = "completed" + receipt_value = { + "id": cell["id"], + "status": status, + "returncode": process.returncode, + "command": cell["command"], + "result_sha256": _hash(result_file) if result_file.exists() else None, + "stdout_sha256": _hash(output / "stdout.log"), + "stderr_sha256": _hash(output / "stderr.log"), + } + _json(receipt, receipt_value, exclusive=True) + statuses.append(receipt_value) + _json(root / "execution.json", {"runs": statuses}) + return {"runs": statuses, "executed": executed} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--spec", type=Path) + parser.add_argument("--candidate", action="store_true") + parser.add_argument("--resume", type=Path) + parser.add_argument("--stop-after", type=int) + parser.add_argument("--prepare", type=Path, help="JSON comparison configuration") + parser.add_argument("--run", type=Path, help="Prepared experiment.json") + parser.add_argument("--trust-external-code", action="store_true") + parser.add_argument("--max-runs", type=int) + args = parser.parse_args() + if args.candidate and args.spec: + value = run_candidate( + json.loads(args.spec.read_text()), + resume=args.resume, + stop_after=args.stop_after, + ) + elif args.prepare: + value = { + "manifest": str(prepare_comparison(json.loads(args.prepare.read_text()))) + } + elif args.run: + value = run_comparison( + args.run, + trust_external_code=args.trust_external_code, + max_runs=args.max_runs, + ) + else: + parser.error("Choose --prepare, --run or --candidate --spec") + print(json.dumps(value, indent=2, allow_nan=False)) + if args.run and any(run["status"] == "failed" for run in value["runs"]): + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/examples/learning/nerd/_support/_evaluation.py b/examples/learning/nerd/_support/_evaluation.py new file mode 100644 index 0000000000..e0aceb53b1 --- /dev/null +++ b/examples/learning/nerd/_support/_evaluation.py @@ -0,0 +1,103 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Shared rollout metrics and teacher-content identity without simulator imports.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +import hashlib +import json +import math +from typing import Any + +import numpy as np +import torch + +__all__: list[str] = [] + + +def _teacher_row_sha256(record: Mapping[str, Any], values: Sequence[np.ndarray]) -> str: + """Hash ordered canonical state/control content and logical trajectory identity.""" + digest = hashlib.sha256() + digest.update( + json.dumps( + { + "schema": "teacher-content-v1", + "id": record["id"], + "sampling_seed": record.get("metadata", {}).get("sampling_seed"), + }, + sort_keys=True, + separators=(",", ":"), + ).encode() + ) + for name, value in zip(("q", "qd", "joint_f"), values): + array = np.ascontiguousarray(value) + digest.update( + json.dumps( + {"name": name, "shape": array.shape, "dtype": array.dtype.str}, + sort_keys=True, + separators=(",", ":"), + ).encode() + ) + digest.update(array.tobytes()) + return digest.hexdigest() + + +def _error_values( + robot: str, + q: torch.Tensor, + qd: torch.Tensor, + target_q: torch.Tensor, + target_qd: torch.Tensor, +) -> torch.Tensor: + # Accumulate errors in float64 so large but finite FP32 states remain + # distinguishable from actual simulator NaN/Inf failures. + q, qd, target_q, target_qd = ( + value.double() for value in (q, qd, target_q, target_qd) + ) + start = 7 if robot == "ant" else 0 + dq = q[..., start:] - target_q[..., start:] + if robot == "ant": + # Ant scalar joints have finite limits, so a full-turn prediction error + # is a real error. Only the Cartpole continuous pole coordinate wraps. + position = (q[..., :3] - target_q[..., :3]).norm(dim=-1) + dot = ( + ( + torch.nn.functional.normalize(q[..., 3:7], dim=-1) + * torch.nn.functional.normalize(target_q[..., 3:7], dim=-1) + ) + .sum(-1) + .abs() + .clamp(0, 1) + ) + orientation = 2 * torch.acos(dot) + else: + dq = torch.stack( + ( + dq[..., 0], + torch.remainder(dq[..., 1] + torch.pi, 2 * torch.pi) - torch.pi, + ), + -1, + ) + position = orientation = torch.zeros_like(q[..., 0]) + return torch.stack( + (dq.abs().mean(-1), (qd - target_qd).abs().mean(-1), position, orientation), -1 + ) + + +def _safe_number(value: float) -> float | None: + return float(value) if math.isfinite(value) else None diff --git a/examples/learning/nerd/_support/_generation.py b/examples/learning/nerd/_support/_generation.py new file mode 100644 index 0000000000..d9d8b6c74c --- /dev/null +++ b/examples/learning/nerd/_support/_generation.py @@ -0,0 +1,352 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Generate trajectory shards using Newton's built-in analytical solvers. + +This module never imports DexSim or installs a custom physics solver. +""" + +from __future__ import annotations + +from collections.abc import Mapping +import hashlib +import json +import math +from pathlib import Path +import time +from typing import Any + +import numpy as np +import torch + +from embodichain_models.dynamics.nerd.codec import quat_rotate +from embodichain_models.dynamics.nerd.data import DynamicsShardWriter +from embodichain.utils import configclass +from ._reference_assets import reference_profile +from ._newton import _NewtonTeacher + +__all__ = ["DynamicsGenerationCfg", "generate_dataset"] + + +def _write_json(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".pending") + temporary.write_text(json.dumps(value, indent=2, allow_nan=False) + "\n") + temporary.replace(path) + + +def _sample( + robot: str, + seed: int, + horizon: int, + motion_mode: str, + initial_states: str, + effort_scale: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Sample one stable trajectory ID independently of collector batch size.""" + profile = reference_profile(robot) + rng = torch.Generator(device="cpu").manual_seed(seed) + q, qd = torch.tensor(profile["default_q"]), torch.tensor(profile["default_qd"]) + if initial_states == "sample": + if robot == "cartpole": + q = (2 * torch.rand(2, generator=rng) - 1) * torch.tensor([3.5, torch.pi]) + qd = (2 * torch.rand(2, generator=rng) - 1) * 10 + else: + q[:3] = torch.tensor([-10.0, -10.0, 1.0]) + torch.rand( + 3, generator=rng + ) * torch.tensor([20.0, 20.0, 0.3]) + angle = (2 * torch.rand((), generator=rng) - 1) * (torch.pi / 2) + axis = 2 * torch.rand(3, generator=rng) - 1 + axis = axis / axis.norm().clamp_min(1e-12) + q[3:7] = torch.cat( + (axis * torch.sin(angle / 2), torch.cos(angle / 2).reshape(1)) + ) + low = torch.tensor( + [-40.0, 30.0, -40.0, -100.0, -40.0, -100.0, -40.0, 30.0] + ) * (torch.pi / 180) + high = torch.tensor( + [40.0, 100.0, 40.0, -30.0, 40.0, -30.0, 40.0, 100.0] + ) * (torch.pi / 180) + q[7:] = low + torch.rand(8, generator=rng) * (high - low) + qd[:3] = quat_rotate(q[3:7], 2 * torch.rand(3, generator=rng) - 1) + qd[3:6] = quat_rotate( + q[3:7], (2 * torch.rand(3, generator=rng) - 1) * torch.pi + ) + qd[6:] = (2 * torch.rand(8, generator=rng) - 1) * (2 * torch.pi) + elif initial_states != "default": + raise ValueError("initial_states must be sample or default.") + efforts = torch.zeros((horizon, profile["qd_dim"])) + if motion_mode == "random_effort": + limits = torch.tensor(profile["effort_limits"]) * effort_scale + sampled = (2 * torch.rand((horizon, len(limits)), generator=rng) - 1) * limits + efforts[:, 6:] = sampled if robot == "ant" else efforts[:, 6:] + if robot == "cartpole": + efforts = sampled + elif motion_mode != "passive": + raise ValueError("motion_mode must be passive or random_effort.") + return q, qd, efforts + + +def _positive_integer(name: str, value: Any, minimum: int = 1) -> int: + if type(value) is not int or value < minimum: + raise ValueError(f"{name} must be an integer >= {minimum}, not {value!r}") + return value + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +@configclass +class DynamicsGenerationCfg: + """Validated analytical sampling settings; paths belong to the caller.""" + + robot: str = "cartpole" + output_dir: str = "outputs/dynamics/data" + device: str = "cpu" + num_envs: int = 64 + trajectories: int = 10000 + horizon: int = 100 + seed: int = 0 + teacher_solver: dict[str, Any] | None = None + num_substeps: int = 10 + trajectories_per_shard: int = 64 + resume: bool = False + motion_mode: str = "random_effort" + split_seed: int = 0 + initial_states: str = "sample" + max_attempts: int | None = None + effort_scale: float = 1.0 + physics_profile: str = "native" + use_cuda_graph: bool = True + + def __post_init__(self) -> None: + for name in ( + "num_envs", + "trajectories", + "horizon", + "num_substeps", + "trajectories_per_shard", + ): + _positive_integer(name, getattr(self, name)) + for name in ("seed", "split_seed"): + _positive_integer(name, getattr(self, name), 0) + if self.max_attempts is not None: + _positive_integer("max_attempts", self.max_attempts, self.trajectories) + if type(self.resume) is not bool or type(self.use_cuda_graph) is not bool: + raise TypeError("resume and use_cuda_graph must be boolean") + if ( + isinstance(self.effort_scale, bool) + or not isinstance(self.effort_scale, (float, int)) + or not math.isfinite(self.effort_scale) + or self.effort_scale < 0 + ): + raise ValueError("effort_scale must be finite and nonnegative") + if ( + self.robot not in {"cartpole", "ant"} + or self.motion_mode not in {"passive", "random_effort"} + or self.initial_states not in {"sample", "default"} + ): + raise ValueError("Unsupported robot, motion_mode or initial_states") + if self.physics_profile not in {"native", "warp_reference"}: + raise ValueError("Unsupported reference physics profile") + if self.teacher_solver is not None and not isinstance( + self.teacher_solver, Mapping + ): + raise TypeError("teacher_solver must be a mapping") + if not isinstance(self.output_dir, (str, Path)) or not str(self.output_dir): + raise ValueError("output_dir must be a nonempty path") + + +def generate_dataset(config: Mapping[str, Any]) -> dict[str, Any]: + """Generate committed finite trajectories with bounded attempts and staging. + + Config validation happens before simulator creation. Each candidate has an + independent seed. Resume advances only past examined candidates, including + a final partial batch; unused padded worlds never acquire trajectory IDs. + + Args: + config: Settings accepted by :class:`DynamicsGenerationCfg`. + + Returns: + Actual committed trajectory/transition counts, rejected candidate count, + rejection evidence path and elapsed wall time. + """ + cfg = DynamicsGenerationCfg(**dict(config)) + cfg.validate() + output = Path(cfg.output_dir) + progress_path = output / "generation-progress.json" + started = time.perf_counter() + with _NewtonTeacher( + cfg.robot, + device=cfg.device, + num_envs=cfg.num_envs, + solver=cfg.teacher_solver, + num_substeps=cfg.num_substeps, + physics_profile=cfg.physics_profile, + use_cuda_graph=cfg.use_cuda_graph, + ) as host: + manifest = dict(host.manifest) + manifest["generation"] = { + "horizon": cfg.horizon, + "seed": cfg.seed, + "motion_mode": cfg.motion_mode, + "initial_states": cfg.initial_states, + "effort_scale": cfg.effort_scale, + "num_envs": cfg.num_envs, + } + with DynamicsShardWriter( + output, + manifest, + cfg.trajectories_per_shard, + resume=cfg.resume, + split_seed=cfg.split_seed, + ) as writer: + progress = writer.manifest.get("generation_progress") + if progress is None: + # Older datasets used a separate progress file. New collections + # recover from the manifest even if that compatibility report + # was not written, or still describes an earlier boundary. + progress = ( + json.loads(progress_path.read_text()) + if cfg.resume and progress_path.exists() + else {"next_candidate": 0, "rejected": 0, "rejections": []} + ) + progress.setdefault("rejections", []) + accepted = len(writer.committed_ids) + if accepted > cfg.trajectories: + raise ValueError( + "Requested trajectory count is smaller than committed data" + ) + first = max( + [int(value.rsplit(":", 1)[-1]) + 1 for value in writer.committed_ids] + or [0] + ) + cursor = max(first, int(progress["next_candidate"])) + progress.update(next_candidate=cursor, accepted=accepted) + # Keep the current candidate boundary beside pending shard state. + # add_trajectory may commit automatically, so update this same + # mapping before adding each accepted trajectory. + writer.manifest["generation_progress"] = progress + writer.commit() + maximum = ( + cfg.max_attempts + if cfg.max_attempts is not None + else cfg.trajectories * 10 + ) + while accepted < cfg.trajectories: + if cursor >= maximum: + raise RuntimeError( + f"Generation attempt limit reached: {accepted}/{cfg.trajectories} valid trajectories, {progress['rejected']} rejected" + ) + candidates = list(range(cursor, min(cursor + cfg.num_envs, maximum))) + samples = [ + _sample( + cfg.robot, + cfg.seed + i, + cfg.horizon, + cfg.motion_mode, + cfg.initial_states, + cfg.effort_scale, + ) + for i in candidates + ] + samples.extend([samples[-1]] * (cfg.num_envs - len(samples))) + q, qd, force = [ + torch.stack([row[i] for row in samples]).to(host.device) + for i in range(3) + ] + host.reset(q, qd) + states_list, contacts_list = [], [] + with torch.no_grad(): + for index in range(cfg.horizon): + state, contacts = host.snapshot() + states_list.append(state) + contacts_list.append(contacts) + host.step(force[:, index]) + state, _ = host.snapshot() + states_list.append(state) + states = { + key: torch.stack([value[key] for value in states_list], dim=1) + .cpu() + .numpy() + for key in states_list[0] + } + transitions = { + key: torch.stack([value[key] for value in contacts_list], dim=1) + .cpu() + .numpy() + for key in contacts_list[0] + } + transitions["joint_f"] = force.cpu().numpy() + transitions["gravity_dir"] = ( + host.gravity_dir[:, None].expand(-1, cfg.horizon, -1).cpu().numpy() + ) + for row, candidate in enumerate(candidates): + if accepted >= cfg.trajectories: + break + cursor = candidate + 1 + progress["next_candidate"] = cursor + finite = all( + np.isfinite(v[row]).all() + for v in (*states.values(), *transitions.values()) + ) + bounded = ( + finite + and np.max(np.abs(states["q"][row])) < 1e5 + and np.max(np.abs(states["qd"][row])) < 1e5 + ) + if not bounded: + progress["rejected"] += 1 + progress["rejections"].append( + { + "id": f"{cfg.seed}:{candidate}", + "sampling_seed": cfg.seed + candidate, + "reason": ( + "nonfinite" + if not finite + else "state_magnitude_limit" + ), + } + ) + continue + progress["accepted"] = accepted + 1 + writer.add_trajectory( + f"{cfg.seed}:{candidate}", + {k: v[row] for k, v in states.items()}, + {k: v[row] for k, v in transitions.items()}, + metadata={"sampling_seed": cfg.seed + candidate}, + ) + accepted += 1 + writer.commit() + # Preserve the historical evidence-file layout as a derived completion + # report. It is no longer a second recovery transaction. + _write_json(progress_path, progress) + result = { + "status": "completed", + "dataset": str(output.resolve()), + "trajectories": accepted, + "transitions": accepted * cfg.horizon, + "rejected": progress["rejected"], + "rejection_evidence": str(progress_path.resolve()), + "elapsed_seconds": time.perf_counter() - started, + } + _write_json(output / "generation-result.json", result) + return result diff --git a/examples/learning/nerd/_support/_historical_physics.py b/examples/learning/nerd/_support/_historical_physics.py new file mode 100644 index 0000000000..6090c5d8d6 --- /dev/null +++ b/examples/learning/nerd/_support/_historical_physics.py @@ -0,0 +1,233 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Standalone historical physics checks without importing EmbodiChain or SDK.""" + +from __future__ import annotations +from collections.abc import Mapping +import hashlib +import json +import math +from pathlib import Path +from typing import Any +from ._official_evaluation_data import _hash + +__all__: list[str] = [] + +OFFICIAL_COMMIT = "370e89d2b8ce68aa3578aa7812859c681b7502bc" + + +def _positive(name: str, value: Any, minimum: int = 1) -> int: + if type(value) is not int or value < minimum: + raise ValueError(f"{name} must be an integer >= {minimum}") + return value + + +def _physical_arrays(model: Any) -> dict[str, Any]: + """Read the same physical fields recorded by original-teacher collection.""" + import numpy as np + + result = {} + for name in ( + "body_mass", + "body_inertia", + "body_com", + "joint_armature", + "joint_limit_lower", + "joint_limit_upper", + "joint_limit_ke", + "joint_limit_kd", + "joint_axis_mode", + "joint_target_ke", + "joint_target_kd", + ): + value = getattr(model, name, None) + if value is not None: + result[name] = ( + value.numpy() if hasattr(value, "numpy") else np.asarray(value) + ) + for field in ("mu", "ke", "kd", "kf"): + value = getattr(model.shape_materials, field, None) + if value is not None: + result["shape_material_" + field] = value.numpy() + result["shape_thickness"] = model.shape_geo.thickness.numpy() + result["shape_body"] = model.shape_body.numpy() + result["ground_plane"] = model.ground_plane.numpy() + result["gravity"] = np.asarray(model.gravity, dtype=np.float64) + return result + + +def _canonical_physics( + arrays: Mapping[str, Any], batch: int, bodies: int +) -> dict[str, Any]: + """Remove replica indexing while proving the recorded replicas homogeneous.""" + import numpy as np + + result = {} + shape_body = np.asarray(arrays["shape_body"]) + # Original Warp always appends one shared ground-plane shape. + if (len(shape_body) - 1) % batch: + raise ValueError("Historical shape replication is not homogeneous") + shapes = (len(shape_body) - 1) // batch + for key, values in arrays.items(): + value = np.asarray(values) + if key in {"ground_plane", "gravity"}: + result[key] = value.tolist() + continue + if key.startswith("shape_"): + if len(value) != batch * shapes + 1: + raise ValueError(f"Unexpected historical shape field {key}") + rows = value[:-1].reshape(batch, shapes, *value.shape[1:]).copy() + if key == "shape_body": + rows = np.where( + rows >= 0, rows - np.arange(batch)[:, None] * bodies, rows + ) + tail = value[-1:].tolist() + else: + if len(value) % batch: + raise ValueError(f"Unexpected historical replication for {key}") + rows = value.reshape(batch, len(value) // batch, *value.shape[1:]) + tail = None + if not np.allclose(rows, rows[:1], rtol=1e-6, atol=1e-7, equal_nan=False): + raise ValueError(f"Historical physical field {key} differs across replicas") + result[key] = {"per_environment": rows[0].tolist(), "shared": tail} + return result + + +def _verify_historical_physics( + environment: Any, + manifest: Mapping[str, Any], + source: Mapping[str, Any], + warp_version: str, +) -> dict[str, Any]: + """Require actual historical physical data, never relabel a Newton manifest.""" + import numpy as np + + recorded = manifest.get("teacher_origin", {}) + if ( + recorded.get("kind") != "official_warp" + or recorded.get("source_commit") != OFFICIAL_COMMIT + or recorded.get("warp_version") != warp_version + ): + raise ValueError( + "Model does not record the required original-Warp teacher identity" + ) + if not isinstance(manifest.get("source_physical_arrays"), Mapping): + raise ValueError( + "Historical source_physical_arrays are required; do not relabel a Newton model" + ) + if ( + recorded.get("up_axis") != int(environment.model.up_axis) + or recorded.get("num_substeps") != environment.sim_substeps_gt + ): + raise ValueError("Original teacher axis/substeps differ from model provenance") + if recorded.get("integrator") != str(environment.integrator_type_gt): + raise ValueError("Original teacher integrator differs from model provenance") + if "gravity" in recorded and not np.allclose( + np.asarray(recorded["gravity"]), + np.asarray(environment.model.gravity, dtype=np.float64), + rtol=1e-6, + atol=1e-7, + ): + raise ValueError("Original teacher gravity differs from recorded provenance") + if not math.isclose( + float(manifest.get("dt", 0)), + float(environment.frame_dt), + rel_tol=0, + abs_tol=1e-12, + ): + raise ValueError("Original teacher timestep differs from model provenance") + robot = manifest["codec"]["robot"] + expected_continuous = [1] if robot == "cartpole" else [] + if ( + list(manifest["codec"].get("continuous_q_indices", expected_continuous)) + != expected_continuous + ): + raise ValueError( + "Historical codec continuous coordinates differ from original runtime" + ) + if not np.allclose( + manifest["codec"].get("base_com", [0.0, 0.0, 0.0]), + environment.model.body_com.numpy()[0], + rtol=1e-6, + atol=1e-7, + ): + raise ValueError("Historical codec base COM differs from actual source model") + asset = ( + Path(source["root"]) + / "envs" + / "warp_sim_envs" + / "assets" + / ("ant.xml" if robot == "ant" else "cartpole_single.urdf") + ) + if manifest.get("asset_sha256") != _hash(asset): + raise ValueError("Historical model asset checksum differs") + batch = manifest.get("sampling", {}).get("batch_size") + split_batch = ( + manifest.get("source_split_identities", {}).get("train", {}).get("num_envs") + ) + if batch is not None and split_batch is not None and batch != split_batch: + raise ValueError("Conflicting source sampling batch sizes") + batch = _positive( + "source sampling batch_size", batch if batch is not None else split_batch + ) + bodies = _positive("source bodies_per_env", recorded.get("bodies_per_env")) + if bodies != environment.bodies_per_env: + raise ValueError("Historical body count differs") + actual = _canonical_physics( + _physical_arrays(environment.model), environment.num_envs, bodies + ) + expected = _canonical_physics(manifest["source_physical_arrays"], batch, bodies) + missing_fields = sorted(set(actual) - set(expected)) + if set(expected) - set(actual) or set(missing_fields) - {"gravity"}: + raise ValueError("Historical physical field set differs") + for key in expected: + if key in {"ground_plane", "gravity"}: + pairs = [(actual[key], expected[key])] + else: + pairs = [ + (actual[key][name], expected[key][name]) + for name in ("per_environment", "shared") + if actual[key][name] is not None or expected[key][name] is not None + ] + for left, right in pairs: + if ( + left is None + or right is None + or np.shape(left) != np.shape(right) + or not np.allclose(left, right, rtol=1e-6, atol=1e-7) + ): + raise ValueError(f"Historical physical mismatch in {key}") + identity = hashlib.sha256( + json.dumps( + actual, sort_keys=True, separators=(",", ":"), allow_nan=False + ).encode() + ).hexdigest() + return { + "schema": "official-warp-physical-arrays-v2", + "sha256": identity, + "matched_fields": sorted(expected), + "missing_recorded_fields": missing_fields, + "matched_recorded_fields": True, + "complete": not missing_fields, + "status": ( + "matched" if not missing_fields else "pending_missing_recorded_physics" + ), + "gravity_live": np.asarray( + environment.model.gravity, dtype=np.float64 + ).tolist(), + "gravity_magnitude_in_training_metadata": "gravity" in expected, + } diff --git a/examples/learning/nerd/_support/_matched_small_budget.py b/examples/learning/nerd/_support/_matched_small_budget.py new file mode 100644 index 0000000000..e53584ae96 --- /dev/null +++ b/examples/learning/nerd/_support/_matched_small_budget.py @@ -0,0 +1,265 @@ +# ---------------------------------------------------------------------------- +# 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 the migrated trainer against a captured, literal original-entry protocol. + +Unlike the general controlled comparison, the original trainer supplies model +initialization, FP32 RunningMeanStd statistics and actual consumed DataLoader +indices. Epoch-zero and final validation use its actual emitted validation order. +The original run retains its independent best-evaluation checkpoint as well. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path +import platform +import time +from typing import Any + +import numpy as np +import torch + +from embodichain_models.dynamics.nerd.codec import NeRDStateCodec +from embodichain_models.dynamics.nerd.data import ( + DynamicsWindowDataset, + dataset_fingerprint, + load_manifest, +) +from examples.learning.nerd._support._trainer import ( + DynamicsTrainCfg, + DynamicsTrainer, + encode_batch, + evaluate_model, + normalized_mse, +) +from examples.learning.nerd._support._workflow import run_export +from embodichain_models.dynamics.nerd import convert_legacy_checkpoint, load_model + +__all__: list[str] = [] + + +def _hash(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 _json(path: Path, value: Any) -> None: + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(value, indent=2, allow_nan=False) + "\n") + temporary.replace(path) + + +def _first_batch_differences( + model: Any, + inputs: dict[str, torch.Tensor], + targets: torch.Tensor, + official: Path, +) -> tuple[dict[str, float], str]: + """Compare features with an explicit pre/post-normalization observation stage.""" + path = official / "first-preprocessed-batch.pt" + sidecar = official / "first-preprocessed-batch.json" + if sidecar.exists(): + observation = json.loads(sidecar.read_text()) + stage = observation.get("feature_stage") + if ( + observation.get("format_version") != 1 + or stage != "physical_before_model_forward" + or observation.get("tensor_sha256") != _hash(path) + ): + raise ValueError("Invalid original first-batch observation metadata") + else: + # Historical observer snapshots were taken after ModelMixedInput.forward + # normalized its input dictionary in place. Preserve those frozen files. + stage = "legacy_unlabeled_normalized_after_model_forward" + original = torch.load(path, map_location="cpu", weights_only=True) + differences = {} + for key, value in inputs.items(): + original_key = "joint_acts" if key == "joint_f" else key + rms = model.input_rms[key] + mean, denominator = rms.mean.cpu(), torch.sqrt(rms.var.cpu() + 1e-5) + normalized = (value.detach().cpu() - mean) / denominator + expected = original[original_key] + if stage == "physical_before_model_forward": + expected = (expected - mean) / denominator + differences[key] = float((normalized - expected).abs().max()) + differences["target"] = float((targets.cpu() - original["target"]).abs().max()) + return differences, stage + + +@torch.no_grad() +def _validation(trainer: DynamicsTrainer, order: np.ndarray) -> float: + trainer.model.eval() + accumulated = torch.zeros((), device=trainer.cfg.device) + for indices in order: + inputs, targets, _ = encode_batch( + trainer.validation_dataset.get_batch(indices), + trainer.codec, + trainer.cfg.device, + ) + accumulated += normalized_mse( + trainer.model(inputs), targets, trainer.model.output_rms.var + ) + return float(accumulated) / len(order) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--official-dir", type=Path, required=True) + parser.add_argument("--dataset", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--wait", action="store_true") + args = parser.parse_args() + official, output = args.official_dir, args.output_dir + while not (official / "completed.json").exists(): + if not args.wait: + raise ValueError("The original entry must finish its 1000-update run first") + time.sleep(2) + completed = json.loads((official / "completed.json").read_text()) + if completed["steps"] != 1000 or not completed["source_files_unchanged"]: + raise ValueError("Invalid original-entry completion receipt") + output.mkdir(parents=True, exist_ok=False) + torch.set_num_threads(1 if args.device.startswith("cuda") else 8) + torch.set_num_interop_threads(1) + torch.set_float32_matmul_precision("highest") + metadata = load_manifest(args.dataset, verify=True)["metadata"] + convert_legacy_checkpoint( + official / "initial.pt", + output / "initial-export", + metadata, + trust_pickle=True, + ) + model, manifest = load_model(output / "initial-export", args.device) + train = DynamicsWindowDataset(args.dataset, "train", 10, verify=False) + validation = DynamicsWindowDataset(args.dataset, "validation", 10, verify=False) + order = np.load(official / "train-indices.npy", allow_pickle=False) + if order.shape != (1000, 512) or len(np.unique(order)) != order.size: + raise ValueError( + "Expected the first 1000 full, nonreplacement official batches" + ) + # This is the exact rate of the first trained official epoch. Keeping both + # endpoints equal preserves it across this truncated 1000-update experiment. + cfg = DynamicsTrainCfg( + total_steps=1000, + batch_size=512, + learning_rate=0.001 * (1 - 1 / 1000) + 0.0001 / 1000, + final_learning_rate=0.001 * (1 - 1 / 1000) + 0.0001 / 1000, + device=args.device, + seed=0, + validation_interval=1000, + checkpoint_interval=250, + max_grad_norm=1.0, + resident_cache_bytes=4 * 1024**3, + ) + trainer = DynamicsTrainer( + model, + NeRDStateCodec(**metadata["codec"]), + train, + validation, + cfg, + output, + manifest=manifest, + initialize_normalization=False, + sample_order=torch.from_numpy(order.copy()), + ) + # Verify initialization and statistics rather than relying on a shared seed. + fresh, _ = load_model(output / "initial-export", "cpu") + if any( + not torch.equal(v.detach().cpu(), fresh.state_dict()[k]) + for k, v in trainer.model.state_dict().items() + ): + raise RuntimeError("Model initialization or normalization changed") + first_inputs, first_targets, _ = encode_batch( + train.get_batch(order[0]), trainer.codec, args.device + ) + differences, observation_stage = _first_batch_differences( + model, first_inputs, first_targets, official + ) + _json( + output / "provenance.json", + { + "python": platform.python_version(), + "torch": torch.__version__, + "numpy": np.__version__, + "device": args.device, + "dataset_fingerprint": dataset_fingerprint(load_manifest(args.dataset)), + "original_completion_sha256": _hash(official / "completed.json"), + "original_initial_sha256": _hash(official / "initial.pt"), + "sample_order_sha256": _hash(official / "train-indices.npy"), + "validation_orders": { + p.name: _hash(p) + for p in sorted(official.glob("validation-*-indices.npy")) + }, + "cfg": cfg.to_dict(), + "first_batch_feature_target_max_abs_difference": differences, + "feature_comparison": "normalized input features; unnormalized physical target", + "original_first_batch_feature_stage": observation_stage, + "known_differences": [ + "official Python3.8/Torch2.2 CPU versus current Python3.11/Torch2.7 and selected candidate device", + "native COM/spatial velocity transport can introduce FP32 roundoff in Ant", + "the same normalized-MSE validation rule selects epoch-zero or update1000; official also retains its original independent rollout-best checkpoint", + ], + }, + ) + selected = [] + best = float("inf") + for step in (0, 1000): + if step: + with (output / "metrics.jsonl").open("w", buffering=1) as metrics: + while trainer.global_step < 1000: + record = trainer.step() + metrics.write(json.dumps(record, allow_nan=False) + "\n") + if trainer.global_step in (1, 10, 100, 250, 500, 750, 1000): + trainer.save(output / ("step-%04d.pt" % trainer.global_step)) + print("MIGRATED_UPDATE", json.dumps(record), flush=True) + loss = _validation( + trainer, + np.load( + official / ("validation-%04d-indices.npy" % step), allow_pickle=False + ), + ) + selected.append({"step": step, "loss": loss}) + if loss < best: + best = loss + trainer.best_validation = best + trainer.save(output / "best.pt") + trainer.save(output / "last.pt") + for name in ("best", "last"): + run_export(output / (name + ".pt"), output / (name + "-export")) + best_model, _ = load_model(output / "best-export", args.device) + test = DynamicsWindowDataset(args.dataset, "test", 10, verify=False) + result = { + "steps": trainer.global_step, + "validation": selected, + "selected_step": min(selected, key=lambda x: x["loss"])["step"], + "test": evaluate_model(best_model, trainer.codec, test, 512, args.device), + "best_weights_sha256": _hash(output / "best-export/weights.pt"), + "last_weights_sha256": _hash(output / "last-export/weights.pt"), + } + _json(output / "completed.json", result) + train.close() + validation.close() + test.close() + + +if __name__ == "__main__": + main() diff --git a/examples/learning/nerd/_support/_newton.py b/examples/learning/nerd/_support/_newton.py new file mode 100644 index 0000000000..2e684d2ea3 --- /dev/null +++ b/examples/learning/nerd/_support/_newton.py @@ -0,0 +1,423 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Batched analytical Newton teacher for offline trajectory collection.""" + +from __future__ import annotations + +from collections.abc import Mapping +import hashlib +from importlib import metadata +from pathlib import Path +import json +from typing import Any + +import numpy as np +import torch +import warp as wp + +from embodichain_models.dynamics.nerd.codec import plane_contact_features, quat_rotate +from ._reference_assets import ( + _ant_reference_friction, + _import_reference, + reference_profile, +) + +__all__: list[str] = [] + + +class _NewtonTeacher: + """Own an independent Newton model without a DexSim registry or Scene.""" + + def __init__( + self, + robot: str, + *, + device: str = "cpu", + num_envs: int = 1, + solver: Mapping[str, Any] | None = None, + num_substeps: int = 10, + physics_profile: str = "native", + use_cuda_graph: bool = True, + ) -> None: + import newton + + self.profile = reference_profile(robot) + self.robot, self.device, self.num_envs = robot, torch.device(device), num_envs + self.dt, self.num_substeps = self.profile["dt"], num_substeps + solver_cfg = dict( + solver + or ( + { + "solver_type": "featherstone", + "update_mass_matrix_interval": num_substeps, + } + if physics_profile == "warp_reference" + else { + "solver_type": "mujoco_warp", + "use_mujoco_cpu": self.device.type == "cpu", + "separate_worlds": True, + } + ) + ) + kind = solver_cfg.pop("solver_type", "mujoco_warp") + if kind not in {"featherstone", "mujoco_warp"}: + raise ValueError("Offline teachers support featherstone or mujoco_warp") + if kind == "featherstone": + interval = solver_cfg.setdefault("update_mass_matrix_interval", 1) + if type(interval) is not int or interval < 1 or num_substeps % interval: + raise ValueError( + "Featherstone mass-matrix interval must divide num_substeps" + ) + friction = None + if ( + robot == "ant" + and physics_profile == "warp_reference" + and kind == "featherstone" + ): + friction = _ant_reference_friction() + delta = solver_cfg.setdefault( + "friction_smoothing", friction["friction_smoothing"] + ) + if isinstance(delta, bool) or delta != friction["friction_smoothing"]: + raise ValueError( + "Ant warp_reference requires the coupled contact mapping" + ) + # CUDA graphs cannot capture MuJoCo's CPU executor. + capture = ( + use_cuda_graph + and self.device.type == "cuda" + and not solver_cfg.get("use_mujoco_cpu", False) + ) + template = _import_reference(robot, physics_profile) + if friction is not None: + template.shape_material_mu[:] = [ + mu * friction["material_mu_scale"] for mu in template.shape_material_mu + ] + builder = newton.ModelBuilder(gravity=(0.0, 0.0, -9.81)) + ground = None + if robot == "ant": + ground = dict( + ke=1e5, + kd=1e3, + kf=1e3, + mu=friction["ground_mu"] if friction else 0.5, + margin=1e-5, + gap=0.05, + mu_torsional=0.0, + mu_rolling=0.0, + ) + builder.add_ground_plane(cfg=newton.ModelBuilder.ShapeConfig(**ground)) + for _ in range(num_envs): + builder.add_world(template) + self.model = builder.finalize(device=str(self.device)) + self.state, self.next_state = self.model.state(), self.model.state() + self.control = self.model.control() + self.pipeline = newton.CollisionPipeline(self.model) + self.contacts = self.pipeline.contacts() + solver_type = ( + newton.solvers.SolverFeatherstone + if kind == "featherstone" + else newton.solvers.SolverMuJoCo + ) + self.solver = solver_type(self.model, **solver_cfg) + self._mujoco = kind == "mujoco_warp" + qdim, vdim = self.profile["q_dim"], self.profile["qd_dim"] + if ( + self.model.joint_coord_count != num_envs * qdim + or self.model.joint_dof_count != num_envs * vdim + ): + raise ValueError( + "Imported reference state dimensions do not match the codec" + ) + names = [name.rsplit("/", 1)[-1] for name in template.joint_label] + scalar = [ + name + for name, start, end in zip( + names, + template.joint_qd_start, + template.joint_qd_start[1:] + [template.joint_dof_count], + ) + if end - start == 1 + ] + if scalar != self.profile["joint_names"]: + raise ValueError("Imported scalar joint order differs from the reference") + self._root_bodies = torch.tensor( + np.asarray(self.model.articulation_start.numpy())[:-1], + device=self.device, + dtype=torch.long, + ) + child = wp.to_torch(self.model.joint_child).long() + self._root_bodies = child[self._root_bodies] + self.gravity_dir = wp.to_torch(self.model.gravity)[:num_envs].clone() + self.gravity_dir /= self.gravity_dir.norm(dim=-1, keepdim=True) + self._prepare_contacts(newton) + arrays = {} + for name in ( + "body_mass", + "body_inertia", + "body_com", + "body_flags", + "joint_type", + "joint_parent", + "joint_child", + "joint_q_start", + "joint_qd_start", + "joint_X_p", + "joint_X_c", + "joint_axis", + "joint_armature", + "joint_damping", + "joint_limit_lower", + "joint_limit_upper", + "joint_limit_ke", + "joint_limit_kd", + "joint_target_ke", + "joint_target_kd", + "shape_body", + "shape_type", + "shape_scale", + "shape_transform", + "shape_margin", + "shape_gap", + "shape_material_mu", + "shape_material_ke", + "shape_material_kd", + "shape_material_kf", + ): + values = getattr(template, name) + arrays[name] = hashlib.sha256(np.asarray(values).tobytes()).hexdigest() + gravity = wp.to_torch(self.model.gravity)[0].tolist() + signature = hashlib.sha256( + json.dumps( + { + "arrays": arrays, + "ground": ground, + "gravity": gravity, + "profile": physics_profile, + "robot": robot, + }, + sort_keys=True, + ).encode() + ).hexdigest() + origin = { + "kind": "current_newton", + "collector": "direct_newton", + "device": str(self.device), + "solver": {"solver_type": kind, **solver_cfg}, + "num_substeps": num_substeps, + "physics_profile": physics_profile, + "use_cuda_graph": capture, + "versions": { + name: metadata.version(name) + for name in ("newton", "warp-lang", "torch") + }, + } + self.manifest = { + **self.profile, + "physics_profile": physics_profile, + "teacher_profile": "direct_newton_reference_import", + "physics_signature": signature, + "physics_signature_schema": "nerd-direct-newton-builder-v1", + "source_physical_array_sha256": arrays, + "ground_material": ground, + "collector_source_sha256": { + logical: hashlib.sha256( + Path(__file__).with_name(filename).read_bytes() + ).hexdigest() + for logical, filename in ( + ("_newton.py", "_newton.py"), + ("_reference_assets.py", "_reference_assets.py"), + ("generation.py", "_generation.py"), + ) + }, + "gravity": gravity, + "codec": { + **self.profile["codec"], + "base_com": wp.to_torch(self.model.body_com)[ + self._root_bodies[0] + ].tolist(), + }, + "contact_schema": "nerd-static-plane-v1", + "teacher": origin, + "teacher_origin": { + **origin, + "dt": self.dt, + "physics_signature": signature, + "physics_signature_schema": "nerd-direct-newton-builder-v1", + }, + } + if friction is not None: + self.manifest["reference_friction_mapping"] = friction + self.graph = None + self.reset() + if capture: + # Warm up outside capture, then restore inputs. A graph must return + # to the same buffer addresses even for an odd substep count. + self._advance() + self.reset() + with wp.ScopedCapture(device=str(self.device)) as captured: + self._advance() + self.graph = captured.graph + self.reset() + + def _prepare_contacts(self, newton: Any) -> None: + bodies, points, thicknesses = [], [], [] + if self.robot == "ant": + shape_body = self.model.shape_body.numpy() + flags = self.model.shape_flags.numpy() + transforms = self.model.shape_transform.numpy() + scales, types = ( + self.model.shape_scale.numpy(), + self.model.shape_type.numpy(), + ) + margins = self.model.shape_margin.numpy() + body_world = self.model.body_world.numpy() + for world in range(self.num_envs): + wb, wpnt, wt = [], [], [] + for index, body in enumerate(shape_body): + if ( + body < 0 + or body_world[body] != world + or not flags[index] & int(newton.ShapeFlags.COLLIDE_SHAPES) + ): + continue + if types[index] == int(newton.GeoType.SPHERE): + centers = torch.zeros(1, 3) + elif types[index] == int(newton.GeoType.CAPSULE): + centers = torch.tensor( + [ + [0.0, 0.0, scales[index, 1]], + [0.0, 0.0, -scales[index, 1]], + ] + ) + else: + raise ValueError( + "Ant contact features require sphere/capsule shapes" + ) + transform = torch.tensor(transforms[index]) + local = transform[:3] + quat_rotate(transform[3:], centers) + wb.extend([int(body)] * len(centers)) + wpnt.extend(local.tolist()) + wt.extend([float(scales[index, 0] + margins[index])] * len(centers)) + if len(wb) != 25: + raise ValueError("Ant requires 25 ordered contact slots") + bodies.append(wb) + points.append(wpnt) + thicknesses.append(wt) + self._contact_bodies = torch.tensor( + bodies, device=self.device, dtype=torch.long + ) + self._contact_points = torch.tensor( + points, device=self.device, dtype=torch.float32 + ) + self._contact_thickness = torch.tensor( + thicknesses, device=self.device, dtype=torch.float32 + ) + + def snapshot(self) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: + q = wp.to_torch(self.state.joint_q).view(self.num_envs, -1).clone() + qd = wp.to_torch(self.state.joint_qd).view(self.num_envs, -1).clone() + body_q = wp.to_torch(self.state.body_q) + state = {"q": q, "qd": qd, "root_body_q": body_q[self._root_bodies].clone()} + if self.robot == "cartpole": + return state, { + name: q.new_empty((self.num_envs, 0)) + for name in ( + "contact_points_0", + "contact_points_1", + "contact_normals", + "contact_depths", + "contact_thicknesses_0", + "contact_thicknesses_1", + ) + } + poses = body_q[self._contact_bodies] + points = poses[..., :3] + quat_rotate(poses[..., 3:], self._contact_points) + contacts = plane_contact_features( + points, + self._contact_thickness, + q.new_tensor([0.0, 0.0, 1.0]), + q.new_zeros(()), + 1e-5, + ) + contacts["contact_points_0"] = self._contact_points.flatten(1).clone() + contacts.pop("contact_masks", None) + return state, contacts + + def reset( + self, q: torch.Tensor | None = None, qd: torch.Tensor | None = None + ) -> None: + import newton + + if q is None: + q = torch.tensor(self.profile["default_q"], device=self.device).expand( + self.num_envs, -1 + ) + if qd is None: + qd = torch.tensor(self.profile["default_qd"], device=self.device).expand( + self.num_envs, -1 + ) + if q.shape != (self.num_envs, self.profile["q_dim"]) or qd.shape != ( + self.num_envs, + self.profile["qd_dim"], + ): + raise ValueError("Reset tensors must match the complete reference batch") + if self._mujoco: + self.solver.reset(self.state) + for state in (self.state, self.next_state): + wp.to_torch(state.joint_q).copy_(q.reshape(-1)) + wp.to_torch(state.joint_qd).copy_(qd.reshape(-1)) + state.clear_forces() + newton.eval_fk(self.model, state.joint_q, state.joint_qd, state) + self.control.joint_f.zero_() + + def _advance(self) -> None: + for _ in range(self.num_substeps): + self.state.clear_forces() + if not self._mujoco: + self.pipeline.collide(self.state, self.contacts) + self.solver.step( + self.state, + self.next_state, + self.control, + self.contacts, + self.dt / self.num_substeps, + ) + self.state, self.next_state = self.next_state, self.state + if self.num_substeps % 2: + # Publish into the stable input buffer used by captured kernels. + self.next_state.assign(self.state) + self.state, self.next_state = self.next_state, self.state + + def step(self, effort: torch.Tensor) -> None: + if effort.shape != (self.num_envs, self.profile["qd_dim"]): + raise ValueError("Effort must use the full reference generalized layout") + if self.robot == "ant" and bool(torch.any(effort[:, :6] != 0)): + raise ValueError("Ant sampling cannot apply free-base forces") + self.control.joint_f.assign(wp.from_torch(effort.contiguous().reshape(-1))) + if self.graph is None: + self._advance() + else: + wp.capture_launch(self.graph) + + def __enter__(self) -> _NewtonTeacher: + return self + + def __exit__(self, *_args: Any) -> None: + # All arrays and solver buffers belong to this instance, with no World + # singleton or native renderer requiring asynchronous destruction. + self.graph = None diff --git a/examples/learning/nerd/_support/_official_evaluation_data.py b/examples/learning/nerd/_support/_official_evaluation_data.py new file mode 100644 index 0000000000..c63c9b01a5 --- /dev/null +++ b/examples/learning/nerd/_support/_official_evaluation_data.py @@ -0,0 +1,555 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Freeze original-Warp evaluation data in Python3.8/Torch2.2, without the SDK. + +Only pure Torch coordinate/observation modules are loaded under an isolated +namespace for canonical transport. Physics and sampling execute unmodified +upstream code. Memory is bounded by one B x (T+1) batch; failures are never filtered. +""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.util +import json +import math +from pathlib import Path +import platform +import subprocess +import sys +import time +import types +from typing import Any + +__all__: list[str] = [] +_COMMIT = "370e89d2b8ce68aa3578aa7812859c681b7502bc" + + +def _hash(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 _default_transport_root() -> Path: + """Find installed or source codec files without importing the model SDK.""" + relative = Path("embodichain_models/dynamics/nerd") + candidates = [ + Path(__file__).resolve().parents[4] + / "extensions/embodichain_models" + / relative, + *(Path(entry) / relative for entry in sys.path), + ] + for candidate in candidates: + if (candidate / "codec.py").is_file(): + return candidate + raise ValueError("Set transport_root to a frozen NeRD codec directory") + + +def _transport_paths(directory: Path) -> dict[str, Path]: + """Locate unchanged logical transport roles in old bundles or the new layout.""" + return { + "codec": directory / "codec.py", + "reference": ( + directory / "reference.py" + if (directory / "reference.py").is_file() + else Path(__file__).with_name("_task_values.py") + ), + "_evaluation": ( + directory / "_evaluation.py" + if (directory / "_evaluation.py").is_file() + else Path(__file__).with_name("_evaluation.py") + ), + } + + +def _transport(directory: Path) -> tuple[Any, Any]: + package = types.ModuleType("_isolated_nerd_transport") + package.__path__ = [str(directory)] + sys.modules[package.__name__] = package + paths = _transport_paths(directory) + loaded = [] + for name in ("codec", "reference"): + path = paths[name] + spec = importlib.util.spec_from_file_location( + package.__name__ + "." + name, path + ) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + if name == "reference": + # Bind the moved task-math import to the isolated codec rather than + # importing the SDK into a historical reference interpreter. + source = path.read_text().replace( + "from embodichain_models.dynamics.nerd.codec import", + "from .codec import", + ) + exec(compile(source, str(path), "exec"), module.__dict__) + else: + spec.loader.exec_module(module) + loaded.append(module) + return tuple(loaded) + + +def _canonical( + raw: Any, codec: Any, rotation: Any, math_module: Any +) -> tuple[Any, Any]: + import torch + + q = raw[..., : codec.q_dim] + qd = codec.legacy_to_native(q, raw[..., codec.q_dim :]) + if codec.robot == "ant": + q = torch.cat( + ( + math_module.quat_rotate(rotation, q[..., :3]), + math_module.quat_multiply(rotation, q[..., 3:7]), + q[..., 7:], + ), + -1, + ) + qd = torch.cat( + ( + math_module.quat_rotate(rotation, qd[..., :3]), + math_module.quat_rotate(rotation, qd[..., 3:6]), + qd[..., 6:], + ), + -1, + ) + return q, qd + + +def _physics(model: Any) -> dict[str, Any]: + import numpy as np + + result = {} + for name in ( + "body_mass", + "body_inertia", + "body_com", + "joint_armature", + "joint_limit_lower", + "joint_limit_upper", + "joint_limit_ke", + "joint_limit_kd", + "joint_axis_mode", + "joint_target_ke", + "joint_target_kd", + ): + value = getattr(model, name, None) + if value is not None: + result[name] = value.numpy().tolist() + for name in ("mu", "ke", "kd", "kf"): + result["shape_material_" + name] = ( + getattr(model.shape_materials, name).numpy().tolist() + ) + result["shape_thickness"] = model.shape_geo.thickness.numpy().tolist() + result["shape_body"] = model.shape_body.numpy().tolist() + result["ground_plane"] = model.ground_plane.numpy().tolist() + result["gravity"] = np.asarray(model.gravity, dtype=np.float64).tolist() + return result + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", type=Path, required=True) + parser.add_argument("--trust-external-code", action="store_true", required=True) + args = parser.parse_args() + cfg = json.loads(args.config.read_text()) + allowed = { + "source", + "output_dir", + "robot", + "num_envs", + "trajectories", + "horizon", + "seed", + "initial_states", + "motion_mode", + "policy", + "transport_root", + } + if set(cfg) - allowed: + raise ValueError("Unknown original evaluation generation settings") + robot, mode = cfg["robot"], cfg["motion_mode"] + if robot not in {"cartpole", "ant"} or mode not in { + "passive", + "random_effort", + "fixed_policy", + }: + raise ValueError("Unsupported robot or motion mode") + if (mode == "fixed_policy") != ("policy" in cfg): + raise ValueError("Only fixed_policy mode requires an explicit PD policy") + batch, count, horizon, seed = ( + cfg[name] for name in ("num_envs", "trajectories", "horizon", "seed") + ) + if ( + any(type(x) is not int or x < 1 for x in (batch, count, horizon)) + or type(seed) is not int + or seed < 0 + ): + raise ValueError( + "Positive integer batch/count/horizon and nonnegative seed required" + ) + if cfg.get("initial_states", "sample") not in {"default", "sample"}: + raise ValueError("Unsupported initial state sampling") + source = Path(cfg["source"]).resolve() + commit = subprocess.check_output( + ["git", "-C", str(source), "rev-parse", "HEAD"], text=True + ).strip() + changed = subprocess.check_output( + ["git", "-C", str(source), "status", "--porcelain", "--untracked-files=no"], + text=True, + ).strip() + if commit != _COMMIT or changed: + raise ValueError("Original source must be the pinned unchanged checkout") + out = Path(cfg["output_dir"]) + if out.exists(): + raise FileExistsError(out) + sys.path.insert(0, str(source)) + import h5py + import numpy as np + import torch + import warp as wp + from envs.neural_environment import NeuralEnvironment + from generate.trajectory_sampler import TrajectorySampler + from generate.trajectory_sampler_ant import TrajectorySamplerAnt + from utils.commons import ( + JOINT_Q_MIN, + JOINT_Q_MAX, + JOINT_QD_MIN, + JOINT_QD_MAX, + JOINT_ACT_SCALE, + ) + from utils.python_utils import set_random_seed + + version = "1.5.1" if robot == "ant" else "1.8.0" + if ( + platform.python_version() != "3.8.20" + or torch.__version__.split("+")[0] != "2.2.2" + or np.__version__ != "1.24.4" + or wp.__version__ != version + ): + raise ValueError( + "Use supported Python3.8.20/Torch2.2.2/NumPy1.24.4 and robot-pinned Warp" + ) + torch.set_num_threads(1) + transport = ( + Path(cfg["transport_root"]) + if "transport_root" in cfg + else _default_transport_root() + ) + math_module, observations = _transport(transport) + set_random_seed(seed) + started = time.monotonic() + with wp.ScopedDevice("cpu"): + env = NeuralEnvironment( + env_name=robot.capitalize(), + num_envs=batch, + neural_model=None, + neural_integrator_cfg={ + "name": "TransformerNeuralIntegrator", + "num_states_history": 10, + "states_frame": "body" if robot == "ant" else "world", + "anchor_frame_step": "every", + "states_embedding_type": "identical", + "prediction_type": "relative", + "orientation_prediction_parameterization": "quaternion", + }, + default_env_mode="ground-truth", + warp_env_cfg={ + "seed": seed, + "random_reset": False, + "env_offset": (0.0, 0.0, 0.0), + }, + device="cpu", + render=False, + ) + try: + sampler_cls = TrajectorySamplerAnt if robot == "ant" else TrajectorySampler + name = env.robot_name + sampler = sampler_cls( + env, + JOINT_Q_MIN[name], + JOINT_Q_MAX[name], + JOINT_QD_MIN[name], + JOINT_QD_MAX[name], + JOINT_ACT_SCALE[name], + contact_prob=0.0, + ) + limits = sampler.joint_act_scale + codec = math_module.NeRDStateCodec( + robot, base_com=env.model.body_com.numpy()[0].tolist() + ) + rotation = torch.tensor( + [math.sqrt(0.5), 0.0, 0.0, math.sqrt(0.5)] + if robot == "ant" + else [0.5, 0.5, 0.5, 0.5] + ) + default = torch.cat( + ( + wp.to_torch(env.model.joint_q).reshape(batch, -1), + wp.to_torch(env.model.joint_qd).reshape(batch, -1), + ), + -1, + )[0].clone() + target = default[7:15].clone() + policy = cfg.get("policy") + identity = None + if policy is not None: + keys = ( + {"kind", "gains"} if robot == "cartpole" else {"kind", "kp", "kd"} + ) + if policy.get("kind") != "pd" or set(policy) - keys: + raise ValueError("Only declared PD baseline settings supported") + identity = { + "kind": "pd", + "implementation": "embodichain-reference-feedback-v1", + "observation_schema": ( + "cartpole-state-4-v1" + if robot == "cartpole" + else "ant-com-state-27-v1" + ), + "deterministic": True, + "scope": "fixed_baseline_not_paper_policy", + "effort_limits": limits.tolist(), + } + if robot == "cartpole": + gains = policy.get("gains", [5.0, 50.0, 5.0, 20.0]) + if len(gains) != 4 or any( + type(x) not in (int, float) or not math.isfinite(x) + for x in gains + ): + raise ValueError("Four finite PD gains required") + identity["gains"] = gains + gains = torch.tensor(gains) + else: + kp, kd = policy.get("kp", 20.0), policy.get("kd", 1.0) + if any( + type(x) not in (int, float) or not math.isfinite(x) or x < 0 + for x in (kp, kd) + ): + raise ValueError("Finite nonnegative PD gains required") + identity.update(kp=float(kp), kd=float(kd), target=target.tolist()) + identity["sha256"] = hashlib.sha256( + json.dumps(identity, sort_keys=True).encode() + ).hexdigest() + runtime_files = [ + "envs/neural_environment.py", + "envs/abstract_contact_environment.py", + "envs/warp_sim_envs/environment.py", + "envs/warp_sim_envs/env_cartpole.py", + "envs/warp_sim_envs/env_ant.py", + "integrators/integrator_neural.py", + "integrators/integrator_neural_stateful.py", + "integrators/integrator_neural_transformer.py", + "generate/trajectory_sampler.py", + "generate/trajectory_sampler_ant.py", + "generate/simulation_sampler.py", + ] + origin = { + "kind": "official_warp", + "source_commit": commit, + "warp_version": wp.__version__, + "device": "cpu", + "integrator": str(env.integrator_type_gt), + "num_substeps": env.sim_substeps_gt, + "up_axis": int(env.model.up_axis), + "bodies_per_env": env.bodies_per_env, + "gravity": np.asarray(env.model.gravity, dtype=np.float64).tolist(), + } + asset = ( + source + / "envs/warp_sim_envs/assets" + / ("ant.xml" if robot == "ant" else "cartpole_single.urdf") + ) + manifest = { + "schema_version": 1, + "trace_type": "evaluation_including_failures", + "split_seed": 0, + "metadata": { + "robot": robot, + "dt": env.frame_dt, + "codec": { + "robot": robot, + "base_com": list(codec.base_com), + "continuous_q_indices": list(codec.continuous_q_indices), + }, + "world_rotation_xyzw": rotation.tolist(), + "asset_sha256": _hash(asset), + "source_physical_arrays": _physics(env.model), + "sampling": {"batch_size": batch}, + "teacher_origin": origin, + "source_runtime_files": { + name: _hash(source / name) for name in runtime_files + }, + "generation_environment": { + "python": platform.python_version(), + "torch": torch.__version__, + "numpy": np.__version__, + "warp": wp.__version__, + "runtime_imports_embodichain": False, + "runtime_imports_dexsim": False, + "runtime_imports_newton": False, + }, + "canonical_transport_modules": { + name: { + "path": str(_transport_paths(transport)[name]), + "sha256": _hash(_transport_paths(transport)[name]), + } + for name in ("codec", "reference") + }, + "generator_sha256": _hash(Path(__file__)), + "config_sha256": _hash(args.config), + }, + "evaluation": { + "seed": seed, + "trajectories": count, + "horizon": horizon, + "num_envs": batch, + "initial_states": cfg.get("initial_states", "sample"), + "motion_mode": mode, + "policy": identity, + }, + "shards": [], + } + out.mkdir(parents=True) + for start in range(0, count, batch): + n = min(batch, count - start) + samples, commands = [], [] + for index in range(start, start + n): + set_random_seed(seed + index) + raw = default[None].clone() + if cfg.get("initial_states", "sample") == "sample": + sampler.sample_initial_states(1, raw) + command = torch.zeros(horizon, env.joint_act_dim) + if mode == "random_effort": + sampler.sampler.sample(horizon, -limits, limits, command) + samples.append(raw[0]) + commands.append(command) + samples.extend([samples[-1]] * (batch - n)) + commands.extend([commands[-1]] * (batch - n)) + command = torch.stack(commands) + env.reset(initial_states=torch.stack(samples).clone()) + raw_rows, controls = [env.states.clone()], [] + with torch.no_grad(): + for step in range(horizon): + force = command[:, step] + if policy is not None: + q, qd = _canonical(env.states, codec, rotation, math_module) + obs, _, _ = observations.reference_task_values(robot, q, qd) + if robot == "cartpole": + force = torch.zeros_like(force) + force[:, 0] = ( + -(torch.nan_to_num(obs) * gains) + .sum(-1) + .clamp(-limits[0], limits[0]) + ) + else: + force = ( + kp * (target - torch.nan_to_num(q[:, 7:])) + - kd * torch.nan_to_num(qd[:, 6:]) + ).clamp(-limits, limits) + force = torch.where( + torch.isfinite(obs).all(-1)[:, None], force, 0 + ) + force = force.contiguous() + controls.append(force.clone()) + raw_rows.append( + env.step_with_joint_act( + force, env_mode="ground-truth" + ).clone() + ) + raw = torch.stack(raw_rows, 1)[:n] + q, qd = _canonical(raw, codec, rotation, math_module) + raw_force = torch.stack(controls, 1)[:n] + canonical_force = ( + torch.cat((raw_force.new_zeros(n, horizon, 6), raw_force), -1) + if robot == "ant" + else raw_force + ) + file = out / ("teacher-%06d.h5" % len(manifest["shards"])) + rows = [] + with h5py.File(file, "w") as stream: + for row in range(n): + group = stream.create_group(str(row)) + for key, value in { + "states/raw": raw, + "states/q": q, + "states/qd": qd, + "transitions/joint_f": canonical_force, + "transitions/raw_joint_f": raw_force, + }.items(): + group.create_dataset( + key, data=value[row].numpy(), compression="lzf" + ) + rows.append( + { + "id": "%d:%d" % (seed, start + row), + "group": str(row), + "length": horizon, + "split": "test", + "metadata": {"sampling_seed": seed + start + row}, + } + ) + manifest["shards"].append( + {"file": file.name, "sha256": _hash(file), "trajectories": rows} + ) + forbidden = {key.split(".")[0] for key in sys.modules} & { + "embodichain", + "dexsim", + "newton", + } + if forbidden: + raise RuntimeError( + "Unsupported packages imported in original environment" + ) + temporary = out / "manifest.pending.json" + temporary.write_text( + json.dumps(manifest, indent=2, allow_nan=False) + "\n" + ) + temporary.replace(out / "manifest.json") + print( + json.dumps( + { + "event": "committed_batch", + "trajectories": start + n, + "requested": count, + } + ), + flush=True, + ) + finally: + env.close() + print( + json.dumps( + { + "status": "completed", + "scope": "Original supported-environment evaluation teacher data only, no model evaluated", + "trajectories": count, + "horizon": horizon, + "manifest_sha256": _hash(out / "manifest.json"), + "elapsed_seconds": time.monotonic() - started, + } + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/learning/nerd/_support/_official_math_rollout.py b/examples/learning/nerd/_support/_official_math_rollout.py new file mode 100644 index 0000000000..e7a192e738 --- /dev/null +++ b/examples/learning/nerd/_support/_official_math_rollout.py @@ -0,0 +1,371 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Replay reference/math portable weights through explicit upstream math attention. + +Runs both trained weight sets in the same original Python3.8/Torch2.2 runtime. +No ground-truth physics is regenerated. The original integrator evolves its own +state, contacts and Transformer history. Fixed-policy actions use its own state. +Pure tensor transport and metrics load independently of the EmbodiChain package. +""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.util +import json +from pathlib import Path +import platform +import sys +from contextlib import nullcontext +from typing import Any + +from ._official_evaluation_data import _canonical, _hash, _transport, _transport_paths +from ._portable_reference import ( + array_sha256, + backend_flags, + load_portable, + verify_source, +) +from ._historical_physics import _verify_historical_physics + +__all__: list[str] = [] + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source", type=Path, required=True) + parser.add_argument("--portable", type=Path, required=True) + parser.add_argument( + "--scope", choices=("smoke", "checkpoint_accuracy"), required=True + ) + parser.add_argument("--teacher", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--transport-root", type=Path, required=True) + parser.add_argument("--count", type=int, default=16) + parser.add_argument("--trust-external-code", action="store_true", required=True) + args = parser.parse_args() + if args.output.exists() or args.output.with_suffix(".npz").exists(): + raise FileExistsError(args.output) + manifest = json.loads((args.teacher / "manifest.json").read_text()) + metadata, evaluation = manifest["metadata"], manifest["evaluation"] + robot = metadata["robot"] + source_identity = verify_source(args.source, metadata, args.transport_root) + sys.dont_write_bytecode = True + sys.path.insert(0, str(args.source)) + import h5py + import numpy as np + import torch + import warp as wp + from envs.neural_environment import NeuralEnvironment + + expected_warp = "1.5.1" if robot == "ant" else "1.8.0" + if ( + platform.python_version(), + torch.__version__.split("+")[0], + np.__version__, + wp.__version__, + ) != ("3.8.20", "2.2.2", "1.24.4", expected_warp): + raise ValueError("Replay requires the original supported environment") + generation = metadata["generation_environment"] + if ( + generation["python"] != "3.8.20" + or generation["torch"].split("+")[0] != "2.2.2" + or generation["warp"] != expected_warp + ): + raise ValueError("Teacher was not generated in its supported environment") + torch.set_num_threads(4) + torch.set_num_interop_threads(1) + math_module, task_module = _transport(args.transport_root) + metric_spec = importlib.util.spec_from_file_location( + "_isolated_nerd_transport._evaluation", + _transport_paths(args.transport_root)["_evaluation"], + ) + metric_module = importlib.util.module_from_spec(metric_spec) + metric_spec.loader.exec_module(metric_module) + codec = math_module.NeRDStateCodec(**metadata["codec"]) + rotation = torch.tensor(metadata["world_rotation_xyzw"], dtype=torch.float32) + horizon, count = 100, args.count + if not 0 < count <= evaluation["trajectories"] or evaluation["horizon"] < horizon: + raise ValueError("Invalid fixed prefix selection") + rows = [] + identities = [] + verified_shards = {} + for shard in manifest["shards"]: + if len(rows) >= count: + break + path = args.teacher / shard["file"] + actual_hash = _hash(path) + if actual_hash != shard["sha256"]: + raise ValueError("Teacher shard checksum mismatch") + verified_shards[shard["file"]] = actual_hash + with h5py.File(str(path), "r") as stream: + for record in shard["trajectories"]: + if len(rows) >= count: + break + group = stream[record["group"]] + row = { + "raw": group["states/raw"][: horizon + 1], + "q": group["states/q"][: horizon + 1], + "qd": group["states/qd"][: horizon + 1], + "raw_joint_f": group["transitions/raw_joint_f"][:horizon], + "joint_f": group["transitions/joint_f"][:horizon], + } + rows.append(row) + identities.append( + { + **record, + "prefix_sha256": metric_module._teacher_row_sha256( + record, [row["q"], row["qd"], row["joint_f"]] + ), + } + ) + teacher = { + key: torch.from_numpy(np.stack([row[key] for row in rows])) for key in rows[0] + } + tq, tqd = _canonical(teacher["raw"], codec, rotation, math_module) + for actual, key in ((tq, "q"), (tqd, "qd")): + torch.testing.assert_close(actual, teacher[key], atol=0, rtol=0) + torch.manual_seed(0) + model, portable_manifest, portable_identity = load_portable( + args.portable, robot, args.scope + ) + flags_before = backend_flags(torch) + env = NeuralEnvironment( + env_name=robot.capitalize(), + num_envs=count, + neural_model=model, + neural_integrator_cfg={ + "name": "TransformerNeuralIntegrator", + "num_states_history": 10, + "states_frame": "body" if robot == "ant" else "world", + "anchor_frame_step": "every", + "states_embedding_type": "identical", + "prediction_type": "relative", + "orientation_prediction_parameterization": "quaternion", + }, + warp_env_cfg={ + "seed": evaluation["seed"], + "random_reset": False, + "env_offset": (0.0, 0.0, 0.0), + "use_graph_capture": False, + }, + default_env_mode="neural", + device="cpu", + render=False, + ) + physical_identity = { + "teacher": _verify_historical_physics( + env, metadata, {"root": str(args.source)}, wp.__version__ + ), + "training": _verify_historical_physics( + env, portable_manifest, {"root": str(args.source)}, wp.__version__ + ), + } + original_stream = wp.stream_to_torch + wp.stream_to_torch = lambda device: ( + None if wp.get_device(device).is_cpu else original_stream(device) + ) + env.reset(initial_states=teacher["raw"][:, 0].contiguous()) + policy = evaluation["policy"] + + def controls(q: Any, qd: Any) -> Any: + observation, _, _ = task_module.reference_task_values(robot, q, qd) + finite = torch.isfinite(observation).all(-1) + limits = torch.tensor(policy["effort_limits"]) + if robot == "cartpole": + effort = torch.zeros_like(qd) + effort[:, 0] = ( + -(torch.nan_to_num(observation) * torch.tensor(policy["gains"])) + .sum(-1) + .clamp(-limits[0], limits[0]) + ) + else: + effort = ( + policy["kp"] + * (torch.tensor(policy["target"]) - torch.nan_to_num(q[:, 7:])) + - policy["kd"] * torch.nan_to_num(qd[:, 6:]) + ).clamp(-limits, limits) + return torch.where(finite[:, None], effort, 0) + + predicted, efforts = [env.states.clone()], [] + attention = torch.backends.cuda.sdp_kernel( + enable_flash=False, enable_math=True, enable_mem_efficient=False + ) + with torch.no_grad(), attention: + flags_inside = backend_flags(torch) + if flags_inside != {"flash": False, "math": True, "memory_efficient": False}: + raise AssertionError("Math attention context was not applied") + for step in range(horizon): + if evaluation["motion_mode"] == "fixed_policy": + torch.testing.assert_close( + controls(tq[:, step], tqd[:, step]), + teacher["raw_joint_f"][:, step], + atol=0, + rtol=0, + ) + q, qd = _canonical(env.states, codec, rotation, math_module) + effort = controls(q, qd) + else: + effort = teacher["raw_joint_f"][:, step] + effort = effort.contiguous() + efforts.append(effort.clone()) + # Profile the real first step, with no extra advance or reset. + with ( + torch.autograd.profiler.profile(use_cuda=False) + if step == 0 + else nullcontext() + ) as profile: + predicted.append( + env.step_with_joint_act(effort, env_mode="neural").clone() + ) + if step == 0: + attention_events = sorted( + { + event.key + for event in profile.key_averages() + if "scaled_dot_product" in event.key + } + ) + if ( + "aten::_scaled_dot_product_attention_math" not in attention_events + or any( + "flash" in key or "efficient_attention" in key + for key in attention_events + ) + ): + raise AssertionError( + "Actual neural forward did not execute math attention" + ) + flags_after = backend_flags(torch) + if flags_after != flags_before: + raise AssertionError("Attention context leaked outside the runner") + raw = torch.stack(predicted, 1) + q, qd = _canonical(raw, codec, rotation, math_module) + errors = metric_module._error_values( + robot, q[:, 1:], qd[:, 1:], tq[:, 1:], tqd[:, 1:] + ) + finite = torch.isfinite(q).all(-1) & torch.isfinite(qd).all(-1) + teacher_finite = torch.isfinite(tq).all(-1) & torch.isfinite(tqd).all(-1) + valid = (finite & teacher_finite).cummin(1).values[:, 1:] + _, _, teacher_failed = task_module.reference_task_values(robot, tq, tqd) + _, _, neural_failed = task_module.reference_task_values(robot, q, qd) + teacher_failed, neural_failed = ( + teacher_failed.cummax(1).values[:, 1:], + neural_failed.cummax(1).values[:, 1:], + ) + names = ("q_mae", "qd_mae", "base_position_error", "base_orientation_error_rad") + metrics = {} + for h in (1, 10, 100): + mask = valid[:, h - 1] + endpoint = errors[:, h - 1][mask] + prefix = errors[:, :h][mask] + metrics[str(h)] = { + "endpoint": { + name: metric_module._safe_number(endpoint[:, i].mean()) + for i, name in enumerate(names) + }, + "prefix_mean": { + name: metric_module._safe_number(prefix[..., i].mean()) + for i, name in enumerate(names) + }, + "valid_trajectories": int(mask.sum()), + "nonfinite_trajectories": int((~mask).sum()), + "teacher_nonfinite_trajectories": int( + (~teacher_finite.cummin(1).values[:, h]).sum() + ), + "teacher_failed_trajectories": int(teacher_failed[:, h - 1].sum()), + "neural_failed_trajectories": int(neural_failed[:, h - 1].sum()), + } + args.output.parent.mkdir(parents=True, exist_ok=True) + evidence = args.output.with_suffix(".npz") + np.savez_compressed( + str(evidence), + q=q.numpy(), + qd=qd.numpy(), + teacher_q=tq.numpy(), + teacher_qd=tqd.numpy(), + raw=raw.numpy(), + efforts=torch.stack(efforts, 1).numpy(), + teacher_efforts=teacher["raw_joint_f"].numpy(), + errors=errors.numpy(), + valid=valid.numpy(), + ) + result = { + "robot": robot, + "motion_mode": evaluation["motion_mode"], + "trajectories": count, + "horizons": [1, 10, 100], + "seed": evaluation["seed"], + "metrics": metrics, + "portable_model": portable_identity, + "source_identity": source_identity, + "physical_identity": physical_identity, + "attention": { + "flags_before": flags_before, + "flags_inside": flags_inside, + "flags_after": flags_after, + "first_step_profile_events": attention_events, + }, + "input_frame": "original source world (Ant Y-up), states/raw; never canonical q/qd", + "output_metric_frame": "frozen canonical transport only after prediction; does not rotate the model input", + "raw_prefixes": [ + { + "id": identity["id"], + "sampling_seed": identity["metadata"]["sampling_seed"], + "raw_states_sha256": array_sha256(row["raw"]), + "raw_actions_sha256": array_sha256(row["raw_joint_f"]), + } + for identity, row in zip(identities, rows) + ], + "runner_sha256": _hash(Path(__file__)), + "helper_sha256": { + name: _hash(Path(__file__).parent / name) + for name in ( + "_portable_reference.py", + "_historical_physics.py", + "_official_evaluation_data.py", + ) + }, + "teacher_manifest_sha256": _hash(args.teacher / "manifest.json"), + "teacher_shards": verified_shards, + "teacher_prefixes": identities, + "evidence_sha256": _hash(evidence), + "runtime": { + "python": platform.python_version(), + "torch": torch.__version__, + "warp": wp.__version__, + "runtime_imports_embodichain": "embodichain" in sys.modules, + "runtime_imports_dexsim": "dexsim" in sys.modules, + "runtime_imports_newton": "newton" in sys.modules, + "device": "cpu", + "threads": torch.get_num_threads(), + "cuda_visible_devices": __import__("os").environ.get( + "CUDA_VISIBLE_DEVICES" + ), + }, + "teacher_generation_environment": generation, + "evaluation_scope": "original neural integration and common original supported runtime; training-weight comparison; not native Newton deployment", + } + args.output.write_text(json.dumps(result, indent=2, allow_nan=False) + "\n") + print(json.dumps(metrics, indent=2), flush=True) + env.close() + wp.stream_to_torch = original_stream + + +if __name__ == "__main__": + main() diff --git a/examples/learning/nerd/_support/_official_model_reference.py b/examples/learning/nerd/_support/_official_model_reference.py new file mode 100644 index 0000000000..7d04d40807 --- /dev/null +++ b/examples/learning/nerd/_support/_official_model_reference.py @@ -0,0 +1,441 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Generate official data and model predictions without importing EmbodiChain. + +Run this worker with Python 3.8.20 / Torch 2.2.2 and the robot's pinned Warp. +Original source and model objects execute unchanged. The CPU stream adapter +only maps a requested CUDA stream on a CPU device to Torch's no-op stream. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path +import platform +import subprocess +import sys +import xml.etree.ElementTree as ET + +__all__: list[str] = [] +_COMMIT = "370e89d2b8ce68aa3578aa7812859c681b7502bc" + + +def _hash(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source", type=Path, required=True) + parser.add_argument("--checkpoint", type=Path) + parser.add_argument("--robot", choices=["cartpole", "ant"], required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--trust-external-code", action="store_true", required=True) + parser.add_argument("--count", type=int, default=32) + parser.add_argument("--horizon", type=int, default=100) + parser.add_argument("--num-envs", type=int) + parser.add_argument("--seed", type=int, default=500) + parser.add_argument("--generate-only", action="store_true") + parser.add_argument( + "--window-inputs", + type=Path, + help="Evaluate immutable common-frame window tensors without generating new physics data.", + ) + args = parser.parse_args() + if args.checkpoint is None and not args.generate_only: + parser.error("Model reference execution requires --checkpoint") + if args.window_inputs is not None and args.generate_only: + parser.error("Choose generation or reference-window evaluation") + source = args.source.resolve() + commit = subprocess.check_output( + ["git", "-C", str(source), "rev-parse", "HEAD"], text=True + ).strip() + changed = subprocess.check_output( + ["git", "-C", str(source), "status", "--porcelain", "--untracked-files=no"], + text=True, + ).strip() + if commit != _COMMIT or changed: + raise ValueError("Reference source must be the pinned, unchanged checkout") + if args.output.exists(): + raise FileExistsError(args.output) + if args.horizon < 10 or args.count < 1: + raise ValueError("horizon >= 10 and count >= 1 required") + sys.path.insert(0, str(source)) + import h5py + import numpy as np + import torch + import warp as wp + from envs.neural_environment import NeuralEnvironment + from generate.trajectory_sampler import TrajectorySampler + from generate.trajectory_sampler_ant import TrajectorySamplerAnt + from utils.commons import ( + JOINT_Q_MIN, + JOINT_Q_MAX, + JOINT_QD_MIN, + JOINT_QD_MAX, + JOINT_ACT_SCALE, + ) + from utils.python_utils import set_random_seed + + version = "1.8.0" if args.robot == "cartpole" else "1.5.1" + if ( + platform.python_version() != "3.8.20" + or torch.__version__.split("+")[0] != "2.2.2" + or wp.__version__ != version + ): + raise ValueError( + "Reference requires Python 3.8.20, Torch 2.2.2 and Warp " + version + ) + torch.set_num_threads(1) + set_random_seed(args.seed) + model = ( + None + if args.generate_only + else torch.load(args.checkpoint, map_location="cpu")[0].eval() + ) + cfg = { + "name": "TransformerNeuralIntegrator", + "num_states_history": 10, + "states_frame": "body" if args.robot == "ant" else "world", + "anchor_frame_step": "every", + "states_embedding_type": "identical", + "prediction_type": "relative", + "orientation_prediction_parameterization": "quaternion", + } + env = NeuralEnvironment( + env_name=args.robot.capitalize(), + num_envs=args.num_envs or min(args.count, 512), + neural_model=model, + neural_integrator_cfg=cfg, + default_env_mode="ground-truth", + warp_env_cfg={"seed": args.seed, "random_reset": False}, + device="cpu", + render=False, + ) + if args.window_inputs is not None: + # Both trainers must see identical floating-point coordinates. Comparing + # the reference in Y-up against its rotated Z-up inputs additionally + # measures the reference's own finite-precision frame non-invariance. + args.output.mkdir(parents=True) + expected = {} + hashes = {} + with np.load(args.window_inputs, allow_pickle=False) as windows: + for length in [1, 9, 10]: + prefix = str(length) + "__raw__" + batch = {} + for key in windows.files: + if key.startswith(prefix): + value = np.ascontiguousarray(windows[key]) + hashes[key] = hashlib.sha256(value.tobytes()).hexdigest() + batch[key[len(prefix) :]] = torch.from_numpy(value.copy()) + batch["contact_masks"] = env.integrator_neural.get_contact_masks( + batch["contact_depths"], batch["contact_thicknesses"] + ) + env.integrator_neural.process_neural_model_inputs(batch) + target = env.integrator_neural.convert_next_states_to_prediction( + batch["states"], batch["next_states"], env.frame_dt + ) + with torch.no_grad(): + prediction = model( + {key: value.clone() for key, value in batch.items()} + ) + for key in model.low_dim_input_names: + expected[str(length) + "__features__" + key] = batch[key].numpy() + expected[str(length) + "__prediction"] = prediction.numpy() + expected[str(length) + "__target"] = target.numpy() + np.savez(args.output / "official-windows.npz", **expected) + identity = { + "scope": "Original model on byte-identical common-frame comparison inputs", + "source_commit": commit, + "checkpoint_sha256": _hash(args.checkpoint), + "python": platform.python_version(), + "torch": torch.__version__, + "warp": wp.__version__, + "numpy": np.__version__, + "window_inputs_sha256": _hash(args.window_inputs), + "window_field_sha256": hashes, + "predictions_sha256": _hash(args.output / "official-windows.npz"), + "runtime_imports_embodichain": any( + k.split(".")[0] == "embodichain" for k in sys.modules + ), + } + (args.output / "reference-identity.json").write_text( + json.dumps(identity, indent=2) + ) + print(json.dumps(identity, indent=2)) + return + name = env.robot_name + sampler_cls = TrajectorySamplerAnt if args.robot == "ant" else TrajectorySampler + sampler = sampler_cls( + env, + JOINT_Q_MIN[name], + JOINT_Q_MAX[name], + JOINT_QD_MIN[name], + JOINT_QD_MAX[name], + JOINT_ACT_SCALE[name], + contact_prob=0.0, + ) + rows = sampler.sample_trajectories_joint_act_mode( + args.count * args.horizon, + args.horizon, + passive=False, + initial_states_source="sample", + ) + data = {} + for key, value in rows.items(): + if isinstance(value, dict): + data.update( + { + k: v.detach() + .cpu() + .numpy() + .reshape( + args.horizon, + -1, + v.shape[-1] if v.ndim == 3 else v.shape[-2] * v.shape[-1], + ) + for k, v in value.items() + } + ) + else: + data[key] = value.detach().cpu().numpy() + args.output.mkdir(parents=True) + with h5py.File(args.output / "test.raw.h5", "w") as h5: + group = h5.create_group("data") + group.attrs["env"] = name + group.attrs["mode"] = "trajectory" + group.attrs["total_trajectories"] = data["states"].shape[1] + group.attrs["total_transitions"] = ( + data["states"].shape[0] * data["states"].shape[1] + ) + for key, value in data.items(): + group.create_dataset(key, data=value) + if args.generate_only: + physical = {} + for key in ( + "body_mass", + "body_inertia", + "body_com", + "joint_armature", + "joint_limit_lower", + "joint_limit_upper", + "joint_limit_ke", + "joint_limit_kd", + "joint_axis_mode", + "joint_target_ke", + "joint_target_kd", + ): + value = getattr(env.model, key, None) + if value is not None: + physical[key] = value.numpy().tolist() + for key in ("mu", "ke", "kd", "kf"): + physical["shape_material_" + key] = ( + getattr(env.model.shape_materials, key).numpy().tolist() + ) + physical["shape_thickness"] = env.model.shape_geo.thickness.numpy().tolist() + physical["shape_body"] = env.model.shape_body.numpy().tolist() + physical["ground_plane"] = env.model.ground_plane.numpy().tolist() + physical["gravity"] = np.asarray(env.model.gravity, dtype=np.float64).tolist() + asset = ( + source + / "envs/warp_sim_envs/assets" + / ("ant.xml" if args.robot == "ant" else "cartpole_single.urdf") + ) + generator_files = [ + "generate/trajectory_sampler.py", + "generate/simulation_sampler.py", + "utils/commons.py", + "utils/env_utils.py", + "utils/warp_utils.py", + "envs/neural_environment.py", + "envs/abstract_contact.py", + "envs/abstract_contact_environment.py", + "envs/warp_sim_envs/environment.py", + "envs/warp_sim_envs/env_" + args.robot + ".py", + ] + if args.robot == "ant": + generator_files.append("generate/trajectory_sampler_ant.py") + file_hashes = {name: _hash(source / name) for name in generator_files} + metadata = { + "robot": args.robot, + "dt": env.frame_dt, + "q_dim": env.dof_q_per_env, + "qd_dim": env.dof_qd_per_env, + "codec": { + "robot": args.robot, + "base_com": env.model.body_com.numpy()[0].tolist(), + "continuous_q_indices": [1] if args.robot == "cartpole" else [], + }, + "joint_names": [ + joint.attrib["name"] + for joint in ET.parse(asset).getroot().iter("joint") + if "name" in joint.attrib and joint.attrib.get("type") != "free" + ], + "asset_sha256": _hash(asset), + "source_physical_arrays": physical, + "physics_contract_version": "official-warp-physical-v2", + "teacher_origin": { + "kind": "official_warp", + "source_commit": commit, + "warp_version": wp.__version__, + "integrator": str(env.integrator_type_gt), + "num_substeps": env.sim_substeps_gt, + "up_axis": int(env.model.up_axis), + "gravity": physical["gravity"], + "gravity_magnitude": float(np.linalg.norm(physical["gravity"])), + "physics_contract_version": "official-warp-physical-v2", + "device": "cpu", + "bodies_per_env": env.env.bodies_per_env, + }, + "source_generator_files": file_hashes, + "source_generator_sha256": hashlib.sha256( + json.dumps(file_hashes, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest(), + "generation_adapter_sha256": _hash(Path(__file__)), + "sampling": { + "seed": args.seed, + "batch_size": env.num_envs, + "initial_states_source": "sample", + "passive": False, + }, + } + (args.output / "source-physics.json").write_text( + json.dumps(metadata, indent=2, allow_nan=False) + ) + identity = { + "scope": "Official analytical teacher generation; no model supplied or evaluated", + "source_commit": commit, + "python": platform.python_version(), + "torch": torch.__version__, + "warp": wp.__version__, + "numpy": np.__version__, + "robot": args.robot, + "trajectories": data["states"].shape[1], + "horizon": args.horizon, + "num_envs": env.num_envs, + "seed": args.seed, + "raw_sha256": _hash(args.output / "test.raw.h5"), + "metadata_sha256": _hash(args.output / "source-physics.json"), + "runtime_imports_embodichain": any( + key.split(".")[0] == "embodichain" for key in sys.modules + ), + } + (args.output / "reference-identity.json").write_text( + json.dumps(identity, indent=2) + ) + print(json.dumps(identity, indent=2)) + return + # Preserve the original uninitialized gravity feature in raw evidence. + # Model comparison uses the measured physical direction for both models. + direction = np.asarray(env.model.gravity, dtype=np.float32) + direction /= np.linalg.norm(direction) + corrected = {key: value.copy() for key, value in data.items()} + corrected["gravity_dir"][:] = direction + selected = np.random.default_rng(17) + windows = {} + integrator = env.integrator_neural + for length in [1, 9, 10]: + ids = np.stack( + ( + selected.integers(0, data["states"].shape[1], 64), + selected.integers(0, args.horizon - length + 1, 64), + ), + axis=-1, + ) + batch = { + key: torch.from_numpy( + np.stack( + [value[start : start + length, episode] for episode, start in ids] + ) + ) + for key, value in corrected.items() + } + batch["contact_masks"] = integrator.get_contact_masks( + batch["contact_depths"], batch["contact_thicknesses"] + ) + integrator.process_neural_model_inputs(batch) + target = integrator.convert_next_states_to_prediction( + batch["states"], batch["next_states"], env.frame_dt + ) + with torch.no_grad(): + # Official forward replaces dictionary entries with normalized + # tensors. Preserve the physical inputs as independent fixtures. + prediction = model({key: value.clone() for key, value in batch.items()}) + windows[str(length) + "__indices"] = ids + for key in model.low_dim_input_names: + windows[str(length) + "__features__" + key] = batch[key].numpy() + windows[str(length) + "__prediction"] = prediction.numpy() + windows[str(length) + "__target"] = target.numpy() + np.savez(args.output / "official-windows.npz", **windows) + # Exercise the untouched original runtime, including history rollover. + original_stream = wp.stream_to_torch + wp.stream_to_torch = lambda device: ( + None if not wp.get_device(device).is_cuda else original_stream(device) + ) + try: + env.set_env_mode("neural") + env.reset(initial_states=torch.from_numpy(data["states"][0].copy())) + states = [env.states.detach().cpu().numpy().copy()] + for step in range(args.horizon): + state = env.step_with_joint_act( + torch.from_numpy(data["joint_acts"][step].copy()), env_mode="neural" + ) + states.append(state.detach().cpu().numpy().copy()) + np.savez( + args.output / "official-rollout.npz", + states=np.stack(states), + joint_acts=data["joint_acts"], + ) + finally: + wp.stream_to_torch = original_stream + identity = { + "source_commit": commit, + "robot": args.robot, + "python": platform.python_version(), + "torch": torch.__version__, + "warp": wp.__version__, + "numpy": np.__version__, + "device": "cpu", + "dtype": "float32", + "count": args.count, + "horizon": args.horizon, + "checkpoint_sha256": _hash(args.checkpoint), + "seed": 500, + "runtime_imports_embodichain": any( + k.split(".")[0] == "embodichain" for k in sys.modules + ), + "cpu_stream_adapter": "CPU stream request mapped to None; numerical source unchanged", + "gravity_transport": { + "direction": direction.tolist(), + "raw_preserved": True, + "reason": "Official sampler leaves non-up-axis gravity components uninitialized", + }, + "files": {p.name: _hash(p) for p in args.output.iterdir() if p.is_file()}, + "model_inputs": list(model.low_dim_input_names), + "state_frame": cfg["states_frame"], + "body_com": env.model.body_com.numpy()[0].tolist(), + } + (args.output / "reference-identity.json").write_text( + json.dumps(identity, indent=2, allow_nan=False) + ) + print(json.dumps(identity, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/examples/learning/nerd/_support/_official_small_budget.py b/examples/learning/nerd/_support/_official_small_budget.py new file mode 100644 index 0000000000..25b7af7b3e --- /dev/null +++ b/examples/learning/nerd/_support/_official_small_budget.py @@ -0,0 +1,345 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Observe the original NeRD training entry in its supported Python environment. + +The source constructor, statistics, samplers, losses, validation, simulation +evaluation and checkpoint selectors execute unchanged. Observers capture actual +consumed indices, initial normalizers/weights and minibatch metrics. A budget +guard stops before the second trained epoch, after the first epoch's original +validation/evaluation/checkpoint selection has finished. This file supports +Python 3.8 and never imports EmbodiChain or the current SDK. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +from pathlib import Path +import platform +import runpy +import subprocess +import sys +import time +from typing import Any + +__all__: list[str] = [] + + +class _BudgetReached(Exception): + pass + + +def _hash(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 _json(path: Path, value: Any) -> None: + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(value, indent=2, allow_nan=False) + "\n") + temporary.replace(path) + + +def _save_first_batch(output: Path, data: dict[str, Any]) -> None: + """Capture physical features before the original model normalizes its input.""" + import torch + + path = output / "first-preprocessed-batch.pt" + torch.save( + { + key: value.detach().cpu().clone() + for key, value in data.items() + if torch.is_tensor(value) + }, + path, + ) + _json( + output / "first-preprocessed-batch.json", + { + "format_version": 1, + "feature_stage": "physical_before_model_forward", + "tensor_sha256": _hash(path), + }, + ) + + +def _observed_loss( + trainer: Any, + data: dict[str, Any], + train: bool, + original_loss: Any, + state: dict[str, Any], + output: Path, +) -> Any: + """Observe the first training batch before an in-place model forward.""" + if train and state["steps"] == 0: + _save_first_batch(output, data) + result = original_loss(trainer, data, train) + if train: + state["last_loss"] = float(result[0].detach()) + return result + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", type=Path, required=True) + parser.add_argument("--trust-external-code", action="store_true", required=True) + args = parser.parse_args() + plan = json.loads(args.config.read_text()) + source, output = Path(plan["source"]), Path(plan["output_dir"]) + output.mkdir(parents=True, exist_ok=False) + import numpy as np + import torch + import warp as wp + + expected_warp = "1.5.1" if plan["robot"] == "ant" else "1.8.0" + if ( + platform.python_version() != "3.8.20" + or torch.__version__.split("+")[0] != "2.2.2" + or np.__version__ != "1.24.4" + or wp.__version__ != expected_warp + ): + raise RuntimeError( + "Official baseline requires the declared supported environment" + ) + commit = subprocess.check_output( + ["git", "-C", str(source), "rev-parse", "HEAD"], text=True + ).strip() + if commit != "370e89d2b8ce68aa3578aa7812859c681b7502bc": + raise ValueError("Unexpected official source commit") + if subprocess.check_output(["git", "-C", str(source), "diff", "--name-only"]): + raise ValueError("Official tracked source must be unmodified") + torch.set_num_threads(plan.get("cpu_threads", 8)) + torch.set_num_interop_threads(1) + torch.set_float32_matmul_precision("highest") + # Upstream simulation unconditionally asks for a CUDA stream even on CPU. + # CPU work is synchronous; the existing torch.cuda.stream(None) is a no-op. + # Keep its numerical integrator and evaluation/selection code unchanged. + original_stream_to_torch = wp.stream_to_torch + + def stream_to_torch(device: Any) -> Any: + if wp.get_device(device).is_cpu: + return None + return original_stream_to_torch(device) + + wp.stream_to_torch = stream_to_torch + sys.path.insert(0, str(source)) + sys.path.insert(0, str(source / "train")) + from torch.utils.data import BatchSampler, DataLoader + from algorithms.vanilla_trainer import VanillaTrainer + + state: dict[str, Any] = { + "steps": 0, + "stage": "initialization", + "trainer": None, + "last_loss": None, + "validation": [], + "saves": [], + "started": time.time(), + } + metrics = (output / "metrics.jsonl").open("w", buffering=1) + original_batch_iter = BatchSampler.__iter__ + original_loader_iter = DataLoader.__iter__ + + def batch_iter(sampler: Any) -> Any: + for indices in original_batch_iter(sampler): + record = getattr(sampler, "_nerd_observation", None) + if record is not None: + record.append(indices) + yield indices + + def loader_iter(loader: Any) -> Any: + # The original iterator still owns workers, RNG, shuffling and prefetch. + # Pair observations at consumption, rather than counting prefetched rows. + recorded: list[Any] = [] + loader.batch_sampler._nerd_observation = recorded + iterator = original_loader_iter(loader) + + def observe() -> Any: + consumed = 0 + while True: + try: + data = next(iterator) + except StopIteration: + return + indices = recorded[consumed] + consumed += 1 + if state["stage"] != "initialization": + state["indices"].append(np.asarray(indices, dtype=np.int64)) + yield data + + # Create the original iterator eagerly: worker prefetch can consume the + # global sampling RNG before the validation iterators are constructed. + return observe() + + BatchSampler.__iter__ = batch_iter + DataLoader.__iter__ = loader_iter + + original_init = VanillaTrainer.__init__ + original_epoch = VanillaTrainer.one_epoch + original_loss = VanillaTrainer.compute_loss + original_save = VanillaTrainer.save_model + + def init(trainer: Any, *positional: Any, **keywords: Any) -> None: + original_init(trainer, *positional, **keywords) + state["trainer"] = trainer + torch.save( + [trainer.neural_model, trainer.neural_env.robot_name], output / "initial.pt" + ) + torch.save( + { + key: {"mean": rms.mean, "var": rms.var, "count": float(rms.count)} + for key, rms in trainer.dataset_rms.items() + }, + output / "original-statistics.pt", + ) + original_step = trainer.optimizer.step + + def step(*pos: Any, **kw: Any) -> Any: + result = original_step(*pos, **kw) + state["steps"] += 1 + record = { + "step": state["steps"], + "loss": state["last_loss"], + "learning_rate": trainer.optimizer.param_groups[0]["lr"], + "elapsed_seconds": time.time() - state["started"], + } + metrics.write(json.dumps(record, allow_nan=False) + "\n") + if state["steps"] in (1, 10, 100, 250, 500, 750, plan["steps"]): + torch.save( + [trainer.neural_model, trainer.neural_env.robot_name], + output / ("step-%04d.pt" % state["steps"]), + ) + print("OBSERVED_UPDATE", json.dumps(record), flush=True) + return result + + trainer.optimizer.step = step + _json( + output / "initial-ready.json", + {"initial_sha256": _hash(output / "initial.pt")}, + ) + + def loss(trainer: Any, data: Any, train: bool) -> Any: + return _observed_loss(trainer, data, train, original_loss, state, output) + + def epoch(trainer: Any, *positional: Any, **keywords: Any) -> Any: + train = keywords.get("train", positional[0] if positional else None) + if train and state["steps"] >= plan["steps"]: + raise _BudgetReached() + state["stage"] = "train" if train else "validation" + state["indices"] = [] + result = original_epoch(trainer, *positional, **keywords) + label = "train" if train else "validation-%04d" % state["steps"] + np.save(str(output / (label + "-indices.npy")), np.stack(state["indices"])) + if not train: + state["validation"].append( + {"step": state["steps"], "loss": result[0], "itemized": result[1]} + ) + _json(output / "validation.json", state["validation"]) + state["stage"] = "initialization" + return result + + def save(trainer: Any, filename: Any = None) -> None: + original_save(trainer, filename) + state["saves"].append({"step": state["steps"], "filename": filename}) + _json(output / "checkpoint-selection.json", state["saves"]) + + VanillaTrainer.__init__ = init + VanillaTrainer.compute_loss = loss + VanillaTrainer.one_epoch = epoch + VanillaTrainer.save_model = save + entry = source / "train" / "train.py" + original_files = { + str(path.relative_to(source)): _hash(path) + for path in ( + entry, + source / "algorithms" / "vanilla_trainer.py", + source / "algorithms" / "sequence_model_trainer.py", + source / "utils" / "datasets.py", + source / "utils" / "running_mean_std.py", + source / "models" / "models.py", + source / "models" / "model_transformer.py", + ) + } + _json( + output / "provenance.json", + { + "plan": plan, + "source_commit": commit, + "original_files": original_files, + "observer_sha256": _hash(Path(__file__)), + "python": platform.python_version(), + "torch": torch.__version__, + "numpy": np.__version__, + "warp": wp.__version__, + "device": "cpu", + "entry": str(entry), + "cpu_compatibility": "wp.stream_to_torch(cpu) returns None; upstream torch.cuda.stream(None) remains a no-op; no numerical or selector changes", + "budget_adaptation": "1000 original updates in first trained epoch; preserve 1000-epoch LR denominator and epoch-zero selection; stop before next trained epoch", + }, + ) + sys.argv = [ + str(entry), + "--cfg", + plan["training_config"], + "--logdir", + str(output / "original-log"), + "--no-time-stamp", + "--device", + "cpu", + "--seed", + str(plan["seed"]), + "--save-interval", + "1", + ] + previous_cwd = Path.cwd() + os.chdir(source / "train") + try: + runpy.run_path(str(entry), run_name="__main__") + except _BudgetReached: + trainer = state["trainer"] + trainer.save_model("final_model") + torch.save(trainer.optimizer.state_dict(), output / "final-optimizer.pt") + trainer.logger.finish() + finally: + metrics.close() + os.chdir(previous_cwd) + if state["steps"] != plan["steps"]: + raise RuntimeError( + "Original entry did not reach the requested optimizer budget" + ) + _json( + output / "completed.json", + { + "steps": state["steps"], + "elapsed_seconds": time.time() - state["started"], + "validation": state["validation"], + "checkpoint_selection": state["saves"], + "source_files_unchanged": all( + _hash(source / p) == h for p, h in original_files.items() + ), + }, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/learning/nerd/_support/_official_small_budget_evaluate.py b/examples/learning/nerd/_support/_official_small_budget_evaluate.py new file mode 100644 index 0000000000..8ab2ac9989 --- /dev/null +++ b/examples/learning/nerd/_support/_official_small_budget_evaluate.py @@ -0,0 +1,166 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Evaluate original checkpoints on fixed held-out data in Python3.8/Torch2.2.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path +import platform +import sys +from typing import Any + +__all__: list[str] = [] + + +def _hash(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 main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source", type=Path, required=True) + parser.add_argument("--official-dir", type=Path, required=True) + parser.add_argument("--robot", choices=("cartpole", "ant"), required=True) + parser.add_argument("--test-data", type=Path, required=True) + parser.add_argument("--validation-data", type=Path, required=True) + parser.add_argument("--trust-external-code", action="store_true", required=True) + args = parser.parse_args() + sys.path.insert(0, str(args.source)) + import numpy as np + import torch + import warp as wp + from torch.utils.data import DataLoader, Subset + from algorithms.vanilla_trainer import VanillaTrainer + from envs.neural_environment import NeuralEnvironment + from utils.datasets import TrajectoryDataset + + if (platform.python_version(), torch.__version__.split("+")[0], wp.__version__) != ( + "3.8.20", + "2.2.2", + "1.5.1" if args.robot == "ant" else "1.8.0", + ): + raise ValueError("Unexpected original evaluation environment") + torch.set_num_threads(8) + torch.set_num_interop_threads(1) + model_path = ( + args.official_dir / "original-log/nn/best_valid_exp_trajectory_model.pt" + ) + model = torch.load(model_path, map_location="cpu")[0].eval() + env = NeuralEnvironment( + env_name=args.robot.capitalize(), + num_envs=512, + neural_model=model, + neural_integrator_cfg={ + "name": "TransformerNeuralIntegrator", + "num_states_history": 10, + "states_frame": "body" if args.robot == "ant" else "world", + "anchor_frame_step": "every", + "states_embedding_type": "identical", + "prediction_type": "relative", + "orientation_prediction_parameterization": "quaternion", + }, + device="cpu", + render=False, + ) + trainer = VanillaTrainer.__new__(VanillaTrainer) + trainer.device, trainer.batch_size = "cpu", 512 + trainer.neural_env, trainer.neural_integrator = env, env.integrator_neural + trainer.neural_model = model + + @torch.no_grad() + def measure(dataset: Any, golden: bool = False) -> dict[str, Any]: + accumulated = torch.zeros(()) + itemized: dict[str, Any] = {} + loader = DataLoader( + dataset, batch_size=512, shuffle=False, drop_last=True, num_workers=4 + ) + for index, data in enumerate(loader): + data = trainer.preprocess_data_batch(data) + if golden and index == 0: + inputs = { + key: data[key].clone() + for key in trainer.neural_model.low_dim_input_names + } + with torch.no_grad(): + prediction = trainer.neural_model( + {key: value.clone() for key, value in inputs.items()} + ) + torch.save( + { + "inputs": inputs, + "prediction": prediction, + "target": data["target"].clone(), + }, + args.official_dir / "test-first-batch-golden.pt", + ) + loss, items = trainer.compute_loss(data, False) + accumulated += loss + for key, value in items.items(): + itemized[key] = itemized.get(key, 0) + value + return { + "normalized_mse": float(accumulated) / len(loader), + "itemized": { + key: float(value) / len(loader) for key, value in itemized.items() + }, + "windows": len(loader) * 512, + "valid_positions": len(loader) * 5120, + } + + test = TrajectoryDataset(str(args.test_data), sample_sequence_length=10) + test_result = measure(test, golden=True) + validation = TrajectoryDataset(str(args.validation_data), sample_sequence_length=10) + fixed_indices = ( + np.load(str(args.official_dir / "validation-0000-indices.npy"))[:8] + .reshape(-1) + .tolist() + ) + fixed = Subset(validation, fixed_indices) + curves = [] + for step in (0, 250, 500, 750, 1000): + path = args.official_dir / ( + "initial.pt" if step == 0 else "step-%04d.pt" % step + ) + trainer.neural_model = torch.load(path, map_location="cpu")[0].eval() + curves.append({"step": step, **measure(fixed)}) + result = { + "python": platform.python_version(), + "torch": torch.__version__, + "warp": wp.__version__, + "selected_model_sha256": _hash(model_path), + "test_data_sha256": _hash(args.test_data), + "test": test_result, + "fixed_validation_curve": curves, + "curve_indices_sha256": hashlib.sha256( + np.asarray(fixed_indices, dtype=np.int64).tobytes() + ).hexdigest(), + "golden_sha256": _hash(args.official_dir / "test-first-batch-golden.pt"), + } + path = args.official_dir / "heldout-evaluation.json" + path.write_text(json.dumps(result, indent=2, allow_nan=False) + "\n") + print(json.dumps(result, indent=2), flush=True) + env.close() + + +if __name__ == "__main__": + main() diff --git a/examples/learning/nerd/_support/_official_small_budget_rollout.py b/examples/learning/nerd/_support/_official_small_budget_rollout.py new file mode 100644 index 0000000000..1e85dfbe54 --- /dev/null +++ b/examples/learning/nerd/_support/_official_small_budget_rollout.py @@ -0,0 +1,276 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Replay fixed supported-environment teacher prefixes through original NeRD. + +Runs both trained weight sets in the same original Python3.8/Torch2.2 runtime. +No ground-truth physics is regenerated. The original integrator evolves its own +state, contacts and Transformer history. Fixed-policy actions use its own state. +Pure tensor transport and metrics load independently of the EmbodiChain package. +""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.util +import json +from pathlib import Path +import platform +import sys +from typing import Any + +from ._official_evaluation_data import _canonical, _hash, _transport, _transport_paths + +__all__: list[str] = [] + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source", type=Path, required=True) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--teacher", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--transport-root", type=Path, required=True) + parser.add_argument("--count", type=int, default=16) + parser.add_argument("--trust-external-code", action="store_true", required=True) + args = parser.parse_args() + if args.output.exists(): + raise FileExistsError(args.output) + manifest = json.loads((args.teacher / "manifest.json").read_text()) + metadata, evaluation = manifest["metadata"], manifest["evaluation"] + robot = metadata["robot"] + sys.path.insert(0, str(args.source)) + import h5py + import numpy as np + import torch + import warp as wp + from envs.neural_environment import NeuralEnvironment + + expected_warp = "1.5.1" if robot == "ant" else "1.8.0" + if ( + platform.python_version(), + torch.__version__.split("+")[0], + np.__version__, + wp.__version__, + ) != ("3.8.20", "2.2.2", "1.24.4", expected_warp): + raise ValueError("Replay requires the original supported environment") + generation = metadata["generation_environment"] + if ( + generation["python"] != "3.8.20" + or generation["torch"].split("+")[0] != "2.2.2" + or generation["warp"] != expected_warp + ): + raise ValueError("Teacher was not generated in its supported environment") + torch.set_num_threads(4) + torch.set_num_interop_threads(1) + math_module, task_module = _transport(args.transport_root) + metric_spec = importlib.util.spec_from_file_location( + "_isolated_nerd_transport._evaluation", + _transport_paths(args.transport_root)["_evaluation"], + ) + metric_module = importlib.util.module_from_spec(metric_spec) + metric_spec.loader.exec_module(metric_module) + codec = math_module.NeRDStateCodec(**metadata["codec"]) + rotation = torch.tensor(metadata["world_rotation_xyzw"], dtype=torch.float32) + horizon, count = 100, args.count + if not 0 < count <= evaluation["trajectories"] or evaluation["horizon"] < horizon: + raise ValueError("Invalid fixed prefix selection") + rows = [] + identities = [] + verified_shards = {} + for shard in manifest["shards"]: + if len(rows) >= count: + break + path = args.teacher / shard["file"] + actual_hash = _hash(path) + if actual_hash != shard["sha256"]: + raise ValueError("Teacher shard checksum mismatch") + verified_shards[shard["file"]] = actual_hash + with h5py.File(str(path), "r") as stream: + for record in shard["trajectories"]: + if len(rows) >= count: + break + group = stream[record["group"]] + row = { + "raw": group["states/raw"][: horizon + 1], + "q": group["states/q"][: horizon + 1], + "qd": group["states/qd"][: horizon + 1], + "raw_joint_f": group["transitions/raw_joint_f"][:horizon], + "joint_f": group["transitions/joint_f"][:horizon], + } + rows.append(row) + identities.append( + { + **record, + "prefix_sha256": metric_module._teacher_row_sha256( + record, [row["q"], row["qd"], row["joint_f"]] + ), + } + ) + teacher = { + key: torch.from_numpy(np.stack([row[key] for row in rows])) for key in rows[0] + } + tq, tqd = _canonical(teacher["raw"], codec, rotation, math_module) + for actual, key in ((tq, "q"), (tqd, "qd")): + torch.testing.assert_close(actual, teacher[key], atol=0, rtol=0) + model = torch.load(args.checkpoint, map_location="cpu")[0].eval() + env = NeuralEnvironment( + env_name=robot.capitalize(), + num_envs=count, + neural_model=model, + neural_integrator_cfg={ + "name": "TransformerNeuralIntegrator", + "num_states_history": 10, + "states_frame": "body" if robot == "ant" else "world", + "anchor_frame_step": "every", + "states_embedding_type": "identical", + "prediction_type": "relative", + "orientation_prediction_parameterization": "quaternion", + }, + warp_env_cfg={ + "seed": evaluation["seed"], + "random_reset": False, + "env_offset": (0.0, 0.0, 0.0), + "use_graph_capture": False, + }, + default_env_mode="neural", + device="cpu", + render=False, + ) + original_stream = wp.stream_to_torch + wp.stream_to_torch = lambda device: ( + None if wp.get_device(device).is_cpu else original_stream(device) + ) + env.reset(initial_states=teacher["raw"][:, 0].contiguous()) + policy = evaluation["policy"] + + def controls(q: Any, qd: Any) -> Any: + observation, _, _ = task_module.reference_task_values(robot, q, qd) + finite = torch.isfinite(observation).all(-1) + limits = torch.tensor(policy["effort_limits"]) + if robot == "cartpole": + effort = torch.zeros_like(qd) + effort[:, 0] = ( + -(torch.nan_to_num(observation) * torch.tensor(policy["gains"])) + .sum(-1) + .clamp(-limits[0], limits[0]) + ) + else: + effort = ( + policy["kp"] + * (torch.tensor(policy["target"]) - torch.nan_to_num(q[:, 7:])) + - policy["kd"] * torch.nan_to_num(qd[:, 6:]) + ).clamp(-limits, limits) + return torch.where(finite[:, None], effort, 0) + + predicted, efforts = [env.states.clone()], [] + with torch.no_grad(): + for step in range(horizon): + if evaluation["motion_mode"] == "fixed_policy": + torch.testing.assert_close( + controls(tq[:, step], tqd[:, step]), + teacher["raw_joint_f"][:, step], + atol=0, + rtol=0, + ) + q, qd = _canonical(env.states, codec, rotation, math_module) + effort = controls(q, qd) + else: + effort = teacher["raw_joint_f"][:, step] + effort = effort.contiguous() + efforts.append(effort.clone()) + predicted.append(env.step_with_joint_act(effort, env_mode="neural").clone()) + raw = torch.stack(predicted, 1) + q, qd = _canonical(raw, codec, rotation, math_module) + errors = metric_module._error_values( + robot, q[:, 1:], qd[:, 1:], tq[:, 1:], tqd[:, 1:] + ) + finite = torch.isfinite(q).all(-1) & torch.isfinite(qd).all(-1) + teacher_finite = torch.isfinite(tq).all(-1) & torch.isfinite(tqd).all(-1) + valid = (finite & teacher_finite).cummin(1).values[:, 1:] + _, _, teacher_failed = task_module.reference_task_values(robot, tq, tqd) + _, _, neural_failed = task_module.reference_task_values(robot, q, qd) + teacher_failed, neural_failed = ( + teacher_failed.cummax(1).values[:, 1:], + neural_failed.cummax(1).values[:, 1:], + ) + names = ("q_mae", "qd_mae", "base_position_error", "base_orientation_error_rad") + metrics = {} + for h in (1, 10, 100): + mask = valid[:, h - 1] + endpoint = errors[:, h - 1][mask] + prefix = errors[:, :h][mask] + metrics[str(h)] = { + "endpoint": { + name: metric_module._safe_number(endpoint[:, i].mean()) + for i, name in enumerate(names) + }, + "prefix_mean": { + name: metric_module._safe_number(prefix[..., i].mean()) + for i, name in enumerate(names) + }, + "valid_trajectories": int(mask.sum()), + "nonfinite_trajectories": int((~mask).sum()), + "teacher_nonfinite_trajectories": int( + (~teacher_finite.cummin(1).values[:, h]).sum() + ), + "teacher_failed_trajectories": int(teacher_failed[:, h - 1].sum()), + "neural_failed_trajectories": int(neural_failed[:, h - 1].sum()), + } + args.output.parent.mkdir(parents=True, exist_ok=True) + evidence = args.output.with_suffix(".npz") + np.savez_compressed( + str(evidence), + q=q.numpy(), + qd=qd.numpy(), + teacher_q=tq.numpy(), + teacher_qd=tqd.numpy(), + raw=raw.numpy(), + efforts=torch.stack(efforts, 1).numpy(), + teacher_efforts=teacher["raw_joint_f"].numpy(), + errors=errors.numpy(), + valid=valid.numpy(), + ) + result = { + "robot": robot, + "motion_mode": evaluation["motion_mode"], + "trajectories": count, + "horizons": [1, 10, 100], + "seed": evaluation["seed"], + "metrics": metrics, + "checkpoint_sha256": _hash(args.checkpoint), + "teacher_manifest_sha256": _hash(args.teacher / "manifest.json"), + "teacher_shards": verified_shards, + "teacher_prefixes": identities, + "evidence_sha256": _hash(evidence), + "runtime": { + "python": platform.python_version(), + "torch": torch.__version__, + "warp": wp.__version__, + "runtime_imports_embodichain": "embodichain" in sys.modules, + "runtime_imports_dexsim": "dexsim" in sys.modules, + }, + "teacher_generation_environment": generation, + "evaluation_scope": "original neural integration and common original supported runtime; training-weight comparison; not native Newton deployment", + } + args.output.write_text(json.dumps(result, indent=2, allow_nan=False) + "\n") + print(json.dumps(metrics, indent=2), flush=True) + env.close() + + +if __name__ == "__main__": + main() diff --git a/examples/learning/nerd/_support/_portable_reference.py b/examples/learning/nerd/_support/_portable_reference.py new file mode 100644 index 0000000000..9deba5b13f --- /dev/null +++ b/examples/learning/nerd/_support/_portable_reference.py @@ -0,0 +1,261 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Restore reference/math exports in the unchanged original model classes.""" + +from __future__ import annotations + +import hashlib +import json +import os +from pathlib import Path +import subprocess +from typing import Any + +from ._official_evaluation_data import _hash, _transport_paths + +__all__: list[str] = [] +COMMIT = "370e89d2b8ce68aa3578aa7812859c681b7502bc" + + +def array_sha256(value: Any) -> str: + """Bind shape, dtype and bytes of one raw state/action array.""" + import numpy as np + + array = np.ascontiguousarray(value) + digest = hashlib.sha256() + digest.update( + json.dumps( + {"shape": array.shape, "dtype": array.dtype.str}, sort_keys=True + ).encode() + ) + digest.update(array.tobytes()) + return digest.hexdigest() + + +def verify_source( + source: Path, metadata: dict[str, Any], transport: Path +) -> dict[str, Any]: + """Verify source and postprocessing identities before external execution.""" + if os.environ.get("CUDA_VISIBLE_DEVICES") != "": + raise ValueError("This CPU-only runner requires CUDA_VISIBLE_DEVICES=''") + commit = subprocess.check_output( + ["git", "-C", str(source), "rev-parse", "HEAD"], text=True + ).strip() + if commit != COMMIT: + raise ValueError("Official source commit mismatch") + dirty = subprocess.check_output( + ["git", "-C", str(source), "diff", "--name-only", "HEAD"], text=True + ) + if dirty.strip(): + raise ValueError("Official tracked source files must remain unchanged") + verified = {} + for name, expected in metadata["source_runtime_files"].items(): + actual = _hash(source / name) + if actual != expected: + raise ValueError("Official runtime file checksum mismatch: " + name) + verified[name] = actual + # These files affect inference but were not all listed in the old teacher. + for name in ( + "models/models.py", + "models/model_transformer.py", + "models/base_models.py", + "utils/running_mean_std.py", + ): + verified[name] = _hash(source / name) + verified_transport = {} + for name, identity in metadata["canonical_transport_modules"].items(): + actual = _hash(_transport_paths(transport)[name]) + if actual != identity["sha256"]: + raise ValueError("Frozen postprocessing module checksum mismatch: " + name) + verified_transport[name] = actual + return { + "commit": commit, + "tracked_clean": True, + "files": verified, + "postprocessing": verified_transport, + } + + +def load_portable( + directory: Path, robot: str, scope: str +) -> tuple[Any, dict[str, Any], dict[str, Any]]: + """Restore FP32 weights/RMS exactly; attention is set by the caller's context. + + The portable JSON must explicitly declare reference/math and source world. + Only the feature spelling joint_f -> joint_acts changes. RMS counts are not + portable, and the frozen inference method reads only the stored mean/var. + """ + import torch + from models.models import ModelMixedInput + from utils.running_mean_std import RunningMeanStd + + path = directory / "manifest.json" + manifest = json.loads(path.read_text()) + cfg = manifest["model_config"] + if ( + manifest.get("robot") != robot + or manifest.get("world_rotation_xyzw", "missing") is not None + ): + raise ValueError( + "Portable model must declare the matching unrotated source world" + ) + if ( + manifest.get("numerical_profile"), + manifest.get("attention_backend"), + cfg.get("attention_backend"), + ) != ("reference", "math", "math"): + raise ValueError( + "Portable model must explicitly declare reference/math at every level" + ) + if cfg["history_length"] != 10: + raise ValueError("This fixed protocol requires the original ten-state history") + training = manifest.get("training", {}) + if training: + if (training.get("numerical_profile"), training.get("attention_backend")) != ( + "reference", + "math", + ): + raise ValueError( + "Training and model numerical policy declarations disagree" + ) + updates = training["global_step"] + else: + initial_path = directory / "identity.json" + if not initial_path.exists(): + raise ValueError("Missing exact update provenance") + updates = json.loads(initial_path.read_text())["source_update_count"] + if type(updates) is not int or updates < 0: + raise ValueError("Invalid checkpoint update count") + if updates == 0 and scope != "smoke": + raise ValueError("An initial model can only be labeled smoke") + weights = directory / "weights.pt" + if _hash(weights) != manifest["weights_sha256"]: + raise ValueError("Portable weights checksum mismatch") + state = torch.load(weights, map_location="cpu", weights_only=True) + if not state or any( + not isinstance(v, torch.Tensor) + or v.dtype != torch.float32 + or not torch.isfinite(v).all() + for v in state.values() + ): + raise ValueError( + "The original integrator requires finite FP32 portable tensors" + ) + rename = lambda key: "joint_acts" if key == "joint_f" else key + model = ModelMixedInput( + input_sample={ + rename(k): torch.zeros((1, 1, v), dtype=torch.float32) + for k, v in cfg["input_dims"].items() + }, + output_dim=cfg["output_dim"], + input_cfg={"low_dim": [rename(k) for k in cfg["input_dims"]]}, + network_cfg={ + "encoder": { + "low_dim": {"activation": "relu", "layer_sizes": [], "layernorm": False} + }, + "model": { + "mlp": { + "activation": "relu", + "layer_sizes": cfg["head_hidden_dims"], + "layernorm": False, + } + }, + "normalize_input": cfg["normalize_input"], + "normalize_output": cfg["normalize_output"], + "transformer": { + k: cfg[k] + for k in ( + "n_embd", + "n_layer", + "n_head", + "block_size", + "bias", + "dropout", + ) + }, + "output_tanh": cfg["output_tanh"], + }, + device="cpu", + ) + torch.nn.Module.to(model, device="cpu", dtype=torch.float32) + parameters = { + k: v + for k, v in state.items() + if not k.startswith(("input_rms.", "output_rms.")) + } + model.load_state_dict(parameters, strict=True) + stats = {} + expected = set(parameters) + for key, width in cfg["input_dims"].items(): + rms = RunningMeanStd(shape=(width,), device="cpu") + for field in ("mean", "var"): + name = "input_rms." + key + "." + field + value = state[name] + if tuple(value.shape) != (width,): + raise ValueError("Unexpected input RMS shape") + setattr(rms, field, value.clone()) + expected.add(name) + stats[rename(key)] = rms + output = RunningMeanStd(shape=(cfg["output_dim"],), device="cpu") + for field in ("mean", "var"): + name = "output_rms." + field + if tuple(state[name].shape) != (cfg["output_dim"],): + raise ValueError("Unexpected output RMS shape") + setattr(output, field, state[name].clone()) + expected.add(name) + if expected != set(state): + raise ValueError("Unexpected portable state fields") + model.set_input_rms(stats) + model.set_output_rms(output) + model.eval() + # Prove restoration didn't silently rename/reshape/drop any parameter value. + if any( + not torch.equal( + model.state_dict()[k].view(torch.int32), value.view(torch.int32) + ) + for k, value in parameters.items() + ): + raise AssertionError("Original model parameter restoration changed bits") + for key in cfg["input_dims"]: + for field in ("mean", "var"): + if not torch.equal( + getattr(model.input_rms[rename(key)], field).view(torch.int32), + state["input_rms." + key + "." + field].view(torch.int32), + ): + raise AssertionError("Original RMS restoration changed bits") + identity = { + "directory": str(directory), + "manifest_sha256": _hash(path), + "weights_sha256": _hash(weights), + "global_step": updates, + "scope": scope, + "numerical_profile": "reference", + "attention_backend": "math", + "strict_parameters_and_input_rms_bit_equal": True, + "rms_counts_not_persisted": True, + "rms_frozen_inference_uses_only_mean_and_variance": True, + "not_a_default_official_training_reproduction": True, + } + return model, manifest, identity + + +def backend_flags(torch: Any) -> dict[str, bool]: + return { + "flash": torch.backends.cuda.flash_sdp_enabled(), + "math": torch.backends.cuda.math_sdp_enabled(), + "memory_efficient": torch.backends.cuda.mem_efficient_sdp_enabled(), + } diff --git a/examples/learning/nerd/_support/_qualification.py b/examples/learning/nerd/_support/_qualification.py new file mode 100644 index 0000000000..e8113e3ae5 --- /dev/null +++ b/examples/learning/nerd/_support/_qualification.py @@ -0,0 +1,179 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Explicit numerical parity and seed-paired accuracy acceptance gates.""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import torch +from torch import Tensor + +from ._trainer import normalized_mse + +__all__ = ["capture_training_step", "compare_training_steps", "paired_noninferiority"] + + +def capture_training_step( + model: torch.nn.Module, + inputs: dict[str, Tensor], + targets: Tensor, + *, + optimizer_state: dict[str, Any] | None = None, + learning_rate: float = 1e-3, + max_grad_norm: float = 1.0, + mask: Tensor | None = None, +) -> dict[str, Any]: + """Capture one real Adam update; caller owns model/RNG initialization. + + This mutates the supplied model. Run the official implementation separately + and serialize the same named tensors; comparison never substitutes a second + invocation of this function for independent official evidence. + """ + model.train() + optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) + if optimizer_state is not None: + optimizer.load_state_dict(optimizer_state) + optimizer.zero_grad(set_to_none=True) + prediction = model(inputs) + loss = normalized_mse( + prediction, + targets, + model.output_rms.var if model.cfg.normalize_output else None, + mask, + ) + loss.backward() + raw = { + key: parameter.grad.detach().cpu().clone() + for key, parameter in model.named_parameters() + if parameter.grad is not None + } + norm = torch.nn.utils.clip_grad_norm_( + model.parameters(), max_grad_norm, error_if_nonfinite=True + ) + clipped = { + key: parameter.grad.detach().cpu().clone() + for key, parameter in model.named_parameters() + if parameter.grad is not None + } + optimizer.step() + return { + "prediction": prediction.detach().cpu(), + "loss": loss.detach().cpu(), + "grad_norm": norm.detach().cpu(), + "raw_gradients": raw, + "clipped_gradients": clipped, + "updated_parameters": { + key: parameter.detach().cpu().clone() + for key, parameter in model.named_parameters() + }, + "optimizer": optimizer.state_dict(), + } + + +def compare_training_steps( + reference: dict[str, Any], + candidate: dict[str, Any], + *, + atol: float = 1e-5, + rtol: float = 1e-4, +) -> dict[str, Any]: + """Compare independently captured forward, gradient and Adam-update tensors.""" + results: dict[str, Any] = {} + + def compare(name: str, expected: Tensor, actual: Tensor) -> None: + same_shape = expected.shape == actual.shape + finite = bool(torch.isfinite(expected).all() and torch.isfinite(actual).all()) + passed = ( + same_shape + and finite + and torch.allclose(expected, actual, atol=atol, rtol=rtol) + ) + results[name] = { + "passed": bool(passed), + "max_abs": ( + float((expected - actual).abs().max()) + if same_shape and finite and expected.numel() + else None + ), + } + + for name in ("prediction", "loss", "grad_norm"): + compare(name, reference[name], candidate[name]) + for group in ("raw_gradients", "clipped_gradients", "updated_parameters"): + if reference[group].keys() != candidate[group].keys(): + results[group] = {"passed": False, "reason": "parameter_names_differ"} + for name in reference[group].keys() & candidate[group].keys(): + compare(f"{group}/{name}", reference[group][name], candidate[group][name]) + return { + "passed": all(value["passed"] for value in results.values()), + "atol": atol, + "rtol": rtol, + "tensors": results, + } + + +def paired_noninferiority( + reference: list[float], + candidate: list[float], + *, + margin: float = 0.05, + absolute_tolerance: float = 1e-5, + resamples: int = 10000, + seed: int = 0, +) -> dict[str, Any]: + """One-sided 95% seed-paired bootstrap gate for lower-is-better metrics. + + Fewer than five seeds is pending qualification. Near-zero references use + absolute error, and non-finite measurements always fail qualification. + """ + baseline, measured = np.asarray(reference, dtype=float), np.asarray( + candidate, dtype=float + ) + if baseline.shape != measured.shape or baseline.ndim != 1 or baseline.size == 0: + raise ValueError("Expected equally sized nonempty paired seed vectors") + if margin < 0 or absolute_tolerance <= 0 or resamples < 1: + raise ValueError("Invalid acceptance thresholds") + if not np.isfinite(baseline).all() or not np.isfinite(measured).all(): + return {"status": "failed", "reason": "nonfinite_metrics"} + if (baseline < 0).any() or (measured < 0).any(): + raise ValueError("Error metrics must be nonnegative") + absolute = bool((baseline <= absolute_tolerance).any()) + differences = measured - baseline if absolute else (measured - baseline) / baseline + rng = np.random.default_rng(seed) + # Bounded chunks also support larger external qualification runs. + estimates = np.empty(resamples) + for start in range(0, resamples, 1024): + count = min(1024, resamples - start) + indices = rng.integers(0, baseline.size, size=(count, baseline.size)) + estimates[start : start + count] = differences[indices].mean(axis=1) + upper = float(np.quantile(estimates, 0.95)) + threshold = absolute_tolerance if absolute else margin + status = ( + "pending" if baseline.size < 5 else "passed" if upper <= threshold else "failed" + ) + return { + "status": status, + "seeds": int(baseline.size), + "mean_difference": float(differences.mean()), + "upper_95": upper, + "threshold": threshold, + "difference_mode": "absolute" if absolute else "relative", + "resamples": resamples, + "seed": seed, + } diff --git a/examples/learning/nerd/_support/_reference.py b/examples/learning/nerd/_support/_reference.py new file mode 100644 index 0000000000..1d889374ab --- /dev/null +++ b/examples/learning/nerd/_support/_reference.py @@ -0,0 +1,1267 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Explicitly trusted, isolated execution of the unmodified official NeRD trainer. + +The adapter binds data, model initialization and normalizers to original +VanillaTrainer methods. It never imports external code until explicitly requested. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +from pathlib import Path +import subprocess +import sys +from types import SimpleNamespace +from typing import Any + +__all__ = [ + "OFFICIAL_COMMIT", + "reference_identity", + "run_reference", + "export_legacy_dataset", + "import_legacy_datasets", + "generate_official_dataset", + "repair_gravity_transport", +] +OFFICIAL_COMMIT = "370e89d2b8ce68aa3578aa7812859c681b7502bc" +_OFFICIAL_PHYSICS_CONTRACT_VERSION = "official-warp-physical-v2" + + +def _hash(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def reference_identity(source: str | Path) -> dict[str, Any]: + """Verify pinned, unmodified tracked source before importing external code.""" + source = Path(source).resolve() + commit = subprocess.check_output( + ["git", "-C", str(source), "rev-parse", "HEAD"], text=True + ).strip() + if commit != OFFICIAL_COMMIT: + raise ValueError(f"Official trainer requires {OFFICIAL_COMMIT}, found {commit}") + changed = subprocess.check_output( + ["git", "-C", str(source), "status", "--porcelain", "--untracked-files=no"], + text=True, + ).strip() + if changed: + raise ValueError("Official reference checkout contains tracked modifications") + paths = ( + "algorithms/vanilla_trainer.py", + "integrators/integrator_neural.py", + "models/models.py", + "models/model_transformer.py", + "models/base_models.py", + "utils/torch_utils.py", + "utils/running_mean_std.py", + ) + return { + "root": str(source), + "commit": commit, + "files": {path: _hash(source / path) for path in paths}, + "source_patches": [], + "adapter": "common-state binding; original preprocess_data_batch, compute_loss, one_epoch and get_scheduled_learning_rate; constructor, logging, data I/O and shared validation evaluation supplied by EmbodiChain", + } + + +def _legacy_states(q: Any, qd: Any, codec: Any) -> Any: + import torch + + # Reversing cross-product operands and subtracting is algebraically equal, + # but its FP32 fused operations can differ. The codec owns this transport + # for both trainers, before the unmodified official preprocessing runs. + return torch.cat((q, codec.native_to_legacy(q, qd)), -1) + + +def _official_batch( + batch: dict[str, Any], codec: Any, *, device: str | None = None +) -> dict[str, Any]: + """Use preserved source twists when present, otherwise transport native qd.""" + import torch + + q = batch["states"]["q"].to(device=device, dtype=torch.float32) + qd = batch["states"]["qd"].to(device=q.device, dtype=torch.float32) + transitions = { + key: value.to(device=q.device, dtype=torch.float32) + for key, value in batch["transitions"].items() + } + legacy_velocity = batch["states"].get("legacy_velocity") + if legacy_velocity is None: + states = _legacy_states(q, qd, codec) + else: + legacy_velocity = legacy_velocity.to(device=q.device, dtype=torch.float32) + codec._check(q, qd) + if legacy_velocity.shape != qd.shape: + raise ValueError("Preserved legacy velocity shape must match native qd") + states = torch.cat((q, legacy_velocity), -1) + shape = q.shape[:2] + shape = (shape[0], shape[1] - 1) + root = batch["states"].get("root_body_q") + if root is not None: + root = root.to(device=q.device, dtype=torch.float32) + if root is None: + root = ( + q[..., :7] + if codec.robot == "ant" + else q.new_tensor([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]).expand( + *q.shape[:2], 7 + ) + ) + result = { + "states": states[:, :-1], + "next_states": states[:, 1:], + "root_body_q": root[:, :-1].float(), + "joint_acts": ( + transitions["joint_f"][..., -8:] + if codec.robot == "ant" + else transitions["joint_f"] + ), + "gravity_dir": transitions.get( + "gravity_dir", q.new_tensor([0.0, 0.0, -1.0]).expand(*shape, 3) + ), + } + slots = transitions.get("contact_depths", q.new_zeros(*shape, 0)).shape[-1] + # Main's preprocessing reshapes contact tensors with an inferred final axis, + # which cannot represent zero slots. A masked dummy is unobserved by the + # contact-free Cartpole network and does not alter its input or loss. + if not slots: + if codec.robot != "cartpole": + raise ValueError("Ant requires contact slots") + slots = 1 + result.update( + contact_depths=q.new_full((*shape, 1), 1e4), + contact_thicknesses=q.new_zeros(*shape, 1), + contact_normals=q.new_zeros(*shape, 3), + contact_points_0=q.new_zeros(*shape, 3), + contact_points_1=q.new_zeros(*shape, 3), + ) + else: + for name in ( + "contact_depths", + "contact_normals", + "contact_points_0", + "contact_points_1", + ): + width = slots if name == "contact_depths" else slots * 3 + result[name] = transitions.get(name, q.new_zeros(*shape, width)) + result["contact_thicknesses"] = transitions.get( + "contact_thicknesses_0", q.new_zeros(*shape, slots) + ) + transitions.get("contact_thicknesses_1", q.new_zeros(*shape, slots)) + return result + + +def _reference_integrator(robot: str, slots: int, device: str) -> Any: + import torch + import warp as wp + from integrators.integrator_neural import NeuralIntegrator + + value = NeuralIntegrator.__new__(NeuralIntegrator) + ant = robot == "ant" + value.torch_device = device + value.dof_q_per_env = 15 if ant else 2 + value.dof_qd_per_env = 14 if ant else 2 + value.state_dim = value.prediction_dim = value.dof_q_per_env + value.dof_qd_per_env + value.num_contacts_per_env = slots + value.joint_types = ( + [wp.sim.JOINT_FREE] + [wp.sim.JOINT_REVOLUTE] * 8 + if ant + else [wp.sim.JOINT_PRISMATIC, wp.sim.JOINT_REVOLUTE] + ) + value.num_joints_per_env = len(value.joint_types) + value.joint_q_start = [0] + list(range(7, 15)) if ant else [0, 1] + value.joint_q_end = [7] + list(range(8, 16)) if ant else [1, 2] + value.is_continuous_dof = torch.zeros( + value.state_dim, dtype=torch.bool, device=device + ) + if not ant: + value.is_continuous_dof[1] = True + value.states_frame = "body" if ant else "world" + value.anchor_frame_step = "every" + value.states_embedding_type = "identical" + value.orientation_prediction_parameterization = "quaternion" + value.prediction_type = "relative" + return value + + +def _restore_official_model( + common: Any, + *, + device: str, + input_sample: dict[str, Any] | None = None, +) -> Any: + """Reconstruct the original network from a portable model without aliasing it. + + The caller must first establish its explicit external-code trust boundary, + verify the pinned checkout and expose its modules on the isolated subprocess + import path. This helper performs no source discovery or simulator startup. + It preserves all parameters, normalization values, ordered features and dtype, + returns evaluation mode, and renames only joint_f to legacy joint_acts. + Physics callers must enforce their own supported inference dtype (the old + Warp integrator uses FP32); standalone network restoration also supports FP64. + """ + import torch + from models.models import ModelMixedInput + from utils.running_mean_std import RunningMeanStd + + c = common.cfg + dtype = next(common.parameters()).dtype + rename = lambda key: "joint_acts" if key == "joint_f" else key + if input_sample is None: + input_sample = { + rename(key): torch.zeros((1, 1, width), dtype=dtype, device=device) + for key, width in c.input_dims.items() + } + model = ModelMixedInput( + input_sample=input_sample, + output_dim=c.output_dim, + input_cfg={"low_dim": [rename(key) for key in c.input_dims]}, + network_cfg={ + "encoder": { + "low_dim": {"activation": "relu", "layer_sizes": [], "layernorm": False} + }, + "model": { + "mlp": { + "activation": "relu", + "layer_sizes": list(c.head_hidden_dims), + "layernorm": False, + } + }, + "normalize_input": c.normalize_input, + "normalize_output": c.normalize_output, + "transformer": { + "n_embd": c.n_embd, + "n_layer": c.n_layer, + "n_head": c.n_head, + "block_size": c.block_size, + "bias": c.bias, + "dropout": c.dropout, + }, + "output_tanh": c.output_tanh, + }, + device=device, + ) + shared = { + key: value + for key, value in common.state_dict().items() + if not key.startswith(("input_rms.", "output_rms.")) + } + # The original ModelMixedInput.to only visits some children; call the base + # implementation so transformer parameters follow the requested dtype/device. + torch.nn.Module.to(model, device=device, dtype=dtype) + model.load_state_dict(shared, strict=True) + stats = {} + for key, value in common.input_rms.items(): + rms = RunningMeanStd(shape=value.mean.shape, device=device) + rms.mean = value.mean.detach().to(device=device, dtype=dtype).clone() + rms.var = value.var.detach().to(device=device, dtype=dtype).clone() + stats[rename(key)] = rms + output = RunningMeanStd(shape=common.output_rms.mean.shape, device=device) + output.mean = common.output_rms.mean.detach().to(device=device, dtype=dtype).clone() + output.var = common.output_rms.var.detach().to(device=device, dtype=dtype).clone() + model.set_input_rms(stats) + model.set_output_rms(output) + return model.eval() + + +def run_reference( + spec: dict[str, Any], + *, + trust_external_code: bool = False, + resume: str | Path | None = None, + stop_after: int | None = None, +) -> dict[str, Any]: + """Run actual official trainer methods against matched immutable artifacts. + + Invoke in an isolated subprocess with legacy Warp dependencies on PYTHONPATH. + This binds an original trainer instance without its constructor's eager + dataset/GPU allocation; no original method is monkeypatched or replaced. + """ + if not trust_external_code: + raise ValueError("Official source execution requires trust_external_code=True") + from embodichain_models.dynamics.nerd._numerics import _validate_numerical_profile + + if ( + _validate_numerical_profile( + spec.get("train", {}).get("numerical_profile", "default") + ) + != "default" + ): + raise ValueError( + "Original-method reference adapter supports only default numerical_profile" + ) + identity = reference_identity(spec["source"]) + sys.path.insert(0, identity["root"]) + import torch + import numpy as np + from algorithms.vanilla_trainer import VanillaTrainer + from utils.time_report import TimeReport + from embodichain_models.dynamics.nerd.codec import NeRDStateCodec + from embodichain_models.dynamics.nerd.data import DynamicsWindowDataset + from embodichain_models.dynamics.nerd import load_model, export_model + from ._trainer import ( + DynamicsTrainCfg, + evaluate_model, + _atomic_save, + _capture_rng_state, + _restore_rng_state, + _PREPROCESSING_EXECUTION, + ) + from ._runtime import _training_runtime_identity + + cfg = DynamicsTrainCfg(**spec["train"]) + torch.set_float32_matmul_precision("highest") + device = cfg.device + torch.set_num_threads(int(spec.get("cpu_threads", 1))) + common, manifest = load_model(spec["initial_model"], device=device) + if common.cfg.attention_backend != "default": + raise ValueError( + "Original-method reference adapter does not support math attention initial_model" + ) + if ( + spec.get("initial_model_sha256") + and _hash(Path(spec["initial_model"]) / "weights.pt") + != spec["initial_model_sha256"] + ): + raise ValueError("Common initialization changed") + codec = NeRDStateCodec(**manifest["codec"]) + training = DynamicsWindowDataset( + spec["dataset"], "train", common.cfg.history_length, verify=True + ) + validation = DynamicsWindowDataset( + spec["dataset"], "validation", common.cfg.history_length + ) + test = DynamicsWindowDataset(spec["dataset"], "test", common.cfg.history_length) + if training.fingerprint != spec["dataset_fingerprint"]: + raise ValueError("Reference dataset fingerprint mismatch") + order_path = Path(spec["sample_order"]) + if _hash(order_path) != spec["sample_order_sha256"]: + raise ValueError("Reference sample order checksum mismatch") + from ._sample_order import _open_order + + order = _open_order(order_path) + if order.shape != (cfg.total_steps, cfg.batch_size): + raise ValueError("Reference sample-order shape mismatch") + contract = { + "dataset_fingerprint": training.fingerprint, + "sample_order_sha256": spec["sample_order_sha256"], + "initial_model_sha256": _hash(Path(spec["initial_model"]) / "weights.pt"), + "train_cfg": cfg.to_dict(), + "source_commit": identity["commit"], + "source_files": identity["files"], + "torch_version": str(torch.__version__), + "adapter_sha256": _hash(Path(__file__)), + "preprocessing_execution": _PREPROCESSING_EXECUTION, + } + restored = None + start_step = 0 + if resume is not None: + restored = torch.load(resume, map_location="cpu", weights_only=True) + if ( + restored.get("format") != "embodichain.nerd.reference_training.v1" + or restored.get("contract") != contract + ): + raise ValueError("Reference resume contract mismatch") + common.load_state_dict(restored["model"]) + start_step = restored["global_step"] + end_step = ( + cfg.total_steps if stop_after is None else min(stop_after, cfg.total_steps) + ) + if type(end_step) is not int or not start_step <= end_step <= cfg.total_steps: + raise ValueError("Invalid reference training stop cursor") + sample = training.get_batch([0]) + raw = _official_batch(sample, codec, device=device) + integrator = _reference_integrator( + codec.robot, raw["contact_depths"].shape[-1], device + ) + original = VanillaTrainer.__new__(VanillaTrainer) + original.device = device + original.neural_env = SimpleNamespace(frame_dt=manifest["dt"]) + original.neural_integrator = integrator + processed = original.preprocess_data_batch({k: v.clone() for k, v in raw.items()}) + original.neural_model = _restore_official_model( + common, device=device, input_sample=processed + ) + integrator.neural_model = original.neural_model + original.batch_size = cfg.batch_size + original.truncate_grad = True + original.grad_norm = cfg.max_grad_norm + original.lr_start = cfg.learning_rate + original.lr_end = cfg.final_learning_rate + original.lr_schedule = "linear" + original.optimizer = torch.optim.Adam( + original.neural_model.parameters(), lr=cfg.learning_rate + ) + if restored is not None: + original.optimizer.load_state_dict(restored["optimizer"]) + original.time_report = TimeReport(cuda_synchronize=False) + original.time_report.add_timers(["dataloader", "compute_loss", "backward", "other"]) + # Initialization and transport above may consume RNG; both runs begin their + # optimizer updates with the configured seed, as DynamicsTrainer does. + torch.manual_seed(cfg.seed) + if device.startswith("cuda"): + torch.cuda.manual_seed_all(cfg.seed) + if cfg.resident_cache_bytes: + used = training.enable_device_cache(device, cfg.resident_cache_bytes) + validation.enable_device_cache(device, cfg.resident_cache_bytes - used) + if torch.device(device).type == "cuda": + warm_batch = training.get_batch([0] * cfg.batch_size) + with torch.enable_grad(): + for _ in range(3): + original.preprocess_data_batch( + _official_batch(warm_batch, codec, device=device) + ) + del warm_batch + destination = Path(spec["output_dir"]) + destination.mkdir(parents=True, exist_ok=True) + best = float("inf") + best_weights = None + if restored is not None: + best, best_weights = restored["best_validation"], restored["best_weights"] + _restore_rng_state(restored["rng"]) + + def transfer() -> None: + current = common.state_dict() + current.update(original.neural_model.state_dict()) + common.load_state_dict(current) + + def checkpoint(step: int) -> None: + transfer() + _atomic_save( + destination / "reference-last.pt", + { + "format": "embodichain.nerd.reference_training.v1", + "contract": contract, + "model": common.state_dict(), + "optimizer": original.optimizer.state_dict(), + "global_step": step, + "best_validation": best, + "best_weights": best_weights, + "rng": _capture_rng_state(device), + }, + ) + + try: + log_path = destination / "metrics.jsonl" + if log_path.exists(): + if resume is None: + raise FileExistsError( + "Reference output already has progress; use explicit resume" + ) + pending_log = log_path.with_suffix(".jsonl.pending") + with log_path.open() as old, pending_log.open("w") as retained: + for line in old: + try: + record = json.loads(line) + except json.JSONDecodeError: + break + if record["step"] <= start_step: + retained.write(line) + os.replace(pending_log, log_path) + with log_path.open("a") as log: + for step in range(start_step, end_step): + if np.any(order[step] < 0) or np.any(order[step] >= len(training)): + raise ValueError("Reference sample order index out of bounds") + batch = training.get_batch(order[step]) + transported = _official_batch(batch, codec, device=device) + rate = original.get_scheduled_learning_rate( + step // cfg.lr_update_interval + cfg.lr_offset, + cfg.total_steps / cfg.lr_update_interval + cfg.lr_offset, + ) + for group in original.optimizer.param_groups: + group["lr"] = rate + loader = [transported] + loss, _, grad = original.one_epoch(True, loader, iter(loader), 1) + record = { + "step": step + 1, + "loss": loss, + "learning_rate": rate, + "grad_norm": float(grad["grad_norm_before_clip"]), + } + if ( + step + 1 + ) % cfg.validation_interval == 0 or step + 1 == cfg.total_steps: + transfer() + metrics = evaluate_model( + common, + codec, + validation, + cfg.batch_size, + device, + max_batches=cfg.validation_batches, + ) + record["validation"] = metrics + if metrics["normalized_mse"] < best: + best = metrics["normalized_mse"] + best_weights = { + key: value.detach().cpu().clone() + for key, value in common.state_dict().items() + } + if ( + (step + 1) % cfg.log_interval == 0 + or step + 1 == end_step + or "validation" in record + ): + log.write(json.dumps(record, allow_nan=False) + "\n") + log.flush() + if (step + 1) % cfg.checkpoint_interval == 0: + checkpoint(step + 1) + checkpoint(end_step) + if end_step != cfg.total_steps: + result = { + "status": "interrupted", + "completed_updates": end_step, + "budget": cfg.total_steps, + "checkpoint": str(destination / "reference-last.pt"), + } + (destination / "result.json").write_text(json.dumps(result, indent=2)) + return result + transfer() + export_model(destination / "last-export", common, manifest) + common.load_state_dict(best_weights) + export_model(destination / "best-export", common, manifest) + metrics = evaluate_model(common, codec, test, cfg.batch_size, device) + result = { + "status": "completed", + "completed_updates": cfg.total_steps, + "budget": cfg.total_steps, + "best_validation": best, + "test": metrics, + "reference": identity, + "dataset_fingerprint": training.fingerprint, + "sample_order_sha256": spec["sample_order_sha256"], + "seed": cfg.seed, + "qualification": "matched_training_run_only", + "torch_version": str(torch.__version__), + "runtime_identity": _training_runtime_identity(original.neural_model), + "device": device, + "adapter_source_sha256": _hash(Path(__file__)), + "best_model_sha256": _hash(destination / "best-export" / "weights.pt"), + "last_model_sha256": _hash(destination / "last-export" / "weights.pt"), + "training_methods": [ + "VanillaTrainer.preprocess_data_batch", + "VanillaTrainer.compute_loss", + "VanillaTrainer.one_epoch", + "VanillaTrainer.get_scheduled_learning_rate", + ], + } + (destination / "result.json").write_text( + json.dumps(result, indent=2, allow_nan=False) + ) + return result + finally: + training.close() + validation.close() + test.close() + + +def export_legacy_dataset( + root: str | Path, output: str | Path, split: str +) -> dict[str, Any]: + """Stream canonical trajectories to official T,B HDF5 transport arrays. + + Variable lengths are retained via traj_lengths; padding is never sampled by + the official TrajectoryDataset. No source dataset or split is modified. + Preserved source velocities bypass native conversion when available. + """ + import h5py + import numpy as np + import torch + from embodichain_models.dynamics.nerd.codec import NeRDStateCodec + from embodichain_models.dynamics.nerd.data import ( + load_manifest, + dataset_fingerprint, + ) + + root = Path(root) + output = Path(output) + if output.exists(): + raise FileExistsError(output) + manifest = load_manifest(root, verify=True) + codec = NeRDStateCodec(**manifest["metadata"]["codec"]) + rows = [ + (shard["file"], row) + for shard in manifest["shards"] + for row in shard["trajectories"] + if row["split"] == split + ] + if not rows: + raise ValueError("No trajectories for requested split") + max_length = max(row["length"] for _, row in rows) + output.parent.mkdir(parents=True, exist_ok=True) + temporary = output.with_suffix(output.suffix + ".pending") + with h5py.File(temporary, "w") as target: + group = target.create_group("data") + group.attrs["mode"] = "trajectory" + group.attrs["total_transitions"] = sum(row["length"] for _, row in rows) + group.create_dataset( + "traj_lengths", + data=np.array([row["length"] for _, row in rows], dtype=np.int32), + ) + for index, (name, row) in enumerate(rows): + with h5py.File(root / name, "r") as source: + trajectory = source[row["group"]] + batch = { + kind: { + key: torch.from_numpy(value[()]).unsqueeze(0) + for key, value in trajectory[kind].items() + } + for kind in ("states", "transitions") + } + raw = _official_batch(batch, codec) + for key, value in raw.items(): + array = value[0].numpy() + if key not in group: + group.create_dataset( + key, + shape=(max_length, len(rows), array.shape[-1]), + dtype=np.float32, + chunks=(min(max_length, 128), 1, array.shape[-1]), + compression="lzf", + fillvalue=0.0, + ) + group[key][: row["length"], index] = array + os.replace(temporary, output) + result = { + "dataset_fingerprint": dataset_fingerprint(manifest), + "split": split, + "trajectories": len(rows), + "path": str(output.resolve()), + "sha256": _hash(output), + "layout": "T,B,features", + "native_to_legacy_qd": "legacy_velocity" not in manifest["fields"]["states"], + "cartpole_dummy_masked_contact": codec.robot == "cartpole", + } + output.with_suffix(output.suffix + ".json").write_text(json.dumps(result, indent=2)) + return result + + +def import_legacy_datasets( + sources: dict[str, str | Path], + output: str | Path, + metadata: dict[str, Any], + *, + trajectories_per_shard: int = 64, + world_rotation_xyzw: tuple[float, float, float, float] | None = None, + gravity_direction: tuple[float, float, float] | None = None, + preserve_legacy_velocity: bool = False, +) -> dict[str, Any]: + """Import official world-frame HDF5 trajectories into canonical shards. + + Source files explicitly name their split; the importer never re-splits + official validation/test trajectories. Only one trajectory is staged at a + time. Adjacent states must be continuous within each declared valid length. + Legacy combined contact thickness is stored in side zero, side one zero; + the combined feature/mask is preserved without inventing separate margins. + + Args: + sources: Original HDF5 paths keyed by explicit dataset split. + output: New canonical shard directory. + metadata: Codec and dataset provenance, copied without modifying it. + trajectories_per_shard: Maximum trajectories per committed shard. + world_rotation_xyzw: Optional unit quaternion rotating world coordinates. + gravity_direction: Optional explicit unit gravity transport correction. + preserve_legacy_velocity: Also retain original FP32 velocities in + ``states/legacy_velocity`` before COM conversion. This requires + unrotated source-world coordinates and bitwise continuous original + states, including signed zero. Native ``q``/``qd`` remain available + to normal runtime consumers; the original input file is read only. + + Returns: + Imported dataset identity, trajectory count and immutable fingerprint. + """ + import h5py + import numpy as np + import torch + from embodichain_models.dynamics.nerd.codec import NeRDStateCodec + from embodichain_models.dynamics.nerd.data import ( + DynamicsShardWriter, + load_manifest, + dataset_fingerprint, + ) + + if not sources or set(sources) - {"train", "validation", "test"}: + raise ValueError( + "Legacy sources must explicitly name train/validation/test splits" + ) + if preserve_legacy_velocity and world_rotation_xyzw is not None: + raise ValueError( + "Preserved legacy velocity requires the source world frame without rotation" + ) + codec = NeRDStateCodec(**metadata["codec"]) + from embodichain_models.dynamics.nerd.codec import quat_rotate, quat_multiply + + rotation = ( + None + if world_rotation_xyzw is None + else torch.tensor(world_rotation_xyzw, dtype=torch.float32) + ) + if rotation is not None and ( + rotation.shape != (4,) + or not torch.isfinite(rotation).all() + or not torch.isclose(rotation.norm(), torch.tensor(1.0), atol=1e-6) + ): + raise ValueError("World rotation must be a unit xyzw quaternion") + source_identity = { + split: {"path": str(Path(path).resolve()), "sha256": _hash(Path(path))} + for split, path in sources.items() + } + manifest_metadata = { + **metadata, + "legacy_sources": source_identity, + "legacy_thickness_transport": "combined_in_side_zero", + "world_rotation_xyzw": None if rotation is None else rotation.tolist(), + "gravity_transport_override": gravity_direction, + } + if preserve_legacy_velocity: + manifest_metadata["legacy_velocity"] = { + "version": 1, + "field": "states/legacy_velocity", + "convention": "official_nerd_raw_velocity", + "layout": ( + "angular_first_world_origin_spatial" + if codec.robot == "ant" + else "joint_velocity" + ), + "world_frame": "source", + "source_field": "data/states[..., q_dim:]", + "terminal_source_field": "data/next_states[valid_length-1, ..., q_dim:]", + "dtype": "float32", + } + else: + # This importer owns the field declaration. Reusing a prior preserved + # manifest cannot label a native-only import as lossless source data. + manifest_metadata.pop("legacy_velocity", None) + total = 0 + with DynamicsShardWriter( + output, manifest_metadata, trajectories_per_shard=trajectories_per_shard + ) as writer: + for split, path in sources.items(): + with h5py.File(path, "r") as file: + data = file["data"] + if data.attrs.get("mode") != "trajectory": + raise ValueError("Expected official trajectory HDF5 layout") + if ( + data["states"].ndim != 3 + or data["states"].shape[-1] != codec.q_dim + codec.qd_dim + ): + raise ValueError("Legacy state layout does not match the codec") + horizon, trajectories, _ = data["states"].shape + lengths = ( + data["traj_lengths"][()] + if "traj_lengths" in data + else np.full(trajectories, horizon) + ) + for index, valid in enumerate(lengths): + length = int(valid) + if not 1 <= length <= horizon: + raise ValueError("Invalid official trajectory valid length") + current = data["states"][:length, index].astype(np.float32) + following = data["next_states"][:length, index].astype(np.float32) + if length > 1: + if preserve_legacy_velocity: + if not np.array_equal( + current[1:].view(np.uint32), + following[:-1].view(np.uint32), + ): + raise ValueError( + "Preserved legacy velocity requires bitwise continuous " + "FP32 states, including signed zero" + ) + elif not np.allclose( + current[1:], following[:-1], atol=1e-6, rtol=1e-5 + ): + raise ValueError( + "Legacy trajectory crosses an undeclared reset/discontinuity" + ) + legacy = torch.from_numpy( + np.concatenate((current, following[-1:]), axis=0) + ) + q = legacy[:, : codec.q_dim] + qd = codec.legacy_to_native(q, legacy[:, codec.q_dim :]) + if rotation is not None and codec.robot == "ant": + q = torch.cat( + ( + quat_rotate(rotation, q[:, :3]), + quat_multiply(rotation, q[:, 3:7]), + q[:, 7:], + ), + -1, + ) + qd = torch.cat( + ( + quat_rotate(rotation, qd[:, :3]), + quat_rotate(rotation, qd[:, 3:6]), + qd[:, 6:], + ), + -1, + ) + states = {"q": q.numpy(), "qd": qd.numpy()} + if preserve_legacy_velocity: + states["legacy_velocity"] = legacy[:, codec.q_dim :].numpy() + if "root_body_q" in data: + root = data["root_body_q"][:length, index].astype(np.float32) + if rotation is not None: + root_tensor = torch.from_numpy(root) + root = torch.cat( + ( + quat_rotate(rotation, root_tensor[:, :3]), + quat_multiply(rotation, root_tensor[:, 3:7]), + ), + -1, + ).numpy() + last = q[-1:, :7].numpy() if codec.robot == "ant" else root[-1:] + states["root_body_q"] = np.concatenate((root, last), axis=0) + effort = data["joint_acts"][:length, index].astype(np.float32) + if codec.robot == "ant" and effort.shape[-1] == 8: + effort = np.concatenate( + (np.zeros((length, 6), np.float32), effort), axis=-1 + ) + transitions = {"joint_f": effort} + if "gravity_dir" in data: + transitions["gravity_dir"] = data["gravity_dir"][ + :length, index + ].astype(np.float32) + for name in ( + "contact_depths", + "contact_normals", + "contact_points_0", + "contact_points_1", + ): + if name in data: + transitions[name] = data[name][:length, index].astype( + np.float32 + ) + if "contact_thicknesses" in data: + thickness = data["contact_thicknesses"][:length, index].astype( + np.float32 + ) + transitions["contact_thicknesses_0"] = thickness + transitions["contact_thicknesses_1"] = np.zeros_like(thickness) + if gravity_direction is not None: + direction = np.asarray(gravity_direction, dtype=np.float32) + if ( + direction.shape != (3,) + or not np.isfinite(direction).all() + or not np.isclose(np.linalg.norm(direction), 1.0, atol=1e-6) + ): + raise ValueError( + "Gravity transport direction must be a finite unit vector" + ) + transitions["gravity_dir"] = np.broadcast_to( + direction, (length, 3) + ).copy() + if rotation is not None: + for field in ( + "gravity_dir", + "contact_normals", + "contact_points_1", + ): + if field in transitions: + array = transitions[field] + transitions[field] = ( + quat_rotate( + rotation, torch.from_numpy(array.reshape(-1, 3)) + ) + .reshape(array.shape) + .numpy() + ) + writer.add_trajectory( + f"{split}:{index}", + states, + transitions, + split=split, + metadata={ + "legacy_source_sha256": source_identity[split]["sha256"], + "legacy_index": index, + }, + ) + total += 1 + manifest = load_manifest(output, verify=True) + return { + "path": str(Path(output).resolve()), + "trajectories": total, + "fingerprint": dataset_fingerprint(manifest), + "sources": source_identity, + } + + +def repair_gravity_transport( + gravity: Any, direction: tuple[float, float, float], *, enabled: bool = False +) -> tuple[Any, dict[str, Any]]: + """Explicitly repair upstream uninitialized gravity components in a copy. + + The original sampler writes only its up-axis into ``torch.empty``. Preserve + raw returned arrays separately; only this declared field is repaired. + """ + import numpy as np + + values = np.asarray(gravity) + expected = np.asarray(direction, dtype=values.dtype) + if ( + values.shape[-1] != 3 + or expected.shape != (3,) + or not np.isfinite(expected).all() + or not np.isclose(np.linalg.norm(expected), 1.0, atol=1e-6) + ): + raise ValueError("Invalid gravity transport shape/direction") + mismatch = ~np.isfinite(values) | ~np.isclose(values, expected, atol=1e-7, rtol=0.0) + evidence = { + "enabled": enabled, + "changed_components": int(mismatch.sum()), + "nonfinite_components": int((~np.isfinite(values)).sum()), + "physical_direction": expected.tolist(), + "source_issue": "TrajectorySampler allocates torch.empty gravity_dir and writes only up-axis", + } + if mismatch.any() and not enabled: + raise ValueError( + "Official sampler returned uninitialized gravity values; explicitly enable gravity repair and retain raw artifacts" + ) + return ( + np.broadcast_to(expected, values.shape).copy() if enabled else values.copy() + ), evidence + + +def generate_official_dataset( + spec: dict[str, Any], *, trust_external_code: bool = False +) -> dict[str, Any]: + """Collect real original Warp trajectories and retain raw HDF5 evidence. + + Uses untouched original NeuralEnvironment, TrajectorySampler[Ant], commons, + sampling, collision and integration. Gravity repair and global frame rotation + are explicit transport operations, recorded separately from raw artifacts. + Use a fresh subprocess with the requested isolated Warp version on PYTHONPATH. + """ + if not trust_external_code: + raise ValueError("Official generation requires trust_external_code=True") + identity = reference_identity(spec["source"]) + sys.path.insert(0, identity["root"]) + import h5py + import numpy as np + import torch + import warp as wp + from envs.neural_environment import NeuralEnvironment + from generate.trajectory_sampler import TrajectorySampler + from generate.trajectory_sampler_ant import TrajectorySamplerAnt + from utils.commons import ( + JOINT_Q_MIN, + JOINT_Q_MAX, + JOINT_QD_MIN, + JOINT_QD_MAX, + JOINT_ACT_SCALE, + ) + from utils.python_utils import set_random_seed + import math + import time + import xml.etree.ElementTree as ET + from embodichain_models.dynamics.nerd.codec import quat_multiply + + robot = str(spec["robot"]).lower() + if robot not in {"cartpole", "ant"}: + raise ValueError("Official generation supports cartpole and ant") + expected_warp = spec.get("warp_version", "1.5.1" if robot == "ant" else "1.8.0") + if wp.__version__ != expected_warp: + raise ValueError( + f"Expected isolated Warp {expected_warp}, loaded {wp.__version__}" + ) + device = spec.get("device", "cpu") + torch.set_num_threads(int(spec.get("cpu_threads", 1))) + batch = int(spec.get("num_envs", 8)) + horizon = int(spec.get("horizon", 100)) + splits = spec.get("splits", {"train": 32, "validation": 8, "test": 8}) + if ( + min(batch, horizon) < 1 + or set(splits) != {"train", "validation", "test"} + or any(int(v) < 1 for v in splits.values()) + ): + raise ValueError( + "Positive batch/horizon and all three explicit split counts are required" + ) + output = Path(spec["output_dir"]).resolve() + output.mkdir(parents=True, exist_ok=True) + if any(output.iterdir()): + raise FileExistsError("Official generation output must be empty") + set_random_seed(int(spec.get("seed", 0))) + started = time.perf_counter() + environment = NeuralEnvironment( + env_name="Ant" if robot == "ant" else "Cartpole", + num_envs=batch, + neural_integrator_cfg={}, + neural_model=None, + default_env_mode="ground-truth", + warp_env_cfg={"seed": int(spec.get("seed", 0))}, + device=device, + render=False, + ) + name = environment.robot_name + sampler_type = TrajectorySamplerAnt if robot == "ant" else TrajectorySampler + sampler = sampler_type( + environment, + JOINT_Q_MIN[name], + JOINT_Q_MAX[name], + JOINT_QD_MIN[name], + JOINT_QD_MAX[name], + JOINT_ACT_SCALE[name], + contact_prob=0.0, + ) + source_dir = Path(identity["root"]) / "envs" / "warp_sim_envs" / "assets" + asset = source_dir / ("ant.xml" if robot == "ant" else "cartpole_single.urdf") + scalar_names = [ + joint.attrib["name"] + for joint in ET.parse(asset).getroot().iter("joint") + if "name" in joint.attrib and joint.attrib.get("type") != "free" + ] + model = environment.model + nbody = environment.env.bodies_per_env + base_com = model.body_com.numpy()[0].tolist() + direction = [0.0, 0.0, 0.0] + direction[int(model.up_axis)] = -1.0 + physical = {} + for field in ( + "body_mass", + "body_inertia", + "body_com", + "joint_armature", + "joint_limit_lower", + "joint_limit_upper", + "joint_limit_ke", + "joint_limit_kd", + "joint_axis_mode", + "joint_target_ke", + "joint_target_kd", + ): + value = getattr(model, field, None) + if value is not None: + array = value.numpy() if hasattr(value, "numpy") else np.asarray(value) + physical[field] = array.tolist() + for field in ("mu", "ke", "kd", "kf"): + value = getattr(model.shape_materials, field, None) + if value is not None: + physical["shape_material_" + field] = value.numpy().tolist() + physical["shape_thickness"] = model.shape_geo.thickness.numpy().tolist() + physical["shape_body"] = model.shape_body.numpy().tolist() + physical["ground_plane"] = model.ground_plane.numpy().tolist() + gravity_vector = np.asarray(model.gravity, dtype=np.float64).reshape(3) + if not np.isfinite(gravity_vector).all(): + raise ValueError("Official teacher gravity must be finite") + physical["gravity"] = gravity_vector.tolist() + generator_paths = ( + "generate/trajectory_sampler.py", + "generate/simulation_sampler.py", + "utils/commons.py", + "utils/env_utils.py", + "utils/warp_utils.py", + "envs/neural_environment.py", + "envs/abstract_contact.py", + "envs/abstract_contact_environment.py", + "envs/warp_sim_envs/environment.py", + f"envs/warp_sim_envs/env_{robot}.py", + ) + if robot == "ant": + generator_paths += ("generate/trajectory_sampler_ant.py",) + generator_files = { + name: _hash(Path(identity["root"]) / name) for name in generator_paths + } + generator_hash = hashlib.sha256( + json.dumps(generator_files, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + metadata = { + "robot": robot, + "dt": environment.frame_dt, + "codec": { + "robot": robot, + "base_com": base_com, + "continuous_q_indices": [1] if robot == "cartpole" else [], + }, + "asset_sha256": _hash(asset), + "joint_names": scalar_names, + "q_dim": environment.dof_q_per_env, + "qd_dim": environment.dof_qd_per_env, + "teacher_origin": { + "kind": "official_warp", + "source_commit": OFFICIAL_COMMIT, + "warp_version": str(wp.__version__), + "integrator": str(environment.env.integrator_type), + "num_substeps": environment.env.sim_substeps, + "up_axis": int(model.up_axis), + "gravity": gravity_vector.tolist(), + "gravity_magnitude": float(np.linalg.norm(gravity_vector)), + "physics_contract_version": _OFFICIAL_PHYSICS_CONTRACT_VERSION, + "bodies_per_env": nbody, + "device": device, + }, + "physics_contract_version": _OFFICIAL_PHYSICS_CONTRACT_VERSION, + "source_physical_arrays": physical, + "source_generator_files": generator_files, + "source_generator_sha256": generator_hash, + "generation_adapter_sha256": _hash(Path(__file__)), + "source": identity, + "sampling": { + "seed": spec.get("seed", 0), + "initial_states_source": spec.get("initial_states_source", "sample"), + "passive": bool(spec.get("passive", False)), + "batch_size": batch, + "horizon": horizon, + }, + } + rotation = None + if spec.get("canonical_world", "source") == "z_up": + rotation_tensor = torch.tensor([math.sqrt(0.5), 0.0, 0.0, math.sqrt(0.5)]) + if robot == "cartpole": + rotation_tensor = quat_multiply( + torch.tensor([0.0, 0.0, math.sqrt(0.5), math.sqrt(0.5)]), + rotation_tensor, + ) + rotation = tuple(rotation_tensor.tolist()) + elif spec.get("canonical_world", "source") != "source": + raise ValueError("canonical_world must be source or z_up") + sources = {} + corrections = {} + source_counts = {} + for split_index, (split, count) in enumerate(splits.items()): + set_random_seed(int(spec.get("seed", 0)) + split_index) + path = output / f"{split}.raw.h5" + written = 0 + correction = None + with h5py.File(path, "w") as target: + data = target.create_group("data") + data.attrs["mode"] = "trajectory" + while written < int(count): + # Invoke the unchanged sampler for at most one environment batch + # and stream it immediately. Its numerical-invalid retry remains + # upstream behavior; no entire split is retained on CPU or GPU. + requested = min(int(count) - written, batch) + rollouts = sampler.sample_trajectories_joint_act_mode( + requested * horizon, + horizon, + passive=bool(spec.get("passive", False)), + initial_states_source=spec.get("initial_states_source", "sample"), + render=False, + ) + arrays = { + key: value[:, :requested].detach().cpu().numpy() + for key, value in rollouts.items() + if not isinstance(value, dict) + } + arrays.update( + { + key: value[:, :requested].detach().cpu().numpy() + for key, value in rollouts["contacts"].items() + } + ) + produced = arrays["states"].shape[1] + for key, value in arrays.items(): + value = value.reshape(*value.shape[:2], -1) + if key not in data: + data.create_dataset( + key, + shape=(horizon, 0, value.shape[-1]), + maxshape=(horizon, None, value.shape[-1]), + chunks=(min(horizon, 128), 1, value.shape[-1]), + dtype=value.dtype, + compression="lzf", + ) + data[key].resize(written + produced, axis=1) + data[key][:, written : written + produced] = value + written += produced + data.attrs["total_transitions"] = horizon * written + data.attrs["total_trajectories"] = written + target.flush() + _, observed = repair_gravity_transport( + arrays["gravity_dir"], + tuple(direction), + enabled=bool(spec.get("repair_uninitialized_gravity", False)), + ) + if correction is None: + correction = {**observed, "sampler_calls": 1} + else: + correction["changed_components"] += observed["changed_components"] + correction["nonfinite_components"] += observed[ + "nonfinite_components" + ] + correction["sampler_calls"] += 1 + del rollouts, arrays + corrections[split] = correction + source_counts[split] = written + sources[split] = path + metadata["gravity_repair"] = corrections + metadata["source_raw_hdf5"] = { + split: {"path": str(path), "sha256": _hash(path)} + for split, path in sources.items() + } + result = import_legacy_datasets( + sources, + output / "canonical", + metadata, + world_rotation_xyzw=rotation, + gravity_direction=( + tuple(direction) + if spec.get("repair_uninitialized_gravity", False) + else None + ), + preserve_legacy_velocity=spec.get("preserve_legacy_velocity", False), + ) + result.update( + status="completed", + source_trajectory_counts=source_counts, + teacher_origin=metadata["teacher_origin"], + gravity_repair=corrections, + source_bug_corrected=bool(spec.get("repair_uninitialized_gravity", False)), + elapsed_seconds=time.perf_counter() - started, + scope="real_official_teacher_data_not_full_training_qualification", + ) + (output / "generation-result.json").write_text( + json.dumps(result, indent=2, allow_nan=False) + ) + (output / "source-physics.json").write_text( + json.dumps(metadata, indent=2, allow_nan=False) + ) + return result + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--spec", type=Path, required=True) + parser.add_argument("--trust-external-code", action="store_true", required=True) + parser.add_argument("--resume", type=Path) + parser.add_argument("--stop-after", type=int) + parser.add_argument( + "--generate", + action="store_true", + help="Run real original teacher collection instead of training", + ) + args = parser.parse_args() + operation = generate_official_dataset if args.generate else run_reference + extra = ( + {} if args.generate else {"resume": args.resume, "stop_after": args.stop_after} + ) + result = operation( + json.loads(args.spec.read_text()), + trust_external_code=args.trust_external_code, + **extra, + ) + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/examples/learning/nerd/_support/_reference_assets.py b/examples/learning/nerd/_support/_reference_assets.py new file mode 100644 index 0000000000..618ec18d92 --- /dev/null +++ b/examples/learning/nerd/_support/_reference_assets.py @@ -0,0 +1,228 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Import pinned robot assets into an independent Newton teacher builder.""" + +from __future__ import annotations + +import hashlib +from pathlib import Path +from typing import Any +import xml.etree.ElementTree as ET + +import numpy as np + +__all__: list[str] = [] + +_SOURCE_COMMIT = "476dc67fccc6a99be293bf611f897b1c06f13e30" +_ASSETS = {"cartpole": "cartpole_single.urdf", "ant": "ant.xml"} +_JOINTS = { + "cartpole": ("slider_to_cart", "cart_to_pole"), + "ant": ( + "hip_1", + "ankle_1", + "hip_2", + "ankle_2", + "hip_3", + "ankle_3", + "hip_4", + "ankle_4", + ), +} + + +def _ant_reference_friction() -> dict[str, Any]: + """Describe the coupled old Euclidean-to-Newton Huber friction mapping. + + Newton uses ``H=delta*(speed-delta/2)`` above ``delta`` in both + normalization and the viscous/Coulomb minimum. Scaling both materials' + Coulomb coefficients by the same ``delta`` approximates the original + Euclidean law without changing the installed solver. This is a parameter + semantics conversion, not equality of the numerical material arrays. + """ + delta = 1e-3 + return { + "schema": "ant-warp-euclidean-to-newton-huber-friction-v1", + "friction_smoothing": delta, + "material_mu_scale": delta, + "reference_shape_mu": 0.75, + "reference_ground_mu": 0.5, + "shape_mu": 0.75 * delta, + "ground_mu": 0.5 * delta, + "interpretation": ( + "Coupled Huber/Coulomb parameter mapping; approximate original " + "Euclidean friction above the smoothing scale. Actual Newton " + "coefficients differ from the original coefficients." + ), + } + + +def _asset(robot: str) -> Path: + if robot not in _ASSETS: + raise ValueError(f"Unsupported NeRD reference robot: {robot}") + from .._resources import prepare_reference_assets + + return prepare_reference_assets() / _ASSETS[robot] + + +def reference_asset_hash(robot: str) -> str: + """Return the SHA256 of the pinned downloaded reference robot source.""" + return hashlib.sha256(_asset(robot).read_bytes()).hexdigest() + + +def reference_joint_names(robot: str) -> tuple[str, ...]: + """Return actuated/scalar joints in the official generalized-state order.""" + _asset(robot) + return _JOINTS[robot] + + +def reference_profile(robot: str) -> dict[str, Any]: + """Return source identity, state dimensions, controls and default reset state.""" + _asset(robot) + return { + "robot": robot, + "dt": 1 / 60, + "source_commit": _SOURCE_COMMIT, + "asset_sha256": reference_asset_hash(robot), + "joint_names": list(reference_joint_names(robot)), + "q_dim": 2 if robot == "cartpole" else 15, + "qd_dim": 2 if robot == "cartpole" else 14, + "default_q": ( + [0.0, 0.0] + if robot == "cartpole" + else [ + 0.0, + 0.0, + 1.0, + 0.0, + 0.0, + 0.0, + 1.0, + 0.0, + 1.0, + 0.0, + -1.0, + 0.0, + -1.0, + 0.0, + 1.0, + ] + ), + "default_qd": [0.0] * (2 if robot == "cartpole" else 14), + "effort_limits": [1500.0, 0.0] if robot == "cartpole" else [400.0] * 8, + "codec": { + "robot": robot, + "base_com": [0.0, 0.0, 0.0], + "continuous_q_indices": [1] if robot == "cartpole" else [], + }, + "teacher_profile": "current_newton_import_generalized_effort", + "source_actuators_imported": False, + "source_solver_options_imported": False, + } + + +def _import_reference(robot: str, physics_profile: str = "native") -> Any: + """Interpret the pinned asset with upstream environment import defaults.""" + import newton + import warp as wp + + if physics_profile not in {"native", "warp_reference"}: + raise ValueError(f"Unsupported reference physics profile: {physics_profile}") + legacy = physics_profile == "warp_reference" + builder = newton.ModelBuilder() + if robot == "cartpole": + builder.default_joint_cfg = newton.ModelBuilder.JointDofConfig( + armature=0.01, limit_ke=1e4, limit_kd=10.0 + ) + if legacy: + builder.default_shape_cfg = newton.ModelBuilder.ShapeConfig( + density=1000.0, mu=0.25 + ) + builder.add_urdf( + str(_asset(robot)), + xform=wp.transform( + (0.0, 0.0, 0.0), + ( + wp.quat_identity() + if legacy + else wp.quat_from_axis_angle(wp.vec3(0.0, 0.0, 1.0), wp.pi * 0.5) + ), + ), + floating=False, + enable_self_collisions=False, + collapse_fixed_joints=True, + ignore_inertial_definitions=legacy, + ) + if legacy: + # Warp 1.8 URDF imports ignore authored inertials by default and + # add armature to BOTH body inertia and generalized joint inertia. + for index, inertia in enumerate(builder.body_inertia): + builder.body_inertia[index] = wp.mat33( + np.asarray(inertia).reshape(3, 3) + 0.01 * np.eye(3) + ) + builder.body_inv_inertia[index] = wp.inverse( + builder.body_inertia[index] + ) + else: + builder.default_joint_cfg = newton.ModelBuilder.JointDofConfig( + armature=0.05, limit_ke=1e3, limit_kd=1e2 + ) + builder.default_shape_cfg = newton.ModelBuilder.ShapeConfig( + ke=4e4, kd=1e2, kf=3e3, mu=0.75 + ) + source = str(_asset(robot)) + if legacy: + # Warp 1.5.1 honors density for spheres but uses the parser's + # density=1000 for capsules, ignoring MJCF geom density=5. + # Express those measured semantics in a temporary XML string; + # the pinned, packaged asset remains byte-for-byte unchanged. + root = ET.parse(source).getroot() + for geom in root.findall(".//body/geom"): + if geom.get("type", "sphere") == "capsule": + geom.set("density", "1000") + # Old Y-axis capsules emit +Y then -Y contact slots. The + # modern fromto importer directs +Z toward the other end. + # Reverse the equivalent capsule frame to preserve the + # network's ordered contact features, not just geometry. + endpoints = geom.get("fromto", "").split() + if len(endpoints) == 6: + geom.set("fromto", " ".join(endpoints[3:] + endpoints[:3])) + geom.set("friction", "0.75 0 0") + geom.set("margin", "0") + source = ET.tostring(root, encoding="unicode") + builder.add_mjcf( + source, + enable_self_collisions=False, + collapse_fixed_joints=True, + ignore_names=["floor", "ground"], + parse_mujoco_options=not legacy, + ) + if legacy: + # The original root authors armature=0; scalar joints author .01. + # Old force mode ignores target_kd=.1. Newton now imports that + # value as passive damping, so leaving it active changes physics. + for index in range(1, builder.body_count): + builder.body_inertia[index] = wp.mat33( + np.asarray(builder.body_inertia[index]).reshape(3, 3) + + 0.01 * np.eye(3) + ) + builder.body_inv_inertia[index] = wp.inverse( + builder.body_inertia[index] + ) + builder.joint_damping[:] = [0.0] * len(builder.joint_damping) + builder.joint_target_ke[:] = [0.0] * len(builder.joint_target_ke) + builder.joint_target_kd[:] = [0.0] * len(builder.joint_target_kd) + return builder diff --git a/examples/learning/nerd/_support/_reference_rollout.py b/examples/learning/nerd/_support/_reference_rollout.py new file mode 100644 index 0000000000..a1f9975a3d --- /dev/null +++ b/examples/learning/nerd/_support/_reference_rollout.py @@ -0,0 +1,1117 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Explicit historical-Warp rollout qualification using original NeRD runtime. + +Run in an isolated subprocess with pinned Warp 1.8.0 (Cartpole) or 1.5.1 (Ant). +This harness does not deploy old physical models into current Newton/DexSim. +""" + +from __future__ import annotations +import argparse +from contextlib import contextmanager +import hashlib +import json +import math +from pathlib import Path +import platform +import sys +import time +from typing import Any +from collections.abc import Mapping, Sequence +from ._reference import OFFICIAL_COMMIT, reference_identity, _hash + +__all__ = ["evaluate_reference_rollouts"] + + +def _positive(name: str, value: Any, minimum: int = 1) -> int: + if type(value) is not int or value < minimum: + raise ValueError(f"{name} must be an integer >= {minimum}") + return value + + +@contextmanager +def _cpu_stream_compatibility(wp: Any, device: str): + """Adapt only the original runtime's CUDA-stream request on CPU. + + NeuralIntegrator.simulate unconditionally requests a Torch CUDA stream. + torch.cuda.stream(None) is the official no-op CPU compatibility boundary; + simulation, preprocessing, network, decoding and FK methods stay unmodified. + """ + original = wp.stream_to_torch + if not wp.get_device(device).is_cuda: + + def compatible(stream_or_device: Any): + candidate = getattr(stream_or_device, "device", stream_or_device) + if not wp.get_device(candidate).is_cuda: + return None + return original(stream_or_device) + + wp.stream_to_torch = compatible + try: + yield + finally: + wp.stream_to_torch = original + + +def _physical_arrays(model: Any) -> dict[str, Any]: + """Read the same physical fields recorded by original-teacher collection.""" + import numpy as np + + result = {} + for name in ( + "body_mass", + "body_inertia", + "body_com", + "joint_armature", + "joint_limit_lower", + "joint_limit_upper", + "joint_limit_ke", + "joint_limit_kd", + "joint_axis_mode", + "joint_target_ke", + "joint_target_kd", + ): + value = getattr(model, name, None) + if value is not None: + result[name] = ( + value.numpy() if hasattr(value, "numpy") else np.asarray(value) + ) + for field in ("mu", "ke", "kd", "kf"): + value = getattr(model.shape_materials, field, None) + if value is not None: + result["shape_material_" + field] = value.numpy() + result["shape_thickness"] = model.shape_geo.thickness.numpy() + result["shape_body"] = model.shape_body.numpy() + result["ground_plane"] = model.ground_plane.numpy() + result["gravity"] = np.asarray(model.gravity, dtype=np.float64) + return result + + +def _canonical_physics( + arrays: Mapping[str, Any], batch: int, bodies: int +) -> dict[str, Any]: + """Remove replica indexing while proving the recorded replicas homogeneous.""" + import numpy as np + + result = {} + shape_body = np.asarray(arrays["shape_body"]) + # Original Warp always appends one shared ground-plane shape. + if (len(shape_body) - 1) % batch: + raise ValueError("Historical shape replication is not homogeneous") + shapes = (len(shape_body) - 1) // batch + for key, values in arrays.items(): + value = np.asarray(values) + if key in {"ground_plane", "gravity"}: + result[key] = value.tolist() + continue + if key.startswith("shape_"): + if len(value) != batch * shapes + 1: + raise ValueError(f"Unexpected historical shape field {key}") + rows = value[:-1].reshape(batch, shapes, *value.shape[1:]).copy() + if key == "shape_body": + rows = np.where( + rows >= 0, rows - np.arange(batch)[:, None] * bodies, rows + ) + tail = value[-1:].tolist() + else: + if len(value) % batch: + raise ValueError(f"Unexpected historical replication for {key}") + rows = value.reshape(batch, len(value) // batch, *value.shape[1:]) + tail = None + if not np.allclose(rows, rows[:1], rtol=1e-6, atol=1e-7, equal_nan=False): + raise ValueError(f"Historical physical field {key} differs across replicas") + result[key] = {"per_environment": rows[0].tolist(), "shared": tail} + return result + + +def _verify_historical_physics( + environment: Any, + manifest: Mapping[str, Any], + source: Mapping[str, Any], + warp_version: str, +) -> dict[str, Any]: + """Require actual historical physical data, never relabel a Newton manifest.""" + import numpy as np + + recorded = manifest.get("teacher_origin", {}) + if ( + recorded.get("kind") != "official_warp" + or recorded.get("source_commit") != OFFICIAL_COMMIT + or recorded.get("warp_version") != warp_version + ): + raise ValueError( + "Model does not record the required original-Warp teacher identity" + ) + if not isinstance(manifest.get("source_physical_arrays"), Mapping): + raise ValueError( + "Historical source_physical_arrays are required; do not relabel a Newton model" + ) + if ( + recorded.get("up_axis") != int(environment.model.up_axis) + or recorded.get("num_substeps") != environment.sim_substeps_gt + ): + raise ValueError("Original teacher axis/substeps differ from model provenance") + if recorded.get("integrator") != str(environment.integrator_type_gt): + raise ValueError("Original teacher integrator differs from model provenance") + if "gravity" in recorded and not np.allclose( + np.asarray(recorded["gravity"]), + np.asarray(environment.model.gravity, dtype=np.float64), + rtol=1e-6, + atol=1e-7, + ): + raise ValueError("Original teacher gravity differs from recorded provenance") + if not math.isclose( + float(manifest.get("dt", 0)), + float(environment.frame_dt), + rel_tol=0, + abs_tol=1e-12, + ): + raise ValueError("Original teacher timestep differs from model provenance") + robot = manifest["codec"]["robot"] + expected_continuous = [1] if robot == "cartpole" else [] + if ( + list(manifest["codec"].get("continuous_q_indices", expected_continuous)) + != expected_continuous + ): + raise ValueError( + "Historical codec continuous coordinates differ from original runtime" + ) + if not np.allclose( + manifest["codec"].get("base_com", [0.0, 0.0, 0.0]), + environment.model.body_com.numpy()[0], + rtol=1e-6, + atol=1e-7, + ): + raise ValueError("Historical codec base COM differs from actual source model") + asset = ( + Path(source["root"]) + / "envs" + / "warp_sim_envs" + / "assets" + / ("ant.xml" if robot == "ant" else "cartpole_single.urdf") + ) + if manifest.get("asset_sha256") != _hash(asset): + raise ValueError("Historical model asset checksum differs") + batch = manifest.get("sampling", {}).get("batch_size") + split_batch = ( + manifest.get("source_split_identities", {}).get("train", {}).get("num_envs") + ) + if batch is not None and split_batch is not None and batch != split_batch: + raise ValueError("Conflicting source sampling batch sizes") + batch = _positive( + "source sampling batch_size", batch if batch is not None else split_batch + ) + bodies = _positive("source bodies_per_env", recorded.get("bodies_per_env")) + if bodies != environment.bodies_per_env: + raise ValueError("Historical body count differs") + actual = _canonical_physics( + _physical_arrays(environment.model), environment.num_envs, bodies + ) + expected = _canonical_physics(manifest["source_physical_arrays"], batch, bodies) + missing_fields = sorted(set(actual) - set(expected)) + if set(expected) - set(actual) or set(missing_fields) - {"gravity"}: + raise ValueError("Historical physical field set differs") + for key in expected: + if key in {"ground_plane", "gravity"}: + pairs = [(actual[key], expected[key])] + else: + pairs = [ + (actual[key][name], expected[key][name]) + for name in ("per_environment", "shared") + if actual[key][name] is not None or expected[key][name] is not None + ] + for left, right in pairs: + if ( + left is None + or right is None + or np.shape(left) != np.shape(right) + or not np.allclose(left, right, rtol=1e-6, atol=1e-7) + ): + raise ValueError(f"Historical physical mismatch in {key}") + identity = hashlib.sha256( + json.dumps( + actual, sort_keys=True, separators=(",", ":"), allow_nan=False + ).encode() + ).hexdigest() + return { + "schema": "official-warp-physical-arrays-v2", + "sha256": identity, + "matched_fields": sorted(expected), + "missing_recorded_fields": missing_fields, + "matched_recorded_fields": True, + "complete": not missing_fields, + "status": ( + "matched" if not missing_fields else "pending_missing_recorded_physics" + ), + "gravity_live": np.asarray( + environment.model.gravity, dtype=np.float64 + ).tolist(), + "gravity_magnitude_in_training_metadata": "gravity" in expected, + } + + +def _canonical_states(states: Any, codec: Any, rotation: Any) -> tuple[Any, Any]: + import torch + from embodichain_models.dynamics.nerd.codec import quat_rotate, quat_multiply + + q = states[..., : codec.q_dim] + qd = codec.legacy_to_native(q, states[..., codec.q_dim :]) + if codec.robot == "ant": + q = torch.cat( + ( + quat_rotate(rotation, q[..., :3]), + quat_multiply(rotation, q[..., 3:7]), + q[..., 7:], + ), + -1, + ) + qd = torch.cat( + ( + quat_rotate(rotation, qd[..., :3]), + quat_rotate(rotation, qd[..., 3:6]), + qd[..., 6:], + ), + -1, + ) + return q, qd + + +class _HistoricalPD: + """Declared fixed PD baseline operating on the shared canonical state schema.""" + + def __init__( + self, robot: str, config: Mapping[str, Any], limits: Any, target: Any + ) -> None: + import torch + + self.robot, self.limits, self.target = robot, limits, target + if config.get("kind") != "pd": + raise ValueError( + "Historical runtime currently supports policy.kind=pd only; actor artifacts are unsupported" + ) + allowed = {"kind", "gains"} if robot == "cartpole" else {"kind", "kp", "kd"} + if set(config) - allowed: + raise ValueError("Unknown historical PD setting") + self.identity = { + "kind": "pd", + "implementation": "embodichain-reference-feedback-v1", + "observation_schema": ( + "cartpole-state-4-v1" if robot == "cartpole" else "ant-com-state-27-v1" + ), + "deterministic": True, + "scope": "fixed_baseline_not_paper_policy", + "effort_limits": limits.detach().cpu().tolist(), + } + if robot == "cartpole": + values = config.get("gains", [5.0, 50.0, 5.0, 20.0]) + if ( + not isinstance(values, Sequence) + or len(values) != 4 + or any( + isinstance(x, bool) + or not isinstance(x, (int, float)) + or not math.isfinite(x) + for x in values + ) + ): + raise ValueError("Cartpole gains must be four finite values") + self.gains = torch.tensor(values, device=limits.device) + self.identity["gains"] = list(values) + else: + self.kp, self.kd = config.get("kp", 20.0), config.get("kd", 1.0) + if any( + isinstance(x, bool) + or not isinstance(x, (int, float)) + or not math.isfinite(x) + or x < 0 + for x in (self.kp, self.kd) + ): + raise ValueError("Ant gains must be finite and nonnegative") + self.identity.update( + kp=float(self.kp), + kd=float(self.kd), + target=target.detach().cpu().tolist(), + ) + self.identity["sha256"] = hashlib.sha256( + json.dumps(self.identity, sort_keys=True).encode() + ).hexdigest() + + def __call__(self, q: Any, qd: Any) -> Any: + import torch + from ._task_values import reference_task_values + + observation, _, _ = reference_task_values(self.robot, q, qd) + finite = torch.isfinite(observation).all(-1) + if self.robot == "cartpole": + force = torch.zeros_like(qd) + force[:, 0] = ( + -(torch.nan_to_num(observation) * self.gains) + .sum(-1) + .clamp(-self.limits[0], self.limits[0]) + ) + else: + force = ( + self.kp * (self.target - torch.nan_to_num(q[:, 7:])) + - self.kd * torch.nan_to_num(qd[:, 6:]) + ).clamp(-self.limits, self.limits) + return torch.where(finite[:, None], force, 0) + + +def evaluate_reference_rollouts( + spec: Mapping[str, Any], *, trust_external_code: bool = False +) -> dict[str, Any]: + """Run original physics and original neural integration in their historical domain. + + Args: + spec: source checkout, portable model export, robot, output JSON path, + device, num_envs, trajectories, horizons, seed, motion_mode and + initial_states(default/sample). Fixed-policy mode requires a PD policy + mapping. verify_portable_steps bounds an extra equivalence probe. + teacher_dataset optionally selects immutable evaluation traces made + by the standalone generator in the original supported environment. + trust_external_code: Explicit permission to execute the pinned external + checkout. Run in a fresh subprocess with its legacy Warp on PYTHONPATH. + + Returns: + Measured error curves, returns, failures and hashed trajectory/source + receipts compatible with qualification report attachment. This provides + no evidence of deployment compatibility with current Newton. + """ + if not trust_external_code: + raise ValueError("Historical runtime requires trust_external_code=True") + allowed = { + "source", + "model", + "robot", + "output", + "device", + "num_envs", + "trajectories", + "horizons", + "seed", + "motion_mode", + "initial_states", + "policy", + "warp_version", + "cpu_threads", + "verify_portable_steps", + "teacher_dataset", + } + if set(spec) - allowed: + raise ValueError( + f"Unknown historical rollout settings: {sorted(set(spec)-allowed)}" + ) + robot = spec["robot"] + if robot not in {"cartpole", "ant"}: + raise ValueError("Historical rollout supports cartpole and ant") + batch = _positive("num_envs", spec.get("num_envs", 8)) + count = _positive("trajectories", spec.get("trajectories", 2048)) + seed = _positive("seed", spec.get("seed", 500), 0) + horizons = sorted( + {_positive("horizon", h) for h in spec.get("horizons", [1, 10, 100, 1000])} + ) + if not horizons: + raise ValueError("horizons must not be empty") + horizon = horizons[-1] + mode = spec.get("motion_mode", "random_effort") + if mode not in {"passive", "random_effort", "fixed_policy"}: + raise ValueError("Unknown motion_mode") + if (mode == "fixed_policy") != ("policy" in spec): + raise ValueError( + "fixed_policy requires an explicit PD policy only in that mode" + ) + initial = spec.get("initial_states", "sample") + if initial not in {"default", "sample"}: + raise ValueError("initial_states must be default or sample") + expected_warp = "1.5.1" if robot == "ant" else "1.8.0" + if spec.get("warp_version", expected_warp) != expected_warp: + raise ValueError("Historical robot requires its pinned Warp version") + identity = reference_identity(spec["source"]) + import warp as wp + + if str(wp.__version__) != expected_warp: + raise ValueError( + f"Run in isolated Warp {expected_warp}; loaded {wp.__version__}" + ) + sys.path.insert(0, identity["root"]) + import torch + import numpy as np + import h5py + from envs.neural_environment import NeuralEnvironment + from generate.trajectory_sampler import TrajectorySampler + from generate.trajectory_sampler_ant import TrajectorySamplerAnt + from utils.commons import ( + JOINT_Q_MIN, + JOINT_Q_MAX, + JOINT_QD_MIN, + JOINT_QD_MAX, + JOINT_ACT_SCALE, + ) + from utils.python_utils import set_random_seed + from embodichain_models.dynamics.nerd.codec import NeRDStateCodec + from ._task_values import reference_task_values + from examples.learning.nerd._support._evaluation import ( + _error_values, + _safe_number, + _teacher_row_sha256, + ) + from embodichain_models.dynamics.nerd import load_model + from ._reference import _restore_official_model + + device = str(spec.get("device", "cpu")) + torch.set_num_threads(_positive("cpu_threads", spec.get("cpu_threads", 1))) + model_path = Path(spec["model"]) + output = Path(spec["output"]) + model_sha, manifest_sha = _hash(model_path / "weights.pt"), _hash( + model_path / "manifest.json" + ) + common, manifest = load_model(model_path, device=device) + if ( + next(common.parameters()).dtype != torch.float32 + or manifest.get("codec", {}).get("robot") != robot + ): + raise ValueError("Historical runtime requires a matching FP32 model") + codec = NeRDStateCodec(**manifest["codec"]) + root_half = math.sqrt(0.5) + expected_rotation = ( + [0.5, 0.5, 0.5, 0.5] + if robot == "cartpole" + else [root_half, 0.0, 0.0, root_half] + ) + rotation = torch.tensor( + manifest.get("world_rotation_xyzw") or expected_rotation, + device=device, + dtype=torch.float32, + ) + if not torch.allclose( + rotation, rotation.new_tensor(expected_rotation), atol=1e-6, rtol=0 + ): + raise ValueError("Unsupported historical-to-canonical world rotation") + verify_steps = min( + horizon, + _positive( + "verify_portable_steps", + spec.get("verify_portable_steps", common.cfg.history_length + 1), + ), + ) + trace_path = output.with_name(output.stem + "-evidence.h5") + rows_path = output.with_name(output.stem + "-trajectories.jsonl") + for path in (output, trace_path, rows_path): + if path.exists(): + raise FileExistsError(path) + output.parent.mkdir(parents=True, exist_ok=True) + runtime_files = ( + "envs/neural_environment.py", + "envs/abstract_contact_environment.py", + "envs/warp_sim_envs/environment.py", + "envs/warp_sim_envs/env_cartpole.py", + "envs/warp_sim_envs/env_ant.py", + "integrators/integrator_neural.py", + "integrators/integrator_neural_stateful.py", + "integrators/integrator_neural_transformer.py", + "generate/trajectory_sampler.py", + "generate/trajectory_sampler_ant.py", + "generate/simulation_sampler.py", + ) + identity["runtime_files"] = { + name: _hash(Path(identity["root"]) / name) for name in runtime_files + } + identity["adapter"] = ( + "Original NeuralEnvironment, TransformerNeuralIntegrator, sampler and torque APIs; portable weights restored into original model; canonical coordinate transport and metric collection only" + ) + integrator_cfg = { + "name": "TransformerNeuralIntegrator", + "num_states_history": common.cfg.history_length, + "states_frame": "body" if robot == "ant" else "world", + "anchor_frame_step": "every", + "states_embedding_type": "identical", + "prediction_type": "relative", + "orientation_prediction_parameterization": "quaternion", + } + + def environment(model=None): + return NeuralEnvironment( + env_name="Ant" if robot == "ant" else "Cartpole", + num_envs=batch, + warp_env_cfg={ + "seed": seed, + "random_reset": False, + "env_offset": (0.0, 0.0, 0.0), + "use_graph_capture": False, + }, + neural_integrator_cfg=integrator_cfg, + neural_model=model, + default_env_mode="ground-truth" if model is None else "neural", + device=device, + render=False, + ) + + class PortableAdapter(torch.nn.Module): + def __init__(self, model): + super().__init__() + self.common = model + + def evaluate(self, inputs): + return self.common.evaluate( + { + "joint_f" if key == "joint_acts" else key: value + for key, value in inputs.items() + } + ) + + def init_rnn(self, batch_size): + del batch_size + + teacher = student = portable = None + started = time.perf_counter() + sums = torch.zeros(horizon, 4, device=device, dtype=torch.float64) + counts = torch.zeros(horizon, 5, device=device, dtype=torch.long) + returns = torch.zeros(horizon, 2, device=device, dtype=torch.float64) + hsum = {h: torch.zeros(4, device=device, dtype=torch.float64) for h in horizons} + hcount = {h: torch.zeros((), device=device, dtype=torch.long) for h in horizons} + selected_digest = hashlib.sha256() + parity_max = 0.0 + parity_nonfinite = False + teacher_dataset = None + teacher_feedback_max_abs_error = 0.0 + teacher_transport_max_abs_error = 0.0 + teacher_trace_physical = None + names = ("q_mae", "qd_mae", "base_position_error", "base_orientation_error") + with wp.ScopedDevice(device): + try: + teacher = environment() + expected_inputs = { + "states_embedding": teacher.state_dim, + "joint_f": teacher.joint_act_dim, + } + if robot == "ant": + slots = teacher.num_contacts_per_env + expected_inputs.update( + contact_normals=3 * slots, + contact_points_1=3 * slots, + contact_depths=slots, + gravity_dir=3, + ) + if ( + dict(common.cfg.input_dims) != expected_inputs + or common.cfg.output_dim != teacher.state_dim + ): + raise ValueError( + "Model feature/target schema does not match original runtime" + ) + physical = _verify_historical_physics( + teacher, manifest, identity, expected_warp + ) + original = _restore_official_model(common, device=device) + student = environment(original) + _verify_historical_physics(student, manifest, identity, expected_warp) + portable = environment(PortableAdapter(common)) + _verify_historical_physics(portable, manifest, identity, expected_warp) + name = teacher.robot_name + sampler_type = TrajectorySamplerAnt if robot == "ant" else TrajectorySampler + sampler = sampler_type( + teacher, + JOINT_Q_MIN[name], + JOINT_Q_MAX[name], + JOINT_QD_MIN[name], + JOINT_QD_MAX[name], + JOINT_ACT_SCALE[name], + contact_prob=0.0, + ) + default = torch.cat( + ( + wp.to_torch(teacher.model.joint_q).reshape(batch, -1), + wp.to_torch(teacher.model.joint_qd).reshape(batch, -1), + ), + -1, + )[0].clone() + limits = sampler.joint_act_scale + canonical_default, _ = _canonical_states(default[None], codec, rotation) + policy = ( + _HistoricalPD(robot, spec["policy"], limits, canonical_default[0, 7:]) + if mode == "fixed_policy" + else None + ) + if spec.get("teacher_dataset") is not None: + from ._reference_traces import ( + _load_reference_trace, + _reference_trace_batches, + ) + + teacher_directory = Path(spec["teacher_dataset"]) + teacher_dataset, teacher_records = _load_reference_trace( + teacher_directory, + spec, + identity, + manifest, + None if policy is None else policy.identity, + count, + horizon, + ) + teacher_trace_physical = _verify_historical_physics( + teacher, teacher_dataset["metadata"], identity, expected_warp + ) + if not teacher_trace_physical["complete"]: + raise ValueError( + "Teacher trace is missing measured physical fields" + ) + batches = _reference_trace_batches( + teacher_directory, teacher_records, batch, horizon, robot, device + ) + else: + batches = (None for _ in range(0, count, batch)) + with ( + _cpu_stream_compatibility(wp, device), + torch.no_grad(), + h5py.File(trace_path, "w") as trace, + rows_path.open("w") as rows_file, + ): + for start, recorded in zip(range(0, count, batch), batches): + n = min(batch, count - start) + cached = None if recorded is None else recorded[1] + if cached is None: + samples = [] + commands = [] + for trajectory in range(start, start + n): + set_random_seed(seed + trajectory) + raw = default[None].clone() + if initial == "sample": + sampler.sample_initial_states(1, raw) + command = torch.zeros( + horizon, teacher.joint_act_dim, device=device + ) + if mode == "random_effort": + sampler.sampler.sample( + horizon, -limits, limits, command + ) + samples.append(raw[0]) + commands.append(command) + samples.extend([samples[-1]] * (batch - n)) + commands.extend([commands[-1]] * (batch - n)) + initial_raw = torch.stack(samples) + command = torch.stack(commands) + teacher.reset(initial_states=initial_raw.clone()) + else: + initial_raw = cached["raw"][:, 0].contiguous() + command = cached["raw_joint_f"] + transported = _canonical_states(cached["raw"], codec, rotation) + for actual, key in zip(transported, ("q", "qd")): + if not torch.allclose( + actual, + cached[key], + atol=1e-5, + rtol=1e-4, + equal_nan=True, + ): + raise ValueError( + "Recorded raw/canonical teacher state transport differs" + ) + error = (actual - cached[key]).abs() + teacher_transport_max_abs_error = max( + teacher_transport_max_abs_error, + float( + torch.where(torch.isfinite(error), error, 0).max() + ), + ) + if mode == "passive" and bool((command != 0).any()): + raise ValueError( + "Passive teacher trace contains nonzero effort" + ) + student.reset(initial_states=initial_raw.clone()) + if start == 0: + portable.reset(initial_states=initial_raw.clone()) + teacher_raw = [teacher.states.clone()] if cached is None else None + student_raw = [student.states.clone()] + teacher_controls = [] + student_controls = [] + for step in range(horizon): + tq, tqd = ( + _canonical_states(teacher.states, codec, rotation) + if cached is None + else (cached["q"][:, step], cached["qd"][:, step]) + ) + nq, nqd = _canonical_states(student.states, codec, rotation) + tf = ( + command[:, step] + if policy is None or cached is not None + else policy(tq, tqd) + ) + nf = command[:, step] if policy is None else policy(nq, nqd) + if cached is not None and policy is not None: + expected = policy(tq, tqd) + if not torch.allclose(tf, expected, atol=1e-5, rtol=1e-4): + raise ValueError( + "Recorded original teacher actions differ from declared feedback" + ) + teacher_feedback_max_abs_error = max( + teacher_feedback_max_abs_error, + float((tf - expected).abs().max()), + ) + # Original torque API uses view(-1), so transport a + # contiguous frame from batch-major command storage. + tf, nf = tf.contiguous(), nf.contiguous() + teacher_controls.append(tf.clone()) + student_controls.append(nf.clone()) + if cached is None: + teacher_raw.append( + teacher.step_with_joint_act( + tf, env_mode="ground-truth" + ).clone() + ) + student_raw.append( + student.step_with_joint_act(nf, env_mode="neural").clone() + ) + if start == 0 and step < verify_steps: + alternate = portable.step_with_joint_act( + nf, env_mode="neural" + ).clone() + torch.testing.assert_close( + alternate, + student_raw[-1], + atol=1e-5, + rtol=1e-4, + equal_nan=True, + ) + difference = (alternate - student_raw[-1]).abs() + parity_nonfinite = parity_nonfinite or not bool( + torch.isfinite(difference).all() + ) + parity_max = max( + parity_max, + float( + torch.nan_to_num( + difference, nan=0.0, posinf=0.0 + ).max() + ), + ) + traw = ( + torch.stack(teacher_raw, 1)[:n] + if cached is None + else cached["raw"][:n] + ) + nraw = torch.stack(student_raw, 1)[:n] + tq, tqd = ( + _canonical_states(traw, codec, rotation) + if cached is None + else (cached["q"][:n], cached["qd"][:n]) + ) + q, qd = _canonical_states(nraw, codec, rotation) + tc = torch.stack(teacher_controls, 1)[:n] + nc = torch.stack(student_controls, 1)[:n] + if robot == "ant": + tc = torch.cat((tc.new_zeros(n, horizon, 6), tc), -1) + nc = torch.cat((nc.new_zeros(n, horizon, 6), nc), -1) + tb = ~(torch.isfinite(tq).all(-1) & torch.isfinite(tqd).all(-1)) + nb = ~(torch.isfinite(q).all(-1) & torch.isfinite(qd).all(-1)) + tn = tb.cummax(1).values[:, 1:] + nn = nb.cummax(1).values[:, 1:] + valid = ~(tn | nn) + _, treward, tfail = reference_task_values(robot, tq, tqd) + _, nreward, nfail = reference_task_values(robot, q, qd) + tfail = (tfail | tb | ~torch.isfinite(treward)).cummax(1).values + nfail = (nfail | nb | ~torch.isfinite(nreward)).cummax(1).values + tret = torch.where( + ~tfail[:, :-1] & torch.isfinite(treward[:, 1:]), + treward[:, 1:], + 0, + ).cumsum(1) + nret = torch.where( + ~nfail[:, :-1] & torch.isfinite(nreward[:, 1:]), + nreward[:, 1:], + 0, + ).cumsum(1) + error = torch.where( + valid[..., None], + _error_values( + robot, q[:, 1:], qd[:, 1:], tq[:, 1:], tqd[:, 1:] + ), + 0, + ) + sums += error.double().sum(0) + counts += torch.stack( + ( + valid.sum(0), + tn.sum(0), + nn.sum(0), + tfail[:, 1:].sum(0), + nfail[:, 1:].sum(0), + ), + -1, + ) + returns += torch.stack( + (tret.double().sum(0), nret.double().sum(0)), -1 + ) + cumulative = error.cumsum(1) + for h in horizons: + hsum[h] += ( + torch.where( + valid[:, h - 1, None], cumulative[:, h - 1] / h, 0 + ) + .double() + .sum(0) + ) + hcount[h] += valid[:, h - 1].sum() + arrays = { + key: value.detach().cpu().numpy() + for key, value in { + "teacher_raw": traw, + "neural_raw": nraw, + "teacher_q": tq, + "teacher_qd": tqd, + "q": q, + "qd": qd, + "teacher_joint_f": tc, + "joint_f": nc, + "errors": error, + "valid": valid, + "teacher_nonfinite": tn, + "neural_nonfinite": nn, + "teacher_failed": tfail[:, 1:], + "neural_failed": nfail[:, 1:], + "teacher_return": tret, + "neural_return": nret, + }.items() + } + group = trace.create_group(str(start // batch)) + for key, value in arrays.items(): + group.create_dataset(key, data=value, compression="lzf") + for row in range(n): + record = { + "id": f"{seed}:{start+row}", + "metadata": {"sampling_seed": seed + start + row}, + } + digest = _teacher_row_sha256( + record, + [ + arrays["teacher_q"][row], + arrays["teacher_qd"][row], + arrays["teacher_joint_f"][row], + ], + ) + selected_digest.update(bytes.fromhex(digest)) + evidence = { + "trajectory_id": record["id"], + "sampling_seed": seed + start + row, + "reference_sha256": digest, + "evidence_group": str(start // batch), + "evidence_row": row, + "horizons": {}, + } + for h in horizons: + mean = ( + arrays["errors"][row, :h].mean(0) + if arrays["valid"][row, h - 1] + else [float("nan")] * 4 + ) + evidence["horizons"][str(h)] = { + **{ + key: _safe_number(value) + for key, value in zip(names, mean) + }, + "teacher_nonfinite": bool( + arrays["teacher_nonfinite"][row, h - 1] + ), + "neural_nonfinite": bool( + arrays["neural_nonfinite"][row, h - 1] + ), + "teacher_failed": bool( + arrays["teacher_failed"][row, h - 1] + ), + "neural_failed": bool( + arrays["neural_failed"][row, h - 1] + ), + "teacher_policy_return": _safe_number( + arrays["teacher_return"][row, h - 1] + ), + "policy_return": _safe_number( + arrays["neural_return"][row, h - 1] + ), + } + rows_file.write(json.dumps(evidence, allow_nan=False) + "\n") + origin = { + "kind": "official_warp", + "source_commit": OFFICIAL_COMMIT, + "warp_version": expected_warp, + "integrator": str(teacher.integrator_type_gt), + "num_substeps": teacher.sim_substeps_gt, + "up_axis": int(teacher.model.up_axis), + "bodies_per_env": teacher.bodies_per_env, + "device": device, + "physical_arrays_sha256": physical["sha256"], + "gravity": physical["gravity_live"], + } + finally: + for value in (portable, student, teacher): + if value is not None: + value.close() + if ( + _hash(model_path / "weights.pt") != model_sha + or _hash(model_path / "manifest.json") != manifest_sha + ): + raise ValueError("Model export changed during historical evaluation") + curve = [] + means = (sums / counts[:, 0].clamp_min(1)[:, None]).cpu().numpy() + count_array = counts.cpu().numpy() + return_array = (returns / count).cpu().numpy() + for step in range(horizon): + row = count_array[step] + curve.append( + { + "step": step + 1, + **{ + key: (_safe_number(v) if row[0] else None) + for key, v in zip(names, means[step]) + }, + "valid_trajectories": int(row[0]), + "teacher_nonfinite_count": int(row[1]), + "neural_nonfinite_count": int(row[2]), + "teacher_failed_count": int(row[3]), + "failed_trajectories": int(row[4]), + "teacher_nonfinite_rate": float(row[1] / count), + "nonfinite_rate": float(row[2] / count), + "teacher_failure_rate": float(row[3] / count), + "failure_rate": float(row[4] / count), + "teacher_policy_return": float(return_array[step, 0]), + "policy_return": float(return_array[step, 1]), + } + ) + metrics = {} + for h in horizons: + number = int(hcount[h].cpu()) + mean = (hsum[h] / max(number, 1)).cpu().tolist() + metrics[str(h)] = { + **curve[h - 1], + **{ + key: (_safe_number(v) if number else None) + for key, v in zip(names, mean) + }, + "valid_trajectories": number, + "numerically_invalid_trajectories": count - number, + } + result = { + "status": "completed", + "qualification_status": ( + "runtime_measurements_only" + if physical["complete"] + else "pending_missing_recorded_physics" + ), + "scope": "historical_official_warp_runtime_training_quality_only_not_current_newton_deployment", + "robot": robot, + "motion_mode": mode, + "seed": seed, + "model": str(model_path.resolve()), + "model_sha256": model_sha, + "model_manifest_sha256": manifest_sha, + "teacher_origin": origin, + "teacher_split": "test", + "teacher_dataset": ( + None + if teacher_dataset is None + else str(Path(spec["teacher_dataset"]).resolve()) + ), + "teacher_dataset_manifest_sha256": ( + None + if teacher_dataset is None + else _hash(Path(spec["teacher_dataset"]) / "manifest.json") + ), + "teacher_generation_environment": ( + None + if teacher_dataset is None + else teacher_dataset["metadata"]["generation_environment"] + ), + "teacher_feedback_verification": ( + "all_selected_actions_verified_against_teacher_states" + if teacher_dataset is not None and policy is not None + else "live_feedback_execution" if policy is not None else None + ), + "teacher_feedback_max_abs_error": teacher_feedback_max_abs_error, + "teacher_transport_max_abs_error": teacher_transport_max_abs_error, + "teacher_physical_verification": teacher_trace_physical, + "trajectories": count, + "num_envs": batch, + "state_dtype": "float32", + "selected_teacher_sha256": selected_digest.hexdigest(), + "selected_teacher_hash_schema": "ordered-teacher-content-v1", + "runtime_identity": { + "engine": "official_warp", + "reference_commit": OFFICIAL_COMMIT, + "device": device, + "state_dtype": "float32", + "python": platform.python_version(), + "versions": { + "torch": str(torch.__version__), + "warp-lang": str(wp.__version__), + }, + "cuda_runtime": torch.version.cuda, + }, + "reference": identity, + "physical_verification": physical, + "world_rotation_xyzw": rotation.cpu().tolist(), + "compatibility_adapters": ( + [] + if wp.get_device(device).is_cuda + else [ + "wp.stream_to_torch(cpu)->None within scoped adapter; original simulate unchanged" + ] + ), + "replica_placement": [0.0, 0.0, 0.0], + "sampling": "original sample_initial_states/UniformSampler with per-trajectory seed; no rejection/resampling", + "portable_runtime_equivalence": { + "steps": verify_steps, + "batch_size": batch, + "atol": 1e-5, + "rtol": 1e-4, + "max_abs_error": parity_max, + "matched_nonfinite": parity_nonfinite, + "status": "matched_nonfinite" if parity_nonfinite else "passed", + }, + "policy": None if policy is None else policy.identity, + "policy_sha256": None if policy is None else policy.identity["sha256"], + "horizons": metrics, + "error_vs_time": curve, + "trajectory_evidence": str(rows_path.resolve()), + "trajectory_evidence_sha256": _hash(rows_path), + "trace_evidence": str(trace_path.resolve()), + "trace_evidence_sha256": _hash(trace_path), + "failure_rate": curve[-1]["failure_rate"], + "nonfinite_rate": curve[-1]["nonfinite_rate"], + "teacher_failure_rate": curve[-1]["teacher_failure_rate"], + "teacher_nonfinite_rate": curve[-1]["teacher_nonfinite_rate"], + "policy_return": curve[-1]["policy_return"] if policy else None, + "teacher_policy_return": curve[-1]["teacher_policy_return"] if policy else None, + "metric_population": "Error means condition on no numerical failure through the horizon; physical failures stay included. Returns stop after first failure. All requested seeds are retained.", + "adapter_source_sha256": _hash(Path(__file__)), + "elapsed_seconds": time.perf_counter() - started, + } + output.write_text(json.dumps(result, indent=2, allow_nan=False) + "\n") + return result + + +def main() -> None: + """Run the isolated historical qualification worker from an explicit spec.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--spec", type=Path, required=True) + parser.add_argument("--trust-external-code", action="store_true", required=True) + args = parser.parse_args() + result = evaluate_reference_rollouts( + json.loads(args.spec.read_text()), trust_external_code=args.trust_external_code + ) + print(json.dumps(result, indent=2, allow_nan=False)) + + +if __name__ == "__main__": + main() diff --git a/examples/learning/nerd/_support/_reference_traces.py b/examples/learning/nerd/_support/_reference_traces.py new file mode 100644 index 0000000000..0fd582e14f --- /dev/null +++ b/examples/learning/nerd/_support/_reference_traces.py @@ -0,0 +1,198 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Bounded replay of supported-environment original-Warp evaluation traces.""" + +from __future__ import annotations + +from collections import OrderedDict +from collections.abc import Iterator, Mapping +from pathlib import Path +from typing import Any + +import h5py +import numpy as np +import torch + +from examples.learning.nerd._support._evaluation import _teacher_row_sha256 +from embodichain_models.dynamics.nerd.data import load_manifest +from ._reference import OFFICIAL_COMMIT + +__all__: list[str] = [] + + +def _load_reference_trace( + directory: Path, + spec: Mapping[str, Any], + source: Mapping[str, Any], + model_metadata: Mapping[str, Any], + policy_identity: Mapping[str, Any] | None, + count: int, + horizon: int, +) -> tuple[dict[str, Any], list[dict[str, Any]]]: + manifest = load_manifest(directory, verify=True) + metadata, evaluation = manifest.get("metadata", {}), manifest.get("evaluation", {}) + environment = metadata.get("generation_environment", {}) + robot = spec["robot"] + origin = metadata.get("teacher_origin", {}) + expected_warp = "1.5.1" if robot == "ant" else "1.8.0" + if ( + manifest.get("trace_type") != "evaluation_including_failures" + or metadata.get("robot") != robot + or origin.get("kind") != "official_warp" + or origin.get("source_commit") != OFFICIAL_COMMIT + or origin.get("warp_version") != expected_warp + or environment.get("python") != "3.8.20" + or str(environment.get("torch", "")).split("+")[0] != "2.2.2" + or environment.get("numpy") != "1.24.4" + or environment.get("warp") != expected_warp + or any( + environment.get("runtime_imports_" + name) is not False + for name in ("embodichain", "dexsim", "newton") + ) + ): + raise ValueError( + "Teacher trace must be original-Warp evaluation data generated in the supported isolated environment" + ) + if ( + metadata.get("asset_sha256") != model_metadata.get("asset_sha256") + or metadata.get("codec") != model_metadata.get("codec") + or metadata.get("dt") != model_metadata.get("dt") + or metadata.get("world_rotation_xyzw") + != model_metadata.get("world_rotation_xyzw") + or metadata.get("source_runtime_files") != source.get("runtime_files") + ): + raise ValueError( + "Teacher source, asset or coordinate contract differs from the model/runtime" + ) + if ( + evaluation.get("motion_mode") != spec.get("motion_mode", "random_effort") + or evaluation.get("seed") != spec.get("seed", 500) + or evaluation.get("initial_states") != spec.get("initial_states", "sample") + or evaluation.get("policy") != policy_identity + or type(evaluation.get("horizon")) is not int + or evaluation["horizon"] < horizon + ): + raise ValueError( + "Teacher evaluation seed, initial state, horizon, motion or policy mismatch" + ) + records = [ + {**row, "file": shard["file"]} + for shard in manifest["shards"] + for row in shard["trajectories"] + ] + if evaluation.get("trajectories") != len(records) or len(records) < count: + raise ValueError("Teacher corpus omits requested evaluation trajectories") + for index, row in enumerate(records): + if ( + row.get("split") != "test" + or row.get("length") != evaluation["horizon"] + or row.get("metadata", {}).get("sampling_seed") + != evaluation["seed"] + index + or row.get("id") != f"{evaluation['seed']}:{index}" + ): + raise ValueError( + "Teacher corpus reorders, filters or mislabels evaluation trajectories" + ) + return manifest, records[:count] + + +def _reference_trace_batches( + directory: Path, + records: list[dict[str, Any]], + batch: int, + horizon: int, + robot: str, + device: str, +) -> Iterator[tuple[list[dict[str, Any]], dict[str, torch.Tensor]]]: + fields = { + "raw": "states/raw", + "q": "states/q", + "qd": "states/qd", + "joint_f": "transitions/joint_f", + "raw_joint_f": "transitions/raw_joint_f", + } + q_dim, qd_dim, force_dim = (15, 14, 8) if robot == "ant" else (2, 2, 2) + widths = { + "raw": q_dim + qd_dim, + "q": q_dim, + "qd": qd_dim, + "joint_f": qd_dim, + "raw_joint_f": force_dim, + } + handles: OrderedDict[str, h5py.File] = OrderedDict() + try: + for start in range(0, len(records), batch): + selected = records[start : start + batch] + values = [] + for row in selected: + filename = row["file"] + if filename not in handles: + if len(handles) >= 4: + _, stream = handles.popitem(last=False) + stream.close() + handles[filename] = h5py.File(directory / filename, "r") + handles.move_to_end(filename) + group = handles[filename][row["group"]] + arrays = {} + for name, key in fields.items(): + if key not in group: + raise ValueError( + "Original teacher trace is missing raw/canonical fields" + ) + steps = horizon + 1 if name in {"raw", "q", "qd"} else horizon + value = group[key][:steps] + if value.dtype != np.dtype("float32") or value.shape != ( + steps, + widths[name], + ): + raise ValueError( + "Original teacher trace has invalid dtype, shape or length" + ) + arrays[name] = value + if not np.isfinite(arrays["raw_joint_f"]).all(): + raise ValueError( + "Original teacher input efforts must remain finite" + ) + expected = ( + np.concatenate( + ( + np.zeros((horizon, 6), dtype=np.float32), + arrays["raw_joint_f"], + ), + -1, + ) + if robot == "ant" + else arrays["raw_joint_f"] + ) + if not np.array_equal(expected, arrays["joint_f"]): + raise ValueError( + "Recorded canonical and original teacher controls differ" + ) + row["reference_sha256"] = _teacher_row_sha256( + row, [arrays["q"], arrays["qd"], arrays["joint_f"]] + ) + values.append(arrays) + values.extend([values[-1]] * (batch - len(values))) + yield selected, { + name: torch.from_numpy(np.stack([value[name] for value in values])).to( + device + ) + for name in fields + } + finally: + for stream in handles.values(): + stream.close() diff --git a/examples/learning/nerd/_support/_report.py b/examples/learning/nerd/_support/_report.py new file mode 100644 index 0000000000..3eeaa3d67c --- /dev/null +++ b/examples/learning/nerd/_support/_report.py @@ -0,0 +1,1148 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Evidence-based comparison reports with seed-level confidence intervals.""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Any + +from ._comparison import _read_plan, _json +from ._qualification import paired_noninferiority +from ._reference import OFFICIAL_COMMIT, _hash + +__all__ = ["write_comparison_report", "annotate_rollout_evidence"] + +_PARITY_REFERENCES = { + OFFICIAL_COMMIT: "official_warp_main", + "476dc67fccc6a99be293bf611f897b1c06f13e30": "newton_preview", +} +_HORIZONS = (1, 10, 100, 1000) + + +def _finite_nonnegative(value: Any, maximum: float | None = None) -> bool: + return ( + isinstance(value, (int, float)) + and not isinstance(value, bool) + and math.isfinite(value) + and value >= 0 + and (maximum is None or value <= maximum) + ) + + +def _sha_identity(value: Any) -> bool: + return ( + isinstance(value, str) + and len(value) == 64 + and all(character in "0123456789abcdef" for character in value) + ) + + +def _numerical_evidence(path: str | Path, robot: str) -> dict[str, Any]: + """Read independent runner measurements instead of accepting bare pass flags.""" + file = Path(path) + result = { + **json.loads(file.read_text()), + "path": str(file.resolve()), + "sha256": _hash(file), + } + groups = { + "preprocessing": ( + "max_codec_feature_error", + "max_codec_target_error", + "max_codec_decode_error", + ), + "forward": ("max_forward_error",), + "parameter_gradients": ("max_parameter_gradient_error",), + "input_gradients": ("max_input_gradient_error",), + "adam_update": ("max_adam_parameter_error",), + } + required = {field for fields in groups.values() for field in fields} + required.update( + { + "robot", + "reference_commit", + "checkpoint_sha256", + "codec_fixtures_sha256", + "history_lengths", + "adam_updates", + "precision", + "atol", + "rtol", + } + ) + if not required.issubset(result): + return { + **result, + "status": "pending", + "reason": "Numerical report lacks measured values or reference identity", + } + identity_matches = ( + str(result["robot"]).lower() == robot + and isinstance(result["reference_commit"], str) + and result["reference_commit"] in _PARITY_REFERENCES + and _sha_identity(result["checkpoint_sha256"]) + and _sha_identity(result["codec_fixtures_sha256"]) + and result["precision"] == "float32" + and isinstance(result["history_lengths"], list) + and {1, 9, 10}.issubset(result["history_lengths"]) + and isinstance(result["adam_updates"], int) + and result["adam_updates"] >= 1 + and _finite_nonnegative(result["atol"], 1e-5) + and _finite_nonnegative(result["rtol"], 1e-4) + ) + if not identity_matches: + return { + **result, + "status": "failed", + "reason": "Numerical reference, robot, fixtures or tolerance differs from qualification contract", + } + checks = {} + for name, fields in groups.items(): + finite = all(_finite_nonnegative(result[field]) for field in fields) + # A maximum absolute error below atol is sufficient for every reference + # magnitude. Larger residuals need per-element relative-error evidence; + # a caller-supplied Boolean is not that evidence. + checks[name] = { + "passed": finite + and all(result[field] <= result["atol"] for field in fields), + "max_abs": max(result[field] for field in fields) if finite else None, + } + if any(item["max_abs"] is None for item in checks.values()): + status, reason = "failed", "Nonfinite or invalid numerical measurements" + elif all(item["passed"] for item in checks.values()): + status, reason = ( + "passed", + "Measured absolute errors satisfy the declared tolerance", + ) + else: + status, reason = ( + "pending", + "Residuals need per-element relative-tolerance evidence", + ) + return { + **result, + "reference_kind": _PARITY_REFERENCES[result["reference_commit"]], + "checks": checks, + "status": status, + "reason": reason, + } + + +def _rollout_evidence( + record: dict[str, Any] | None, + cell: dict[str, Any], + robot: str, + *, + minimum_trajectories: int = 2048, + required_horizons: tuple[int, ...] = _HORIZONS, +) -> tuple[str, str]: + """Validate an evidence scope; formal reports always use the fixed defaults.""" + required = { + "robot", + "motion_mode", + "state_dtype", + "teacher_origin", + "teacher_split", + "selected_teacher_sha256", + "trajectories", + "horizons", + "failure_rate", + "nonfinite_rate", + "errors", + "model_sha256", + "model_manifest_sha256", + "source_report", + "source_sha256", + "trace_evidence", + "trace_evidence_sha256", + "trajectory_evidence", + "trajectory_evidence_sha256", + "selected_teacher_hash_schema", + "runtime_identity", + } + if record is None or not required.issubset(record): + return "pending", "Missing rollout measurements or provenance" + if "output_dir" not in cell: + return "pending", "Training cell has no exported model" + if ( + record["robot"] != robot + or record["state_dtype"] != "float32" + or not isinstance(record["teacher_origin"], dict) + or record["teacher_origin"].get("kind") != cell["teacher"] + or record["teacher_split"] != "test" + ): + return "failed", "Rollout robot, teacher domain or test split differs" + if ( + not isinstance(record["trajectories"], int) + or isinstance(record["trajectories"], bool) + or record["trajectories"] < minimum_trajectories + or not isinstance(record["horizons"], list) + or not set(required_horizons).issubset(record["horizons"]) + ): + return "pending", "Required rollout population or horizons incomplete" + if not _sha_identity(record["selected_teacher_sha256"]): + return "failed", "Missing selected-trajectory content identity" + if record["selected_teacher_hash_schema"] != "ordered-teacher-content-v1": + return "failed", "Unsupported selected-trajectory content schema" + runtime = record["runtime_identity"] + if ( + not isinstance(runtime, dict) + or not isinstance(runtime.get("device"), str) + or runtime.get("state_dtype") != "float32" + or not runtime.get("python") + or not isinstance(runtime.get("versions"), dict) + or not all(runtime["versions"].get(name) for name in ("torch", "warp-lang")) + ): + return "failed", "Actual neural execution identity is incomplete" + if cell["teacher"] == "official_warp": + expected_warp = "1.5.1" if robot == "ant" else "1.8.0" + if ( + runtime.get("engine") != "official_warp" + or runtime.get("reference_commit") != OFFICIAL_COMMIT + or runtime["versions"]["warp-lang"] != expected_warp + or record["teacher_origin"].get("source_commit") != OFFICIAL_COMMIT + or record["teacher_origin"].get("warp_version") != expected_warp + ): + return ( + "failed", + "Historical rollout requires the pinned original Warp runtime", + ) + physical = record.get("physical_verification") + if ( + not isinstance(physical, dict) + or physical.get("schema") != "official-warp-physical-arrays-v2" + or physical.get("complete") is not True + or physical.get("matched_recorded_fields") is not True + or physical.get("missing_recorded_fields") != [] + or "gravity" not in physical.get("matched_fields", []) + or not _sha_identity(physical.get("sha256")) + ): + return ( + "pending", + "Historical training physics must include verified gravity and recorded parameters", + ) + elif runtime.get("engine") != "current_newton" or not runtime["versions"].get( + "newton" + ): + return "failed", "Current-domain rollout requires recorded Newton execution" + if not all( + _finite_nonnegative(record[name], 1.0) + for name in ("failure_rate", "nonfinite_rate") + ): + return "failed", "Invalid failure/nonfinite rates" + metrics = ["q_mae", "qd_mae"] + if robot == "ant": + metrics.extend(("base_position_error", "base_orientation_error")) + if not isinstance(record["errors"], dict): + return "failed", "Invalid horizon error mapping" + for horizon in required_horizons: + values = record["errors"].get(str(horizon), {}) + if not isinstance(values, dict) or any(name not in values for name in metrics): + return "pending", "Missing predeclared horizon errors" + if not all(_finite_nonnegative(values[name]) for name in metrics): + return "failed", "Undefined or nonfinite horizon errors" + if record["motion_mode"] == "fixed_policy": + if not _sha_identity(record.get("policy_sha256")): + return "pending", "Fixed-policy checkpoint/config identity missing" + policy_return = record.get("policy_return") + if ( + not isinstance(policy_return, (int, float)) + or isinstance(policy_return, bool) + or not math.isfinite(policy_return) + ): + return "failed", "Undefined or nonfinite policy return" + model = Path(cell["output_dir"]) / "best-export" + artifacts = ( + (model / "weights.pt", record["model_sha256"]), + (model / "manifest.json", record["model_manifest_sha256"]), + (Path(record["source_report"]), record["source_sha256"]), + (Path(record["trace_evidence"]), record["trace_evidence_sha256"]), + (Path(record["trajectory_evidence"]), record["trajectory_evidence_sha256"]), + ) + if any(not file.is_file() or _hash(file) != digest for file, digest in artifacts): + return "failed", "Rollout source, model or trajectory evidence changed" + raw = json.loads(Path(record["source_report"]).read_text()) + if raw.get("horizons") != record["errors"] or any( + raw.get(name) != record.get(name) + for name in ( + "robot", + "motion_mode", + "state_dtype", + "trajectories", + "selected_teacher_sha256", + "teacher_origin", + "teacher_split", + "model_sha256", + "model_manifest_sha256", + "failure_rate", + "nonfinite_rate", + "policy_return", + "policy_sha256", + "trace_evidence", + "trace_evidence_sha256", + "trajectory_evidence", + "trajectory_evidence_sha256", + "selected_teacher_hash_schema", + "runtime_identity", + "physical_verification", + ) + ): + return "failed", "Annotated rollout differs from measured source" + return "available", "Measured rollout and source identities verified" + + +def _matched_rollouts(left: dict[str, Any], right: dict[str, Any]) -> bool: + """Paired errors require the same actual teacher traces and fixed policy.""" + return all( + left.get(name) == right.get(name) + for name in ( + "selected_teacher_sha256", + "teacher_origin", + "trajectories", + "horizons", + "policy_sha256", + "state_dtype", + ) + ) + + +def _supported_split_sources(metadata: dict[str, Any]) -> bool: + """Verify raw split receipts generated by the isolated supported environment.""" + splits = {"train", "validation", "test"} + sources = metadata.get("legacy_sources", {}) + identities = metadata.get("source_split_identities", {}) + if ( + not isinstance(sources, dict) + or not isinstance(identities, dict) + or set(sources) != splits + or set(identities) != splits + ): + return False + robot = metadata.get("codec", {}).get("robot") + if robot not in {"cartpole", "ant"}: + return False + warp_version = "1.5.1" if robot == "ant" else "1.8.0" + for split in sorted(splits): + source, identity = sources[split], identities[split] + if not isinstance(source, dict) or not isinstance(identity, dict): + return False + torch_version = identity.get("torch") + if ( + identity.get("source_commit") != OFFICIAL_COMMIT + or identity.get("robot") != robot + or identity.get("python") != "3.8.20" + or identity.get("numpy") != "1.24.4" + or not isinstance(torch_version, str) + or torch_version.split("+", 1)[0] != "2.2.2" + or identity.get("warp") != warp_version + or identity.get("runtime_imports_embodichain") is not False + or not _sha_identity(source.get("sha256")) + or identity.get("raw_sha256") != source["sha256"] + or not isinstance(source.get("path"), str) + ): + return False + raw = Path(source["path"]) + if not raw.is_file() or _hash(raw) != source["sha256"]: + return False + return True + + +def _dataset_evidence(teacher: str, declaration: dict[str, Any]) -> dict[str, Any]: + """Check provenance in immutable data, rather than trusting a matrix row label.""" + if declaration.get("status") != "available": + return {"status": "pending", "reason": "Teacher dataset unavailable"} + from embodichain_models.dynamics.nerd.data import ( + dataset_fingerprint, + load_manifest, + ) + + manifest = load_manifest(declaration["path"], verify=True) + if dataset_fingerprint(manifest) != declaration.get("fingerprint"): + return { + "status": "failed", + "reason": "Teacher dataset changed after preparation", + } + metadata = manifest["metadata"] + if teacher == "official_warp": + origin = metadata.get("teacher_origin", {}) + verified = ( + origin.get("kind") == teacher + and origin.get("source_commit") == OFFICIAL_COMMIT + and bool(metadata.get("source_physical_arrays")) + and ( + ( + metadata.get("source", {}).get("commit") == OFFICIAL_COMMIT + and bool(metadata.get("source_raw_hdf5")) + ) + or _supported_split_sources(metadata) + ) + ) + else: + native = metadata.get("teacher", {}) + verified = ( + _sha_identity(metadata.get("physics_signature")) + and bool(native.get("versions", {}).get("newton")) + and bool(native.get("versions", {}).get("warp-lang")) + and bool(native.get("solver", {}).get("solver_type")) + and isinstance(native.get("num_substeps"), int) + and native["num_substeps"] > 0 + ) + origin = {"kind": teacher, **native} if verified else {} + counts = {split: 0 for split in ("train", "validation", "test")} + transitions = counts.copy() + for shard in manifest["shards"]: + for row in shard["trajectories"]: + counts[row["split"]] += 1 + transitions[row["split"]] += row["length"] + return { + "status": "available" if verified else "pending", + "reason": ( + "Recorded teacher provenance verified" + if verified + else "Matrix label lacks recorded teacher provenance" + ), + "origin": origin, + "fingerprint": declaration["fingerprint"], + "trajectories": counts, + "transitions": transitions, + } + + +def _full_profile_issues(plan: dict[str, Any]) -> list[str]: + """Check effective initial artifacts, preserving feature order lost in sorted JSON.""" + import torch + + robot = "ant" if plan["model"]["output_dim"] == 29 else "cartpole" + inputs = ( + [ + ("states_embedding", 29), + ("contact_normals", 75), + ("contact_points_1", 75), + ("contact_depths", 25), + ("joint_f", 8), + ("gravity_dir", 3), + ] + if robot == "ant" + else [("states_embedding", 4), ("joint_f", 2)] + ) + expected_model = { + "output_dim": 29 if robot == "ant" else 4, + "n_embd": 384 if robot == "ant" else 192, + "n_layer": 6, + "n_head": 12, + "history_length": 10, + "block_size": 32, + "head_hidden_dims": [64], + "bias": False, + "dropout": 0.0, + "normalize_input": True, + "normalize_output": True, + "output_tanh": False, + } + expected_train = { + "total_steps": 4995000, + "batch_size": 512, + "learning_rate": 1e-3, + "final_learning_rate": 1e-4, + "max_grad_norm": 1.0, + "lr_update_interval": 5000, + "lr_offset": 1, + "validation_interval": 5000, + "validation_batches": 100, + } + issues = [ + f"train.{name} differs from the matched official profile" + for name, value in expected_train.items() + if plan["train"].get(name) != value + ] + seen = set() + for cell in plan["cells"]: + if cell.get("status") != "prepared": + continue + spec = json.loads(Path(cell["spec"]).read_text()) + path = Path(spec["initial_model"]) / "manifest.json" + if path in seen: + continue + seen.add(path) + model = json.loads(path.read_text())["model_config"] + weights = torch.load( + path.with_name("weights.pt"), map_location="cpu", weights_only=True + ) + floating = [value for value in weights.values() if value.is_floating_point()] + if not floating or any(value.dtype != torch.float32 for value in floating): + issues.append( + f"{cell['id']}: actual initial-model weights/normalizers are not FP32" + ) + if list(model.get("input_dims", {}).items()) != inputs: + issues.append( + f"{cell['id']}: actual initial-model feature order/dimensions differ" + ) + issues.extend( + f"{cell['id']}: actual initial-model {name} differs" + for name, value in expected_model.items() + if model.get(name) != value + ) + if not seen: + issues.append( + "No hash-verified initial model is available to verify the full profile" + ) + return issues + + +def _curve( + path: Path, budget: int, max_points: int = 1000 +) -> tuple[list[int], list[float]]: + """Keep the first/final observation and bounded update-spaced samples. + + Sparse logging need not divide the plotting stride. Selecting observations + by modulo can otherwise discard nearly the entire measured learning curve. + """ + if max_points < 2: + raise ValueError("Learning curves require at least two plot points") + steps = [] + losses = [] + stride = max(1, (budget + max_points - 2) // (max_points - 1)) + next_step = 0 + last = None + if path.exists(): + with path.open() as stream: + for line in stream: + record = json.loads(line) + last = record + if not steps or record["step"] >= next_step: + steps.append(record["step"]) + losses.append(record["loss"]) + next_step = record["step"] + stride + if last is not None and steps[-1] != last["step"]: + steps.append(last["step"]) + losses.append(last["loss"]) + return steps, losses + + +def annotate_rollout_evidence( + source: str | Path, run_id: str, output: str | Path +) -> dict[str, Any]: + """Attach a measured simulation report to a matrix cell without changing errors. + + The simulation report must already identify the exact model weights and + motion scenario. This adapter never infers a scenario or model hash from a + filename and never turns missing stability/return measurements into zeros. + """ + source = Path(source) + raw = json.loads(source.read_text()) + for key in ("model_sha256", "motion_mode", "trajectories", "horizons"): + if key not in raw: + raise ValueError(f"Rollout source lacks measured identity/metadata: {key}") + errors = raw["horizons"] + if not isinstance(errors, dict) or not errors: + raise ValueError( + "Expected per-horizon metric mapping from simulation evaluator" + ) + last = errors[str(max(map(int, errors)))] + result = { + "run_id": run_id, + "model_sha256": raw["model_sha256"], + "motion_mode": raw["motion_mode"], + "trajectories": raw["trajectories"], + "horizons": sorted(map(int, errors)), + "errors": errors, + "source_report": str(source.resolve()), + "source_sha256": _hash(source), + "teacher_dataset_fingerprint": raw.get("teacher_dataset_fingerprint"), + "teacher_origin": raw.get("teacher_origin"), + } + for name in ( + "robot", + "seed", + "state_dtype", + "teacher_split", + "selected_teacher_sha256", + "model_manifest_sha256", + "policy_sha256", + "trace_evidence", + "trace_evidence_sha256", + "trajectory_evidence", + "trajectory_evidence_sha256", + "selected_teacher_hash_schema", + "teacher_source_motion_mode", + "runtime_identity", + "physical_verification", + ): + if name in raw: + result[name] = raw[name] + for name in ("failure_rate", "nonfinite_rate", "policy_return"): + if name in raw: + result[name] = raw[name] + elif name in last: + result[name] = last[name] + output = Path(output) + output.parent.mkdir(parents=True, exist_ok=True) + _json(output, result, exclusive=True) + return result + + +def write_comparison_report( + experiment: str | Path, + output: str | Path | None = None, + *, + numerical_parity_report: str | Path | None = None, + rollout_reports: list[str | Path] | None = None, +) -> dict[str, Any]: + """Report verified run receipts; missing evidence always remains pending. + + Full rollout attachments use one JSON per cell/scenario with run_id, + model_sha256, trajectories, horizons, motion_mode, failure_rate, + nonfinite_rate and (for fixed_policy) policy_return. Attachments are measured + evidence, not substitutes for the required training matrix or numerical gate. + Their source/model/trace hashes, teacher domain, selected trajectory identity + and common fixed-policy identity must also verify before paired statistics. + """ + root, plan = _read_plan(experiment) + destination = Path(output).resolve() if output else root / "report" + destination.mkdir(parents=True, exist_ok=True) + cells = [] + completed = {} + pending = [] + for cell in plan["cells"]: + receipt_path = root / (cell["id"] + ".receipt.json") + if not receipt_path.exists(): + cells.append( + { + "id": cell["id"], + "status": "pending", + "reason": cell.get("reason", "not executed"), + } + ) + continue + receipt = json.loads(receipt_path.read_text()) + result_path = Path(cell["output_dir"]) / "result.json" + if ( + receipt["status"] != "completed" + or not result_path.exists() + or _hash(result_path) != receipt["result_sha256"] + ): + cells.append( + { + "id": cell["id"], + "status": "failed", + "reason": "run failed or result receipt mismatch", + } + ) + continue + result = json.loads(result_path.read_text()) + spec = json.loads(Path(cell["spec"]).read_text()) + if ( + result.get("completed_updates") != spec["train"]["total_steps"] + or result.get("dataset_fingerprint") != spec["dataset_fingerprint"] + or result.get("sample_order_sha256") != spec["sample_order_sha256"] + ): + cells.append( + { + "id": cell["id"], + "status": "failed", + "reason": "budget/data/order evidence mismatch", + } + ) + continue + weights = Path(cell["output_dir"]) / "best-export" / "weights.pt" + if not weights.exists() or _hash(weights) != result.get("best_model_sha256"): + cells.append( + { + "id": cell["id"], + "status": "failed", + "reason": "exported model identity missing or changed", + } + ) + continue + completed[cell["id"]] = result + cells.append( + { + "id": cell["id"], + "status": "completed", + "completed_updates": result["completed_updates"], + "test": result["test"], + "result": str(result_path), + } + ) + acceptance = plan["acceptance"] + dataset_evidence = { + teacher: _dataset_evidence(teacher, declaration) + for teacher, declaration in plan["datasets"].items() + } + comparisons = [] + metrics = list(acceptance["primary_metrics"]) + if any( + "base_orientation_error_rad" in value["test"] for value in completed.values() + ): + metrics.extend(["base_position_error", "base_orientation_error_rad"]) + for teacher in ("official_warp", "current_newton"): + for metric in metrics: + reference = [] + candidate = [] + paired_seeds = [] + for seed in plan["seeds"]: + left = completed.get(f"{teacher}-official-{seed}") + right = completed.get(f"{teacher}-embodichain-{seed}") + if ( + left + and right + and metric in left["test"] + and metric in right["test"] + ): + reference.append(float(left["test"][metric])) + candidate.append(float(right["test"][metric])) + paired_seeds.append(seed) + gate = ( + paired_noninferiority( + reference, + candidate, + margin=acceptance["relative_margin"], + absolute_tolerance=acceptance["absolute_tolerance"], + resamples=acceptance["resamples"], + seed=acceptance["bootstrap_seed"], + ) + if reference + else {"status": "pending", "reason": "no completed seed pairs"} + ) + comparisons.append( + { + "teacher": teacher, + "metric": metric, + "seeds": paired_seeds, + "official": reference, + "embodichain": candidate, + "gate": gate, + } + ) + if len(plan["seeds"]) != acceptance["required_seeds"]: + pending.append("Five predeclared paired training seeds are required.") + if ( + plan["scope"] != "full" + or plan["train"]["total_steps"] != acceptance["official_updates"] + ): + pending.append( + "This is a shortened/smoke training budget, not full official-budget qualification." + ) + profile_issues = _full_profile_issues(plan) if plan["scope"] == "full" else [] + if profile_issues: + pending.append( + "Effective architecture/optimizer/validation settings differ from the predeclared full matched profile." + ) + if any( + value["status"] != "available" + or (value.get("origin") or {}).get("kind") != teacher + for teacher, value in plan["datasets"].items() + ): + pending.append( + "Both teacher datasets require recorded origin evidence; a row label is not teacher provenance." + ) + if any(value["status"] != "available" for value in dataset_evidence.values()): + pending.append( + "Teacher data provenance must be verified from committed dataset metadata." + ) + if any(cell["status"] == "pending" for cell in cells): + pending.append("Some cells have not completed verified training.") + numerical = { + "status": "pending", + "reason": "No independent numerical/Adam parity report attached", + } + robot = "ant" if plan["model"]["output_dim"] == 29 else "cartpole" + if numerical_parity_report: + numerical = _numerical_evidence(numerical_parity_report, robot) + if numerical["status"] != "passed": + pending.append( + "Independent preprocessing, forward, gradient and Adam numerical parity must pass." + ) + rollout = {} + valid_keys = { + (cell["id"], scenario) + for cell in plan["cells"] + for scenario in ("passive", "random_effort", "fixed_policy") + } + for path in rollout_reports or []: + path = Path(path) + record = json.loads(path.read_text()) + key = (record.get("run_id"), record.get("motion_mode")) + if key not in valid_keys: + raise ValueError(f"Unknown rollout cell/scenario: {key}") + if key in rollout: + raise ValueError(f"Duplicate rollout attachment: {key}") + rollout[key] = {**record, "source": str(path.resolve()), "sha256": _hash(path)} + qualified_rollouts = [] + verified_rollouts = {} + for cell in plan["cells"]: + for scenario in ("passive", "random_effort", "fixed_policy"): + record = rollout.get((cell["id"], scenario)) + status, reason = _rollout_evidence(record, cell, robot) + if status == "available": + verified_rollouts[(cell["id"], scenario)] = record + qualified_rollouts.append( + { + "id": cell["id"], + "scenario": scenario, + "status": status, + "reason": reason, + "evidence": record, + } + ) + if any(record["status"] != "available" for record in qualified_rollouts): + pending.append( + "Matched passive/random-effort/fixed-policy rollouts, stability and return evidence are incomplete." + ) + rollout_gates = [] + pair_mismatches = [] + rollout_metrics = list( + acceptance.get("rollout_primary_metrics", ["q_mae", "qd_mae"]) + ) + if plan["model"]["output_dim"] == 29: + rollout_metrics.extend(["base_position_error", "base_orientation_error"]) + for teacher in ("official_warp", "current_newton"): + for scenario in ("passive", "random_effort", "fixed_policy"): + pairs = [] + for seed in plan["seeds"]: + left = verified_rollouts.get((f"{teacher}-official-{seed}", scenario)) + right = verified_rollouts.get( + (f"{teacher}-embodichain-{seed}", scenario) + ) + if left and right: + if _matched_rollouts(left, right): + pairs.append((left, right)) + else: + pair_mismatches.append( + { + "teacher": teacher, + "seed": seed, + "scenario": scenario, + "reason": "Paired rollout teacher trajectories, population or policy differ", + } + ) + for horizon in (1, 10, 100, 1000): + for metric in rollout_metrics: + values = [ + ( + left["errors"][str(horizon)][metric], + right["errors"][str(horizon)][metric], + ) + for left, right in pairs + if metric in left.get("errors", {}).get(str(horizon), {}) + and metric in right.get("errors", {}).get(str(horizon), {}) + ] + gate = ( + paired_noninferiority( + [x[0] for x in values], + [x[1] for x in values], + margin=acceptance["relative_margin"], + absolute_tolerance=acceptance["absolute_tolerance"], + resamples=acceptance["resamples"], + seed=acceptance["bootstrap_seed"], + ) + if values + else {"status": "pending"} + ) + rollout_gates.append( + { + "teacher": teacher, + "scenario": scenario, + "horizon": horizon, + "metric": metric, + "gate": gate, + } + ) + for metric, threshold_name in ( + ("failure_rate", "failure_rate_increase"), + ("nonfinite_rate", "nonfinite_rate_increase"), + ): + values = [ + (left[metric], right[metric]) + for left, right in pairs + if metric in left and metric in right + ] + gate = ( + {"status": "pending"} + if len(values) < 5 + else { + "status": ( + "passed" + if all( + candidate - reference <= acceptance[threshold_name] + for reference, candidate in values + ) + else "failed" + ), + "threshold": acceptance[threshold_name], + } + ) + rollout_gates.append( + { + "teacher": teacher, + "scenario": scenario, + "metric": metric, + "gate": gate, + } + ) + if scenario == "fixed_policy": + import numpy as np + + values = np.asarray( + [ + (left["policy_return"], right["policy_return"]) + for left, right in pairs + if "policy_return" in left and "policy_return" in right + ], + dtype=float, + ) + gate = {"status": "pending"} + if len(values) >= 5: + if not np.isfinite(values).all(): + gate = {"status": "failed", "reason": "nonfinite_returns"} + else: + absolute = bool( + ( + np.abs(values[:, 0]) <= acceptance["absolute_tolerance"] + ).any() + ) + degradation = values[:, 0] - values[:, 1] + if not absolute: + degradation = degradation / np.abs(values[:, 0]) + threshold = ( + acceptance["absolute_tolerance"] + if absolute + else acceptance["policy_return_relative_margin"] + ) + rng = np.random.default_rng(acceptance["bootstrap_seed"]) + upper = float( + np.quantile( + degradation[ + rng.integers( + 0, + len(values), + size=(acceptance["resamples"], len(values)), + ) + ].mean(1), + 0.95, + ) + ) + gate = { + "status": ("passed" if upper <= threshold else "failed"), + "upper_95": upper, + "threshold": threshold, + "difference_mode": "absolute" if absolute else "relative", + } + rollout_gates.append( + { + "teacher": teacher, + "scenario": scenario, + "metric": "policy_return", + "gate": gate, + } + ) + if any(item["gate"]["status"] == "pending" for item in rollout_gates): + pending.append( + "Paired seed-level rollout error/stability/policy-return gates lack measurements." + ) + failed = ( + any(cell["status"] == "failed" for cell in cells) + or any(row["gate"]["status"] == "failed" for row in comparisons) + or numerical["status"] == "failed" + or any(row["status"] == "failed" for row in qualified_rollouts) + or bool(pair_mismatches) + or any(value["status"] == "failed" for value in dataset_evidence.values()) + or any(item["gate"]["status"] == "failed" for item in rollout_gates) + ) + result = { + "status": "failed" if failed else "pending" if pending else "passed", + "scope": plan["scope"], + "qualification_kind": "controlled_matched_data_noninferiority", + "paper_reproduction": { + "status": "not_established", + "reason": "Shared initialization, statistics, batch order and validation selection control this comparison; literal upstream data/validation recipe reproduction is separate.", + }, + "full_profile_issues": profile_issues, + "comparison_protocol": plan.get( + "comparison_protocol", + { + "kind": "controlled_matched_data_noninferiority", + "note": "Older manifest: shared statistics/order/checkpoint controls are described in dataset and run artifacts; this is not literal paper-recipe reproduction.", + }, + ), + "source_commit": plan["source"]["commit"], + "cells": cells, + "dataset_evidence": dataset_evidence, + "comparisons": comparisons, + "numerical_parity": numerical, + "rollouts": qualified_rollouts, + "rollout_gates": rollout_gates, + "rollout_pair_mismatches": pair_mismatches, + "pending": pending, + "experiment_sha256": _hash(Path(experiment)), + "training_matrix_status": ( + "completed" if len(completed) == len(plan["cells"]) else "pending" + ), + } + _json(destination / "comparison.json", result) + _plots(plan, completed, comparisons, destination) + lines = [ + "# NeRD official implementation and training comparison", + "", + f"Overall qualification: **{result['status']}**. Training matrix: **{result['training_matrix_status']}**.", + "", + f"Source: `{plan['source']['commit']}`. Scope: `{plan['scope']}`; configured updates per run: {plan['train']['total_steps']}.", + "", + "This qualification measures controlled matched-data noninferiority. Shared statistics, sampling order and checkpoint selection are experiment controls; literal paper data/validation-recipe reproduction is not established.", + "", + "| Run | Status | Completed updates |", + "|---|---|---|", + ] + lines += [ + f"| {cell['id']} | {cell['status']} | {cell.get('completed_updates','—')} |" + for cell in cells + ] + lines += ["", "## Dataset origins", "", "| Row | Recorded origin |", "|---|---|"] + lines += [ + f"| {teacher} | {json.dumps(value.get('origin', {'status': 'unavailable'}), sort_keys=True)} |" + for teacher, value in plan["datasets"].items() + ] + lines += [ + "", + "## Paired seed-level error gates", + "", + "| Teacher | Metric | Paired seeds | Upper 95% difference | Status |", + "|---|---|---|---|---|", + ] + lines += [ + f"| {row['teacher']} | {row['metric']} | {len(row['seeds'])} | {row['gate'].get('upper_95','—')} | {row['gate']['status']} |" + for row in comparisons + ] + lines += [ + "", + "## Outstanding evidence", + "", + *[f"- {reason}" for reason in pending], + "", + "![Training curves](training-curves.png)", + "", + "![Paired offline errors](paired-errors.png)", + "", + "Commands and stdout/stderr are retained in each run directory. Source files remain unmodified; the transport adapter binds common initialization, stats and data to original trainer methods. Offline errors alone do not establish long-rollout or paper equivalence.", + ] + (destination / "REPORT.md").write_text("\n".join(lines) + "\n") + return result + + +def _plots( + plan: dict[str, Any], + completed: dict[str, Any], + comparisons: list[dict[str, Any]], + output: Path, +) -> None: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + synthetic = any( + (value.get("origin") or {}).get("kind") == "synthetic_fixture_only" + for value in plan["datasets"].values() + ) + label = "synthetic fixtures; infrastructure smoke" if synthetic else plan["scope"] + figure, axis = plt.subplots(figsize=(10, 5)) + for cell in plan["cells"]: + if cell["id"] in completed: + steps, loss = _curve( + Path(cell["output_dir"]) / "metrics.jsonl", plan["train"]["total_steps"] + ) + axis.plot(steps, loss, label=cell["id"], alpha=0.65) + axis.set( + xlabel="Committed optimizer update", + ylabel="Normalized training MSE", + title=f"NeRD matched-data training ({label})", + ) + if completed: + axis.legend(fontsize=6, ncol=2) + else: + axis.text( + 0.5, + 0.5, + "No completed training runs", + ha="center", + transform=axis.transAxes, + ) + figure.tight_layout() + figure.savefig(output / "training-curves.png", dpi=140) + plt.close(figure) + figure, axis = plt.subplots(figsize=(10, 5)) + available = [row for row in comparisons if row["official"]] + for index, row in enumerate(available): + axis.scatter( + [index - 0.12] * len(row["official"]), + row["official"], + color="tab:blue", + label="Official" if index == 0 else None, + ) + axis.scatter( + [index + 0.12] * len(row["embodichain"]), + row["embodichain"], + color="tab:orange", + label="EmbodiChain" if index == 0 else None, + ) + axis.set_xticks( + range(len(available)), + [row["teacher"] + "\n" + row["metric"] for row in available], + rotation=15, + fontsize=7, + ) + axis.set( + ylabel="Per-seed held-out error", + title=f"Paired offline metrics ({label})", + ) + if available: + axis.legend() + else: + axis.text( + 0.5, + 0.5, + "No completed matched seed pairs", + ha="center", + transform=axis.transAxes, + ) + figure.tight_layout() + figure.savefig(output / "paired-errors.png", dpi=140) + plt.close(figure) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--experiment", type=Path, required=True) + parser.add_argument("--output", type=Path) + parser.add_argument("--numerical-parity-report", type=Path) + parser.add_argument("--rollout-report", type=Path, action="append", default=[]) + args = parser.parse_args() + result = write_comparison_report( + args.experiment, + args.output, + numerical_parity_report=args.numerical_parity_report, + rollout_reports=args.rollout_report, + ) + print(json.dumps(result, indent=2, allow_nan=False)) + + +if __name__ == "__main__": + main() diff --git a/examples/learning/nerd/_support/_runtime.py b/examples/learning/nerd/_support/_runtime.py new file mode 100644 index 0000000000..ff0c3d1645 --- /dev/null +++ b/examples/learning/nerd/_support/_runtime.py @@ -0,0 +1,154 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Read-only execution provenance without importing any simulation runtime.""" + +from __future__ import annotations + +import csv +import io +import platform +import shutil +import subprocess +import sys +from typing import Any + +import torch + +__all__: list[str] = [] + + +def _driver_inventory() -> dict[str, Any]: + """Query physical GPUs independently of Torch's logical device numbering.""" + executable = shutil.which("nvidia-smi") + if executable is None: + return {"status": "unavailable", "reason": "nvidia-smi not found"} + try: + process = subprocess.run( + [ + executable, + "--query-gpu=index,uuid,name,driver_version", + "--format=csv,noheader,nounits", + ], + capture_output=True, + text=True, + timeout=2.0, + check=False, + ) + if process.returncode: + return { + "status": "error", + "reason": "nonzero_exit", + "returncode": process.returncode, + "stderr": process.stderr[:1024], + } + if len(process.stdout) > 65536: + return {"status": "error", "reason": "inventory_too_large"} + rows = list(csv.reader(io.StringIO(process.stdout), skipinitialspace=True)) + if ( + not rows + or len(rows) > 128 + or any( + len(row) != 4 or any(not value.strip() for value in row) for row in rows + ) + ): + return {"status": "error", "reason": "invalid_inventory"} + return { + "status": "available", + "physical_gpus": [ + dict( + zip( + ("index", "uuid", "name", "driver_version"), + (value.strip() for value in row), + ) + ) + for row in rows + ], + } + except subprocess.TimeoutExpired: + return {"status": "error", "reason": "timeout", "timeout_seconds": 2.0} + except (OSError, UnicodeError) as error: + return { + "status": "error", + "reason": type(error).__name__, + "message": str(error)[:1024], + } + + +def _training_runtime_identity(model: torch.nn.Module) -> dict[str, Any]: + """Describe actual parameter placement without moving tensors or consuming RNG. + + CUDA driver discovery is best effort and limited to two seconds. The driver + inventory uses physical GPU identifiers; it is deliberately not equated with + Torch logical indices, which may be remapped by the process environment. + """ + placements = sorted( + { + (str(parameter.device), str(parameter.dtype)) + for parameter in model.parameters() + } + ) + cuda_devices = sorted( + {device for device, _ in placements if torch.device(device).type == "cuda"} + ) + devices = [] + for device in cuda_devices: + try: + devices.append( + { + "device": device, + "status": "available", + "name": torch.cuda.get_device_name(torch.device(device)), + } + ) + except (RuntimeError, AssertionError, ValueError, OSError) as error: + devices.append( + { + "device": device, + "status": "error", + "reason": type(error).__name__, + "message": str(error)[:1024], + } + ) + result = { + "schema_version": 1, + "python": platform.python_version(), + "python_implementation": platform.python_implementation(), + "platform": platform.platform(), + "torch_version": str(torch.__version__), + "parameter_placements": [ + {"device": device, "dtype": dtype} for device, dtype in placements + ], + "parameter_placement_status": "available" if placements else "no_parameters", + "cuda": { + "used": bool(cuda_devices), + "runtime_version": torch.version.cuda, + "devices": devices, + "driver_inventory": ( + _driver_inventory() if cuda_devices else {"status": "not_used"} + ), + }, + } + warp = sys.modules.get("warp") + result["warp"] = ( + { + "status": "already_loaded", + "version": str(getattr(warp, "__version__", "unknown")), + } + if warp is not None + else {"status": "not_loaded"} + ) + return result diff --git a/examples/learning/nerd/_support/_sample_order.py b/examples/learning/nerd/_support/_sample_order.py new file mode 100644 index 0000000000..593ae1b989 --- /dev/null +++ b/examples/learning/nerd/_support/_sample_order.py @@ -0,0 +1,376 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Bounded storage for exact shared Torch minibatch sampling streams. + +Store RNG states at fixed block boundaries instead of materializing every +sample index. The emitted stream remains identical to the original matrix. +""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path +import sys +from collections.abc import Sequence +from typing import Any + +import numpy as np +import torch + +__all__: list[str] = [] +_FORMAT = "embodichain.nerd.sample_order.v1" +_OFFICIAL_FORMAT = "embodichain.nerd.official_shuffle_order.v1" +_MAX_PERMUTATION_BYTES = 32 * 1024**2 + + +def _official_dimensions( + dataset_size: int, steps: int, batch_size: int, epoch_steps: int +) -> tuple[int, int, int]: + """Bound the source iterator schedule and return pass dimensions.""" + for value in (dataset_size, steps, batch_size, epoch_steps): + if type(value) is not int or value < 1: + raise ValueError("Official shuffle dimensions must be positive integers") + if dataset_size * 8 > _MAX_PERMUTATION_BYTES: + raise ValueError("Official shuffle permutation exceeds 32 MiB") + batches = dataset_size // batch_size + if not batches or epoch_steps < batches: + raise ValueError( + "Official shuffle requires an exhausted source iterator each epoch" + ) + passes_per_epoch = (epoch_steps + batches - 1) // batches + epochs, tail = divmod(steps, epoch_steps) + passes = epochs * passes_per_epoch + ((tail + batches - 1) // batches) + return batches, passes_per_epoch, passes + + +def _permutation_sha256(value: np.ndarray) -> str: + """Hash native little-endian contiguous int64 indices without a full copy.""" + if ( + sys.byteorder != "little" + or value.dtype != np.int64 + or not value.flags.c_contiguous + ): + raise ValueError( + "Official permutation hash requires contiguous little-endian int64" + ) + return hashlib.sha256(memoryview(value).cast("B")).hexdigest() + + +def _write_official_shuffle_order( + path: Path, + *, + dataset_size: int, + steps: int, + batch_size: int, + epoch_steps: int, + pass_seeds: np.ndarray | Sequence[int], + permutation_sha256: Sequence[str], + provenance: dict[str, Any], +) -> tuple[Path, Path, Path]: + """Persist captured original sampler seeds and full permutation digests. + + The caller must establish original global-RNG event provenance. This writer + never invents seeds or equates an independent shuffle with official replay. + Epoch boundaries discard the unfinished final permutation, matching the + pinned source's exhausted outer iterator. Hashes bind each actual CPU + randperm so another Torch version must prove its output before returning a + minibatch. Old with-replacement descriptors retain their original format. + """ + _, _, passes = _official_dimensions(dataset_size, steps, batch_size, epoch_steps) + seeds = np.asarray(pass_seeds) + if seeds.dtype != np.int64 or seeds.shape != (passes,) or np.any(seeds < 0): + raise ValueError( + "Official shuffle requires one nonnegative int64 seed per pass" + ) + hashes = list(permutation_sha256) + if len(hashes) != passes or any( + not isinstance(value, str) + or len(value) != 64 + or any(char not in "0123456789abcdef" for char in value) + for value in hashes + ): + raise ValueError("Official shuffle requires one complete SHA256 per pass") + if not isinstance(provenance, dict) or not provenance: + raise ValueError("Official shuffle requires source RNG provenance") + provenance = json.loads(json.dumps(provenance, allow_nan=False)) + path = Path(path) + if path.suffix != ".json": + raise ValueError("Official shuffle descriptor must use .json") + seeds_path = path.with_name(path.stem + "-seeds.npy") + hashes_path = path.with_name(path.stem + "-permutation-sha256.npy") + if any(value.exists() for value in (path, seeds_path, hashes_path)): + raise FileExistsError(path) + with seeds_path.open("xb") as stream: + np.save(stream, seeds, allow_pickle=False) + with hashes_path.open("xb") as stream: + np.save(stream, np.asarray(hashes, dtype="S64"), allow_pickle=False) + descriptor = { + "format": _OFFICIAL_FORMAT, + "dataset_size": dataset_size, + "steps": steps, + "batch_size": batch_size, + "epoch_steps": epoch_steps, + "drop_last": True, + "discard_epoch_tail": True, + "permutation_algorithm": "torch.randperm-cpu-int64-little-endian", + "seeds_file": seeds_path.name, + "seeds_sha256": _digest(seeds_path), + "permutation_hashes_file": hashes_path.name, + "permutation_hashes_sha256": _digest(hashes_path), + "provenance": provenance, + } + with path.open("x") as stream: + json.dump(descriptor, stream, sort_keys=True, indent=2, allow_nan=False) + return path, seeds_path, hashes_path + + +def _digest(path: Path) -> str: + value = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + value.update(chunk) + return value.hexdigest() + + +def _write_checkpointed_order( + path: Path, + *, + dataset_size: int, + steps: int, + batch_size: int, + seed: int, + block_steps: int = 1024, +) -> tuple[Path, Path]: + """Freeze the same continuous randint stream using block RNG checkpoints.""" + for name, value in ( + ("dataset_size", dataset_size), + ("steps", steps), + ("batch_size", batch_size), + ("block_steps", block_steps), + ): + if type(value) is not int or value < 1: + raise ValueError(f"{name} must be a positive integer") + if type(seed) is not int or seed < 0: + raise ValueError("seed must be a nonnegative integer") + if block_steps * batch_size > 4_194_304: + raise ValueError("sample-order block exceeds 32 MiB") + path = Path(path) + states_path = path.with_name(path.stem + "-rng.npy") + if path.exists() or states_path.exists(): + raise FileExistsError(path) + generator = torch.Generator(device="cpu").manual_seed(seed) + state_width = generator.get_state().numel() + blocks = (steps + block_steps - 1) // block_steps + states = np.lib.format.open_memmap( + states_path, mode="w+", dtype=np.uint8, shape=(blocks, state_width) + ) + for block, start in enumerate(range(0, steps, block_steps)): + states[block] = generator.get_state().numpy() + # Exactly the same call shape/order as the previous matrix writer. + torch.randint( + dataset_size, + (min(block_steps, steps - start), batch_size), + generator=generator, + ) + states.flush() + del states + descriptor = { + "format": _FORMAT, + "dataset_size": dataset_size, + "steps": steps, + "batch_size": batch_size, + "seed": seed, + "block_steps": block_steps, + "torch_version": str(torch.__version__), + "states_file": states_path.name, + "states_sha256": _digest(states_path), + } + with path.open("x") as stream: + json.dump(descriptor, stream, sort_keys=True, indent=2) + return path, states_path + + +class _CheckpointedOrder: + """Read-only random access with at most one materialized index block.""" + + def __init__(self, path: Path) -> None: + path = Path(path).resolve() + spec = json.loads(path.read_text()) + expected_keys = { + "format", + "dataset_size", + "steps", + "batch_size", + "seed", + "block_steps", + "torch_version", + "states_file", + "states_sha256", + } + if set(spec) != expected_keys or spec["format"] != _FORMAT: + raise ValueError("Unsupported checkpointed sample order") + if spec["torch_version"] != str(torch.__version__): + raise ValueError("Sample-order Torch version mismatch") + for key in ("dataset_size", "steps", "batch_size", "block_steps"): + if type(spec[key]) is not int or spec[key] < 1: + raise ValueError("Invalid sample-order dimensions") + if type(spec["seed"]) is not int or spec["seed"] < 0: + raise ValueError("Invalid sample-order seed") + if spec["block_steps"] * spec["batch_size"] > 4_194_304: + raise ValueError("sample-order block exceeds 32 MiB") + states_path = (path.parent / spec["states_file"]).resolve() + if states_path.parent != path.parent or not states_path.is_file(): + raise ValueError("Sample-order RNG file must be beside its descriptor") + if _digest(states_path) != spec["states_sha256"]: + raise ValueError("Sample-order RNG checksum mismatch") + self.states = np.load(states_path, mmap_mode="r", allow_pickle=False) + generator = torch.Generator(device="cpu").manual_seed(spec["seed"]) + blocks = (spec["steps"] + spec["block_steps"] - 1) // spec["block_steps"] + if ( + self.states.dtype != np.uint8 + or self.states.shape != (blocks, generator.get_state().numel()) + or not np.array_equal(self.states[0], generator.get_state().numpy()) + ): + raise ValueError("Invalid sample-order RNG checkpoints") + self.spec = spec + self.shape = (spec["steps"], spec["batch_size"]) + self.dataset_size = spec["dataset_size"] + self._block_id: int | None = None + self._block: np.ndarray | None = None + + def __getitem__(self, index: int) -> np.ndarray: + if not 0 <= index < self.shape[0]: + raise IndexError(index) + block_id, offset = divmod(index, self.spec["block_steps"]) + if self._block_id != block_id: + generator = torch.Generator(device="cpu") + generator.set_state(torch.from_numpy(self.states[block_id].copy())) + start = block_id * self.spec["block_steps"] + count = min(self.spec["block_steps"], self.shape[0] - start) + self._block = torch.randint( + self.dataset_size, (count, self.shape[1]), generator=generator + ).numpy() + self._block_id = block_id + # Returned rows must survive later seeks and may be changed by callers. + return self._block[offset].copy() + + +class _OfficialShuffleOrder: + """Replay captured source passes with one bounded, hash-checked permutation.""" + + def __init__(self, path: Path) -> None: + path = Path(path).resolve() + spec = json.loads(path.read_text()) + keys = { + "format", + "dataset_size", + "steps", + "batch_size", + "epoch_steps", + "drop_last", + "discard_epoch_tail", + "permutation_algorithm", + "seeds_file", + "seeds_sha256", + "permutation_hashes_file", + "permutation_hashes_sha256", + "provenance", + } + if set(spec) != keys or spec["format"] != _OFFICIAL_FORMAT: + raise ValueError("Unsupported official shuffle order") + if ( + spec["drop_last"] is not True + or spec["discard_epoch_tail"] is not True + or spec["permutation_algorithm"] != "torch.randperm-cpu-int64-little-endian" + or not isinstance(spec["provenance"], dict) + or not spec["provenance"] + or sys.byteorder != "little" + ): + raise ValueError("Invalid official shuffle contract") + batches, per_epoch, passes = _official_dimensions( + spec["dataset_size"], spec["steps"], spec["batch_size"], spec["epoch_steps"] + ) + arrays = {} + for key, hash_key in ( + ("seeds_file", "seeds_sha256"), + ("permutation_hashes_file", "permutation_hashes_sha256"), + ): + if not isinstance(spec[key], str): + raise ValueError("Official shuffle payload must be a sibling file") + payload = (path.parent / spec[key]).resolve() + if payload.parent != path.parent or not payload.is_file(): + raise ValueError("Official shuffle payload must be a sibling file") + if _digest(payload) != spec[hash_key]: + raise ValueError("Official shuffle payload checksum mismatch") + arrays[key] = np.load(payload, mmap_mode="r", allow_pickle=False) + seeds, hashes = arrays["seeds_file"], arrays["permutation_hashes_file"] + if seeds.dtype != np.int64 or seeds.shape != (passes,) or np.any(seeds < 0): + raise ValueError("Invalid official shuffle pass seeds") + if hashes.dtype != np.dtype("S64") or hashes.shape != (passes,): + raise ValueError("Invalid official shuffle permutation hashes") + if any( + len(value) != 64 or any(char not in "0123456789abcdef" for char in value) + for value in (bytes(item).decode("ascii") for item in hashes) + ): + raise ValueError("Invalid official shuffle permutation hashes") + self.spec, self.seeds, self.permutation_hashes = spec, seeds, hashes + self.shape = (spec["steps"], spec["batch_size"]) + self.dataset_size = spec["dataset_size"] + self._batches_per_pass, self._passes_per_epoch = batches, per_epoch + self._pass_id: int | None = None + self._permutation: np.ndarray | None = None + + def __getitem__(self, index: int) -> np.ndarray: + if not 0 <= index < self.shape[0]: + raise IndexError(index) + epoch, within_epoch = divmod(index, self.spec["epoch_steps"]) + within_pass, row = divmod(within_epoch, self._batches_per_pass) + pass_id = epoch * self._passes_per_epoch + within_pass + if self._pass_id != pass_id: + # Rows are owned copies; dropping the previous cache cannot + # invalidate a previously returned minibatch or double live storage. + self._permutation, self._pass_id = None, None + generator = torch.Generator(device="cpu").manual_seed( + int(self.seeds[pass_id]) + ) + permutation = torch.randperm( + self.dataset_size, generator=generator, dtype=torch.int64 + ).numpy() + expected = bytes(self.permutation_hashes[pass_id]).decode("ascii") + if _permutation_sha256(permutation) != expected: + raise ValueError( + "Official shuffle permutation checksum mismatch; Torch RNG algorithm differs" + ) + self._permutation, self._pass_id = permutation, pass_id + start = row * self.shape[1] + return self._permutation[start : start + self.shape[1]].copy() + + +def _open_order( + path: str | Path, +) -> np.ndarray | _CheckpointedOrder | _OfficialShuffleOrder: + """Open a new RNG-checkpoint stream or a backward-compatible index matrix.""" + path = Path(path) + if path.suffix == ".json": + if json.loads(path.read_text()).get("format") == _OFFICIAL_FORMAT: + return _OfficialShuffleOrder(path) + return _CheckpointedOrder(path) + order = np.load(path, mmap_mode="r", allow_pickle=False) + if order.dtype.kind not in "iu": + raise ValueError("Memory-mapped sample order requires integer indices") + return order diff --git a/examples/learning/nerd/_support/_task_values.py b/examples/learning/nerd/_support/_task_values.py new file mode 100644 index 0000000000..33bcbf5ac6 --- /dev/null +++ b/examples/learning/nerd/_support/_task_values.py @@ -0,0 +1,70 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Pure reference task observations and rewards for learned robot dynamics. + +Costs follow the Cartpole and forward-Ant definitions in NVIDIA NeRD. Inputs use +Newton world COM linear/angular velocities; legacy twist conversion belongs to +the state codec. These functions do not implement a simulator or reset policy. +""" + +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations +import torch +from torch import Tensor +from embodichain_models.dynamics.nerd.codec import quat_rotate + +__all__ = ["reference_task_values"] + + +def reference_task_values( + robot: str, q: Tensor, qd: Tensor +) -> tuple[Tensor, Tensor, Tensor]: + """Return graph-connected observations/reward and a discrete failure mask. + + Args: + robot: ``cartpole`` or ``ant``. + q: Coordinates with trailing width 2 or 15, arbitrary batch/time axes. + qd: Native velocities with trailing width 2 or 14 and matching axes. + + Returns: + Observation (width 4 for Cartpole, 27 for Ant), scalar reward per row, + and boolean task/numerical failure per row. Cartpole observes wrapped + pole angle; Ant omits root x/y and observes COM velocity. Ant rewards + forward motion along +y plus upright/heading bonuses. Non-finite states + and invalid zero quaternions fail; reward values remain unsanitized so + evaluators can account explicitly for numerical failures. + """ + if robot not in {"cartpole", "ant"}: + raise ValueError(f"Unsupported reference robot {robot!r}") + q_dim, qd_dim = (2, 2) if robot == "cartpole" else (15, 14) + if q.shape[-1] != q_dim or qd.shape[-1] != qd_dim or q.shape[:-1] != qd.shape[:-1]: + raise ValueError("Reference state dimensions or batch axes do not match") + if robot == "cartpole": + angle = torch.remainder(q[..., 1] + torch.pi, 2 * torch.pi) - torch.pi + observation = torch.cat((q[..., :1], angle[..., None], qd), -1) + reward = -( + angle.square() + 0.05 * q[..., 0].square() + 0.1 * qd.square().sum(-1) + ) + failed = (q[..., 0].abs() > 4) | (qd.abs() > 10).any(-1) + else: + observation = torch.cat((q[..., 2:], qd), -1) + up = quat_rotate(q[..., 3:7], q.new_tensor((0.0, 0.0, 1.0))) + heading = quat_rotate(q[..., 3:7], q.new_tensor((0.0, 1.0, 0.0))) + reward = qd[..., 1] + 0.1 * up[..., 2] + heading[..., 1] + failed = (q[..., 2] < 0.3) | (q[..., 3:7].norm(dim=-1) < 1e-6) + failed = failed | ~torch.isfinite(q).all(-1) | ~torch.isfinite(qd).all(-1) + return observation, reward, failed diff --git a/examples/learning/nerd/_support/_trainer.py b/examples/learning/nerd/_support/_trainer.py new file mode 100644 index 0000000000..44c6b93d9f --- /dev/null +++ b/examples/learning/nerd/_support/_trainer.py @@ -0,0 +1,663 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Single-device supervised NeRD training with committed-step resume.""" + +from __future__ import annotations + +import json +import hashlib +import math +from pathlib import Path +import random +from typing import Any + +import numpy as np +import torch +from torch import Tensor + +from embodichain.utils import configclass +from embodichain_models.dynamics.nerd.data import DynamicsWindowDataset +from embodichain.learning.sl import SupervisedLearner +from embodichain.learning.sl._artifacts import _RunArtifacts +from embodichain_models.dynamics.nerd.supervised import NeRDObjective, make_nerd_batch +from ._sample_order import _CheckpointedOrder, _OfficialShuffleOrder, _open_order +from embodichain_models.dynamics.nerd._numerics import ( + _attention_context, + _checkpoint_numerical_profile, + _clip_grad_norm_, + _validate_numerical_profile, +) + +from embodichain_models.dynamics.nerd.integration import ( + _PREPROCESSING_EXECUTION, + encode_batch, + normalized_mse, + compute_statistics, + save_statistics, + load_statistics, + _atomic_save, + _model_variance, + evaluate_model, +) + +__all__ = [ + "DynamicsTrainCfg", + "DynamicsTrainer", + "encode_batch", + "normalized_mse", + "compute_statistics", + "load_statistics", + "save_statistics", + "evaluate_model", +] + + +@configclass +class DynamicsTrainCfg: + """Optimizer and deterministic sampling configuration. + + ``lr_update_interval`` holds the rate for a fixed number of updates, allowing + an official epoch schedule. ``lr_offset`` is measured in schedule intervals; + the official first trained epoch uses offset one because epoch zero evaluates. + The durable cursor is the count of completed optimizer updates. + + ``numerical_profile="reference"`` selects math attention and an ordered + FP64 reduction of FP32 per-parameter gradient norms. The default retains + Torch's gradient clipping and the model's configured attention backend. + """ + + total_steps: int = 10000 + batch_size: int = 64 + learning_rate: float = 1e-3 + final_learning_rate: float = 1e-4 + max_grad_norm: float = 1.0 + seed: int = 0 + device: str = "cpu" + validation_interval: int = 100 + validation_batches: int | None = None + checkpoint_interval: int = 100 + lr_update_interval: int = 1 + lr_offset: int = 0 + log_interval: int = 1 + resident_cache_bytes: int = 0 + numerical_profile: str = "default" + + +def _capture_rng_state(device: str) -> dict[str, Any]: + numpy_state = np.random.get_state() + return { + "torch": torch.get_rng_state(), + "python": random.getstate(), + "numpy": ( + numpy_state[0], + torch.from_numpy(numpy_state[1].astype(np.int64)), + numpy_state[2], + numpy_state[3], + numpy_state[4], + ), + "cuda": ( + torch.cuda.get_rng_state_all() + if torch.device(device).type == "cuda" + else [] + ), + } + + +def _restore_rng_state(state: dict[str, Any]) -> None: + torch.set_rng_state(state["torch"]) + random.setstate(state["python"]) + name, keys, pos, gaussian, cached = state["numpy"] + np.random.set_state((name, keys.numpy().astype(np.uint32), pos, gaussian, cached)) + if state["cuda"]: + torch.cuda.set_rng_state_all(state["cuda"]) + + +def _model_cfg_identity(cfg: dict[str, Any]) -> dict[str, Any]: + """Compare legal head-width sequences without changing feature order.""" + identity = dict(cfg) + if isinstance(identity.get("head_hidden_dims"), (tuple, list)): + identity["head_hidden_dims"] = list(identity["head_hidden_dims"]) + return identity + + +class DynamicsTrainer: + """Train from immutable windows using an explicit committed sample cursor. + + Batches are sampled with replacement by a dedicated CPU generator. Only a + completed optimizer step advances its durable state. External worker prefetch + therefore cannot silently skip examples on resume. ``sample_order`` supplies + a shared explicit index matrix for official-reference comparisons. + + Explicit statistics must identify the training dataset even when + ``initialize_normalization=False``. That flag leaves existing model + normalizers untouched; it does not waive statistics provenance. + """ + + def __init__( + self, + model: torch.nn.Module, + codec: Any, + train_dataset: DynamicsWindowDataset, + validation_dataset: DynamicsWindowDataset, + cfg: DynamicsTrainCfg, + output_dir: str | Path, + *, + manifest: dict[str, Any], + statistics: dict[str, Any] | None = None, + initialize_normalization: bool = True, + sample_order: Tensor | str | Path | None = None, + ) -> None: + cfg.validate() + _validate_numerical_profile(cfg.numerical_profile) + if ( + min( + cfg.total_steps, + cfg.batch_size, + cfg.validation_interval, + cfg.checkpoint_interval, + cfg.lr_update_interval, + cfg.log_interval, + ) + < 1 + ): + raise ValueError("Training sizes and intervals must be positive") + if cfg.validation_batches is not None and cfg.validation_batches < 1: + raise ValueError("validation_batches must be positive") + if cfg.lr_offset < 0: + raise ValueError("lr_offset must be nonnegative") + if type(cfg.resident_cache_bytes) is not int or cfg.resident_cache_bytes < 0: + raise ValueError("resident_cache_bytes must be nonnegative integer bytes") + if ( + cfg.max_grad_norm <= 0 + or min(cfg.learning_rate, cfg.final_learning_rate) < 0 + ): + raise ValueError("Invalid learning rate or gradient clipping norm") + if train_dataset.split != "train" or validation_dataset.split != "validation": + raise ValueError("Trainer requires train and validation splits") + if train_dataset.fingerprint != validation_dataset.fingerprint: + raise ValueError("Training and validation must share an immutable dataset") + if not len(train_dataset) or not len(validation_dataset): + raise ValueError("Training and validation need eligible windows") + if statistics is not None and ( + statistics.get("dataset_fingerprint") != train_dataset.fingerprint + ): + raise ValueError("NeRD statistics/dataset fingerprint mismatch") + self.cfg, self.codec = cfg, codec + self.model = model.to(cfg.device) + if cfg.numerical_profile == "reference": + self.model.cfg.attention_backend = "math" + self.train_dataset, self.validation_dataset = train_dataset, validation_dataset + self.output_dir = Path(output_dir) + self.output_dir.mkdir(parents=True, exist_ok=True) + self._artifacts = _RunArtifacts(self.output_dir) + self.manifest = manifest + self.global_step = 0 + self.best_validation = float("inf") + self._best_model: dict[str, Tensor] | None = None + self._best_step: int | None = None + self.sampler = torch.Generator(device="cpu").manual_seed(cfg.seed) + self._sample_order_identity: dict[str, Any] | None = None + self.sample_order = self._prepare_order(sample_order) + if train_dataset.history_length != model.cfg.history_length: + raise ValueError("Dataset/model history length mismatch") + if initialize_normalization: + self.statistics = statistics or compute_statistics( + train_dataset, codec, cfg.batch_size + ) + self.model.set_normalization( + { + key: (value["mean"], value["var"]) + for key, value in self.statistics["inputs"].items() + }, + self.statistics["output"]["mean"], + self.statistics["output"]["var"], + ) + else: + self.statistics = statistics + # The original trainer preprocesses one real window while constructing + # its integrator/model metadata. Match that JIT warmup before updates; + # cold and optimized CUDA coordinate graphs can round differently. + encode_batch(train_dataset.get_batch([0]), self.codec, cfg.device) + if cfg.resident_cache_bytes: + used = train_dataset.enable_device_cache( + cfg.device, cfg.resident_cache_bytes + ) + validation_dataset.enable_device_cache( + cfg.device, cfg.resident_cache_bytes - used + ) + if torch.device(cfg.device).type == "cuda": + # Specialize the real batch layout before committing updates. This + # makes the first step after resume use the same CUDA graph as an + # uninterrupted run, without consuming sampling/RNG or model state. + warm_batch = train_dataset.get_batch([0] * cfg.batch_size) + with torch.enable_grad(): + for _ in range(3): + encode_batch(warm_batch, self.codec, cfg.device) + del warm_batch + self.optimizer = torch.optim.Adam(self.model.parameters(), lr=cfg.learning_rate) + self.learner = SupervisedLearner( + self.model, + NeRDObjective(include_prefix=train_dataset.include_prefix), + self.optimizer, + max_grad_norm=cfg.max_grad_norm, + clip_grad=lambda parameters, max_norm: _clip_grad_norm_( + parameters, + max_norm, + profile=cfg.numerical_profile, + error_if_nonfinite=cfg.numerical_profile == "default", + ), + numerical_context=lambda: _attention_context(cfg.numerical_profile), + ) + random.seed(cfg.seed) + np.random.seed(cfg.seed) + torch.manual_seed(cfg.seed) + if torch.device(cfg.device).type == "cuda": + torch.cuda.manual_seed_all(cfg.seed) + + def _prepare_order( + self, source: Tensor | str | Path | None + ) -> Tensor | np.ndarray | _CheckpointedOrder | _OfficialShuffleOrder | None: + order, identity = self._load_order(source) + self._sample_order_identity = identity + return order + + def _load_order(self, source: Tensor | str | Path | None) -> tuple[ + Tensor | np.ndarray | _CheckpointedOrder | _OfficialShuffleOrder | None, + dict[str, Any] | None, + ]: + if source is None: + return None, None + identity = None + if isinstance(source, (str, Path)): + path = Path(source).resolve() + digest = hashlib.sha256() + with path.open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + order = _open_order(path) + if isinstance( + order, (_CheckpointedOrder, _OfficialShuffleOrder) + ) and order.dataset_size != len(self.train_dataset): + raise ValueError("Sample-order dataset size mismatch") + identity = { + "path": str(path), + "sha256": digest.hexdigest(), + } + else: + order = source.to(device="cpu", dtype=torch.long) + if tuple(order.shape) != (self.cfg.total_steps, self.cfg.batch_size): + raise ValueError("sample_order must have total_steps x batch_size indices") + return order, identity + + def _order_snapshot(self) -> Any: + return ( + self._sample_order_identity + if self._sample_order_identity is not None + else self.sample_order + ) + + def _lr(self) -> float: + interval = self.cfg.lr_update_interval + ratio = (self.global_step // interval + self.cfg.lr_offset) / ( + self.cfg.total_steps / interval + self.cfg.lr_offset + ) + return ( + self.cfg.learning_rate * (1 - ratio) + self.cfg.final_learning_rate * ratio + ) + + def step(self) -> dict[str, float | int]: + """Commit one Adam update, returning loss, rate and unclipped norm.""" + if self.global_step >= self.cfg.total_steps: + raise RuntimeError("Configured training budget exhausted") + self._artifacts.claim() + candidate_sampler = torch.Generator(device="cpu") + candidate_sampler.set_state(self.sampler.get_state()) + indices = ( + self.sample_order[self.global_step] + if self.sample_order is not None + else torch.randint( + len(self.train_dataset), + (self.cfg.batch_size,), + generator=candidate_sampler, + ) + ) + if any(not 0 <= int(index) < len(self.train_dataset) for index in indices): + raise ValueError("sample_order contains out-of-range dataset indices") + batch = self.train_dataset.get_batch(indices) + supervised_batch = make_nerd_batch(batch, self.codec, self.cfg.device) + lr = self._lr() + for group in self.optimizer.param_groups: + group["lr"] = lr + metrics = self.learner.update(supervised_batch) + self.sampler.set_state(candidate_sampler.get_state()) + self.global_step += 1 + return { + "step": self.global_step, + "loss": metrics["loss"], + "learning_rate": lr, + "grad_norm": metrics["grad_norm"], + } + + def save(self, path: str | Path) -> None: + """Atomically save weights and exact committed training/RNG progress.""" + self._artifacts.claim() + self._ensure_best_snapshot() + state = { + "format": "embodichain.nerd.training.v1", + "run_id": self._artifacts.run_id, + "preprocessing_execution": _PREPROCESSING_EXECUTION, + "model_cfg": self.model.cfg.to_dict(), + "model": self.model.state_dict(), + "manifest": self.manifest, + "train_cfg": self.cfg.to_dict(), + "numerical_profile": self.cfg.numerical_profile, + "attention_backend": self.model.cfg.attention_backend, + "optimizer": self.optimizer.state_dict(), + "global_step": self.global_step, + "scheduler": {"next_step": self.global_step}, + "best_validation": self.best_validation, + "best_model": self._best_model, + "best_step": self._best_step, + "sampler_state": self.sampler.get_state(), + "sample_order": self._order_snapshot(), + "dataset_fingerprint": self.train_dataset.fingerprint, + "window_contract": { + "history_length": self.train_dataset.history_length, + "include_prefix": self.train_dataset.include_prefix, + }, + "statistics": self.statistics, + "rng": _capture_rng_state(self.cfg.device), + } + _atomic_save(Path(path), state) + + def resume(self, path: str | Path) -> None: + """Restore an identical run; changing its budget or data requires fine-tune.""" + path = Path(path) + state = torch.load(path, map_location="cpu", weights_only=True) + if state.get("format") != "embodichain.nerd.training.v1": + raise ValueError("Unsupported dynamics training checkpoint") + if state.get("resumable") is False: + raise ValueError( + "Best-model snapshots support export only; resume from last.pt" + ) + if state.get("preprocessing_execution") != _PREPROCESSING_EXECUTION: + raise ValueError( + "Resume preprocessing execution mismatch; use matching code or fine-tuning" + ) + profile = _checkpoint_numerical_profile(state) + if profile != self.cfg.numerical_profile: + raise ValueError("Resume numerical_profile mismatch; use fine-tuning") + saved_train_cfg = {**state["train_cfg"], "numerical_profile": profile} + saved_model_cfg = {"attention_backend": "default", **state["model_cfg"]} + for key, expected in ( + ("dataset_fingerprint", self.train_dataset.fingerprint), + ("train_cfg", self.cfg.to_dict()), + ("model_cfg", _model_cfg_identity(self.model.cfg.to_dict())), + ("manifest", self.manifest), + ( + "window_contract", + { + "history_length": self.train_dataset.history_length, + "include_prefix": self.train_dataset.include_prefix, + }, + ), + ): + actual = ( + saved_train_cfg + if key == "train_cfg" + else ( + _model_cfg_identity(saved_model_cfg) + if key == "model_cfg" + else state[key] + ) + ) + if actual != expected: + raise ValueError( + f"Resume {key} mismatch; use fine-tuning for a new run" + ) + if list(saved_model_cfg["input_dims"].items()) != list( + self.model.cfg.input_dims.items() + ): + raise ValueError("Resume model_cfg input_dims order mismatch") + saved_order = state["sample_order"] + if self.sample_order is not None: + same = ( + self._sample_order_identity == saved_order + if isinstance(saved_order, dict) + else isinstance(saved_order, Tensor) + and isinstance(self.sample_order, Tensor) + and torch.equal(self.sample_order, saved_order) + ) + if not same: + raise ValueError("Resume sample order mismatch") + if isinstance(saved_order, dict): + order, order_identity = self._load_order(saved_order["path"]) + if order_identity != saved_order: + raise ValueError("Resume sample order checksum mismatch") + else: + order, order_identity = self._load_order(saved_order) + step = state["global_step"] + if ( + type(step) is not int + or not 0 <= step <= self.cfg.total_steps + or state["scheduler"] != {"next_step": step} + ): + raise ValueError("Resume committed step or scheduler mismatch") + self._validate_model_state(state["model"], "model") + groups = state["optimizer"].get("param_groups", []) + if len(groups) != len(self.optimizer.param_groups) or any( + len(saved.get("params", [])) != len(current["params"]) + for saved, current in zip(groups, self.optimizer.param_groups) + ): + raise ValueError("Resume optimizer parameter groups mismatch") + sampler = torch.Generator(device="cpu") + sampler.set_state(state["sampler_state"]) + best_model, best_step = self._checkpoint_best(state, path) + run_id = self._artifacts.validate_resume(state.get("run_id")) + self._artifacts.validate_history(step, key="step") + self._validate_destination_best(state, best_model, best_step, run_id) + # Directory identity, order checksums and every semantic contract are + # admitted before changing model, optimizer, sampling or global RNG. + self.model.load_state_dict(state["model"]) + self.optimizer.load_state_dict(state["optimizer"]) + self.global_step, self.best_validation = ( + state["global_step"], + state["best_validation"], + ) + self.statistics = state["statistics"] + self.sample_order, self._sample_order_identity = order, order_identity + self._best_model, self._best_step = best_model, best_step + self.learner.num_updates = self.global_step + self.sampler.set_state(state["sampler_state"]) + _restore_rng_state(state["rng"]) + self._artifacts.adopt_resume(run_id) + + def _validate_model_state(self, saved: dict[str, Tensor], label: str) -> None: + current = self.model.state_dict() + if current.keys() != saved.keys() or any( + not isinstance(saved[key], Tensor) + or saved[key].shape != value.shape + or saved[key].dtype != value.dtype + for key, value in current.items() + ): + raise ValueError(f"Resume {label} tensor layout mismatch") + + def _checkpoint_best( + self, state: dict[str, Any], path: Path + ) -> tuple[dict[str, Tensor] | None, int | None]: + best_model, best_step = state.get("best_model"), state.get("best_step") + score = state["best_validation"] + if not math.isfinite(score): + if score != float("inf") or best_model is not None or best_step is not None: + raise ValueError("Resume best-model metadata mismatch") + return None, None + if best_model is None: + # Legacy checkpoints stored the score but kept best weights only in + # the sibling file. Read and validate that source without touching + # the destination; an absent best model cannot be invented. + best_path = path.with_name("best.pt") + if not best_path.is_file(): + raise ValueError( + "Resume needs matching best.pt for legacy best weights" + ) + best = torch.load(best_path, map_location="cpu", weights_only=True) + best_cfg = {"attention_backend": "default", **best["model_cfg"]} + if ( + best.get("format") != "embodichain.nerd.training.v1" + or _model_cfg_identity(best_cfg) + != _model_cfg_identity(self.model.cfg.to_dict()) + or list(best_cfg["input_dims"].items()) + != list(self.model.cfg.input_dims.items()) + or best["manifest"] != state["manifest"] + or best["dataset_fingerprint"] != state["dataset_fingerprint"] + or best["best_validation"] != score + or best.get("run_id") != state.get("run_id") + ): + raise ValueError("Resume legacy best.pt identity mismatch") + best_model, best_step = best["model"], best["global_step"] + if type(best_step) is not int or not 0 <= best_step <= state["global_step"]: + raise ValueError("Resume best-model step mismatch") + self._validate_model_state(best_model, "best model") + return { + key: value.detach().to(device="cpu", copy=True) + for key, value in best_model.items() + }, best_step + + def _write_best_snapshot(self) -> None: + assert self._best_model is not None and self._best_step is not None + _atomic_save( + self.output_dir / "best.pt", + { + "format": "embodichain.nerd.training.v1", + "resumable": False, + "run_id": self._artifacts.run_id, + "model_cfg": self.model.cfg.to_dict(), + "model": self._best_model, + "manifest": self.manifest, + "global_step": self._best_step, + "dataset_fingerprint": self.train_dataset.fingerprint, + "numerical_profile": self.cfg.numerical_profile, + "attention_backend": self.model.cfg.attention_backend, + "best_validation": self.best_validation, + }, + ) + + def _validate_destination_best( + self, + state: dict[str, Any], + best_model: dict[str, Tensor] | None, + best_step: int | None, + run_id: str, + ) -> None: + path = self.output_dir / "best.pt" + if not path.exists(): + return + try: + best = torch.load(path, map_location="cpu", weights_only=True) + best_cfg = {"attention_backend": "default", **best["model_cfg"]} + if ( + best_model is None + or best.get("format") != "embodichain.nerd.training.v1" + or best.get("run_id") != run_id + or best["global_step"] != best_step + or best["best_validation"] != state["best_validation"] + or _model_cfg_identity(best_cfg) + != _model_cfg_identity(self.model.cfg.to_dict()) + or list(best_cfg["input_dims"].items()) + != list(self.model.cfg.input_dims.items()) + or best["manifest"] != state["manifest"] + or best["dataset_fingerprint"] != state["dataset_fingerprint"] + ): + raise ValueError("Best selection boundary differs") + self._validate_model_state(best["model"], "destination best model") + if any( + not torch.equal(best["model"][key], value) + for key, value in best_model.items() + ): + raise ValueError("Best selection weights differ") + except (OSError, RuntimeError, KeyError, TypeError, ValueError) as error: + raise ValueError( + "Resume destination best.pt differs from this checkpoint; " + "use a fresh output directory to preserve existing best weights" + ) from error + + def _ensure_best_snapshot(self) -> None: + if self._best_model is not None and not (self.output_dir / "best.pt").exists(): + self._write_best_snapshot() + + def train(self, *, stop_after: int | None = None) -> dict[str, Any]: + """Run to the configured budget or an explicit interruption checkpoint.""" + end = ( + self.cfg.total_steps + if stop_after is None + else min(stop_after, self.cfg.total_steps) + ) + if end < self.global_step: + raise ValueError("stop_after precedes the committed cursor") + self._artifacts.claim() + self._artifacts.validate_history(self.global_step, key="step") + self._ensure_best_snapshot() + last: dict[str, Any] = {} + log_path = self.output_dir / "metrics.jsonl" + with log_path.open("a") as log: + while self.global_step < end: + last = self.step() + if ( + self.global_step % self.cfg.validation_interval == 0 + or self.global_step == self.cfg.total_steps + ): + validation = evaluate_model( + self.model, + self.codec, + self.validation_dataset, + self.cfg.batch_size, + self.cfg.device, + max_batches=self.cfg.validation_batches, + ) + last["validation"] = validation + score = float(validation["normalized_mse"]) + if score < self.best_validation: + self.best_validation = score + self._best_model = { + key: value.detach().to(device="cpu", copy=True) + for key, value in self.model.state_dict().items() + } + self._best_step = self.global_step + self._write_best_snapshot() + if ( + self.global_step % self.cfg.log_interval == 0 + or self.global_step == end + or "validation" in last + ): + log.write(json.dumps(last, allow_nan=False) + "\n") + log.flush() + if self.global_step % self.cfg.checkpoint_interval == 0: + self.save(self.output_dir / "last.pt") + self.save(self.output_dir / "last.pt") + return { + "global_step": self.global_step, + "budget": self.cfg.total_steps, + "status": ( + "completed" + if self.global_step == self.cfg.total_steps + else "interrupted" + ), + "best_validation": ( + self.best_validation if math.isfinite(self.best_validation) else None + ), + "last": last, + } diff --git a/examples/learning/nerd/_support/_worker.py b/examples/learning/nerd/_support/_worker.py new file mode 100644 index 0000000000..9177059336 --- /dev/null +++ b/examples/learning/nerd/_support/_worker.py @@ -0,0 +1,61 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + + +"""Bootstrap one private showcase worker without importing the SDK CLI.""" + +from __future__ import annotations + +import argparse +from pathlib import Path +import runpy +import sys +from types import ModuleType + +__all__: list[str] = [] + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("worker", help="Private module stem in this support directory.") + parser.add_argument("arguments", nargs=argparse.REMAINDER) + args = parser.parse_args() + directory = Path(__file__).resolve().parent + name = args.worker + if ( + not name.isidentifier() + or not name.startswith("_") + or name.startswith("__") + or name == "_worker" + or not (directory / (name + ".py")).is_file() + ): + parser.error("worker must name a private module in the NeRD support directory") + # Examples are wheel data, so their installed parent need not be on sys.path. + # Keep this bootstrap compatible with isolated historical Python 3.8 workers. + examples = directory.parents[2] + package = ModuleType("examples") + package.__path__ = [str(examples)] + sys.modules["examples"] = package + sys.argv = [str(directory / (name + ".py")), *args.arguments] + runpy.run_module( + "examples.learning.nerd._support." + name, + run_name="__main__", + alter_sys=True, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/learning/nerd/_support/_workflow.py b/examples/learning/nerd/_support/_workflow.py new file mode 100644 index 0000000000..c741ded2f9 --- /dev/null +++ b/examples/learning/nerd/_support/_workflow.py @@ -0,0 +1,348 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Lazy offline training, evaluation and model export entry points.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +__all__ = ["run_train", "run_eval", "run_export"] + +# These fields may only come from the dataset's measured deployment contract. +# Even an absent field cannot be invented through a training annotation override. +_DEPLOYMENT_MANIFEST_FIELDS = frozenset( + { + "robot", + "codec", + "dt", + "joint_names", + "q_dim", + "qd_dim", + "asset_sha256", + "assets", + "asset_hashes", + "physics_signature", + "physics_signature_schema", + "gravity", + "contact_schema", + "world_rotation_xyzw", + "gravity_transport_override", + "source_physical_arrays", + "legacy_velocity", + "reference_friction_mapping", + } +) + + +def _arguments(value: dict[str, Any], allowed: set[str]) -> None: + unknown = value.keys() - allowed + if unknown: + raise ValueError(f"Unknown dynamics workflow fields: {sorted(unknown)}") + + +def _load_training_checkpoint(path: str | Path) -> dict[str, Any]: + """Read a tensor-only training payload without constructing a default model.""" + import torch + + state = torch.load(path, map_location="cpu", weights_only=True) + if ( + not isinstance(state, dict) + or state.get("format") != "embodichain.nerd.training.v1" + ): + raise ValueError("Expected a dynamics training checkpoint") + return state + + +def run_train(config: dict[str, Any]) -> dict[str, Any]: + """Train from scratch, resume a committed run, or fine-tune an export. + + Configuration keys are ``dataset``, ``output_dir``, ``model``, ``train``, + ``manifest``, optional ``statistics``, ``resume``, ``fine_tune``, + ``include_prefix``, ``sample_order`` and ``stop_after``. Relative paths should + be resolved by the configuration-file entry point before this call. + """ + _arguments( + config, + { + "dataset", + "output_dir", + "model", + "train", + "manifest", + "statistics", + "resume", + "fine_tune", + "include_prefix", + "sample_order", + "stop_after", + }, + ) + import torch + from embodichain_models.dynamics.nerd.codec import NeRDStateCodec + from embodichain_models.dynamics.nerd import NeRDModel, NeRDModelCfg, load_model + from embodichain_models.dynamics.nerd.data import ( + DynamicsWindowDataset, + load_manifest, + ) + from ._trainer import DynamicsTrainCfg, DynamicsTrainer, load_statistics + from ._runtime import _training_runtime_identity + + if config.get("resume") and config.get("statistics"): + raise ValueError("Resume restores frozen statistics; omit statistics override") + if config.get("resume") and config.get("fine_tune"): + raise ValueError("resume and fine_tune are mutually exclusive") + cfg = DynamicsTrainCfg(**config.get("train", {})) + torch.manual_seed(cfg.seed) + dataset_metadata = load_manifest(config["dataset"])["metadata"] + supplied_manifest = config.get("manifest") + manifest = dataset_metadata if supplied_manifest is None else supplied_manifest + if not isinstance(manifest, dict): + raise ValueError("Training manifest must be a mapping") + # Canonical JSON data is also the export contract; copy before training so an + # external dictionary mutation cannot relabel a run after validation. + manifest = json.loads(json.dumps(manifest, allow_nan=False)) + missing = object() + for key in set(dataset_metadata) | _DEPLOYMENT_MANIFEST_FIELDS: + if manifest.get(key, missing) != dataset_metadata.get(key, missing): + raise ValueError( + f"Training manifest/data {key} mismatch; annotations may add fields but cannot replace or drop dataset provenance" + ) + codec = NeRDStateCodec(**manifest["codec"]) + if config.get("fine_tune"): + model, pretrained_manifest = load_model(config["fine_tune"], device=cfg.device) + if pretrained_manifest["codec"] != manifest["codec"] or pretrained_manifest.get( + "dt" + ) != manifest.get("dt"): + raise ValueError("Fine-tuning requires compatible codec and timestep") + elif config.get("resume"): + from embodichain_models.dynamics.nerd.checkpoint import _model_from_state + + saved = _load_training_checkpoint(config["resume"]) + # Keep the requested architecture/feature order for resume preflight, + # while restoring precision before trainer construction. Release the + # full optimizer payload before resume reads its committed state again. + model = _model_from_state(NeRDModelCfg(**config["model"]), saved["model"]) + del saved + else: + model = NeRDModel(NeRDModelCfg(**config["model"])) + train = DynamicsWindowDataset( + config["dataset"], + "train", + model.cfg.history_length, + include_prefix=config.get("include_prefix", False), + verify=True, + ) + validation = DynamicsWindowDataset( + config["dataset"], + "validation", + model.cfg.history_length, + include_prefix=config.get("include_prefix", False), + ) + statistics = ( + load_statistics(config["statistics"]) if config.get("statistics") else None + ) + order = None + if config.get("sample_order"): + order_path = Path(config["sample_order"]) + order = ( + order_path + if order_path.suffix in {".npy", ".json"} + else torch.load(order_path, map_location="cpu", weights_only=True) + ) + try: + trainer = DynamicsTrainer( + model, + codec, + train, + validation, + cfg, + config["output_dir"], + manifest=manifest, + statistics=statistics, + initialize_normalization=not config.get("resume") + and (not config.get("fine_tune") or statistics is not None), + sample_order=order, + ) + if config.get("resume"): + trainer.resume(config["resume"]) + result = trainer.train(stop_after=config.get("stop_after")) + result["runtime_identity"] = _training_runtime_identity(trainer.model) + Path(config["output_dir"], "result.json").write_text( + json.dumps(result, indent=2, allow_nan=False) + ) + return result + finally: + train.close() + validation.close() + + +def run_eval(config: dict[str, Any]) -> dict[str, Any]: + """Evaluate an exported model on a fixed offline split, without simulation.""" + _arguments( + config, + { + "dataset", + "model", + "split", + "batch_size", + "device", + "output", + "include_prefix", + }, + ) + from embodichain_models.dynamics.nerd.codec import NeRDStateCodec + from embodichain_models.dynamics.nerd import load_model + from embodichain_models.dynamics.nerd.data import DynamicsWindowDataset + from ._trainer import evaluate_model + + device = config.get("device", "cpu") + model, manifest = load_model(config["model"], device=device) + codec = NeRDStateCodec(**manifest["codec"]) + dataset = DynamicsWindowDataset( + config["dataset"], + config.get("split", "test"), + model.cfg.history_length, + include_prefix=config.get("include_prefix", False), + verify=True, + ) + try: + metadata = dataset.manifest["metadata"] + if metadata["codec"] != manifest["codec"] or metadata.get("dt") != manifest.get( + "dt" + ): + raise ValueError("Evaluation data/model codec or timestep mismatch") + result = { + **evaluate_model( + model, codec, dataset, config.get("batch_size", 64), device + ), + "dataset_fingerprint": dataset.fingerprint, + "qualification": "offline_only", + "numerical_profile": manifest.get("numerical_profile", "default"), + "attention_backend": model.cfg.attention_backend, + } + if config.get("output"): + output = Path(config["output"]) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(result, indent=2, allow_nan=False)) + return result + finally: + dataset.close() + + +def run_export(checkpoint: str | Path, output: str | Path) -> dict[str, Any]: + """Export NeRD or shared-SL training tensors without arbitrary unpickling.""" + import importlib + import numpy as np + import torch + + from embodichain_models.dynamics.nerd import NeRDModelCfg, export_model + from embodichain_models.dynamics.nerd.checkpoint import _model_from_state + from embodichain_models.dynamics.nerd._numerics import _checkpoint_numerical_profile + + # Generic SL records NumPy's MT19937 RNG state, unlike the historical + # tensor-only NeRD checkpoint. Restrict loading to those known array types. + array_module = importlib.import_module( + "numpy._core.multiarray" if hasattr(np, "_core") else "numpy.core.multiarray" + ) + reconstruct = array_module._reconstruct + with torch.serialization.safe_globals( + [ + np.ndarray, + np.dtype, + type(np.dtype("uint32")), + (reconstruct, "numpy.core.multiarray._reconstruct"), + (reconstruct, "numpy._core.multiarray._reconstruct"), + ] + ): + state = torch.load(checkpoint, map_location="cpu", weights_only=True) + if state.get("version") == 1 and "learner" in state: + signature = state["signature"] + model_cfg = dict(signature["model"]) + weights = state["learner"]["model"] + layout = signature.get("model_input_layout") + if layout is None: + # Earlier SL checkpoints canonicalized signature dictionaries, but + # state_dict retains the model's ordered normalizer registration. + names = [ + key[len("input_rms.") : -len(".mean")] + for key in weights + if key.startswith("input_rms.") and key.endswith(".mean") + ] + layout = [(name, model_cfg["input_dims"][name]) for name in names] + if ( + len(layout) != len(model_cfg["input_dims"]) + or dict(layout) != model_cfg["input_dims"] + ): + raise ValueError("SL checkpoint has an invalid ordered NeRD feature layout") + model_cfg["input_dims"] = dict(layout) + profile = signature["objective"]["numerical_profile"] + model = _model_from_state(NeRDModelCfg(**model_cfg), weights) + manifest = { + **signature["physics"], + "numerical_profile": profile, + "attention_backend": model.cfg.attention_backend, + "training": { + "global_step": state["learner"]["num_updates"], + "qualification": "unqualified", + "numerical_profile": profile, + "attention_backend": model.cfg.attention_backend, + }, + } + export_model(output, model, manifest) + return { + "output": str(output), + "global_step": state["learner"]["num_updates"], + "qualification": "unqualified", + "numerical_profile": profile, + "attention_backend": model.cfg.attention_backend, + } + if state.get("format") != "embodichain.nerd.training.v1": + raise ValueError("Unsupported NeRD or shared-SL training checkpoint") + profile_state = state + if state.get("resumable") is False and "train_cfg" not in state: + # A selected-model snapshot has no optimizer/training state. Retain the + # same execution-label validation without inventing resumable state. + profile_state = { + **state, + "train_cfg": { + "numerical_profile": state.get("numerical_profile", "default") + }, + } + profile = _checkpoint_numerical_profile(profile_state) + model = _model_from_state(NeRDModelCfg(**state["model_cfg"]), state["model"]) + manifest = { + **state["manifest"], + "numerical_profile": profile, + "attention_backend": model.cfg.attention_backend, + "training": { + "global_step": state["global_step"], + "dataset_fingerprint": state["dataset_fingerprint"], + "qualification": "unqualified", + "numerical_profile": profile, + "attention_backend": model.cfg.attention_backend, + }, + } + export_model(output, model, manifest) + return { + "output": str(output), + "global_step": state["global_step"], + "qualification": "unqualified", + "numerical_profile": profile, + "attention_backend": model.cfg.attention_backend, + } diff --git a/examples/learning/nerd/configs/ant/README.md b/examples/learning/nerd/configs/ant/README.md new file mode 100644 index 0000000000..b174536645 --- /dev/null +++ b/examples/learning/nerd/configs/ant/README.md @@ -0,0 +1,40 @@ +# Ant NeRD configurations + +These files belong to the [NeRD showcase](../../README.md), not runnable Gym +task deployments. The default `train` command uses `sl.train.yaml` and the +common supervised trainer: + +```bash +embodichain run-example learning/nerd prepare --robot ant --resource-root ./outputs/nerd-resources +embodichain run-example learning/nerd train --robot ant --resource-root ./outputs/nerd-resources --device cuda:0 --output-dir ./outputs/nerd/ant/train +embodichain run-example learning/nerd export --checkpoint ./outputs/nerd/ant/train/last.pt --output ./outputs/nerd/ant/export +embodichain run-example learning/nerd eval --robot ant --resource-root ./outputs/nerd-resources --device cpu --model ./outputs/nerd/ant/export +``` + +The SL recipe is bounded to 1,000 updates, with an explicit resident-cache +budget. Use `--device cpu` for CPU execution. To use `embodichain train-sl`, copy +`sl.train.yaml`, set `data.path` and `trainer.output_dir`, and resolve relative +paths from the working directory. The factory is installed under +`embodichain_models.dynamics.nerd.sl_recipe`, so no example imports are needed. +The final SL checkpoint does not imply validation-selected accuracy. + +Teacher collection defaults to `dynamics.generate.smoke.yaml`. Prepare assets +with `prepare --assets-only`, then run `generate --robot ant` with the same +`--resource-root` and an explicit `--output-dir`. Newton's analytical solver +owns integration; no DexSim custom solver or Scene extension is required. +The `full` and `warp-reference` presets retain their recorded scales and +physical profiles, which do not promise identical dynamics across versions. + +An explicitly supplied `dynamics.train.*.yaml` retains the historical trainer +for comparison. Those configuration paths resolve relative to the file; +CLI path overrides resolve from the working directory. Historical `best.pt` is +export-only; `last.pt` restores optimizer/RNG/sampling and the selected best +weights. The `full` historical preset disables resident caching rather than +assuming the entire dataset fits a fixed device budget. + +The pretrained evaluation default uses the pinned published dataset/model from +`prepare`. Downloaded data retains its Scene teacher provenance; new direct +Newton data has a distinct physical signature. See the +[feature guide](../../../../../docs/source/features/neural_dynamics.md) and +[research acceptance](../../research/acceptance.md) for format and evidence +boundaries. Preset names do not establish accuracy qualification. diff --git a/examples/learning/nerd/configs/ant/dynamics.eval.full.yaml b/examples/learning/nerd/configs/ant/dynamics.eval.full.yaml new file mode 100644 index 0000000000..47b0d32aa1 --- /dev/null +++ b/examples/learning/nerd/configs/ant/dynamics.eval.full.yaml @@ -0,0 +1,8 @@ +# Dynamics workflow configuration; not a runnable Gym deployment. +mode: offline +dataset: ../../../../../outputs/dynamics/ant/full/data +model: ../../../../../outputs/dynamics/ant/full/export +split: test +batch_size: 512 +device: cuda:0 +output: ../../../../../outputs/dynamics/ant/full/offline.json diff --git a/examples/learning/nerd/configs/ant/dynamics.eval.pretrained.yaml b/examples/learning/nerd/configs/ant/dynamics.eval.pretrained.yaml new file mode 100644 index 0000000000..772002febb --- /dev/null +++ b/examples/learning/nerd/configs/ant/dynamics.eval.pretrained.yaml @@ -0,0 +1,8 @@ +# Offline evaluation of the published NeRD model and dataset. +mode: offline +dataset: ../../../../../outputs/datasets/neural_dynamics/nerd/ant/v1 +model: ../../../../../outputs/models/neural_dynamics/nerd/ant/v1 +split: test +batch_size: 256 +device: cpu +output: ../../../../../outputs/dynamics/ant/pretrained/offline.json diff --git a/examples/learning/nerd/configs/ant/dynamics.eval.smoke.yaml b/examples/learning/nerd/configs/ant/dynamics.eval.smoke.yaml new file mode 100644 index 0000000000..fd6d951988 --- /dev/null +++ b/examples/learning/nerd/configs/ant/dynamics.eval.smoke.yaml @@ -0,0 +1,8 @@ +# Dynamics workflow configuration; not a runnable Gym deployment. +mode: offline +dataset: ../../../../../outputs/dynamics/ant/smoke/data +model: ../../../../../outputs/dynamics/ant/smoke/export +split: test +batch_size: 64 +device: cpu +output: ../../../../../outputs/dynamics/ant/smoke/offline.json diff --git a/examples/learning/nerd/configs/ant/dynamics.generate.full.yaml b/examples/learning/nerd/configs/ant/dynamics.generate.full.yaml new file mode 100644 index 0000000000..1da63e9fb5 --- /dev/null +++ b/examples/learning/nerd/configs/ant/dynamics.generate.full.yaml @@ -0,0 +1,18 @@ +# Dynamics workflow configuration; not a runnable Gym deployment. +robot: ant +output_dir: ../../../../../outputs/dynamics/ant/full/data +device: cuda:0 +num_envs: 512 +trajectories: 125000 +horizon: 100 +seed: 0 +teacher_solver: + solver_type: mujoco_warp + use_mujoco_cpu: false +num_substeps: 10 +trajectories_per_shard: 256 +resume: false +motion_mode: random_effort +initial_states: sample +effort_scale: 1.0 +split_seed: 0 diff --git a/examples/learning/nerd/configs/ant/dynamics.generate.smoke.yaml b/examples/learning/nerd/configs/ant/dynamics.generate.smoke.yaml new file mode 100644 index 0000000000..06394bd6af --- /dev/null +++ b/examples/learning/nerd/configs/ant/dynamics.generate.smoke.yaml @@ -0,0 +1,19 @@ +# Dynamics workflow configuration; not a runnable Gym deployment. +robot: ant +output_dir: ../../../../../outputs/dynamics/ant/smoke/data +device: cpu +num_envs: 4 +trajectories: 64 +horizon: 32 +seed: 0 +teacher_solver: + solver_type: mujoco_warp + use_mujoco_cpu: true + separate_worlds: true +num_substeps: 10 +trajectories_per_shard: 16 +resume: false +motion_mode: random_effort +initial_states: default +effort_scale: 0.05 +split_seed: 0 diff --git a/examples/learning/nerd/configs/ant/dynamics.generate.warp-reference.yaml b/examples/learning/nerd/configs/ant/dynamics.generate.warp-reference.yaml new file mode 100644 index 0000000000..6db1478a8c --- /dev/null +++ b/examples/learning/nerd/configs/ant/dynamics.generate.warp-reference.yaml @@ -0,0 +1,19 @@ +# Bounded parameter-alignment collection; not full accuracy qualification. +robot: ant +output_dir: ../../../../../outputs/dynamics/ant/warp-reference/data +device: cpu +num_envs: 64 +trajectories: 512 +horizon: 100 +seed: 100 +physics_profile: warp_reference +teacher_solver: + solver_type: featherstone + update_mass_matrix_interval: 16 + friction_smoothing: 0.001 +num_substeps: 16 +trajectories_per_shard: 64 +motion_mode: random_effort +initial_states: sample +effort_scale: 1.0 +split_seed: 0 diff --git a/examples/learning/nerd/configs/ant/dynamics.train.full.yaml b/examples/learning/nerd/configs/ant/dynamics.train.full.yaml new file mode 100644 index 0000000000..cf9d1df029 --- /dev/null +++ b/examples/learning/nerd/configs/ant/dynamics.train.full.yaml @@ -0,0 +1,39 @@ +# Dynamics workflow configuration; not a runnable Gym deployment. +dataset: ../../../../../outputs/dynamics/ant/full/data +output_dir: ../../../../../outputs/dynamics/ant/full/train +model: + input_dims: + states_embedding: 29 + contact_normals: 75 + contact_points_1: 75 + contact_depths: 25 + joint_f: 8 + gravity_dir: 3 + output_dim: 29 + n_embd: 384 + n_layer: 6 + n_head: 12 + block_size: 32 + history_length: 10 + head_hidden_dims: + - 64 + bias: false + dropout: 0.0 + normalize_input: true + normalize_output: true +train: + total_steps: 4995000 + batch_size: 512 + learning_rate: 0.001 + final_learning_rate: 0.0001 + max_grad_norm: 1.0 + seed: 0 + device: cuda:0 + validation_interval: 5000 + validation_batches: 100 + checkpoint_interval: 5000 + lr_update_interval: 5000 + lr_offset: 1 + log_interval: 500 + # Full Ant data exceed the former 4 GiB hard cache budget; use bounded reads. + resident_cache_bytes: 0 diff --git a/examples/learning/nerd/configs/ant/dynamics.train.smoke.yaml b/examples/learning/nerd/configs/ant/dynamics.train.smoke.yaml new file mode 100644 index 0000000000..7e8b697d89 --- /dev/null +++ b/examples/learning/nerd/configs/ant/dynamics.train.smoke.yaml @@ -0,0 +1,36 @@ +# Dynamics workflow configuration; not a runnable Gym deployment. +dataset: ../../../../../outputs/dynamics/ant/smoke/data +output_dir: ../../../../../outputs/dynamics/ant/smoke/train +model: + input_dims: + states_embedding: 29 + contact_normals: 75 + contact_points_1: 75 + contact_depths: 25 + joint_f: 8 + gravity_dir: 3 + output_dim: 29 + n_embd: 24 + n_layer: 1 + n_head: 3 + block_size: 32 + history_length: 10 + head_hidden_dims: + - 16 + bias: false + dropout: 0.0 + normalize_input: true + normalize_output: true +train: + total_steps: 100 + batch_size: 16 + learning_rate: 0.001 + final_learning_rate: 0.0001 + max_grad_norm: 1.0 + seed: 0 + device: cpu + validation_interval: 25 + validation_batches: null + checkpoint_interval: 25 + lr_update_interval: 1 + lr_offset: 0 diff --git a/examples/learning/nerd/configs/ant/sl.train.yaml b/examples/learning/nerd/configs/ant/sl.train.yaml new file mode 100644 index 0000000000..f794bc4c26 --- /dev/null +++ b/examples/learning/nerd/configs/ant/sl.train.yaml @@ -0,0 +1,46 @@ +# Standalone SL recipe; run from the repository root after downloading data. +factory: embodichain_models.dynamics.nerd.sl_recipe:build_supervised_trainer +model: + input_dims: + states_embedding: 29 + contact_normals: 75 + contact_points_1: 75 + contact_depths: 25 + joint_f: 8 + gravity_dir: 3 + output_dim: 29 + n_embd: 384 + n_layer: 6 + n_head: 12 + block_size: 32 + history_length: 10 + head_hidden_dims: + - 64 + bias: false + dropout: 0.0 + normalize_input: true + normalize_output: true + attention_backend: math +data: + path: outputs/datasets/neural_dynamics/nerd/ant/v1 + # Keep full-corpus training on bounded reads unless a fitting budget is set. + resident_cache_bytes: 0 +objective: + include_prefix: false + numerical_profile: reference +optimizer: + learning_rate: 0.001 + final_learning_rate: 0.0001 + lr_update_interval: 500 + lr_offset: 1 + max_grad_norm: 1.0 +trainer: + total_updates: 1000 + batch_size: 512 + device: cuda:0 + seed: 0 + output_dir: outputs/sl/nerd/ant + checkpoint_interval: 500 + validation_interval: 500 +evaluation: + max_batches: 100 diff --git a/examples/learning/nerd/configs/cartpole/README.md b/examples/learning/nerd/configs/cartpole/README.md new file mode 100644 index 0000000000..0772649a0d --- /dev/null +++ b/examples/learning/nerd/configs/cartpole/README.md @@ -0,0 +1,41 @@ +# Cartpole NeRD configurations + +These files belong to the [NeRD showcase](../../README.md), not runnable Gym +task deployments. The default `train` command uses `sl.train.yaml` and the +common supervised trainer: + +```bash +embodichain run-example learning/nerd prepare --robot cartpole --resource-root ./outputs/nerd-resources +embodichain run-example learning/nerd train --robot cartpole --resource-root ./outputs/nerd-resources --device cuda:0 --output-dir ./outputs/nerd/cartpole/train +embodichain run-example learning/nerd export --checkpoint ./outputs/nerd/cartpole/train/last.pt --output ./outputs/nerd/cartpole/export +embodichain run-example learning/nerd eval --robot cartpole --resource-root ./outputs/nerd-resources --device cpu --model ./outputs/nerd/cartpole/export +``` + +The SL recipe is bounded to 1,000 updates, with an explicit resident-cache +budget. Use `--device cpu` for CPU execution. To use `embodichain train-sl`, copy +`sl.train.yaml`, set `data.path` and `trainer.output_dir`, and resolve relative +paths from the working directory. The factory is installed under +`embodichain_models.dynamics.nerd.sl_recipe`, so no example imports are needed. +The final SL checkpoint does not imply validation-selected accuracy. + +Teacher collection defaults to `dynamics.generate.smoke.yaml`. Prepare assets +with `prepare --assets-only`, then run `generate --robot cartpole` with the same +`--resource-root` and an explicit `--output-dir`. Newton's analytical solver +owns integration; no DexSim custom solver or Scene extension is required. +The `full` and `warp-reference` presets retain their recorded scales and +physical profiles, which do not promise identical dynamics across versions. + +An explicitly supplied `dynamics.train.*.yaml` retains the historical trainer +for comparison. Those configuration paths resolve relative to the file; +CLI path overrides resolve from the working directory. Historical `best.pt` is +export-only; `last.pt` restores optimizer/RNG/sampling and the selected best +weights. The full historical Cartpole preset declares a 4 GiB resident-cache +budget; datasets exceeding that explicit limit fail admission rather than +allocating an unbounded cache. Ant's full preset instead disables resident caching. + +The pretrained evaluation default uses the pinned published dataset/model from +`prepare`. Downloaded data retains its Scene teacher provenance; new direct +Newton data has a distinct physical signature. See the +[feature guide](../../../../../docs/source/features/neural_dynamics.md) and +[research acceptance](../../research/acceptance.md) for format and evidence +boundaries. Preset names do not establish accuracy qualification. diff --git a/examples/learning/nerd/configs/cartpole/dynamics.eval.full.yaml b/examples/learning/nerd/configs/cartpole/dynamics.eval.full.yaml new file mode 100644 index 0000000000..e09a098017 --- /dev/null +++ b/examples/learning/nerd/configs/cartpole/dynamics.eval.full.yaml @@ -0,0 +1,8 @@ +# Dynamics workflow configuration; not a runnable Gym deployment. +mode: offline +dataset: ../../../../../outputs/dynamics/cartpole/full/data +model: ../../../../../outputs/dynamics/cartpole/full/export +split: test +batch_size: 512 +device: cuda:0 +output: ../../../../../outputs/dynamics/cartpole/full/offline.json diff --git a/examples/learning/nerd/configs/cartpole/dynamics.eval.pretrained.yaml b/examples/learning/nerd/configs/cartpole/dynamics.eval.pretrained.yaml new file mode 100644 index 0000000000..ced76765d1 --- /dev/null +++ b/examples/learning/nerd/configs/cartpole/dynamics.eval.pretrained.yaml @@ -0,0 +1,8 @@ +# Offline evaluation of the published NeRD model and dataset. +mode: offline +dataset: ../../../../../outputs/datasets/neural_dynamics/nerd/cartpole/v1 +model: ../../../../../outputs/models/neural_dynamics/nerd/cartpole/v1 +split: test +batch_size: 256 +device: cpu +output: ../../../../../outputs/dynamics/cartpole/pretrained/offline.json diff --git a/examples/learning/nerd/configs/cartpole/dynamics.eval.smoke.yaml b/examples/learning/nerd/configs/cartpole/dynamics.eval.smoke.yaml new file mode 100644 index 0000000000..1e4f164d67 --- /dev/null +++ b/examples/learning/nerd/configs/cartpole/dynamics.eval.smoke.yaml @@ -0,0 +1,8 @@ +# Dynamics workflow configuration; not a runnable Gym deployment. +mode: offline +dataset: ../../../../../outputs/dynamics/cartpole/smoke/data +model: ../../../../../outputs/dynamics/cartpole/smoke/export +split: test +batch_size: 64 +device: cpu +output: ../../../../../outputs/dynamics/cartpole/smoke/offline.json diff --git a/examples/learning/nerd/configs/cartpole/dynamics.generate.full.yaml b/examples/learning/nerd/configs/cartpole/dynamics.generate.full.yaml new file mode 100644 index 0000000000..02b2b20ea0 --- /dev/null +++ b/examples/learning/nerd/configs/cartpole/dynamics.generate.full.yaml @@ -0,0 +1,18 @@ +# Dynamics workflow configuration; not a runnable Gym deployment. +robot: cartpole +output_dir: ../../../../../outputs/dynamics/cartpole/full/data +device: cuda:0 +num_envs: 512 +trajectories: 62500 +horizon: 100 +seed: 0 +teacher_solver: + solver_type: mujoco_warp + use_mujoco_cpu: false +num_substeps: 10 +trajectories_per_shard: 256 +resume: false +motion_mode: random_effort +initial_states: sample +effort_scale: 1.0 +split_seed: 0 diff --git a/examples/learning/nerd/configs/cartpole/dynamics.generate.smoke.yaml b/examples/learning/nerd/configs/cartpole/dynamics.generate.smoke.yaml new file mode 100644 index 0000000000..9efc7d69b3 --- /dev/null +++ b/examples/learning/nerd/configs/cartpole/dynamics.generate.smoke.yaml @@ -0,0 +1,19 @@ +# Dynamics workflow configuration; not a runnable Gym deployment. +robot: cartpole +output_dir: ../../../../../outputs/dynamics/cartpole/smoke/data +device: cpu +num_envs: 4 +trajectories: 64 +horizon: 32 +seed: 0 +teacher_solver: + solver_type: mujoco_warp + use_mujoco_cpu: true + separate_worlds: true +num_substeps: 10 +trajectories_per_shard: 16 +resume: false +motion_mode: random_effort +initial_states: default +effort_scale: 1.0 +split_seed: 0 diff --git a/examples/learning/nerd/configs/cartpole/dynamics.generate.warp-reference.yaml b/examples/learning/nerd/configs/cartpole/dynamics.generate.warp-reference.yaml new file mode 100644 index 0000000000..d15857ec8a --- /dev/null +++ b/examples/learning/nerd/configs/cartpole/dynamics.generate.warp-reference.yaml @@ -0,0 +1,18 @@ +# Bounded parameter-alignment collection; not full accuracy qualification. +robot: cartpole +output_dir: ../../../../../outputs/dynamics/cartpole/warp-reference/data +device: cpu +num_envs: 64 +trajectories: 512 +horizon: 100 +seed: 100 +physics_profile: warp_reference +teacher_solver: + solver_type: featherstone + update_mass_matrix_interval: 5 +num_substeps: 5 +trajectories_per_shard: 64 +motion_mode: random_effort +initial_states: sample +effort_scale: 1.0 +split_seed: 0 diff --git a/examples/learning/nerd/configs/cartpole/dynamics.train.full.yaml b/examples/learning/nerd/configs/cartpole/dynamics.train.full.yaml new file mode 100644 index 0000000000..1a1d8f6f93 --- /dev/null +++ b/examples/learning/nerd/configs/cartpole/dynamics.train.full.yaml @@ -0,0 +1,34 @@ +# Dynamics workflow configuration; not a runnable Gym deployment. +dataset: ../../../../../outputs/dynamics/cartpole/full/data +output_dir: ../../../../../outputs/dynamics/cartpole/full/train +model: + input_dims: + states_embedding: 4 + joint_f: 2 + output_dim: 4 + n_embd: 192 + n_layer: 6 + n_head: 12 + block_size: 32 + history_length: 10 + head_hidden_dims: + - 64 + bias: false + dropout: 0.0 + normalize_input: true + normalize_output: true +train: + total_steps: 4995000 + batch_size: 512 + learning_rate: 0.001 + final_learning_rate: 0.0001 + max_grad_norm: 1.0 + seed: 0 + device: cuda:0 + validation_interval: 5000 + validation_batches: 100 + checkpoint_interval: 5000 + lr_update_interval: 5000 + lr_offset: 1 + log_interval: 500 + resident_cache_bytes: 4294967296 diff --git a/examples/learning/nerd/configs/cartpole/dynamics.train.smoke.yaml b/examples/learning/nerd/configs/cartpole/dynamics.train.smoke.yaml new file mode 100644 index 0000000000..a9d1543bf9 --- /dev/null +++ b/examples/learning/nerd/configs/cartpole/dynamics.train.smoke.yaml @@ -0,0 +1,32 @@ +# Dynamics workflow configuration; not a runnable Gym deployment. +dataset: ../../../../../outputs/dynamics/cartpole/smoke/data +output_dir: ../../../../../outputs/dynamics/cartpole/smoke/train +model: + input_dims: + states_embedding: 4 + joint_f: 2 + output_dim: 4 + n_embd: 24 + n_layer: 1 + n_head: 3 + block_size: 32 + history_length: 10 + head_hidden_dims: + - 16 + bias: false + dropout: 0.0 + normalize_input: true + normalize_output: true +train: + total_steps: 100 + batch_size: 16 + learning_rate: 0.001 + final_learning_rate: 0.0001 + max_grad_norm: 1.0 + seed: 0 + device: cpu + validation_interval: 25 + validation_batches: null + checkpoint_interval: 25 + lr_update_interval: 1 + lr_offset: 0 diff --git a/examples/learning/nerd/configs/cartpole/sl.train.yaml b/examples/learning/nerd/configs/cartpole/sl.train.yaml new file mode 100644 index 0000000000..8fa869c445 --- /dev/null +++ b/examples/learning/nerd/configs/cartpole/sl.train.yaml @@ -0,0 +1,41 @@ +# Standalone SL recipe; run from the repository root after downloading data. +factory: embodichain_models.dynamics.nerd.sl_recipe:build_supervised_trainer +model: + input_dims: + states_embedding: 4 + joint_f: 2 + output_dim: 4 + n_embd: 192 + n_layer: 6 + n_head: 12 + block_size: 32 + history_length: 10 + head_hidden_dims: + - 64 + bias: false + dropout: 0.0 + normalize_input: true + normalize_output: true + attention_backend: math +data: + path: outputs/datasets/neural_dynamics/nerd/cartpole/v1 + resident_cache_bytes: 4294967296 +objective: + include_prefix: false + numerical_profile: reference +optimizer: + learning_rate: 0.001 + final_learning_rate: 0.0001 + lr_update_interval: 500 + lr_offset: 1 + max_grad_norm: 1.0 +trainer: + total_updates: 1000 + batch_size: 512 + device: cuda:0 + seed: 0 + output_dir: outputs/sl/nerd/cartpole + checkpoint_interval: 500 + validation_interval: 500 +evaluation: + max_batches: 100 diff --git a/examples/learning/nerd/research/README.md b/examples/learning/nerd/research/README.md new file mode 100644 index 0000000000..c3da68e1fa --- /dev/null +++ b/examples/learning/nerd/research/README.md @@ -0,0 +1,65 @@ +# NeRD reference experiments + +The [showcase](../README.md) owns resource preparation, default SL training, +export and offline evaluation. This directory documents controlled comparisons +with the original NeRD implementation and historical training workflow. Private +runners live in `../_support/`; logs, checkpoints, reports and machine-specific +plans belong in caller-owned output directories. + +Run a support tool through the example launcher so source and installed-wheel +example packages receive the same bootstrap: + +```bash +embodichain run-example learning/nerd/_support/_comparison --help +embodichain run-example learning/nerd/_support/_reference --help +embodichain run-example learning/nerd/_support/_official_math_rollout --help +``` + +These research comparisons require explicit source/runtime paths. Source-bound +acceptance uses source checkouts even though the model package and ordinary +showcase commands are installable. The reference workers use their selected +interpreter and the private stdlib bootstrap; an installed example directory +is not a globally importable `examples` distribution. + +## Environments and frozen sources + +Run unchanged upstream source in its supported Python/Torch/Warp environment. +Use a separate EmbodiChain environment for current Newton collection and model +training. Original-runtime helpers do not import the installed simulation SDK. +Keep virtualenv executable paths intact rather than resolving their symlinks +to a different base interpreter. Trusted external source/checkpoint execution +requires each tool's explicit trust option. + +The [acceptance protocol](acceptance.md) pins reference source +`a58de92a82b4e1c1c260083b5b1930173322aceb`, dataset revision +`b9ec58f36da99cc99a93f97adfe74ebc9fd66586` and model revision +`78e581526570069410bf269b38c0c52c800b24a7`. Reference asset revision +`97e202e41fe7fb504656a6b163291305e400780d` is prepared by the showcase; raw +URDF/XML and license/attribution are external Hub payloads. Keep every existing +report bound to its actual source commit and environment. The extension/example +migration does not constitute a rerun of a historical passing matrix. + +`_portable_reference.py` and `_historical_physics.py` adapt only the isolated +original runtime. Preserve the original numerical implementation separately from +compatibility transport; do not call a migrated implementation an upstream oracle. + +## Comparison contracts + +- Separate model/trainer changes, runtime kernels, coordinate transport and + analytical-solver versions when interpreting differences. +- Preserve sampling, normalization, loss and model selection for trainer-port + parity. Equal seeds alone do not prove the same consumed batches. +- Replay identical realized open-loop efforts for predictive accuracy. A + controller using each model's own state measures closed-loop behavior. +- Keep current Newton and historical Warp in their recorded physical domains; + do not replace signatures or normalizers to make incompatible exports load. +- Retain every requested trajectory and report nonfinite outputs and failures. + Distinguish endpoint error from trajectory-time average error. +- Use the matching analytical solver as the accuracy reference. Distance to + another neural model measures similarity, not ground-truth error. +- State dataset size, update budget, population and horizon independently. + A complete bounded budget or README example does not reproduce paper-scale + training or establish live solver rollout accuracy. + +Tools write to fresh caller-supplied directories. They do not change the model's +production numerical policy or provide a blanket accuracy acceptance. diff --git a/examples/learning/nerd/research/acceptance.md b/examples/learning/nerd/research/acceptance.md new file mode 100644 index 0000000000..1a02b8a37a --- /dev/null +++ b/examples/learning/nerd/research/acceptance.md @@ -0,0 +1,118 @@ +# NeRD supervised-learning acceptance + +The default [showcase](../README.md) uses the installable NeRD model/objective +and the common SL trainer. Its 1,000-update recipes demonstrate the public entry +point, not full-budget convergence. The historical trainer under +`../_support/_trainer.py` also delegates optimizer updates to `SupervisedLearner` +while preserving its source-specific sampling, normalization, numerical schedule +and checkpoint format for the frozen comparisons below. + +```bash +embodichain run-example learning/nerd prepare --robot cartpole --resource-root ./outputs/nerd-resources +embodichain run-example learning/nerd train --robot cartpole --resource-root ./outputs/nerd-resources --device cuda:0 --output-dir ./outputs/nerd/cartpole/train +``` + +Use `--robot ant` for Ant. For direct `embodichain train-sl`, copy the robot's +`../configs//sl.train.yaml`, set `data.path` and `trainer.output_dir`, and +use the installable factory +`embodichain_models.dynamics.nerd.sl_recipe:build_supervised_trainer`. Common SL +checkpoints resume with the same recipe through `train-sl --resume` or the +showcase's default `train --resume`. Export its committed `last.pt` through +`embodichain run-example learning/nerd export`; common SL does not automatically +select a best model. Explicit `dynamics.train.*` configurations select the +historical trainer instead. + +## Run identity and historical checkpoint contract + +Each new training run claims its output with `run.json`. Starting fresh or +fine-tuning requires a new destination; neither may replace another run's +metrics or checkpoints. Resume preserves the checkpoint's run identity and +rejects a different run's directory or metric history newer than that checkpoint. +Metric history is append-only. An invalid or incomplete log record requires a +fresh resume destination so the original evidence is retained. +If a crash published a later best model before its checkpoint/log boundary, +same-directory recovery also requires a fresh destination rather than replacing +that best file with an older selection. +Historical checkpoints without a run identity can resume into an artifact-free +destination. Input feature order is part of recovery identity for both trainer +formats. + +In the historical trainer, `last.pt` is the complete resumable checkpoint. +Automatically selected `best.pt` files are model-only snapshots marked +`resumable=False`, suitable for `embodichain run-example learning/nerd export`. +They retain the best model's own step and frozen normalizers; resume from +`last.pt`. The full checkpoint retains one owned CPU +copy of the best model and its step, so a new destination can recreate `best.pt` +even when subsequent validation does not improve. This adds one model's parameter +and buffer storage on CPU after the first selected best, plus that payload in +`last.pt`; it does not retain historical optimizers or change update arithmetic. +For the frozen published Cartpole/Ant models below, those tensor payloads are +10.35/41.55 MiB respectively; this is storage accounting, not a transient peak +measurement. Resume into an occupied owned directory may read one best snapshot +to validate its selection boundary, outside the optimizer-update path. +Legacy complete best checkpoints remain readable. If a historical checkpoint +has a finite best score but no embedded best weights, its matching source-side +`best.pt` must be present rather than silently substituting the latest weights. + +## Frozen comparison + +The source IDs below name the historical protocol and its independent reference. +Reports retain the exact candidate source used for each run. Moving the model, +data and research code into extension/example packages does not rerun or relabel +those reports. A current-source comparison must write a new output and record +its new candidate identity; a short installed-package smoke cannot substitute +for these quality gates. + +Use PR #855 revision `a58de92a82b4e1c1c260083b5b1930173322aceb` as the +independent reference source. Use published dataset revision +`b9ec58f36da99cc99a93f97adfe74ebc9fd66586` and model revision +`78e581526570069410bf269b38c0c52c800b24a7`. Downloaded roots must contain +`cartpole/v1` and `ant/v1`. The runner verifies all data-shard checksums and +records model/dataset manifest hashes, weights hashes and source identities. + +```bash +CUDA_VISIBLE_DEVICES=1 embodichain run-example learning/nerd/_support/_acceptance \ + --reference-source /path/to/frozen-a58-worktree \ + --candidate-source /path/to/candidate-worktree \ + --datasets /path/to/downloaded/neural_dynamics/nerd \ + --models /path/to/downloaded/neural_dynamics/nerd \ + --device cuda:0 --output /path/to/fresh/results +``` + +The default matrix performs: + +- Full test-split offline evaluation of the same published Cartpole/Ant weights + through independent reference and candidate processes. +- Five paired seeds (`0..4`) per robot, 1,000 updates each, batch size 512. Model + architecture and train-only normalization match the published artifacts; + weights start from scratch. The schedule retains the original 4,995,000-update + horizon, 5,000-update interval and offset rather than shortening it to the + acceptance interruption point. +- First-update comparison of predictions, loss, raw/clipped gradients and + updated parameters using `atol=1e-5`, `rtol=1e-4`. +- Exact final-weight comparison and a five-seed, one-sided 95% paired bootstrap + gate allowing at most 5% normalized-validation-MSE degradation. Near-zero + reference metrics use the established absolute-tolerance mode. +- A process-restarted 500+500 update candidate run for seed zero of each robot, + compared with uninterrupted weights, optimizer, sampler, counters and RNG. + +`final_gradients.pt` is an additional diagnostic artifact, separate from these +26 frozen gates. After training it copies the existing clipped parameter +gradients to CPU and records the reported final pre-clipping norm. It performs +no extra forward/backward/update and adds no training hooks or changes to the +acceptance thresholds or budget. The common-trainer orchestration comparison +can reuse it alongside the original first-update capture and final checkpoint. + +`summary.json` is the machine-readable gate; the command exits nonzero if a gate +fails or fewer than five paired seeds are supplied. Seed lists must contain +distinct integers; duplicate or malformed lists are rejected before output +creation or worker startup. `report.md` contains time, +memory, quality and trainer-ranking tables. Per-process requests, logs, +checkpoints, curves and first-update captures remain under the supplied output +directory. Use a fresh directory for every matrix. + +Time includes worker setup, fitting and validation. Process peak RSS is a Linux +high-water measurement; GPU peaks cover the Torch allocator. The reference +model uses its original physical data signature. No Scene-generated model is +relabeled as a newly generated direct-Newton domain, and these offline checks +do not establish native solver rollout accuracy or paper-scale reproduction. diff --git a/extensions/embodichain_models/README.md b/extensions/embodichain_models/README.md new file mode 100644 index 0000000000..7509d4f05d --- /dev/null +++ b/extensions/embodichain_models/README.md @@ -0,0 +1,65 @@ +# Official models + +`embodichain_models` contains official model implementations grouped by model +family, including `embodichain_models.dynamics.nerd`. It ships in the same +`embodichain` wheel and uses that distribution's version and dependencies. There +is no separate model distribution, version file or model registry. + +Import a named model package explicitly. The top-level package and family +packages do not eagerly import individual models or initialize a simulator. +Shared optimizer/trainer contracts remain in `embodichain.learning.sl`; shared +MLP and normalization building blocks remain in `embodichain.learning.models`. +The core does not import official model implementations. + +NeRD datasets, teacher assets and learned weights are prepared by +`examples/learning/nerd/_resources.py` from pinned Hugging Face revisions. +They remain outside Git and the wheel; the showcase owns their download and +verification lifecycle. Model-specific supervised factories live in the +installable model package so `embodichain train-sl` can resolve them from any +working directory. Example helpers remain owned by their example workflows. + +## Domains and public imports + +Organize implementations as `embodichain_models..`. +The first level describes the model's prediction contract: + +| Domain | Prediction contract | +| --- | --- | +| `policy` | Observations/history to actions, action distributions or action chunks. | +| `dynamics` | State/history, with optional controls and physical context, to future robot/environment states, geometry, observations or latent transitions. | + +Use the following canonical packages for the four current model directions: + +| Model family | Public package | Landing status | +| --- | --- | --- | +| Ordinary state/visual BC | `embodichain_models.policy.bc` | Target for the BC follow-up; not created yet. | +| NeRD | `embodichain_models.dynamics.nerd` | Implemented in this change. | +| Robot world model (RWM) | `embodichain_models.dynamics.rwm` | Target for the RWM follow-up; not created yet. | +| RigidFormer | `embodichain_models.dynamics.rigidformer` | Target for the RigidFormer follow-up; not created yet. | + +RWM predicts action-conditioned robot state transitions, so it belongs to +`dynamics`. World models also describe future-state/observation prediction; +`world_models` is not a parallel top-level domain or an alternative import path. +Keep each implementation in one canonical package. Robots/tasks, state/RGB/point +modalities, MLP/Transformer backbones and SL/RL training methods do not determine +the first-level directory. Add another domain only for a concrete responsibility +that neither existing prediction contract covers. + +The directory does not imply capabilities such as action conditioning, +state/action gradients or solver deployment. Those remain explicit model +contracts: the current RigidFormer adapter predicts rigid geometry without an +action input and does not qualify action gradients. + +## Ownership inside a model family + +A family owns its model/configuration, portable-checkpoint support where +available, codecs, model-format data adapters and model-specific objectives. +Its `sl_recipe.py` owns the installable trainer factory. Shared training, +checkpoint/recovery and distributed orchestration belong to `learning.sl`; +model packages do not introduce parallel trainer frameworks. + +Collection, task/controller integration, reference research and downloadable +resources remain in the corresponding example. Shared MLPs and normalizers +remain in core. All four directions use the same SL contracts regardless of +their domain. Create a family only when its implementation lands, without empty +placeholder directories, eager cross-family imports or a new registry. diff --git a/extensions/embodichain_models/embodichain_models/__init__.py b/extensions/embodichain_models/embodichain_models/__init__.py new file mode 100644 index 0000000000..addbcd4951 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/__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. +# ---------------------------------------------------------------------------- + +"""Official model implementations bundled in the EmbodiChain distribution.""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/extensions/embodichain_models/embodichain_models/dynamics/__init__.py b/extensions/embodichain_models/embodichain_models/dynamics/__init__.py new file mode 100644 index 0000000000..6fd00181de --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/__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. +# ---------------------------------------------------------------------------- + +"""Official learned dynamics models, imported through their named subpackages.""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/extensions/embodichain_models/embodichain_models/dynamics/nerd/THIRD_PARTY_LICENSES b/extensions/embodichain_models/embodichain_models/dynamics/nerd/THIRD_PARTY_LICENSES new file mode 100644 index 0000000000..e30e79e719 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/nerd/THIRD_PARTY_LICENSES @@ -0,0 +1,30 @@ +NeRD model and state conventions originate from: +https://github.com/NVlabs/neural-robot-dynamics +Reference commit: 476dc67fccc6a99be293bf611f897b1c06f13e30 +Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +Licensed under Apache License 2.0 (repository LICENSE). + +The Transformer architecture derives from nanoGPT: +https://github.com/karpathy/nanoGPT + +MIT License + +Copyright (c) 2022 Andrej Karpathy + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/extensions/embodichain_models/embodichain_models/dynamics/nerd/__init__.py b/extensions/embodichain_models/embodichain_models/dynamics/nerd/__init__.py new file mode 100644 index 0000000000..31dd9c6e2d --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/nerd/__init__.py @@ -0,0 +1,58 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Lazy NeRD networks, physical codecs and supervised training integration.""" + +from __future__ import annotations + +from importlib import import_module +from typing import Any + +__all__ = [ + "NeRDModelCfg", + "NeRDModel", + "FeatureNormalizer", + "export_model", + "load_model", + "convert_legacy_checkpoint", + "NeRDStateCodec", + "NeRDObjective", + "make_nerd_batch", + "NeRDWindowSource", + "build_supervised_trainer", +] + +_EXPORT_MODULES = { + "NeRDModelCfg": "model", + "NeRDModel": "model", + "FeatureNormalizer": "model", + "export_model": "checkpoint", + "load_model": "checkpoint", + "convert_legacy_checkpoint": "checkpoint", + "NeRDStateCodec": "codec", + "NeRDObjective": "supervised", + "make_nerd_batch": "supervised", + "NeRDWindowSource": "sl_recipe", + "build_supervised_trainer": "sl_recipe", +} + + +def __getattr__(name: str) -> Any: + if name not in _EXPORT_MODULES: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + value = getattr(import_module(f".{_EXPORT_MODULES[name]}", __name__), name) + globals()[name] = value + return value diff --git a/extensions/embodichain_models/embodichain_models/dynamics/nerd/_numerics.py b/extensions/embodichain_models/embodichain_models/dynamics/nerd/_numerics.py new file mode 100644 index 0000000000..85d127da0b --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/nerd/_numerics.py @@ -0,0 +1,121 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Explicit numerical execution for reference-aligned dynamics training. + +This module depends only on Torch and the standard library so an isolated +historical training worker can load the same policy without importing the SDK. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from contextlib import AbstractContextManager, nullcontext + +import torch +from torch import Tensor + +__all__: list[str] = [] + + +def _validate_numerical_profile(profile: str) -> str: + if not isinstance(profile, str) or profile not in {"default", "reference"}: + raise ValueError("numerical_profile must be 'default' or 'reference'") + return profile + + +def _checkpoint_numerical_profile(state: dict) -> str: + """Resolve legacy default checkpoints and reject conflicting annotations.""" + configured = _validate_numerical_profile( + state["train_cfg"].get("numerical_profile", "default") + ) + declared = _validate_numerical_profile(state.get("numerical_profile", configured)) + if configured != declared: + raise ValueError("Checkpoint numerical_profile metadata mismatch") + backend = state["model_cfg"].get("attention_backend", "default") + if state.get("attention_backend", backend) != backend: + raise ValueError("Checkpoint attention_backend metadata mismatch") + if declared == "reference" and backend != "math": + raise ValueError("Reference checkpoint requires math attention_backend") + return declared + + +def _attention_context(profile: str) -> AbstractContextManager: + """Select math SDPA explicitly, restoring all backend flags on exit.""" + _validate_numerical_profile(profile) + if profile == "default": + return nullcontext() + if hasattr(torch.nn, "attention") and hasattr(torch.nn.attention, "sdpa_kernel"): + return torch.nn.attention.sdpa_kernel([torch.nn.attention.SDPBackend.MATH]) + return torch.backends.cuda.sdp_kernel( + enable_flash=False, enable_math=True, enable_mem_efficient=False + ) + + +@torch.jit.script +def _ordered_sum(values: Tensor) -> Tensor: + # Do not replace this left fold with a parallel/tree sum. The squared + # inputs are computed separately, preserving multiply/add rounding order. + total = values.new_zeros(()) + for index in range(values.numel()): + total = total + values[index] + return total + + +@torch.no_grad() +def _clip_grad_norm_( + parameters: Iterable[Tensor], + max_norm: float, + *, + profile: str, + error_if_nonfinite: bool = True, +) -> Tensor: + """Clip FP32 single-device gradients with a fixed reference total norm. + + Each tensor retains its FP32 L2 norm. Only the final ordered square/sum/sqrt + uses FP64, matching the validated CPU scalar policy. CUDA scalars stay on + device; callers can combine their finite check with step logging by setting + ``error_if_nonfinite=False`` and checking the returned norm before Adam. + """ + _validate_numerical_profile(profile) + if profile == "default": + return torch.nn.utils.clip_grad_norm_( + parameters, max_norm, error_if_nonfinite=error_if_nonfinite + ) + parameters = list(parameters) + grads = [ + parameter.grad.detach() + for parameter in parameters + if parameter.grad is not None + ] + if not grads: + return parameters[0].new_zeros(()) if parameters else torch.zeros(()) + device = grads[0].device + if device.type not in {"cpu", "cuda"} or any( + grad.device != device or grad.dtype != torch.float32 or grad.is_sparse + for grad in grads + ): + raise ValueError( + "reference numerical_profile requires dense FP32 gradients on one CPU/CUDA device" + ) + norms = torch.stack(torch._foreach_norm(grads, 2.0)) + squared = norms.to(dtype=torch.float64).square() + total = _ordered_sum(squared).sqrt().to(dtype=torch.float32) + if error_if_nonfinite and not bool(torch.isfinite(total)): + raise FloatingPointError("Non-finite reference gradient norm") + coefficient = torch.clamp(float(max_norm) / (total + 1e-6), max=1.0) + torch._foreach_mul_(grads, coefficient) + return total diff --git a/extensions/embodichain_models/embodichain_models/dynamics/nerd/_resident.py b/extensions/embodichain_models/embodichain_models/dynamics/nerd/_resident.py new file mode 100644 index 0000000000..9663efcaba --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/nerd/_resident.py @@ -0,0 +1,131 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Budgeted immutable transition storage for long single-process training runs.""" + +from __future__ import annotations + +import math +import os +from typing import Any + +import numpy as np +import torch + +__all__: list[str] = [] + + +class _ResidentWindows: + """Own contiguous raw fields; gather fresh bounded windows without HDF5 I/O.""" + + def __init__( + self, dataset: Any, device: str | torch.device, max_bytes: int + ) -> None: + if type(max_bytes) is not int or max_bytes < 1: + raise ValueError("Resident cache budget must be a positive byte count") + if not len(dataset): + raise ValueError("Resident cache requires eligible trajectory windows") + self.device = torch.device(device) + self.pid = os.getpid() + self.rows = dataset._rows + self.ends = np.asarray(dataset._ends, dtype=np.int64) + self.history = dataset.history_length + self.include_prefix = dataset.include_prefix + self.offsets = {} + self.fields: dict[str, dict[str, torch.Tensor]] = {} + required = 0 + for kind, fields in dataset.manifest["fields"].items(): + sizes = [row["length"] + int(kind == "states") for row in self.rows] + self.offsets[kind] = np.concatenate(([0], np.cumsum(sizes, dtype=np.int64))) + for spec in fields.values(): + required += ( + int(self.offsets[kind][-1]) + * math.prod(spec["shape"]) + * np.dtype(spec["dtype"]).itemsize + ) + if required > max_bytes: + raise ValueError( + f"Resident fields require {required} bytes, exceeding cache budget {max_bytes}" + ) + self.nbytes = required + for kind, fields in dataset.manifest["fields"].items(): + self.fields[kind] = {} + for name, spec in fields.items(): + dtype = torch.from_numpy( + np.empty((), dtype=np.dtype(spec["dtype"])) + ).dtype + self.fields[kind][name] = torch.empty( + (int(self.offsets[kind][-1]), *spec["shape"]), + dtype=dtype, + device=self.device, + ) + # Host staging remains bounded by the existing read-page size even for + # very long trajectories. Copying never changes the persisted precision. + for row_index, row in enumerate(self.rows): + source = dataset._handle(row["file"])[row["group"]] + for kind, fields in dataset.manifest["fields"].items(): + count = row["length"] + int(kind == "states") + base = int(self.offsets[kind][row_index]) + for name, spec in fields.items(): + for start in range(0, count, dataset.read_page_steps): + end = min(count, start + dataset.read_page_steps) + block = dataset._read_block( + source[kind], + name, + start, + end, + tuple(spec["shape"]), + np.dtype(spec["dtype"]), + ) + self.fields[kind][name][base + start : base + end].copy_( + torch.from_numpy(block) + ) + + def get_batch(self, requests: list[int]) -> dict[str, Any]: + if os.getpid() != self.pid: + raise RuntimeError("Resident dynamics caches are owned by one process") + indices = np.asarray(requests, dtype=np.int64) + indices = np.where(indices < 0, indices + self.ends[-1], indices) + if np.any(indices < 0) or np.any(indices >= self.ends[-1]): + raise IndexError("Resident dynamics window index out of range") + rows = np.searchsorted(self.ends, indices, side="right") + first = np.concatenate(([0], self.ends[:-1])) + offsets = indices - first[rows] + ends = offsets + (1 if self.include_prefix else self.history) + starts = np.maximum(0, ends - self.history) + lengths = ends - starts + result = { + "trajectory_id": [self.rows[index]["id"] for index in rows], + "start": torch.from_numpy(starts), + "length": torch.from_numpy(lengths), + "valid_mask": torch.arange(self.history)[None, :] + < torch.from_numpy(lengths)[:, None], + } + for kind, fields in self.fields.items(): + extra = int(kind == "states") + take = ( + self.offsets[kind][rows, None] + + starts[:, None] + + np.minimum( + np.arange(self.history + extra)[None, :], + (lengths - 1 + extra)[:, None], + ) + ) + take = torch.as_tensor(take, device=self.device) + # Advanced indexing returns independent tensors, including repeated + # rows. Later seeks/close cannot mutate a returned learner batch. + result[kind] = {name: value[take] for name, value in fields.items()} + return result diff --git a/extensions/embodichain_models/embodichain_models/dynamics/nerd/checkpoint.py b/extensions/embodichain_models/embodichain_models/dynamics/nerd/checkpoint.py new file mode 100644 index 0000000000..893a582e0f --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/nerd/checkpoint.py @@ -0,0 +1,233 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Versioned portable NeRD exports and explicitly trusted legacy conversion.""" + +from __future__ import annotations +from collections.abc import Mapping +import hashlib +import json +import pickle +from pathlib import Path +from typing import Any +import torch +from torch import nn +from .model import NeRDModel, NeRDModelCfg + +__all__ = ["export_model", "load_model", "convert_legacy_checkpoint"] +_FORMAT_VERSION = 1 + + +def _execution_metadata(manifest: dict[str, Any], cfg: NeRDModelCfg) -> str: + """Validate execution labels without importing training into model loading.""" + training = manifest.get("training", {}) + training = training if isinstance(training, dict) else {} + profile = manifest.get( + "numerical_profile", training.get("numerical_profile", "default") + ) + if not isinstance(profile, str) or profile not in {"default", "reference"}: + raise ValueError("Unknown exported numerical_profile") + if training.get("numerical_profile", profile) != profile: + raise ValueError("Export numerical_profile metadata mismatch") + if ( + manifest.get("attention_backend", cfg.attention_backend) + != cfg.attention_backend + or training.get("attention_backend", cfg.attention_backend) + != cfg.attention_backend + ): + raise ValueError("Export attention_backend metadata mismatch") + if profile == "reference" and cfg.attention_backend != "math": + raise ValueError("Reference exports require math attention_backend") + return profile + + +def _digest(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _model_from_state( + cfg: NeRDModelCfg, + state: Mapping[str, torch.Tensor], + device: str | torch.device = "cpu", +) -> NeRDModel: + """Restore saved floating precision before copying parameters and statistics.""" + dtype = next(value.dtype for value in state.values() if value.is_floating_point()) + model = NeRDModel(cfg).to(dtype=dtype) + model.load_state_dict(state, strict=True) + return model.to(device) + + +def export_model(path: str | Path, model: NeRDModel, manifest: dict[str, Any]) -> None: + """Write tensor weights and JSON schema without executable Python objects. + + Args: + path: Output directory, which must not already contain an export. + model: Model whose parameters and statistics are exported on CPU. + manifest: Deployment metadata, including codec and timestep contracts. + """ + path = Path(path) + if (path / "manifest.json").exists() or (path / "weights.pt").exists(): + raise FileExistsError(f"Model export already exists: {path}") + metadata = dict(manifest) + profile = _execution_metadata(metadata, model.cfg) + metadata.update( + numerical_profile=profile, attention_backend=model.cfg.attention_backend + ) + metadata.update(format_version=_FORMAT_VERSION, model_config=model.cfg.to_dict()) + json.dumps(metadata, allow_nan=False) + path.mkdir(parents=True, exist_ok=True) + weights = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()} + temporary = path / "weights.pt.tmp" + torch.save(weights, temporary) + temporary.replace(path / "weights.pt") + metadata["weights_sha256"] = _digest(path / "weights.pt") + temporary = path / "manifest.json.tmp" + temporary.write_text(json.dumps(metadata, indent=2, allow_nan=False) + "\n") + temporary.replace(path / "manifest.json") + + +def load_model( + path: str | Path, device: str | torch.device = "cpu" +) -> tuple[NeRDModel, dict[str, Any]]: + """Verify and load an exported model; returned model is in evaluation mode.""" + path = Path(path) + manifest = json.loads((path / "manifest.json").read_text()) + if manifest.get("format_version") != _FORMAT_VERSION: + raise ValueError("Unsupported NeRD export format_version") + if _digest(path / "weights.pt") != manifest.get("weights_sha256"): + raise ValueError("NeRD weights checksum mismatch") + cfg = NeRDModelCfg(**manifest["model_config"]) + _execution_metadata(manifest, cfg) + state = torch.load(path / "weights.pt", map_location="cpu", weights_only=True) + return _model_from_state(cfg, state, device).eval(), manifest + + +class _LegacyModule(nn.Module): + pass + + +class _LegacyObject: + pass + + +class _LegacyUnpickler(pickle.Unpickler): + def find_class(self, module: str, name: str) -> Any: + # Remap historical source locations without importing a checkout into runtime. + if module in { + "models.models", + "models.base_models", + "models.model_transformer", + }: + return _LegacyObject if name == "GPTConfig" else _LegacyModule + if module == "utils.running_mean_std" and name == "RunningMeanStd": + return _LegacyObject + return super().find_class(module, name) + + +class _LegacyPickle: + __name__ = "nerd_legacy_pickle" + Unpickler = _LegacyUnpickler + load = staticmethod(pickle.load) + + +def convert_legacy_checkpoint( + checkpoint: str | Path, + destination: str | Path, + manifest: dict[str, Any], + *, + trust_pickle: bool = False, +) -> None: + """Convert a trusted official Transformer checkpoint into a portable export. + + Args: + checkpoint: Official model.pt, containing a serialized Python model. + destination: New export directory. + manifest: Deployment metadata validated by the eventual provider. + trust_pickle: Must be explicitly true: legacy pickle can execute code. + + Only the official identity encoder/ReLU head/Transformer architecture is + accepted. Arbitrary architectures must be migrated explicitly. + """ + if not trust_pickle: + raise ValueError("Legacy checkpoints require trust_pickle=True") + old = torch.load( + checkpoint, map_location="cpu", weights_only=False, pickle_module=_LegacyPickle + ) + if isinstance(old, (list, tuple)): + old = old[0] + if not getattr(old, "is_transformer", False) or getattr(old, "is_rnn", False): + raise ValueError("Only official Transformer checkpoints are supported") + if list(old.encoders) != ["low_dim"] or len(old.encoders["low_dim"].body): + raise ValueError("Only identity low-dimensional encoders are supported") + layers = list(old.model.feature_net.body) + if len(layers) % 2 or any( + not isinstance(m, nn.Linear) if i % 2 == 0 else not isinstance(m, nn.ReLU) + for i, m in enumerate(layers) + ): + raise ValueError("Only Linear/ReLU prediction heads are supported") + + def rename(name: str) -> str: + return "joint_f" if name == "joint_acts" else name + + c = old.transformer_model.config + dims = { + rename(k): int(old.input_rms[k].mean.numel()) for k in old.low_dim_input_names + } + training = manifest.get("training", {}) + training = training if isinstance(training, dict) else {} + profile = manifest.get( + "numerical_profile", training.get("numerical_profile", "default") + ) + cfg = NeRDModelCfg( + input_dims=dims, + output_dim=int(old.output_rms.mean.numel()), + n_embd=c.n_embd, + n_layer=c.n_layer, + n_head=c.n_head, + block_size=c.block_size, + history_length=int(manifest.get("history_length", 10)), + head_hidden_dims=tuple( + m.out_features for m in layers if isinstance(m, nn.Linear) + ), + bias=c.bias, + dropout=c.dropout, + normalize_input=old.normalize_input, + normalize_output=old.normalize_output, + output_tanh=old.output_tanh, + attention_backend=manifest.get( + "attention_backend", "math" if profile == "reference" else "default" + ), + ) + model = NeRDModel(cfg) + state = old.state_dict() + result = model.load_state_dict(state, strict=False) + expected_missing = { + k for k in model.state_dict() if k.startswith(("input_rms.", "output_rms.")) + } + if set(result.missing_keys) != expected_missing or result.unexpected_keys: + raise ValueError(f"Unexpected legacy parameter schema: {result}") + model.set_normalization( + {rename(k): (v.mean, v.var) for k, v in old.input_rms.items()}, + old.output_rms.mean, + old.output_rms.var, + ) + metadata = dict(manifest) + metadata["legacy_checkpoint_sha256"] = _digest(Path(checkpoint)) + export_model(destination, model, metadata) diff --git a/extensions/embodichain_models/embodichain_models/dynamics/nerd/codec.py b/extensions/embodichain_models/embodichain_models/dynamics/nerd/codec.py new file mode 100644 index 0000000000..57a7236f07 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/nerd/codec.py @@ -0,0 +1,533 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Differentiable NeRD state and contact conventions, independent of simulation. + +Coordinate/target conventions follow NVlabs/neural-robot-dynamics (Apache-2.0). +The native-COM adapter and functional history ownership are EmbodiChain additions. +""" + +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from __future__ import annotations +from collections.abc import Mapping, Sequence +import torch +from torch import Tensor + +__all__ = [ + "NeRDStateCodec", + "quat_multiply", + "quat_rotate", + "quat_inverse", + "plane_contact_features", + "append_history", + "reset_history", +] + + +def quat_inverse(q: Tensor) -> Tensor: + """Return the inverse of unit xyzw quaternions.""" + return torch.cat((-q[..., :3], q[..., 3:]), -1) + + +@torch.jit.script +def _quat_product_flat(a: Tensor, b: Tensor) -> Tensor: + """Share the reference product between eager and scripted coordinate math.""" + assert a.shape == b.shape + shape = a.shape + a = a.reshape(-1, 4) + b = b.reshape(-1, 4) + x1, y1, z1, w1 = a[:, 0], a[:, 1], a[:, 2], a[:, 3] + x2, y2, z2, w2 = b[:, 0], b[:, 1], b[:, 2], b[:, 3] + ww = (z1 + x1) * (x2 + y2) + yy = (w1 - y1) * (w2 + z2) + zz = (w1 + y1) * (w2 - z2) + xx = ww + yy + zz + qq = 0.5 * (xx + (z1 - x1) * (x2 - y2)) + w = qq - ww + (z1 - y1) * (y2 - z2) + x = qq - xx + (x1 + w1) * (x2 + w2) + y = qq - yy + (w1 - x1) * (y2 + z2) + z = qq - zz + (z1 + y1) * (w2 - x2) + return torch.stack([x, y, z, w], dim=-1).view(shape) + + +def quat_multiply(a: Tensor, b: Tensor) -> Tensor: + """Multiply broadcastable xyzw quaternions in official NeRD FP32 order. + + Algebraically equivalent Hamilton-product formulas round differently. The + identity orientation in Ant's body-frame features is sensitive to these + differences after normalization, so preserve the reference's operation order. + """ + a, b = torch.broadcast_tensors(a, b) + return _quat_product_flat(a.reshape(-1, 4), b.reshape(-1, 4)).reshape(a.shape) + + +@torch.jit.script +def _inverse_rotate_flat(q: Tensor, value: Tensor) -> Tensor: + """Keep positive-quaternion inverse rotation inside the reference JIT graph.""" + shape = q.shape + scalar, xyz = q[:, -1], q[:, :3] + a = value * (2.0 * scalar**2 - 1.0).unsqueeze(-1) + b = torch.cross(xyz, value, dim=-1) * scalar.unsqueeze(-1) * 2.0 + c = ( + xyz + * torch.bmm(xyz.view(shape[0], 1, 3), value.view(shape[0], 3, 1)).squeeze(-1) + * 2.0 + ) + return a - b + c + + +@torch.jit.script +def _quat_inverse_flat(value: Tensor) -> Tensor: + shape = value.shape + value = value.reshape(-1, 4) + return torch.cat((-value[:, :3], value[:, -1:]), dim=-1).view(shape) + + +@torch.jit.script +def _point_inverse_flat(pos: Tensor, rotation: Tensor, point: Tensor) -> Tensor: + return _inverse_rotate_flat(rotation, point - pos) + + +@torch.jit.script +def _body_coordinates_flat( + frame_pos: Tensor, + frame_quat: Tensor, + pos: Tensor, + rotation: Tensor, + omega: Tensor, + nu: Tensor, +) -> tuple[Tensor, Tensor, Tensor, Tensor]: + """Match the official scripted world-to-body execution, including CUDA rounding.""" + pos_body = _point_inverse_flat(frame_pos, frame_quat, pos) + rotation_body = _quat_product_flat(_quat_inverse_flat(frame_quat), rotation) + omega_body = _inverse_rotate_flat(frame_quat, omega) + nu_body = _inverse_rotate_flat( + frame_quat, nu - torch.cross(frame_pos, omega, dim=-1) + ) + return pos_body, rotation_body, omega_body, nu_body + + +def quat_rotate(q: Tensor, v: Tensor) -> Tensor: + """Rotate vectors by unit xyzw quaternions in official NeRD FP32 order.""" + xyz, v = torch.broadcast_tensors(q[..., :3], v) + scalar = q[..., 3:] + a = v * (2.0 * scalar.square() - 1.0) + b = torch.linalg.cross(xyz, v) * scalar * 2.0 + dot = torch.bmm(xyz.reshape(-1, 1, 3), v.reshape(-1, 3, 1)) + c = xyz * dot.reshape(*v.shape[:-1], 1) * 2.0 + return a + b + c + + +def _canonical(q: Tensor) -> Tensor: + return torch.where(q[..., 3:] < 0, -q, q) + + +def _normalize(q: Tensor) -> Tensor: + return q / torch.linalg.vector_norm(q, dim=-1, keepdim=True).clamp_min(1e-9) + + +def _wrap(x: Tensor) -> Tensor: + return x - torch.floor((x + torch.pi) / (2 * torch.pi)) * (2 * torch.pi) + + +class NeRDStateCodec: + """Encode Newton child-COM velocities as the official NeRD spatial twist. + + Args: + robot: ``cartpole`` or ``ant``; fixes q/qd layout to 2/2 or 15/14. + base_com: Root body's COM offset in its local body frame. + continuous_q_indices: Coordinates wrapped to [-pi, pi). Defaults to the + Cartpole pole coordinate and no Ant coordinates. + + State tensors may have any leading batch/time dimensions. Ant q uses xyz, + xyzw, eight joint coordinates. Native qd uses world COM linear velocity, + world angular velocity, eight joint speeds. The legacy spatial twist uses + angular velocity then linear velocity referenced to the world origin. + """ + + def __init__( + self, + robot: str = "cartpole", + base_com: Sequence[float] = (0.0, 0.0, 0.0), + continuous_q_indices: Sequence[int] | None = None, + ) -> None: + self.robot = robot.lower() + if self.robot not in {"cartpole", "ant"}: + raise ValueError(f"Unsupported NeRD robot: {robot}") + self.q_dim, self.qd_dim = (2, 2) if self.robot == "cartpole" else (15, 14) + if len(base_com) != 3: + raise ValueError("base_com must contain three coordinates") + self.base_com = tuple(base_com) + self.continuous_q_indices = tuple( + continuous_q_indices + if continuous_q_indices is not None + else ((1,) if self.robot == "cartpole" else ()) + ) + if any( + i < 0 or i >= self.q_dim or (self.robot == "ant" and i < 7) + for i in self.continuous_q_indices + ): + raise ValueError("continuous coordinates must be scalar joint coordinates") + + def _check(self, q: Tensor, qd: Tensor) -> None: + if ( + q.shape[-1] != self.q_dim + or qd.shape[-1] != self.qd_dim + or q.shape[:-1] != qd.shape[:-1] + ): + raise ValueError( + f"Expected matching q/qd batch shapes with dimensions {self.q_dim}/{self.qd_dim}" + ) + + def _wrap_q(self, q: Tensor) -> Tensor: + if not self.continuous_q_indices: + return q + return torch.stack( + [ + _wrap(q[..., i]) if i in self.continuous_q_indices else q[..., i] + for i in range(self.q_dim) + ], + -1, + ) + + def native_to_legacy(self, q: Tensor, qd: Tensor) -> Tensor: + """Convert native qd to angular-first world-origin spatial velocity.""" + self._check(q, qd) + if self.robot == "cartpole": + return qd + com = q[..., :3] + quat_rotate(q[..., 3:7], q.new_tensor(self.base_com)) + omega = qd[..., 3:6] + return torch.cat( + (omega, qd[..., :3] + torch.linalg.cross(com, omega), qd[..., 6:]), -1 + ) + + def legacy_to_native(self, q: Tensor, velocity: Tensor) -> Tensor: + """Convert angular-first world-origin velocity to native COM velocity.""" + self._check(q, velocity) + if self.robot == "cartpole": + return velocity + com = q[..., :3] + quat_rotate(q[..., 3:7], q.new_tensor(self.base_com)) + omega = velocity[..., :3] + return torch.cat( + ( + velocity[..., 3:6] - torch.linalg.cross(com, omega), + omega, + velocity[..., 6:], + ), + -1, + ) + + def _resolve_legacy_velocity( + self, q: Tensor, qd: Tensor, legacy_velocity: Tensor | None + ) -> Tensor: + """Keep explicit same-frame source values without a lossy COM roundtrip.""" + if legacy_velocity is None: + return self.native_to_legacy(q, qd) + self._check(q, qd) + if ( + legacy_velocity.shape != qd.shape + or legacy_velocity.dtype != qd.dtype + or legacy_velocity.device != qd.device + ): + raise ValueError( + "Legacy velocity override must match native qd shape, dtype and device" + ) + return legacy_velocity + + def _body( + self, anchor: Tensor, q: Tensor, velocity: Tensor + ) -> tuple[Tensor, Tensor]: + shape = q.shape[:-1] + frame = anchor[..., :7].reshape(-1, 7).contiguous() + state = torch.cat((q, velocity), -1).reshape(-1, self.q_dim + self.qd_dim) + pos, rot, omega, nu = _body_coordinates_flat( + frame[:, :3], + frame[:, 3:7], + state[:, :3], + state[:, 3:7], + state[:, self.q_dim : self.q_dim + 3], + state[:, self.q_dim + 3 : self.q_dim + 6], + ) + return torch.cat( + (pos.reshape(*shape, 3), rot.reshape(*shape, 4), q[..., 7:]), -1 + ), torch.cat( + (omega.reshape(*shape, 3), nu.reshape(*shape, 3), velocity[..., 6:]), -1 + ) + + def encode( + self, + q: Tensor, + qd: Tensor, + effort: Tensor, + contacts: Mapping[str, Tensor] | None = None, + gravity_dir: Tensor | None = None, + *, + legacy_velocity: Tensor | None = None, + ) -> dict[str, Tensor]: + """Encode inputs at t; contact points/normals are world-frame flattened slots. + + Args: + q: Current native joint positions. + qd: Native joint velocities, using COM velocity for the Ant free base. + effort: Generalized efforts, or eight actuated efforts for Ant. + contacts: Fixed ordered world-frame contact features for Ant. + gravity_dir: Optional world-frame unit gravity direction. + legacy_velocity: Optional original NeRD velocity in the same world + frame as ``q``. It must match ``qd`` shape, dtype and device. + This preserves imported FP32 values without a COM roundtrip; + omitting it retains native velocity conversion. It is an + independent input and is never modified or detached. + + Returns: + Named unnormalized model features, retaining their input graphs. + """ + velocity = self._resolve_legacy_velocity(q, qd, legacy_velocity) + if self.robot == "cartpole": + if effort.shape != qd.shape: + raise ValueError( + "Cartpole effort must have two generalized coordinates" + ) + return { + "states_embedding": torch.cat((self._wrap_q(q), velocity), -1), + "joint_f": effort, + } + bq, bv = self._body(q, q, velocity) + bq = torch.cat((bq[..., :3], _canonical(bq[..., 3:7]), bq[..., 7:]), -1) + if effort.shape[-1] not in (8, 14) or effort.shape[:-1] != q.shape[:-1]: + raise ValueError( + "Ant effort must have 8 actuated or 14 generalized coordinates" + ) + if contacts is None: + raise ValueError("Ant requires explicitly ordered fixed contact slots") + depth = contacts["contact_depths"] + slots = depth.shape[-1] + if depth.shape[:-1] != q.shape[:-1]: + raise ValueError("Contact batch/time shape must match state") + inv = quat_inverse(q[..., 3:7]).unsqueeze(-2) + points = contacts["contact_points_1"].reshape(*q.shape[:-1], slots, 3) + normals = contacts["contact_normals"].reshape_as(points) + points = quat_rotate(inv, points - q[..., None, :3]) + normals = quat_rotate(inv, normals) + if "contact_masks" in contacts: + mask = contacts["contact_masks"].bool() + else: + t0 = contacts.get("contact_thicknesses_0", torch.zeros_like(depth)) + t1 = contacts.get("contact_thicknesses_1", torch.zeros_like(depth)) + mask = depth < (4 * (t0 + t1)).clamp_min(0.12) + if gravity_dir is None: + gravity_dir = q.new_tensor((0.0, 0.0, -1.0)).expand(*q.shape[:-1], 3) + return { + "states_embedding": torch.cat((self._wrap_q(bq), bv), -1), + "contact_normals": torch.where(mask[..., None], normals, 0).flatten(-2), + "contact_points_1": torch.where(mask[..., None], points, 0).flatten(-2), + "contact_depths": torch.where(mask, depth, 0), + "joint_f": effort[..., -8:], + "gravity_dir": quat_rotate(inv.squeeze(-2), gravity_dir), + } + + def targets( + self, + q: Tensor, + qd: Tensor, + next_q: Tensor, + next_qd: Tensor, + *, + current_embedding: Tensor | None = None, + legacy_velocity: Tensor | None = None, + next_legacy_velocity: Tensor | None = None, + ) -> Tensor: + """Encode relative next state in the frame at t, including quaternion delta. + + Args: + q: Current native joint positions. + qd: Current native joint velocities. + next_q: Next native joint positions. + next_qd: Next native joint velocities. + current_embedding: Optional unnormalized ``states_embedding`` from + :meth:`encode` on exactly the same current state. Reusing it + preserves the reference trainer's single current-frame transform. + legacy_velocity: Optional original NeRD velocity at t, in the same + world frame as ``q`` and matching ``qd`` shape/dtype/device. + Omission retains native conversion. + next_legacy_velocity: Optional original NeRD velocity at t+1, in + the same world frame as ``next_q`` and matching ``next_qd``. + It is transformed into the frame at t like the native path. + + Returns: + Relative position/orientation and velocity targets. + """ + if current_embedding is not None and current_embedding.shape != ( + *q.shape[:-1], + self.q_dim + self.qd_dim, + ): + raise ValueError("Current embedding shape does not match state") + velocity = self._resolve_legacy_velocity(q, qd, legacy_velocity) + next_velocity = self._resolve_legacy_velocity( + next_q, next_qd, next_legacy_velocity + ) + if self.robot == "ant": + anchor = q + if current_embedding is None: + q, velocity = self._body(anchor, q, velocity) + else: + q, velocity = ( + current_embedding[..., : self.q_dim], + current_embedding[..., self.q_dim :], + ) + next_q, next_velocity = self._body(anchor, next_q, next_velocity) + q, next_q = self._wrap_q(q), self._wrap_q(next_q) + delta = _canonical( + quat_multiply(next_q[..., 3:7], quat_inverse(q[..., 3:7])) + ) + diff = self._wrap_q(next_q - q) + return torch.cat( + (diff[..., :3], delta, diff[..., 7:], next_velocity - velocity), -1 + ) + return torch.cat( + ( + self._wrap_q(self._wrap_q(next_q) - self._wrap_q(q)), + next_velocity - velocity, + ), + -1, + ) + + def decode( + self, + q: Tensor, + qd: Tensor, + prediction: Tensor, + *, + current_embedding: Tensor | None = None, + ) -> tuple[Tensor, Tensor]: + """Decode a relative prediction into native Newton q/qd without detaching. + + Args: + q: Current native joint positions. + qd: Current native joint velocities. + prediction: Relative network prediction in physical target units. + current_embedding: Optional unnormalized ``states_embedding`` from + :meth:`encode` for the same state, shared with model inputs. + + Returns: + Next native joint positions and velocities. + """ + if prediction.shape != (*q.shape[:-1], self.q_dim + self.qd_dim): + raise ValueError("Prediction shape does not match state") + if ( + current_embedding is not None + and current_embedding.shape != prediction.shape + ): + raise ValueError("Current embedding shape does not match state") + velocity = self.native_to_legacy(q, qd) + if self.robot == "cartpole": + return ( + self._wrap_q(self._wrap_q(q) + prediction[..., :2]), + qd + prediction[..., 2:], + ) + if current_embedding is None: + bq, bv = self._body(q, q, velocity) + else: + bq, bv = ( + current_embedding[..., : self.q_dim], + current_embedding[..., self.q_dim :], + ) + predicted_rot = _normalize( + quat_multiply(_normalize(prediction[..., 3:7]), bq[..., 3:7]) + ) + local_q = torch.cat( + ( + bq[..., :3] + prediction[..., :3], + predicted_rot, + bq[..., 7:] + prediction[..., 7 : self.q_dim], + ), + -1, + ) + local_v = bv + prediction[..., self.q_dim :] + rot = q[..., 3:7] + pos = q[..., :3] + quat_rotate(rot, local_q[..., :3]) + next_rot = quat_multiply(rot, local_q[..., 3:7]) + omega = quat_rotate(rot, local_v[..., :3]) + nu = torch.linalg.cross(q[..., :3], omega) + quat_rotate(rot, local_v[..., 3:6]) + next_q = self._wrap_q(torch.cat((pos, next_rot, local_q[..., 7:]), -1)) + next_velocity = torch.cat((omega, nu, local_v[..., 6:]), -1) + return next_q, self.legacy_to_native(next_q, next_velocity) + + +def plane_contact_features( + points_world: Tensor, + thickness0: Tensor, + plane_normal: Tensor, + plane_offset: Tensor | float = 0.0, + thickness1: Tensor | float = 0.0, +) -> dict[str, Tensor]: + """Project fixed slots onto n dot x = offset; normals must be unit length. + + Args: + points_world: Candidate center points, shape (..., slots, 3). + thickness0: Shape radius plus margin, shape (..., slots). + plane_normal: Unit normal, broadcastable to candidate points. + plane_offset: Plane displacement along the normal. + thickness1: Plane margin. + + Returns: + Official contact fields; depth is center-to-plane distance, before thickness. + """ + normals = torch.broadcast_to(plane_normal, points_world.shape) + depth = (points_world * normals).sum(-1) - plane_offset + projected = points_world - depth[..., None] * normals + t0 = torch.broadcast_to(thickness0, depth.shape) + t1 = torch.zeros_like(depth) + thickness1 + return { + "contact_points_1": projected.flatten(-2), + "contact_normals": normals.flatten(-2), + "contact_depths": depth, + "contact_thicknesses_0": t0, + "contact_thicknesses_1": t1, + "contact_masks": depth < (4 * (t0 + t1)).clamp_min(0.12), + } + + +def append_history( + history: Mapping[str, Tensor], lengths: Tensor, features: Mapping[str, Tensor] +) -> tuple[dict[str, Tensor], Tensor]: + """Append one token to left-aligned fixed capacity histories without mutation. + + Full rows slide left. Short rows retain position indices starting at zero. + Storage is O(batch * capacity * feature_width), with no host synchronization. + """ + if not history or set(history) != set(features): + raise ValueError("History and feature keys must be identical and nonempty") + capacity = next(iter(history.values())).shape[1] + pos = torch.arange(capacity, device=lengths.device)[None, :] + full = lengths >= capacity + result = {} + for name, values in history.items(): + shifted = torch.cat((values[:, 1:], torch.zeros_like(values[:, :1])), 1) + base = torch.where(full[:, None, None], shifted, values) + at = pos == lengths.clamp_max(capacity - 1)[:, None] + result[name] = torch.where(at[..., None], features[name][:, None, :], base) + return result, (lengths + 1).clamp_max(capacity) + + +def reset_history( + history: Mapping[str, Tensor], lengths: Tensor, mask: Tensor +) -> tuple[dict[str, Tensor], Tensor]: + """Functionally clear selected rows while retaining other rows' gradients.""" + return { + k: torch.where(mask[:, None, None], torch.zeros_like(v), v) + for k, v in history.items() + }, torch.where(mask, 0, lengths) diff --git a/extensions/embodichain_models/embodichain_models/dynamics/nerd/data.py b/extensions/embodichain_models/embodichain_models/dynamics/nerd/data.py new file mode 100644 index 0000000000..713953ed26 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/nerd/data.py @@ -0,0 +1,573 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Committed HDF5 trajectories and bounded, worker-local window reads.""" + +from __future__ import annotations + +from bisect import bisect_right +from collections import OrderedDict +import hashlib +import json +import os +from pathlib import Path +from typing import Any +from collections.abc import Sequence +import operator + +import h5py +import numpy as np +import torch +from torch.utils.data import Dataset + +__all__ = [ + "DynamicsShardWriter", + "DynamicsWindowDataset", + "trajectory_split", + "load_manifest", + "dataset_fingerprint", +] + + +def _json_bytes(value: Any) -> bytes: + return json.dumps( + value, sort_keys=True, separators=(",", ":"), allow_nan=False + ).encode() + + +def _atomic_json(path: Path, value: Any) -> None: + temporary = path.with_suffix(".json.pending") + with temporary.open("wb") as stream: + stream.write(_json_bytes(value)) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + descriptor = os.open(path.parent, os.O_RDONLY) + try: + os.fsync(descriptor) + finally: + os.close(descriptor) + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def trajectory_split(trajectory_id: str, seed: int = 0) -> str: + """Assign a stable trajectory-level 80/10/10 split independent of order.""" + fraction = ( + int.from_bytes( + hashlib.sha256(f"{seed}:{trajectory_id}".encode()).digest()[:8], "big" + ) + / 2**64 + ) + return "train" if fraction < 0.8 else "validation" if fraction < 0.9 else "test" + + +def load_manifest(root: str | Path, *, verify: bool = False) -> dict[str, Any]: + """Read the committed dataset snapshot, optionally checking shard hashes.""" + root = Path(root) + manifest = json.loads((root / "manifest.json").read_text()) + if manifest.get("schema_version") != 1: + raise ValueError("Unsupported dynamics dataset schema") + for shard in manifest["shards"]: + name = shard["file"] + if Path(name).name != name or not (root / name).is_file(): + raise ValueError(f"Missing or unsafe committed shard: {name}") + if verify and _sha256(root / name) != shard["sha256"]: + raise ValueError(f"Shard checksum mismatch: {name}") + return manifest + + +def dataset_fingerprint(manifest: dict[str, Any]) -> str: + """Hash immutable data identity, excluding operational generation progress. + + Manifests written before generation progress was included retain their + original fingerprint. Rejection-only commits do not change model inputs. + """ + identity = { + key: value for key, value in manifest.items() if key != "generation_progress" + } + return hashlib.sha256(_json_bytes(identity)).hexdigest() + + +class DynamicsShardWriter: + """Write trajectories directly to disk; only complete shards are visible. + + This is a single-writer API. Each caller-owned trajectory is written before + ``add_trajectory`` returns; no unbounded staging queue is retained. Resume + ignores files absent from the manifest, including an interrupted final shard. + """ + + def __init__( + self, + root: str | Path, + metadata: dict[str, Any], + trajectories_per_shard: int = 64, + *, + resume: bool = False, + split_seed: int = 0, + ) -> None: + if trajectories_per_shard < 1: + raise ValueError("trajectories_per_shard must be positive") + self.root = Path(root) + self.root.mkdir(parents=True, exist_ok=True) + self.capacity = trajectories_per_shard + self._file: h5py.File | None = None + self._pending: list[dict[str, Any]] = [] + self._closed = False + path = self.root / "manifest.json" + metadata = json.loads(_json_bytes(metadata)) + if path.exists(): + if not resume: + raise FileExistsError(f"Dataset already exists: {path}") + self.manifest = load_manifest(self.root, verify=True) + if ( + self.manifest["metadata"] != metadata + or self.manifest["split_seed"] != split_seed + ): + raise ValueError("Resume metadata or split seed does not match") + else: + if resume: + raise FileNotFoundError(path) + self.manifest = { + "schema_version": 1, + "metadata": metadata, + "split_seed": split_seed, + "fields": None, + "shards": [], + } + _atomic_json(path, self.manifest) + self._ids = { + row["id"] + for shard in self.manifest["shards"] + for row in shard["trajectories"] + } + + @property + def committed_ids(self) -> frozenset[str]: + """IDs persisted in the current manifest, excluding pending writes.""" + return frozenset( + row["id"] + for shard in self.manifest["shards"] + for row in shard["trajectories"] + ) + + def add_trajectory( + self, + trajectory_id: str, + states: dict[str, np.ndarray], + transitions: dict[str, np.ndarray], + *, + split: str | None = None, + metadata: dict[str, Any] | None = None, + ) -> None: + """Append one finite trajectory, with T+1 states and T transitions.""" + if self._closed: + raise RuntimeError("Writer is closed") + if not trajectory_id or trajectory_id in self._ids: + raise ValueError(f"Empty or duplicate trajectory ID: {trajectory_id}") + if not {"q", "qd"}.issubset(states) or not transitions: + raise ValueError("states requires q/qd and transitions must be nonempty") + arrays = { + "states": {k: np.asarray(v) for k, v in states.items()}, + "transitions": {k: np.asarray(v) for k, v in transitions.items()}, + } + length = len(next(iter(arrays["transitions"].values()))) + if length < 1: + raise ValueError("Trajectory must contain at least one transition") + schema: dict[str, Any] = {} + for group, fields in arrays.items(): + schema[group] = {} + for name, value in fields.items(): + if ( + not name + or "/" in name + or value.ndim < 2 + or value.shape[0] != length + (group == "states") + ): + raise ValueError(f"Invalid trajectory shape/name: {group}/{name}") + if ( + value.dtype.kind not in "fiu b".replace(" ", "") + or not np.isfinite(value).all() + ): + raise ValueError(f"Non-finite/non-numeric trajectory field: {name}") + schema[group][name] = { + "shape": list(value.shape[1:]), + "dtype": value.dtype.str, + } + if self.manifest["fields"] is not None and schema != self.manifest["fields"]: + raise ValueError("Trajectory fields differ from committed schema") + split = split or trajectory_split(trajectory_id, self.manifest["split_seed"]) + if split not in {"train", "validation", "test"}: + raise ValueError(f"Unknown split: {split}") + trajectory_metadata = json.loads(_json_bytes(metadata or {})) + if self._file is None: + self._name = f"shard-{len(self.manifest['shards']):06d}.h5" + self._temporary = self.root / (self._name + ".pending") + self._file = h5py.File(self._temporary, "w") + group_name = str(len(self._pending)) + target = self._file.create_group(group_name) + for group, fields in arrays.items(): + sub = target.create_group(group) + for name, value in fields.items(): + # Zero-width contact features cannot be compressed/chunked. + opts = {"chunks": True, "compression": "lzf"} if value.size else {} + sub.create_dataset(name, data=value, **opts) + self.manifest["fields"] = schema + self._pending.append( + { + "id": trajectory_id, + "group": group_name, + "length": length, + "split": split, + "metadata": trajectory_metadata, + } + ) + self._ids.add(trajectory_id) + if len(self._pending) >= self.capacity: + self.commit() + + def commit(self) -> None: + """Publish pending data and optional generation progress atomically. + + A caller may maintain ``manifest["generation_progress"]`` alongside + pending trajectories. It is published with each automatic shard commit; + a rejection-only update also publishes when no shard is pending. + """ + if self._file is None: + if "generation_progress" in self.manifest: + _atomic_json(self.root / "manifest.json", self.manifest) + return + self._file.flush() + self._file.close() + self._file = None + with self._temporary.open("rb") as stream: + os.fsync(stream.fileno()) + destination = self.root / self._name + os.replace(self._temporary, destination) + shard = { + "file": self._name, + "sha256": _sha256(destination), + "trajectories": self._pending, + } + self.manifest["shards"].append(shard) + _atomic_json(self.root / "manifest.json", self.manifest) + self._pending = [] + + def close(self) -> None: + """Commit the final partial shard and close the writer.""" + if not self._closed: + self.commit() + self._closed = True + + def abort(self) -> None: + """Discard uncommitted work, retaining every previously committed shard.""" + if self._file is not None: + self._file.close() + self._file = None + self._closed = True + + def __enter__(self) -> DynamicsShardWriter: + return self + + def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> None: + self.close() if exc_type is None else self.abort() + + +class DynamicsWindowDataset(Dataset): + """Read bounded windows with O(trajectories) indexing and LRU file handles. + + Default full windows match the official sequence dataset. ``include_prefix`` + also samples startup windows, padded on the right with a validity mask. + No window ever spans a trajectory, reset or split boundary. + """ + + def __init__( + self, + root: str | Path, + split: str, + history_length: int, + *, + include_prefix: bool = False, + max_open_shards: int = 4, + read_page_steps: int = 256, + verify: bool = False, + ) -> None: + if ( + history_length < 1 + or max_open_shards < 1 + or read_page_steps < 1 + or split not in {"train", "validation", "test"} + ): + raise ValueError("Invalid split, history length or handle limit") + self.root = Path(root) + self.manifest = load_manifest(self.root, verify=verify) + self.fingerprint = dataset_fingerprint(self.manifest) + self.split = split + self.history_length = history_length + self.include_prefix = include_prefix + self.max_open_shards = max_open_shards + self.read_page_steps = read_page_steps + self._rows: list[dict[str, Any]] = [] + self._ends: list[int] = [] + total = 0 + for shard in self.manifest["shards"]: + for row in shard["trajectories"]: + count = ( + row["length"] + if include_prefix + else max(0, row["length"] - history_length + 1) + ) + if row["split"] == split and count: + self._rows.append({**row, "file": shard["file"]}) + total += count + self._ends.append(total) + self._pid = os.getpid() + self._handles: OrderedDict[str, h5py.File] = OrderedDict() + self._resident = None + + def enable_device_cache(self, device: str | torch.device, max_bytes: int) -> int: + """Cache raw transitions within an explicit single-process byte budget. + + Args: + device: CPU or CUDA device used by this learner. + max_bytes: Maximum resident raw-field allocation. Insufficient + budgets fail before allocating; no data are dropped or quantized. + + Returns: + Actual resident field bytes. The original shard fingerprint and + statistical/window weighting remain unchanged. + """ + from ._resident import _ResidentWindows + + if self._resident is not None: + raise ValueError("A resident dynamics cache is already enabled") + cache = _ResidentWindows(self, device, max_bytes) + self._resident = cache + return cache.nbytes + + def __len__(self) -> int: + return self._ends[-1] if self._ends else 0 + + def _handle(self, name: str) -> h5py.File: + if os.getpid() != self._pid: + self.close() + self._pid = os.getpid() + if name not in self._handles: + if len(self._handles) >= self.max_open_shards: + _, handle = self._handles.popitem(last=False) + handle.close() + self._handles[name] = h5py.File(self.root / name, "r") + self._handles.move_to_end(name) + return self._handles[name] + + def __getitem__(self, index: int) -> dict[str, Any]: + if index < 0: + index += len(self) + if not 0 <= index < len(self): + raise IndexError(index) + row_index = bisect_right(self._ends, index) + row = self._rows[row_index] + offset = index - (self._ends[row_index - 1] if row_index else 0) + end = offset + 1 if self.include_prefix else offset + self.history_length + start = max(0, end - self.history_length) + length = end - start + source = self._handle(row["file"])[row["group"]] + result: dict[str, Any] = { + "trajectory_id": row["id"], + "start": start, + "length": length, + "valid_mask": torch.arange(self.history_length) < length, + } + for group in ("states", "transitions"): + result[group] = {} + extra = int(group == "states") + for name in self.manifest["fields"][group]: + value = source[group][name][start : end + extra] + padding = self.history_length - length + if padding: + value = np.concatenate( + (value, np.repeat(value[-1:], padding, axis=0)), axis=0 + ) + result[group][name] = torch.from_numpy(value.copy()) + return result + + @staticmethod + def _read_block( + group: h5py.Group, + name: str, + start: int, + stop: int, + shape: tuple[int, ...], + dtype: np.dtype, + ) -> np.ndarray: + """Read a contiguous, exactly sized hyperslab without Dataset wrappers.""" + block = np.empty((stop - start, *shape), dtype=dtype) + dataset = h5py.h5d.open(group.id, name.encode("utf-8")) + source_space = memory_space = None + try: + source_space = dataset.get_space() + actual = source_space.get_simple_extent_dims() + if tuple(actual[1:]) != shape or not 0 <= start < stop <= actual[0]: + raise ValueError(f"Committed field shape/range mismatch: {name}") + if not block.size: + return block + source_space.select_hyperslab((start, *([0] * len(shape))), block.shape) + # The memory space is constructed from this owned C-contiguous + # allocation, never from untrusted file dimensions. H5Dread therefore + # cannot write beyond the allocated destination buffer. + memory_space = h5py.h5s.create_simple(block.shape) + dataset.read(memory_space, source_space, block) + finally: + if memory_space is not None: + memory_space.close() + if source_space is not None: + source_space.close() + dataset.close() + return block + + def get_batch( + self, indices: Sequence[int] | np.ndarray | torch.Tensor + ) -> dict[str, Any]: + """Read an owned batch in caller order using bounded grouped HDF5 reads. + + Requests are grouped by shard, trajectory and a fixed time page. The + largest source read is ``read_page_steps + history_length`` frames; + no trajectory/data cache grows with dataset size. Storage consists of the + output batch, O(batch size) indices, one bounded source block and a + per-field batched gather at a time. Repeated requests occupy independent rows in the returned tensors. + Sampling and prefetch state are not touched by this synchronous method. + + Args: + indices: Nonempty CPU integer window indices. Negative indices have + the same meaning as ``__getitem__``. + + Returns: + The same nested tensors/metadata as PyTorch's default collation of + individual windows, including right-padded startup masks. + """ + if isinstance(indices, torch.Tensor) and indices.device.type != "cpu": + raise ValueError("Dynamics batch indices must be on CPU") + requests = [operator.index(index) for index in indices] + if not requests: + raise ValueError("A dynamics batch must contain at least one window") + if self._resident is not None: + return self._resident.get_batch(requests) + size = len(requests) + grouped: dict[tuple[int, int], list[tuple[int, int, int]]] = {} + identifiers: list[str] = [] + starts = np.empty(size, dtype=np.int64) + lengths = np.empty(size, dtype=np.int64) + for position, index in enumerate(requests): + if index < 0: + index += len(self) + if not 0 <= index < len(self): + raise IndexError(index) + row_index = bisect_right(self._ends, index) + row = self._rows[row_index] + offset = index - (self._ends[row_index - 1] if row_index else 0) + end = offset + 1 if self.include_prefix else offset + self.history_length + start = max(0, end - self.history_length) + starts[position], lengths[position] = start, end - start + identifiers.append(row["id"]) + grouped.setdefault((row_index, start // self.read_page_steps), []).append( + (position, start, end) + ) + output = { + kind: { + name: np.empty( + ( + size, + self.history_length + int(kind == "states"), + *field["shape"], + ), + dtype=np.dtype(field["dtype"]), + ) + for name, field in fields.items() + } + for kind, fields in self.manifest["fields"].items() + } + # Sort only the I/O schedule, never the sample sequence presented to the + # learner. Processing all requests for one shard prevents LRU thrashing. + keys = sorted(grouped, key=lambda key: (self._rows[key[0]]["file"], key)) + last_row = None + source_groups = None + for key in keys: + row_index, _ = key + row = self._rows[row_index] + if row_index != last_row: + source = self._handle(row["file"])[row["group"]] + source_groups = {kind: source[kind] for kind in output} + last_row = row_index + entries = np.asarray(grouped[key], dtype=np.int64) + positions, first, end = entries[:, 0], entries[:, 1], entries[:, 2] + low, high = int(first.min()), int(end.max()) + relative = first - low + for kind, fields in output.items(): + extra = int(kind == "states") + offsets = np.arange(self.history_length + extra) + take = relative[:, None] + np.minimum( + offsets[None, :], (end - first - 1 + extra)[:, None] + ) + for name, array in fields.items(): + block = self._read_block( + source_groups[kind], + name, + low, + high + extra, + tuple(array.shape[2:]), + array.dtype, + ) + array[positions] = block[take] + del block + return { + "trajectory_id": identifiers, + "start": torch.from_numpy(starts), + "length": torch.from_numpy(lengths), + "valid_mask": torch.arange(self.history_length)[None, :] + < torch.from_numpy(lengths)[:, None], + **{ + kind: {name: torch.from_numpy(value) for name, value in fields.items()} + for kind, fields in output.items() + }, + } + + def close(self) -> None: + """Release resident fields and file handles; returned batches stay owned.""" + self._resident = None + for handle in self._handles.values(): + handle.close() + self._handles.clear() + + def __getstate__(self) -> dict[str, Any]: + state = self.__dict__.copy() + state["_handles"] = OrderedDict() + state["_resident"] = None + return state + + def __del__(self) -> None: + # h5py internals may already be finalized during interpreter shutdown. + # Explicit close() remains strict so operational failures are visible. + try: + if hasattr(self, "_handles"): + self.close() + except Exception: + pass diff --git a/extensions/embodichain_models/embodichain_models/dynamics/nerd/integration.py b/extensions/embodichain_models/embodichain_models/dynamics/nerd/integration.py new file mode 100644 index 0000000000..172148fce4 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/nerd/integration.py @@ -0,0 +1,299 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""NeRD encoding, target losses, frozen statistics and offline evaluation.""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any + +import torch +from torch import Tensor + +from .data import DynamicsWindowDataset + +__all__ = [ + "encode_batch", + "normalized_mse", + "compute_statistics", + "save_statistics", + "load_statistics", + "evaluate_model", +] + +_PREPROCESSING_EXECUTION = "nerd-shared-frame-scripted-cuda-warmup-v2" + + +def encode_batch( + batch: dict[str, Any], codec: Any, device: str | torch.device +) -> tuple[dict[str, Tensor], Tensor, Tensor]: + """Encode a collated raw trajectory window through the shared codec.""" + states = { + key: value.to(device=device, dtype=torch.float32) + for key, value in batch["states"].items() + } + transitions = { + key: value.to(device=device, dtype=torch.float32) + for key, value in batch["transitions"].items() + } + effort = transitions.get("joint_f", transitions.get("effort")) + if effort is None: + raise ValueError("Dataset requires transitions/joint_f") + contacts = { + key: value for key, value in transitions.items() if key.startswith("contact_") + } + q, qd = states["q"], states["qd"] + legacy_velocity = states.get("legacy_velocity") + current_velocity = None if legacy_velocity is None else legacy_velocity[:, :-1] + inputs = codec.encode( + q[:, :-1], + qd[:, :-1], + effort, + contacts or None, + transitions.get("gravity_dir"), + legacy_velocity=current_velocity, + ) + target = codec.targets( + q[:, :-1], + qd[:, :-1], + q[:, 1:], + qd[:, 1:], + current_embedding=inputs["states_embedding"], + legacy_velocity=current_velocity, + next_legacy_velocity=( + None if legacy_velocity is None else legacy_velocity[:, 1:] + ), + ) + return inputs, target, batch["valid_mask"].to(device) + + +def normalized_mse( + prediction: Tensor, + target: Tensor, + variance: Tensor | None, + mask: Tensor | None = None, +) -> Tensor: + """Official MSE over normalized physical targets, excluding padded steps.""" + dtype = torch.promote_types(prediction.dtype, target.dtype) + prediction = prediction.to(dtype=dtype) + target = target.to(dtype=dtype) + if variance is None: + scaled_prediction, scaled_target = prediction, target + else: + weights = 1.0 / torch.sqrt(variance + 1e-5) + # Preserve official arithmetic ordering for numerical parity. + scaled_prediction, scaled_target = prediction * weights, target * weights + if mask is None: + return torch.nn.functional.mse_loss(scaled_prediction, scaled_target) + squared = torch.nn.functional.mse_loss( + scaled_prediction, scaled_target, reduction="none" + ) + valid = mask.to(squared.dtype).unsqueeze(-1) + return (squared * valid).sum() / (valid.sum() * squared.shape[-1]).clamp_min(1) + + +class _Moments: + def __init__(self) -> None: + self.count = 0 + self.mean: Tensor | None = None + self.m2: Tensor | None = None + + def add(self, values: Tensor) -> None: + values = values.detach().to(dtype=torch.float64, device="cpu") + count = values.shape[0] + if not count: + return + mean = values.mean(0) + m2 = (values - mean).square().sum(0) + if self.mean is None: + self.count, self.mean, self.m2 = count, mean, m2 + else: + delta = mean - self.mean + total = self.count + count + self.m2 += m2 + delta.square() * self.count * count / total + self.mean += delta * count / total + self.count = total + + def result(self) -> tuple[Tensor, Tensor]: + if self.mean is None: + raise ValueError("Cannot normalize an empty dataset") + return self.mean.float(), (self.m2 / self.count).float() + + +def compute_statistics( + dataset: DynamicsWindowDataset, codec: Any, batch_size: int = 64 +) -> dict[str, Any]: + """Compute bounded training-only, overlapping-window-weighted moments. + + Statistical weighting is explicit. Official comparison runs can instead load + the exact shared reference artifact; no claim of matching an upstream running + estimator's epsilon initialization is implied by these population moments. + """ + if dataset.split != "train": + raise ValueError("Normalization statistics may only use the training split") + if batch_size < 1: + raise ValueError("batch_size must be positive") + moments: dict[str, _Moments] = {} + targets = _Moments() + for start in range(0, len(dataset), batch_size): + batch = dataset.get_batch(range(start, min(len(dataset), start + batch_size))) + inputs, target, mask = encode_batch(batch, codec, "cpu") + for key, value in inputs.items(): + moments.setdefault(key, _Moments()).add(value[mask]) + targets.add(target[mask]) + mean, variance = targets.result() + return { + "schema_version": 1, + "weighting": "valid_window_occurrences", + "dataset_fingerprint": dataset.fingerprint, + "inputs": { + key: {"mean": item.result()[0], "var": item.result()[1]} + for key, item in moments.items() + }, + "output": {"mean": mean, "var": variance}, + } + + +def save_statistics(path: str | Path, statistics: dict[str, Any]) -> None: + """Write a tensor-only statistics artifact usable by both trainers.""" + _atomic_save(Path(path), statistics) + + +def load_statistics(path: str | Path) -> dict[str, Any]: + """Load a tensor-only shared reference statistics artifact.""" + value = torch.load(path, map_location="cpu", weights_only=True) + if value.get("schema_version") != 1: + raise ValueError("Unsupported statistics schema") + for item in [*value["inputs"].values(), value["output"]]: + if ( + not torch.isfinite(item["mean"]).all() + or not torch.isfinite(item["var"]).all() + or (item["var"] < 0).any() + ): + raise ValueError("Invalid normalizer moments") + return value + + +def _atomic_save(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".pending") + with temporary.open("wb") as stream: + torch.save(value, stream) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + + +def _model_variance(model: torch.nn.Module) -> Tensor | None: + return model.output_rms.var if model.cfg.normalize_output else None + + +@torch.no_grad() +def evaluate_model( + model: torch.nn.Module, + codec: Any, + dataset: DynamicsWindowDataset, + batch_size: int = 64, + device: str = "cpu", + max_batches: int | None = None, +) -> dict[str, float | int | str]: + """Evaluate all valid held-out positions without model selection on test.""" + if not len(dataset): + raise ValueError(f"No eligible windows in {dataset.split} split") + if batch_size < 1 or (max_batches is not None and max_batches < 1): + raise ValueError("Evaluation batch limits must be positive") + window_count = ( + len(dataset) + if max_batches is None + else min(len(dataset), max_batches * batch_size) + ) + was_training = model.training + model.eval() + accumulated: dict[str, float] = { + "normalized_mse": 0.0, + "target_mse": 0.0, + "q_error_norm": 0.0, + "qd_error_norm": 0.0, + } + count = 0 + try: + for start in range(0, window_count, batch_size): + batch = dataset.get_batch( + range(start, min(window_count, start + batch_size)) + ) + inputs, target, mask = encode_batch(batch, codec, device) + prediction = model(inputs) + if not torch.isfinite(prediction).all(): + raise FloatingPointError("Non-finite evaluation prediction") + n = int(mask.sum()) + accumulated["normalized_mse"] += ( + float( + normalized_mse( + prediction, + target, + _model_variance(model), + mask if dataset.include_prefix else None, + ) + ) + * n + ) + accumulated["target_mse"] += ( + float((prediction - target).square()[mask].mean()) * n + ) + q = batch["states"]["q"].to(device=device, dtype=prediction.dtype) + qd = batch["states"]["qd"].to(device=device, dtype=prediction.dtype) + next_q, next_qd = codec.decode(q[:, :-1], qd[:, :-1], prediction) + q_error = next_q - q[:, 1:] + periodic = codec.continuous_q_indices + if periodic: + q_error[..., list(periodic)] = torch.atan2( + torch.sin(q_error[..., list(periodic)]), + torch.cos(q_error[..., list(periodic)]), + ) + if codec.robot == "ant": + position = q_error[..., :3].norm(dim=-1)[mask].sum() + dot = ( + ( + torch.nn.functional.normalize(next_q[..., 3:7], dim=-1) + * torch.nn.functional.normalize(q[:, 1:, 3:7], dim=-1) + ) + .sum(-1) + .abs() + .clamp(max=1) + ) + angle = (2 * torch.acos(dot))[mask].sum() + accumulated["base_position_error"] = accumulated.get( + "base_position_error", 0.0 + ) + float(position) + accumulated["base_orientation_error_rad"] = accumulated.get( + "base_orientation_error_rad", 0.0 + ) + float(angle) + q_error = q_error[..., 7:] + accumulated["q_error_norm"] += float(q_error.norm(dim=-1)[mask].sum()) + accumulated["qd_error_norm"] += float( + (next_qd - qd[:, 1:]).norm(dim=-1)[mask].sum() + ) + count += n + finally: + model.train(was_training) + return { + **{key: value / count for key, value in accumulated.items()}, + "valid_positions": count, + "windows": window_count, + "split": dataset.split, + } diff --git a/extensions/embodichain_models/embodichain_models/dynamics/nerd/model.py b/extensions/embodichain_models/embodichain_models/dynamics/nerd/model.py new file mode 100644 index 0000000000..4a36843da5 --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/nerd/model.py @@ -0,0 +1,324 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""NeRD causal Transformer compatible with the official network parameter names. + +Architecture adapted from NVlabs/neural-robot-dynamics (Apache-2.0), whose +Transformer derives from Andrej Karpathy's nanoGPT (MIT). See THIRD_PARTY_LICENSES. +""" + +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations +import math +from collections.abc import Mapping +from contextlib import AbstractContextManager, nullcontext +import torch +from torch import Tensor, nn +from torch.nn import functional as F +from embodichain.utils import configclass + +__all__ = ["NeRDModelCfg", "NeRDModel", "FeatureNormalizer"] + + +@configclass +class NeRDModelCfg: + """Network architecture and ordered input schema for NeRD. + + The default architecture matches the official Cartpole checkpoint; input_dims + and output_dim must be supplied. Ant uses n_embd=384 and its contact schema. + ``attention_backend="math"`` persists an explicit SDPA choice in portable + exports; ``default`` preserves Torch's backend selection. + """ + + input_dims: dict[str, int] = {} + output_dim: int = 4 + n_embd: int = 192 + n_layer: int = 6 + n_head: int = 12 + block_size: int = 32 + history_length: int = 10 + head_hidden_dims: tuple[int, ...] = (64,) + bias: bool = False + dropout: float = 0.0 + normalize_input: bool = True + normalize_output: bool = True + output_tanh: bool = False + attention_backend: str = "default" + + +def _attention_backend_context(backend: str) -> AbstractContextManager: + if backend == "default": + return nullcontext() + if backend != "math": + raise ValueError("attention_backend must be 'default' or 'math'") + if hasattr(torch.nn, "attention") and hasattr(torch.nn.attention, "sdpa_kernel"): + return torch.nn.attention.sdpa_kernel([torch.nn.attention.SDPBackend.MATH]) + return torch.backends.cuda.sdp_kernel( + enable_flash=False, enable_math=True, enable_mem_efficient=False + ) + + +class FeatureNormalizer(nn.Module): + """Frozen mean/variance buffers with the exact official sqrt(var + 1e-5).""" + + def __init__(self, width: int) -> None: + super().__init__() + self.register_buffer("mean", torch.zeros(width)) + self.register_buffer("var", torch.ones(width)) + + def forward(self, value: Tensor, inverse: bool = False) -> Tensor: + """Normalize or denormalize without modifying the caller's tensor.""" + scale = torch.sqrt(self.var + 1e-5) + return value * scale + self.mean if inverse else (value - self.mean) / scale + + +class _LayerNorm(nn.Module): + def __init__(self, width: int, bias: bool) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(width)) + self.bias = nn.Parameter(torch.zeros(width)) if bias else None + + def forward(self, value: Tensor) -> Tensor: + return F.layer_norm(value, self.weight.shape, self.weight, self.bias, 1e-5) + + +class _Attention(nn.Module): + def __init__(self, cfg: NeRDModelCfg) -> None: + super().__init__() + self.c_attn = nn.Linear(cfg.n_embd, 3 * cfg.n_embd, bias=cfg.bias) + self.c_proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=cfg.bias) + self.attn_dropout = nn.Dropout(cfg.dropout) + self.resid_dropout = nn.Dropout(cfg.dropout) + self.n_head, self.n_embd, self.dropout = cfg.n_head, cfg.n_embd, cfg.dropout + + def forward(self, value: Tensor) -> Tensor: + batch, length, width = value.shape + q, k, v = self.c_attn(value).split(self.n_embd, -1) + q, k, v = [ + x.view(batch, length, self.n_head, width // self.n_head).transpose(1, 2) + for x in (q, k, v) + ] + out = F.scaled_dot_product_attention( + q, k, v, dropout_p=self.dropout if self.training else 0.0, is_causal=True + ) + out = out.transpose(1, 2).contiguous().view(batch, length, width) + return self.resid_dropout(self.c_proj(out)) + + +class _MLP(nn.Module): + def __init__(self, cfg: NeRDModelCfg) -> None: + super().__init__() + self.c_fc = nn.Linear(cfg.n_embd, 4 * cfg.n_embd, bias=cfg.bias) + self.gelu = nn.GELU() + self.c_proj = nn.Linear(4 * cfg.n_embd, cfg.n_embd, bias=cfg.bias) + self.dropout = nn.Dropout(cfg.dropout) + + def forward(self, value: Tensor) -> Tensor: + return self.dropout(self.c_proj(self.gelu(self.c_fc(value)))) + + +class _Block(nn.Module): + def __init__(self, cfg: NeRDModelCfg) -> None: + super().__init__() + self.ln_1 = _LayerNorm(cfg.n_embd, cfg.bias) + self.attn = _Attention(cfg) + self.ln_2 = _LayerNorm(cfg.n_embd, cfg.bias) + self.mlp = _MLP(cfg) + + def forward(self, value: Tensor) -> Tensor: + value = value + self.attn(self.ln_1(value)) + return value + self.mlp(self.ln_2(value)) + + +class _GPT(nn.Module): + def __init__(self, cfg: NeRDModelCfg) -> None: + super().__init__() + self.transformer = nn.ModuleDict( + dict( + wte=nn.Linear(sum(cfg.input_dims.values()), cfg.n_embd), + wpe=nn.Embedding(cfg.block_size, cfg.n_embd), + drop=nn.Dropout(cfg.dropout), + h=nn.ModuleList([_Block(cfg) for _ in range(cfg.n_layer)]), + ln_f=_LayerNorm(cfg.n_embd, cfg.bias), + ) + ) + self.lm_head = nn.Linear(cfg.n_embd, cfg.n_embd, bias=False) + self.apply(self._init_weights) + for name, parameter in self.named_parameters(): + if name.endswith("c_proj.weight"): + nn.init.normal_( + parameter, mean=0.0, std=0.02 / math.sqrt(2 * cfg.n_layer) + ) + + @staticmethod + def _init_weights(module: nn.Module) -> None: + if isinstance(module, nn.Linear): + nn.init.normal_(module.weight, mean=0.0, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=0.02) + + def forward(self, value: Tensor) -> Tensor: + pos = torch.arange(value.shape[1], device=value.device) + value = self.transformer.drop( + self.transformer.wte(value) + self.transformer.wpe(pos) + ) + for block in self.transformer.h: + value = block(value) + return self.lm_head(self.transformer.ln_f(value)) + + +class _Head(nn.Module): + def __init__(self, cfg: NeRDModelCfg) -> None: + super().__init__() + modules = [] + width = cfg.n_embd + for next_width in cfg.head_hidden_dims: + modules.extend([nn.Linear(width, next_width), nn.ReLU()]) + width = next_width + self.feature_net = nn.Module() + self.feature_net.body = nn.Sequential(*modules) + self.output_net = nn.Linear(width, cfg.output_dim) + + def forward(self, value: Tensor) -> Tensor: + return self.output_net(self.feature_net.body(value)) + + +class NeRDModel(nn.Module): + """Causal NeRD network with immutable named inputs and portable statistics. + + Args: + cfg: Explicit input layout and Transformer configuration. + + All inputs have shape (batch, time, width). Right-padding is permitted: + causal attention preserves prefix outputs, and the caller masks padded loss. + Features follow the network's floating dtype on their existing device; + dtype conversion preserves input gradients independently of normalization. + ``evaluate`` gathers the final valid prefix token independently per row. + Neither training nor evaluation implicitly detaches gradients. + """ + + def __init__(self, cfg: NeRDModelCfg) -> None: + super().__init__() + if not isinstance(cfg.attention_backend, str) or cfg.attention_backend not in { + "default", + "math", + }: + raise ValueError("attention_backend must be 'default' or 'math'") + if not cfg.input_dims or any(v <= 0 for v in cfg.input_dims.values()): + raise ValueError("input_dims must specify positive ordered feature widths") + if any("." in k or not k for k in cfg.input_dims): + raise ValueError("Feature names must be nonempty and cannot contain dots") + if ( + cfg.n_head <= 0 + or cfg.n_embd <= 0 + or cfg.n_embd % cfg.n_head + or cfg.n_layer <= 0 + ): + raise ValueError( + "Positive embedding must be divisible by positive attention heads" + ) + if ( + not 1 <= cfg.history_length <= cfg.block_size + or cfg.output_dim <= 0 + or not 0 <= cfg.dropout < 1 + ): + raise ValueError("Invalid history, output dimension or dropout") + self.cfg = cfg.copy() + self.encoders = nn.ModuleDict({"low_dim": nn.Module()}) + self.encoders["low_dim"].body = nn.Sequential() + self.transformer_model = _GPT(cfg) + self.model = _Head(cfg) + self.input_rms = nn.ModuleDict( + {name: FeatureNormalizer(width) for name, width in cfg.input_dims.items()} + ) + self.output_rms = FeatureNormalizer(cfg.output_dim) + + def set_normalization( + self, + input_stats: Mapping[str, tuple[Tensor, Tensor]], + output_mean: Tensor, + output_var: Tensor, + ) -> None: + """Set frozen statistics; missing features and invalid variances fail early.""" + if set(input_stats) != set(self.cfg.input_dims): + raise ValueError("Normalization keys must exactly match the feature schema") + entries = [(self.input_rms[k], *input_stats[k]) for k in input_stats] + entries.append((self.output_rms, output_mean, output_var)) + for rms, mean, var in entries: + if ( + mean.shape != rms.mean.shape + or var.shape != rms.var.shape + or not torch.isfinite(mean).all() + or not torch.isfinite(var).all() + or (var < 0).any() + ): + raise ValueError("Invalid normalizer mean/variance") + with torch.no_grad(): + for rms, mean, var in entries: + rms.mean.copy_(mean) + rms.var.copy_(var) + + def forward(self, inputs: Mapping[str, Tensor]) -> Tensor: + """Predict every sequence token in physical target units.""" + values = [] + shape = None + dtype = self.transformer_model.transformer.wte.weight.dtype + for name, width in self.cfg.input_dims.items(): + value = inputs[name] + if value.ndim != 3 or value.shape[-1] != width: + raise ValueError(f"{name} must have shape (batch, time, {width})") + if shape is not None and value.shape[:2] != shape: + raise ValueError("All model features must share batch/time axes") + shape = value.shape[:2] + value = value.to(dtype=dtype) + values.append( + self.input_rms[name](value) if self.cfg.normalize_input else value + ) + if not 1 <= shape[1] <= self.cfg.block_size: + raise ValueError("Sequence length is outside the model block_size") + with _attention_backend_context(self.cfg.attention_backend): + value = self.model(self.transformer_model(torch.cat(values, -1))) + if self.cfg.output_tanh: + value = torch.tanh(value) + return ( + self.output_rms(value, inverse=True) if self.cfg.normalize_output else value + ) + + def evaluate( + self, inputs: Mapping[str, Tensor], valid_lengths: Tensor | None = None + ) -> Tensor: + """Predict the last valid token, shape (batch, 1, output_dim).""" + value = self(inputs) + if valid_lengths is None: + return value[:, -1:] + if ( + valid_lengths.shape != (value.shape[0],) + or valid_lengths.dtype != torch.long + ): + raise ValueError( + "valid_lengths must be long values with one length per row" + ) + valid = ((valid_lengths >= 1) & (valid_lengths <= value.shape[1])).all() + if valid_lengths.device.type == "cpu": + if not valid: + raise ValueError("valid_lengths must be in [1, sequence_length]") + else: + # Keep per-step history validation on device; do not synchronize CUDA. + torch._assert_async(valid, "valid_lengths must be in [1, sequence_length]") + index = (valid_lengths - 1)[:, None, None].expand(-1, 1, value.shape[-1]) + return value.gather(1, index) diff --git a/extensions/embodichain_models/embodichain_models/dynamics/nerd/sl_recipe.py b/extensions/embodichain_models/embodichain_models/dynamics/nerd/sl_recipe.py new file mode 100644 index 0000000000..8616ff480f --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/nerd/sl_recipe.py @@ -0,0 +1,223 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""NeRD window sampling and standalone recipes for ``train-sl``.""" + +from __future__ import annotations + +import random +import hashlib +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +from .codec import NeRDStateCodec +from .data import DynamicsWindowDataset +from .model import NeRDModel, NeRDModelCfg +from embodichain.learning.sl import SupervisedBatch, SupervisedLearner +from embodichain.learning.sl.trainer import SupervisedTrainer, SupervisedTrainerCfg + +from ._numerics import _attention_context, _clip_grad_norm_, _validate_numerical_profile +from .supervised import NeRDObjective, make_nerd_batch +from .integration import ( + _PREPROCESSING_EXECUTION, + compute_statistics, + encode_batch, + evaluate_model, + load_statistics, +) + +__all__ = ["NeRDWindowSource", "build_supervised_trainer"] + + +class NeRDWindowSource: + """Adapt an immutable NeRD window reader without owning sampling RNG. + + Args: + dataset: Caller-owned window dataset, including optional resident cache. + codec: Physical codec matching that dataset. + device: Encoding device, preserving the reference preprocessing policy. + """ + + def __init__( + self, + dataset: DynamicsWindowDataset, + codec: NeRDStateCodec, + device: str | torch.device = "cpu", + ) -> None: + self.dataset, self.codec, self.device = dataset, codec, device + self.source_id = ( + f"{dataset.fingerprint}:{dataset.split}:" + f"history={dataset.history_length}:prefix={dataset.include_prefix}" + ) + + def sample(self, batch_size: int, *, generator: torch.Generator) -> SupervisedBatch: + """Sample copied windows using an explicit CPU generator. + + Args: + batch_size: Positive number of windows sampled with replacement. + generator: Caller-owned sampling generator, committed after updates. + + Returns: + Encoded inputs, physical targets and valid-position mask. + """ + indices = torch.randint(len(self.dataset), (batch_size,), generator=generator) + return make_nerd_batch(self.dataset.get_batch(indices), self.codec, self.device) + + +def build_supervised_trainer(config: dict[str, Any]) -> SupervisedTrainer: + """Build a simulation-independent NeRD recipe for the shared trainer. + + Args: + config: ``train-sl`` mapping. ``model`` is a NeRD model configuration; + ``data.path`` identifies committed shards and optional + ``data.statistics`` supplies shared frozen training statistics. + ``objective`` accepts ``include_prefix`` and ``numerical_profile``. + ``optimizer`` accepts Adam rates, clipping and interval/offset LR + scheduling. ``trainer`` uses ``SupervisedTrainerCfg``. + + Returns: + Common trainer with a NeRD objective and owned window reader references. + Its checkpoints use the common SL format. The NeRD showcase historical + trainer continues reading and writing NeRD training checkpoints. + """ + cfg = SupervisedTrainerCfg(**config.get("trainer", {})) + cfg.validate() + if cfg.total_updates < 1: + raise ValueError("NeRD recipes require a positive update budget") + torch.manual_seed(cfg.seed) + model = NeRDModel(NeRDModelCfg(**config["model"])) + data = config["data"] + objective_cfg = config.get("objective", {}) + include_prefix = objective_cfg.get("include_prefix", False) + profile = _validate_numerical_profile( + objective_cfg.get("numerical_profile", "default") + ) + if profile == "reference": + model.cfg.attention_backend = "math" + train = DynamicsWindowDataset( + data["path"], + "train", + model.cfg.history_length, + include_prefix=include_prefix, + verify=True, + ) + validation = DynamicsWindowDataset( + data["path"], + "validation", + model.cfg.history_length, + include_prefix=include_prefix, + ) + codec = NeRDStateCodec(**train.manifest["metadata"]["codec"]) + statistics = ( + load_statistics(data["statistics"]) + if data.get("statistics") + else compute_statistics(train, codec, cfg.batch_size) + ) + if statistics["dataset_fingerprint"] != train.fingerprint: + raise ValueError("NeRD statistics/dataset fingerprint mismatch") + model.set_normalization( + { + key: (value["mean"], value["var"]) + for key, value in statistics["inputs"].items() + }, + statistics["output"]["mean"], + statistics["output"]["var"], + ) + model.to(cfg.device) + encode_batch(train.get_batch([0]), codec, cfg.device) + cache_bytes = data.get("resident_cache_bytes", 0) + if cache_bytes: + used = train.enable_device_cache(cfg.device, cache_bytes) + validation.enable_device_cache(cfg.device, cache_bytes - used) + if torch.device(cfg.device).type == "cuda": + raw = train.get_batch([0] * cfg.batch_size) + for _ in range(3): + encode_batch(raw, codec, cfg.device) + optimizer_cfg = config.get("optimizer", {}) + max_grad_norm = float(optimizer_cfg.get("max_grad_norm", 1.0)) + rate = float(optimizer_cfg.get("learning_rate", 1e-3)) + final_rate = float(optimizer_cfg.get("final_learning_rate", 1e-4)) + interval = int(optimizer_cfg.get("lr_update_interval", 1)) + offset = int(optimizer_cfg.get("lr_offset", 0)) + if interval < 1 or offset < 0 or rate <= 0 or final_rate < 0 or max_grad_norm <= 0: + raise ValueError("Invalid NeRD optimizer/schedule settings") + optimizer = torch.optim.Adam(model.parameters(), lr=rate) + + def schedule(step: int) -> float: + ratio = (step // interval + offset) / (cfg.total_updates / interval + offset) + return ((1 - ratio) * rate + ratio * final_rate) / rate + + scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, schedule) + learner = SupervisedLearner( + model, + NeRDObjective(include_prefix=include_prefix), + optimizer, + max_grad_norm=max_grad_norm, + clip_grad=lambda parameters, maximum: _clip_grad_norm_( + parameters, + maximum, + profile=profile, + error_if_nonfinite=profile == "default", + ), + numerical_context=lambda: _attention_context(profile), + ) + evaluation = config.get("evaluation", {}) + + def evaluate(current: SupervisedLearner) -> dict[str, float]: + with _attention_context(profile): + report = evaluate_model( + current.model, + codec, + validation, + cfg.batch_size, + cfg.device, + max_batches=evaluation.get("max_batches"), + ) + return { + key: float(value) + for key, value in report.items() + if isinstance(value, (int, float)) + } + + random.seed(cfg.seed) + np.random.seed(cfg.seed) + torch.manual_seed(cfg.seed) + return SupervisedTrainer( + learner, + NeRDWindowSource(train, codec, cfg.device), + cfg, + signature={ + "model": model.cfg.to_dict(), + "model_input_layout": list(model.cfg.input_dims.items()), + "objective": { + "include_prefix": include_prefix, + "numerical_profile": profile, + }, + "optimizer": optimizer_cfg, + "physics": train.manifest["metadata"], + "preprocessing_execution": _PREPROCESSING_EXECUTION, + "statistics_sha256": ( + hashlib.sha256(Path(data["statistics"]).read_bytes()).hexdigest() + if data.get("statistics") + else None + ), + }, + scheduler=scheduler, + evaluate=evaluate, + ) diff --git a/extensions/embodichain_models/embodichain_models/dynamics/nerd/supervised.py b/extensions/embodichain_models/embodichain_models/dynamics/nerd/supervised.py new file mode 100644 index 0000000000..80acc8879c --- /dev/null +++ b/extensions/embodichain_models/embodichain_models/dynamics/nerd/supervised.py @@ -0,0 +1,101 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""NeRD objectives and batch adaptation for the shared supervised learner.""" + +from __future__ import annotations + +from typing import Any + +import torch + +from embodichain.learning.sl import LossResult, SupervisedBatch + +__all__ = ["NeRDObjective", "make_nerd_batch"] + + +class NeRDObjective: + """Preserve the NeRD normalized physical-target loss. + + Args: + include_prefix: Include valid startup positions using the batch mask. + False retains the reference full-window MSE reduction order. + + Input and target normalization remain owned by the NeRD model and codec. + Computing this objective never updates their statistics. + """ + + def __init__(self, *, include_prefix: bool = False) -> None: + self.include_prefix = include_prefix + + def compute_loss( + self, model: torch.nn.Module, batch: SupervisedBatch + ) -> LossResult: + """Compute the normalized state-increment error without detaching it. + + Args: + model: NeRD model with its frozen normalization buffers. + batch: Encoded inputs, physical targets and valid-position mask. + + Returns: + Graph-connected loss for one shared supervised update. + """ + # Share the physical-target reduction with offline evaluation. + from .integration import _model_variance, normalized_mse + + prediction = model(batch.inputs) + loss = normalized_mse( + prediction, + batch.targets, + _model_variance(model), + batch.masks["valid"] if self.include_prefix else None, + ) + count = batch.metadata.get("valid_positions") + if count is None: + count = ( + int(batch.masks["valid"].sum()) + if self.include_prefix + else prediction.numel() // prediction.shape[-1] + ) + return LossResult(loss, weight=float(count)) + + +def make_nerd_batch( + batch: dict[str, Any], codec: Any, device: str | torch.device +) -> SupervisedBatch: + """Encode a raw NeRD sequence without changing its sampling contract. + + Args: + batch: Owned trajectory windows from ``DynamicsWindowDataset``. + codec: The matching physical state/contact/effort codec. + device: Device on which encoding and model execution take place. + + Returns: + Encoded history inputs, physical increment targets and validity mask. + """ + from .integration import encode_batch + + inputs, targets, valid = encode_batch(batch, codec, device) + return SupervisedBatch( + inputs, + targets, + masks={"valid": valid}, + metadata={ + "valid_positions": int(batch["valid_mask"].sum()), + "trajectory_id": tuple(batch["trajectory_id"]), + "start": batch["start"], + }, + ) diff --git a/setup.py b/setup.py index 367259147a..6ea584a60b 100644 --- a/setup.py +++ b/setup.py @@ -33,9 +33,11 @@ CORE_PACKAGE_PATTERNS = ["embodichain", "embodichain.*"] EXTENSIONS_DIR = THIS_DIR / "extensions" TASKS_PROJECT_DIR = EXTENSIONS_DIR / "embodichain_tasks" +MODELS_PROJECT_DIR = EXTENSIONS_DIR / "embodichain_models" TASKS_PACKAGE_PATTERNS = ["embodichain_tasks", "embodichain_tasks.*"] ROBOTS_PACKAGE_PATTERNS = ["embodichain_robots", "embodichain_robots.*"] BENCHMARKS_PACKAGE_PATTERNS = ["embodichain_benchmarks", "embodichain_benchmarks.*"] +MODELS_PACKAGE_PATTERNS = ["embodichain_models", "embodichain_models.*"] EXAMPLES_DIR = THIS_DIR / "examples" # Defer importing torch until it's actually needed (when building extensions). @@ -99,6 +101,9 @@ def get_packages() -> list[str]: benchmark_packages = find_namespace_packages( where=str(EXTENSIONS_DIR), include=BENCHMARKS_PACKAGE_PATTERNS ) + model_packages = find_namespace_packages( + where=str(MODELS_PROJECT_DIR), include=MODELS_PACKAGE_PATTERNS + ) return sorted( set( [ @@ -106,6 +111,7 @@ def get_packages() -> list[str]: *task_packages, *robot_packages, *benchmark_packages, + *model_packages, "embodichain_tasks.configs", ] ) @@ -119,6 +125,7 @@ def get_package_dir() -> dict[str, str]: "embodichain_tasks.configs": "extensions/embodichain_tasks/configs", "embodichain_robots": "extensions/embodichain_robots", "embodichain_benchmarks": "extensions/embodichain_benchmarks", + "embodichain_models": "extensions/embodichain_models/embodichain_models", } @@ -177,6 +184,7 @@ def main(): package_dir=get_package_dir(), package_data={ "embodichain": ["VERSION"], + "embodichain_models.dynamics.nerd": ["THIRD_PARTY_LICENSES"], "embodichain.data": ["default_materials.json"], "embodichain_benchmarks.motion_generation": ["suites/*.yaml"], "embodichain.gen_sim.simready_pipeline.configs": ["*.json"], diff --git a/tests/ci/test_workflows.py b/tests/ci/test_workflows.py index 4e438148ae..8777f81e6a 100644 --- a/tests/ci/test_workflows.py +++ b/tests/ci/test_workflows.py @@ -343,13 +343,16 @@ def test_test_job_restores_git_access_before_running_selected_tests( [ None, "missing_robots", + "missing_models", "missing_entry_point", "wrong_entry_point_group", + "missing_asset_entry_point", + "retired_robot_module", "foreign_package", "bytecode", ], ) -def test_release_guard_validates_bundled_robot_catalog( +def test_release_guard_validates_bundled_runtime_packages( tmp_path: Path, invalid_contents: str | None ) -> None: """Execute the release guard against wheels, including invalid artifacts.""" @@ -369,11 +372,18 @@ def test_release_guard_validates_bundled_robot_catalog( entry_points = ( "[embodichain.tasks]\nembodichain_tasks = embodichain_tasks\n\n" "[embodichain.robots]\nofficial = embodichain_robots._catalog:ROBOT_CONFIGS\n" + "\n[embodichain.assets]\n" + "official_robots = embodichain_robots._asset_catalog:ROBOT_ASSETS\n" ) if invalid_contents == "missing_entry_point": - entry_points = entry_points.split("[embodichain.robots]", 1)[0] + entry_points = entry_points.replace( + "[embodichain.robots]\nofficial = embodichain_robots._catalog:ROBOT_CONFIGS\n", + "", + ) elif invalid_contents == "wrong_entry_point_group": entry_points = entry_points.replace("[embodichain.robots]", "[other.robots]") + elif invalid_contents == "missing_asset_entry_point": + entry_points = entry_points.split("[embodichain.assets]", 1)[0] members = { "embodichain/VERSION": version, "embodichain/gen_sim/simready_pipeline/configs/gen_config.json": "{}", @@ -381,7 +391,12 @@ def test_release_guard_validates_bundled_robot_catalog( task_config: source.read_text(), "embodichain_robots/__init__.py": "", "embodichain_robots/_catalog.py": "ROBOT_CONFIGS = {}\n", + "embodichain_robots/_asset_catalog.py": "ROBOT_ASSETS = {}\n", + "embodichain/data/_asset_registry.py": "", + "embodichain_robots/franka/assets.py": "", + "embodichain_robots/dexforce/w1/assets.py": "", "embodichain_benchmarks/__init__.py": "", + "embodichain_models/__init__.py": "", f"embodichain-{version}.data/data/share/embodichain/examples/__init__.py": "", f"embodichain-{version}.dist-info/METADATA": f"Version: {version}\n", f"embodichain-{version}.dist-info/entry_points.txt": entry_points, @@ -392,6 +407,14 @@ def test_release_guard_validates_bundled_robot_catalog( for path, data in members.items() if not path.startswith("embodichain_robots/") } + elif invalid_contents == "missing_models": + members = { + path: data + for path, data in members.items() + if not path.startswith("embodichain_models/") + } + elif invalid_contents == "retired_robot_module": + members["embodichain_robots/franka_panda.py"] = "" elif invalid_contents == "foreign_package": members["foreign/__init__.py"] = "" elif invalid_contents == "bytecode": diff --git a/tests/docs/test_build_architecture.py b/tests/docs/test_build_architecture.py index e744057884..d539e46810 100644 --- a/tests/docs/test_build_architecture.py +++ b/tests/docs/test_build_architecture.py @@ -99,13 +99,32 @@ def test_dirty_check_checks_index_and_worktree_independently(seed): build_snapshot(root, path) -def test_extracts_only_unambiguous_absolute_module_imports_and_inheritance(seed): +@pytest.mark.parametrize( + "package_path,module_root", + [ + ("pkg", "pkg"), + ( + "extensions/embodichain_models/embodichain_models/dynamics/nerd", + "embodichain_models.dynamics.nerd", + ), + ], +) +def test_extracts_only_unambiguous_absolute_module_imports_and_inheritance( + seed, package_path, module_root +): data, root, path, schema = seed + source_a = f"{package_path}/a.py" + source_b = f"{package_path}/b.py" + write(root, source_b, (root / "pkg/b.py").read_text()) write( root, - "pkg/a.py", - "from pkg.b import Beta as B\nfrom typing import TYPE_CHECKING\nif TYPE_CHECKING:\n import pkg.b\nclass Alpha(B):\n def run(self):\n return 1\n def hidden(self):\n import pkg.b\nraise RuntimeError('must never execute')\n", + source_a, + f"from {module_root}.b import Beta as B\nfrom typing import TYPE_CHECKING\nif TYPE_CHECKING:\n import {module_root}.b\nclass Alpha(B):\n def run(self):\n return 1\n def hidden(self):\n import {module_root}.b\nraise RuntimeError('must never execute')\n", ) + locations = {"pkg/a.py": source_a, "pkg/b.py": source_b} + for item in [*data["nodes"], *data["edges"]]: + for evidence in item["evidence"]: + evidence["path"] = locations[evidence["path"]] data["nodes"][0]["evidence"][0]["excerpt"] = "class Alpha(B):" path.write_text(json.dumps(data)) commit(root) @@ -116,7 +135,7 @@ def test_extracts_only_unambiguous_absolute_module_imports_and_inheritance(seed) assert not any(p["start_line"] == 9 for e in extracted for p in e["evidence"]) assert any("TYPE_CHECKING" in e["scope"] for e in extracted) assert all(e["id"] in result["views"][0]["edge_ids"] for e in extracted) - assert "pkg.a" not in sys.modules + assert f"{module_root}.a" not in sys.modules validate_snapshot(result, repo_root=root, schema_path=schema) diff --git a/tests/docs/test_check_api_docs.py b/tests/docs/test_check_api_docs.py index 67fd341181..a8e462b3c5 100644 --- a/tests/docs/test_check_api_docs.py +++ b/tests/docs/test_check_api_docs.py @@ -263,15 +263,26 @@ def test_main_writes_missing_json_to_stdout( assert captured.err == "" -def test_default_inventory_includes_bundled_benchmark_public_exports(): +@pytest.mark.parametrize( + "package,module,exports", + [ + ( + "embodichain_benchmarks", + "embodichain_benchmarks.motion_generation", + {"run_benchmark", "PlannerAdapter"}, + ), + ( + "embodichain_models", + "embodichain_models.dynamics.nerd.model", + {"NeRDModel", "NeRDModelCfg"}, + ), + ], +) +def test_default_inventory_includes_bundled_extension_public_exports( + package, module, exports +): """Moving evaluation to an extension must not remove it from API coverage.""" - roots = [ - root - for root in _checker.DEFAULT_PACKAGE_ROOTS - if root.module == "embodichain_benchmarks" - ] + roots = [root for root in _checker.DEFAULT_PACKAGE_ROOTS if root.module == package] assert len(roots) == 1 modules = {module.name: module for module in discover_public_modules(roots)} - public = modules["embodichain_benchmarks.motion_generation"] - assert "run_benchmark" in public.exports - assert "PlannerAdapter" in public.exports + assert exports <= set(modules[module].exports) diff --git a/tests/learning/sl/__init__.py b/tests/learning/sl/__init__.py new file mode 100644 index 0000000000..2140aca6cd --- /dev/null +++ b/tests/learning/sl/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Tests for shared supervised learning contracts.""" + +from __future__ import annotations diff --git a/tests/learning/sl/test_learner.py b/tests/learning/sl/test_learner.py new file mode 100644 index 0000000000..b84b3c601b --- /dev/null +++ b/tests/learning/sl/test_learner.py @@ -0,0 +1,107 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Optimizer and graph contracts shared by BC and transition fitting.""" + +from __future__ import annotations + +from contextlib import contextmanager +import copy +import torch +import pytest + +from embodichain.learning.sl import LossResult, SupervisedBatch, SupervisedLearner + + +class Regression: + def compute_loss(self, model, batch): + error = (model(batch.inputs) - batch.targets).square() + mask = batch.masks.get("valid", torch.ones_like(error, dtype=torch.bool)) + return LossResult(error[mask].mean(), weight=int(mask.sum())) + + +def test_update_matches_independent_masked_sgd_and_retains_context(): + torch.manual_seed(11) + model = torch.nn.Linear(2, 1) + reference = copy.deepcopy(model) + x = torch.tensor([[1.0, 2.0], [3.0, -1.0]]) + y = torch.tensor([[0.5], [1000.0]]) + mask = torch.tensor([[True], [False]]) + expected_optimizer = torch.optim.SGD(reference.parameters(), lr=0.1) + expected_loss = (reference(x[:1]) - y[:1]).square().mean() + expected_loss.backward() + expected_optimizer.step() + inside = [] + + @contextmanager + def numerical_context(): + inside.append(True) + try: + yield + finally: + inside.pop() + + model.weight.register_hook( + lambda grad: grad if inside else pytest.fail("backward escaped context") + ) + learner = SupervisedLearner( + model, + Regression(), + torch.optim.SGD(model.parameters(), lr=0.1), + numerical_context=numerical_context, + ) + report = learner.update(SupervisedBatch(x, y, {"valid": mask})) + assert report["loss"] == expected_loss.item() + assert report["num_updates"] == 1 + for actual, expected in zip(model.parameters(), reference.parameters()): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +def test_multistep_objective_preserves_action_graph_and_batch_transfer(): + class Dynamics: + def compute_loss(self, model, batch): + state = batch.inputs["state"] + for action in batch.inputs["actions"].unbind(1): + state = state + model(action) + return LossResult((state - batch.targets).square().mean()) + + model = torch.nn.Linear(1, 1, bias=False) + model.weight.data.fill_(0.5) + actions = torch.ones(2, 3, 1, requires_grad=True) + batch = SupervisedBatch( + {"state": torch.zeros(2, 1), "actions": actions}, torch.zeros(2, 1) + ).to("cpu") + learner = SupervisedLearner( + model, Dynamics(), torch.optim.SGD(model.parameters(), lr=0.01) + ) + learner.update(batch) + torch.testing.assert_close(actions.grad, torch.full_like(actions, 0.75)) + + +def test_nonfinite_loss_does_not_commit_update(): + class Invalid: + def compute_loss(self, model, batch): + return LossResult(model(batch.inputs).sum() * float("inf")) + + model = torch.nn.Linear(1, 1) + before = copy.deepcopy(model.state_dict()) + learner = SupervisedLearner(model, Invalid(), torch.optim.Adam(model.parameters())) + with pytest.raises((RuntimeError, FloatingPointError)): + learner.update(SupervisedBatch(torch.ones(1, 1), None)) + assert learner.num_updates == 0 + for key, value in model.state_dict().items(): + torch.testing.assert_close(value, before[key], rtol=0, atol=0) + assert not learner.optimizer.state diff --git a/tests/learning/sl/test_trainer.py b/tests/learning/sl/test_trainer.py new file mode 100644 index 0000000000..c1681efbad --- /dev/null +++ b/tests/learning/sl/test_trainer.py @@ -0,0 +1,373 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Offline sampling ownership, semantic identity, and exact run recovery.""" + +from __future__ import annotations + +import copy +import random +import numpy as np +import json +import sys +import subprocess +import torch +import pytest +from embodichain.learning.sl import ( + LossResult, + SupervisedBatch, + SupervisedLearner, + SupervisedTrainer, + SupervisedTrainerCfg, + TensorBatchSource, +) + + +class Regression: + def compute_loss(self, model, batch): + return LossResult((model(batch.inputs) - batch.targets).square().mean()) + + +def build(path): + torch.manual_seed(42) + model = torch.nn.Sequential( + torch.nn.Linear(2, 8), torch.nn.Dropout(0.2), torch.nn.Linear(8, 1) + ) + x = torch.arange(40, dtype=torch.float32).view(20, 2) / 40 + data = TensorBatchSource(SupervisedBatch(x, x.sum(-1, keepdim=True))) + learner = SupervisedLearner( + model, Regression(), torch.optim.Adam(model.parameters(), lr=0.01) + ) + cfg = SupervisedTrainerCfg( + total_updates=20, batch_size=4, output_dir=str(path), checkpoint_interval=0 + ) + return SupervisedTrainer( + learner, data, cfg, signature={"target": "sum", "units": "dimensionless"} + ) + + +def test_resume_reproduces_dropout_optimizer_and_sampling(tmp_path): + full = build(tmp_path / "full") + full.fit() + expected = copy.deepcopy(full.learner.state_dict()) + split = build(tmp_path / "split") + split.fit(stop_after=7) + resumed = build(tmp_path / "resumed") + resumed.resume(tmp_path / "split/last.pt") + resumed.fit() + for key, value in expected["model"].items(): + torch.testing.assert_close( + value, resumed.learner.model.state_dict()[key], rtol=0, atol=0 + ) + combined = (tmp_path / "split/metrics.jsonl").read_text() + ( + tmp_path / "resumed/metrics.jsonl" + ).read_text() + assert combined == (tmp_path / "full/metrics.jsonl").read_text() + assert resumed.learner.num_updates == 20 + + +def test_dataset_change_rejected_before_state_mutation(tmp_path): + original = build(tmp_path / "first") + original.fit(stop_after=2) + candidate = build(tmp_path / "second") + candidate.signature["units"] = "metres" + before = copy.deepcopy(candidate.learner.model.state_dict()) + with pytest.raises(ValueError, match="identity"): + candidate.resume(tmp_path / "first/last.pt") + for key, value in before.items(): + torch.testing.assert_close( + value, candidate.learner.model.state_dict()[key], rtol=0, atol=0 + ) + + +def test_source_snapshots_and_sampled_batches_own_storage(): + x = torch.arange(12, dtype=torch.float32).view(4, 3) + source = TensorBatchSource(SupervisedBatch({"measured": x}, x + 1)) + same = TensorBatchSource(SupervisedBatch({"other_semantics": x}, x + 1)) + assert source.source_id != same.source_id + x.fill_(100) + first = source.sample(3, generator=torch.Generator().manual_seed(2)) + expected = first.inputs["measured"].clone() + source.sample(3, generator=torch.Generator().manual_seed(7)).inputs[ + "measured" + ].fill_(200) + torch.testing.assert_close(first.inputs["measured"], expected) + assert expected.max() < 100 + assert [batch.inputs["measured"].shape[0] for batch in source.iter_batches(3)] == [ + 3, + 1, + ] + + +def test_sl_import_does_not_load_simulation(): + code = "import sys; import embodichain.learning.sl; assert not any(k.startswith(('dexsim', 'embodichain.lab', 'embodichain.learning.rl')) for k in sys.modules)" + subprocess.run([sys.executable, "-c", code], check=True) + + +def test_unified_cli_routes_sl_help(): + result = subprocess.run( + [sys.executable, "-m", "embodichain", "train-sl", "--help"], + text=True, + capture_output=True, + check=True, + ) + assert "--config" in result.stdout + assert "--resume" in result.stdout + + +def test_sample_metadata_is_owned(): + source = TensorBatchSource( + SupervisedBatch( + torch.ones(2, 1), torch.zeros(2, 1), metadata={"units": {"position": "m"}} + ) + ) + first = source.sample(1, generator=torch.Generator()) + first.metadata["units"]["position"] = "cm" + second = source.sample(1, generator=torch.Generator()) + assert second.metadata["units"]["position"] == "m" + + +def test_optimizer_incompatibility_rejected_before_model_mutation(): + first = torch.nn.Linear(2, 1) + saved = SupervisedLearner( + first, Regression(), torch.optim.Adam(first.parameters()) + ).state_dict() + for optimizer in (torch.optim.SGD, torch.optim.Adam): + candidate = torch.nn.Linear(2, 1) + parameters = ( + [{"params": [candidate.weight]}, {"params": [candidate.bias]}] + if optimizer is torch.optim.Adam + else candidate.parameters() + ) + learner = SupervisedLearner( + candidate, Regression(), optimizer(parameters, lr=0.01) + ) + before = copy.deepcopy(candidate.state_dict()) + with pytest.raises(ValueError, match="optimizer"): + learner.load_state_dict(saved) + for key, value in before.items(): + torch.testing.assert_close( + candidate.state_dict()[key], value, rtol=0, atol=0 + ) + + +def test_source_identity_retains_key_and_container_types(): + value = torch.ones(2, 1) + trees = [{1: value}, {"1": value}, (value,), [value]] + identifiers = [ + TensorBatchSource(SupervisedBatch(tree, value)).source_id for tree in trees + ] + assert len(set(identifiers)) == len(trees) + + +def test_scheduler_type_rejected_before_learner_mutation(tmp_path): + first = build(tmp_path / "first") + first.scheduler = torch.optim.lr_scheduler.StepLR( + first.learner.optimizer, step_size=2, gamma=0.5 + ) + first.fit(stop_after=2) + second = build(tmp_path / "second") + second.scheduler = torch.optim.lr_scheduler.ExponentialLR( + second.learner.optimizer, gamma=0.5 + ) + before = copy.deepcopy(second.learner.model.state_dict()) + with pytest.raises(ValueError, match="scheduler type"): + second.resume(tmp_path / "first/last.pt") + for key, value in before.items(): + torch.testing.assert_close( + second.learner.model.state_dict()[key], value, rtol=0, atol=0 + ) + + +@pytest.mark.parametrize("callback_fails", [False, True]) +def test_validation_restores_rng_and_mode_even_on_failure(tmp_path, callback_fails): + def reset_rng(): + random.seed(17) + np.random.seed(17) + + def count_forwards(trainer): + trainer.learner.model.register_buffer( + "seen", torch.zeros((), dtype=torch.int64) + ) + + def record_forward(model, inputs): + model.seen.add_(1) + + trainer.learner.model.register_forward_pre_hook(record_forward) + + reset_rng() + baseline = build(tmp_path / "baseline") + count_forwards(baseline) + count = 1 if callback_fails else 3 + baseline.fit(stop_after=count) + expected = copy.deepcopy(baseline.learner.model.state_dict()) + expected_random = (random.random(), np.random.rand(), torch.rand(4)) + reset_rng() + candidate = build(tmp_path / "candidate") + count_forwards(candidate) + candidate.cfg.validation_interval = 1 + + def evaluate(learner): + learner.model(torch.ones(1, 2)) + random.random() + np.random.rand(5) + torch.rand(9) + candidate.save(candidate.output / "selected.pt") + if callback_fails: + raise RuntimeError("validation failed") + return {"probe": 0.0} + + candidate.evaluate = evaluate + if callback_fails: + with pytest.raises(RuntimeError, match="validation failed"): + candidate.fit(stop_after=count) + else: + candidate.fit(stop_after=count) + assert candidate.learner.model.training + for key, value in expected.items(): + torch.testing.assert_close( + candidate.learner.model.state_dict()[key], value, rtol=0, atol=0 + ) + assert random.random() == expected_random[0] + assert np.random.rand() == expected_random[1] + torch.testing.assert_close(torch.rand(4), expected_random[2], rtol=0, atol=0) + resumed = build(tmp_path / "resumed") + count_forwards(resumed) + resumed.cfg.validation_interval = 1 + resumed.resume(candidate.output / "selected.pt") + for key, value in expected.items(): + torch.testing.assert_close( + resumed.learner.model.state_dict()[key], value, rtol=0, atol=0 + ) + assert random.random() == expected_random[0] + assert np.random.rand() == expected_random[1] + torch.testing.assert_close(torch.rand(4), expected_random[2], rtol=0, atol=0) + + +def test_fresh_run_cannot_overwrite_existing_artifacts(tmp_path): + first = build(tmp_path / "run") + first.fit(stop_after=2) + log = (tmp_path / "run/metrics.jsonl").read_bytes() + checkpoint = (tmp_path / "run/last.pt").read_bytes() + with pytest.raises(FileExistsError, match="fresh output"): + build(tmp_path / "run").fit() + assert (tmp_path / "run/metrics.jsonl").read_bytes() == log + assert (tmp_path / "run/last.pt").read_bytes() == checkpoint + resumed = build(tmp_path / "run") + resumed.resume(tmp_path / "run/last.pt") + resumed.fit(stop_after=3) + assert resumed.learner.num_updates == 3 + + +@pytest.mark.parametrize("history", ["newer", "incomplete"]) +def test_checkpoint_recovery_preserves_incompatible_output_history(tmp_path, history): + first = build(tmp_path / "run") + first.fit(stop_after=2) + first.save(tmp_path / "step2.pt") + if history == "newer": + first.fit(stop_after=3) + else: + with (first.output / "metrics.jsonl").open("a") as stream: + stream.write('{"num_updates":') + original_log = (first.output / "metrics.jsonl").read_bytes() + candidate = build(tmp_path / "run") + before = copy.deepcopy(candidate.learner.model.state_dict()) + with pytest.raises(ValueError, match="history"): + candidate.resume(tmp_path / "step2.pt") + assert (first.output / "metrics.jsonl").read_bytes() == original_log + for key, value in before.items(): + torch.testing.assert_close( + candidate.learner.model.state_dict()[key], value, rtol=0, atol=0 + ) + branch = build(tmp_path / "branch") + branch.resume(tmp_path / "step2.pt") + assert branch.fit(stop_after=3)["num_updates"] == 3 + + +def test_same_trainer_may_save_an_initial_policy_before_fitting(tmp_path): + trainer = build(tmp_path / "run") + trainer.save(trainer.output / "step_000000.pt") + assert trainer.fit(stop_after=1)["num_updates"] == 1 + + +def test_preconstructed_trainer_rechecks_output_ownership_at_fit(tmp_path): + first = build(tmp_path / "shared") + second = build(tmp_path / "shared") + first.fit(stop_after=2) + before = (tmp_path / "shared/last.pt").read_bytes() + with pytest.raises(FileExistsError, match="another run"): + second.fit(stop_after=1) + assert (tmp_path / "shared/last.pt").read_bytes() == before + assert second.learner.num_updates == 0 + + +def test_foreign_writer_cannot_replace_an_owned_initial_snapshot(tmp_path): + first = build(tmp_path / "shared") + second = build(tmp_path / "shared") + path = first.output / "step_000000.pt" + first.save(path) + original = path.read_bytes() + with pytest.raises(FileExistsError, match="another run"): + second.save(path) + assert path.read_bytes() == original + assert first.fit(stop_after=1)["num_updates"] == 1 + + +@pytest.mark.parametrize("destination_step", [1, 3]) +def test_resume_rejects_another_run_even_with_older_or_equal_history( + tmp_path, destination_step +): + source = build(tmp_path / "source") + source.fit(stop_after=3) + destination = build(tmp_path / "destination") + destination.fit(stop_after=destination_step) + original = { + name: (destination.output / name).read_bytes() + for name in ("metrics.jsonl", "last.pt", "run.json") + } + candidate = build(destination.output) + before = copy.deepcopy(candidate.learner.state_dict()) + sampler = candidate.sampler.get_state().clone() + rng = torch.get_rng_state().clone() + with pytest.raises(FileExistsError, match="another run"): + candidate.resume(source.output / "last.pt") + for key, value in before["model"].items(): + torch.testing.assert_close( + candidate.learner.model.state_dict()[key], value, rtol=0, atol=0 + ) + assert candidate.learner.optimizer.state_dict() == before["optimizer"] + assert candidate.learner.num_updates == before["num_updates"] + torch.testing.assert_close(candidate.sampler.get_state(), sampler, rtol=0, atol=0) + torch.testing.assert_close(torch.get_rng_state(), rng, rtol=0, atol=0) + for name, content in original.items(): + assert (destination.output / name).read_bytes() == content + + +def test_checkpoint_without_run_identity_requires_a_fresh_destination(tmp_path): + source = build(tmp_path / "source") + source.fit(stop_after=2) + state = torch.load(source.output / "last.pt", weights_only=False) + state.pop("run_id") + legacy_path = tmp_path / "legacy.pt" + torch.save(state, legacy_path) + occupied = build(source.output) + with pytest.raises(FileExistsError, match="fresh output"): + occupied.resume(legacy_path) + fresh = build(tmp_path / "fresh") + fresh.resume(legacy_path) + assert fresh.fit(stop_after=3)["num_updates"] == 3 + saved = torch.load(fresh.output / "last.pt", weights_only=False) + assert saved["run_id"] diff --git a/tests/test_agent_context_map.py b/tests/test_agent_context_map.py index 6e51206086..4e1472bae5 100644 --- a/tests/test_agent_context_map.py +++ b/tests/test_agent_context_map.py @@ -62,6 +62,8 @@ def test_map_registers_the_supported_context_domains() -> None: "motion-planning", "rl-learning", "learning-layers", + "supervised-learning", + "neural-dynamics", "configclass-pattern", "randomization", "atomic-actions", diff --git a/tests/test_benchmark_package.py b/tests/test_benchmark_package.py index c4e67b94ab..35f70ccbf5 100644 --- a/tests/test_benchmark_package.py +++ b/tests/test_benchmark_package.py @@ -53,6 +53,7 @@ def installed_benchmark(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, "extensions/embodichain_robots", "extensions/embodichain_tasks", "extensions/embodichain_benchmarks", + "extensions/embodichain_models", ): shutil.copytree(root / name, source / name, ignore=ignore) built = subprocess.run( diff --git a/tests/test_example_package.py b/tests/test_example_package.py index 8ff228794e..d7a429629e 100644 --- a/tests/test_example_package.py +++ b/tests/test_example_package.py @@ -64,6 +64,7 @@ def test_distribution_preserves_the_ordinary_examples_tree(tmp_path: Path) -> No "embodichain_tasks", "embodichain_robots", "embodichain_benchmarks", + "embodichain_models", } for name in get_packages() ) diff --git a/tests/test_main.py b/tests/test_main.py index 3a97dac048..1a58033c74 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -54,6 +54,7 @@ "scene-randomization", "simready", "train-rl", + "train-sl", "preview-scene", "workspace-cache", } diff --git a/tests/test_release_metadata.py b/tests/test_release_metadata.py index abef9a9a53..dc7bf76eb4 100644 --- a/tests/test_release_metadata.py +++ b/tests/test_release_metadata.py @@ -76,6 +76,7 @@ def test_package_discovery_includes_only_runtime_trees() -> None: "embodichain_tasks", "embodichain_robots", "embodichain_benchmarks", + "embodichain_models", } assert runtime_roots | { "embodichain_tasks.configs", @@ -89,6 +90,7 @@ def test_package_discovery_includes_only_runtime_trees() -> None: "embodichain_tasks.configs": "extensions/embodichain_tasks/configs", "embodichain_robots": "extensions/embodichain_robots", "embodichain_benchmarks": "extensions/embodichain_benchmarks", + "embodichain_models": "extensions/embodichain_models/embodichain_models", }