From 63beff70553546ed3d6694647113d93e59a4aa96 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 24 Oct 2026 14:33:45 +0000 Subject: [PATCH 01/21] feat(dynamics): separate NeRD models, data generation and training --- agent_context/MAP.yaml | 19 + .../topics/data-pipeline/data-pipeline.md | 1 + .../topics/learning-layers/learning-layers.md | 5 +- .../topics/neural-dynamics/neural-dynamics.md | 64 + .../embodichain/embodichain.data_pipeline.rst | 1 + .../embodichain/embodichain.learning.rst | 1 + docs/source/api_reference/public_api.rst | 237 +++ docs/source/features/neural_dynamics.md | 170 +++ docs/source/guides/neural_dynamics.md | 8 + docs/source/index.rst | 1 + embodichain/cli/dynamics.py | 239 ++++ embodichain/cli/main.py | 5 + embodichain/compute/dynamics/__init__.py | 41 + embodichain/compute/dynamics/_evaluation.py | 103 ++ embodichain/compute/dynamics/codec.py | 533 +++++++ embodichain/compute/dynamics/reference.py | 70 + embodichain/data/nerd/ATTRIBUTION.md | 17 + embodichain/data/nerd/LICENSE.md | 176 +++ embodichain/data/nerd/ant.xml | 92 ++ embodichain/data/nerd/cartpole_single.urdf | 82 ++ embodichain/data_pipeline/__init__.py | 22 +- .../data_pipeline/dynamics/__init__.py | 21 + embodichain/data_pipeline/dynamics/_newton.py | 419 ++++++ .../dynamics/_reference_assets.py | 228 +++ .../data_pipeline/dynamics/_resident.py | 131 ++ .../data_pipeline/dynamics/generation.py | 338 +++++ embodichain/data_pipeline/dynamics/shards.py | 559 ++++++++ embodichain/learning/__init__.py | 2 +- embodichain/learning/dynamics/__init__.py | 21 + embodichain/learning/dynamics/_numerics.py | 121 ++ .../learning/dynamics/_reference_traces.py | 198 +++ embodichain/learning/dynamics/_runtime.py | 154 ++ .../learning/dynamics/_sample_order.py | 376 +++++ embodichain/learning/dynamics/comparison.py | 525 +++++++ .../learning/dynamics/qualification.py | 179 +++ embodichain/learning/dynamics/reference.py | 1267 +++++++++++++++++ .../learning/dynamics/reference_rollout.py | 1116 +++++++++++++++ embodichain/learning/dynamics/report.py | 1148 +++++++++++++++ embodichain/learning/dynamics/trainer.py | 741 ++++++++++ embodichain/learning/dynamics/workflow.py | 257 ++++ .../learning/models/nerd/THIRD_PARTY_LICENSES | 30 + embodichain/learning/models/nerd/__init__.py | 30 + .../learning/models/nerd/checkpoint.py | 225 +++ embodichain/learning/models/nerd/model.py | 320 +++++ .../tasks/classic_control/ant/README.md | 25 + .../ant/dynamics.eval.full.yaml | 8 + .../ant/dynamics.eval.pretrained.yaml | 8 + .../ant/dynamics.eval.smoke.yaml | 8 + .../ant/dynamics.generate.full.yaml | 18 + .../ant/dynamics.generate.smoke.yaml | 19 + .../ant/dynamics.generate.warp-reference.yaml | 19 + .../ant/dynamics.train.full.yaml | 38 + .../ant/dynamics.train.smoke.yaml | 36 + .../tasks/classic_control/cartpole/README.md | 25 + .../cartpole/dynamics.eval.full.yaml | 8 + .../cartpole/dynamics.eval.pretrained.yaml | 8 + .../cartpole/dynamics.eval.smoke.yaml | 8 + .../cartpole/dynamics.generate.full.yaml | 18 + .../cartpole/dynamics.generate.smoke.yaml | 19 + .../dynamics.generate.warp-reference.yaml | 18 + .../cartpole/dynamics.train.full.yaml | 34 + .../cartpole/dynamics.train.smoke.yaml | 32 + scripts/benchmark/learning/dynamics_input.py | 320 +++++ .../benchmark/learning/dynamics_sampling.py | 193 +++ .../benchmark/learning/dynamics_training.py | 309 ++++ scripts/experiments/neural_dynamics/README.md | 51 + .../neural_dynamics/_historical_physics.py | 233 +++ .../neural_dynamics/_portable_reference.py | 261 ++++ .../neural_dynamics/compare_official_model.py | 278 ++++ .../neural_dynamics/matched_small_budget.py | 265 ++++ .../official_evaluation_data.py | 513 +++++++ .../neural_dynamics/official_math_rollout.py | 370 +++++ .../official_model_reference.py | 441 ++++++ .../neural_dynamics/official_small_budget.py | 345 +++++ .../official_small_budget_evaluate.py | 166 +++ .../official_small_budget_rollout.py | 275 ++++ setup.py | 2 + tests/cli/test_dynamics.py | 298 ++++ .../dynamics/data/ant-official-frame.json | 64 + .../dynamics/data/ant-official-frame.npz | Bin 0 -> 34898 bytes tests/compute/dynamics/test_codec.py | 331 +++++ .../compute/dynamics/test_dynamics_metrics.py | 99 ++ .../dynamics/test_dynamics_task_reference.py | 70 + .../data_pipeline/dynamics/test_generation.py | 206 +++ tests/data_pipeline/dynamics/test_shards.py | 343 +++++ tests/learning/dynamics/test_comparison.py | 190 +++ .../test_dynamics_runtime_identity.py | 183 +++ tests/learning/dynamics/test_reference.py | 935 ++++++++++++ .../dynamics/test_reference_rollout.py | 533 +++++++ .../dynamics/test_reference_traces.py | 247 ++++ tests/learning/dynamics/test_report.py | 414 ++++++ tests/learning/dynamics/test_sample_order.py | 328 +++++ .../dynamics/test_small_budget_observation.py | 169 +++ tests/learning/dynamics/test_trainer.py | 901 ++++++++++++ .../learning/models/_nerd_reference_check.py | 348 +++++ tests/learning/models/test_nerd.py | 261 ++++ tests/test_agent_context_map.py | 1 + 97 files changed, 20580 insertions(+), 5 deletions(-) create mode 100644 agent_context/topics/neural-dynamics/neural-dynamics.md create mode 100644 docs/source/features/neural_dynamics.md create mode 100644 docs/source/guides/neural_dynamics.md create mode 100644 embodichain/cli/dynamics.py create mode 100644 embodichain/compute/dynamics/__init__.py create mode 100644 embodichain/compute/dynamics/_evaluation.py create mode 100644 embodichain/compute/dynamics/codec.py create mode 100644 embodichain/compute/dynamics/reference.py create mode 100644 embodichain/data/nerd/ATTRIBUTION.md create mode 100644 embodichain/data/nerd/LICENSE.md create mode 100644 embodichain/data/nerd/ant.xml create mode 100644 embodichain/data/nerd/cartpole_single.urdf create mode 100644 embodichain/data_pipeline/dynamics/__init__.py create mode 100644 embodichain/data_pipeline/dynamics/_newton.py create mode 100644 embodichain/data_pipeline/dynamics/_reference_assets.py create mode 100644 embodichain/data_pipeline/dynamics/_resident.py create mode 100644 embodichain/data_pipeline/dynamics/generation.py create mode 100644 embodichain/data_pipeline/dynamics/shards.py create mode 100644 embodichain/learning/dynamics/__init__.py create mode 100644 embodichain/learning/dynamics/_numerics.py create mode 100644 embodichain/learning/dynamics/_reference_traces.py create mode 100644 embodichain/learning/dynamics/_runtime.py create mode 100644 embodichain/learning/dynamics/_sample_order.py create mode 100644 embodichain/learning/dynamics/comparison.py create mode 100644 embodichain/learning/dynamics/qualification.py create mode 100644 embodichain/learning/dynamics/reference.py create mode 100644 embodichain/learning/dynamics/reference_rollout.py create mode 100644 embodichain/learning/dynamics/report.py create mode 100644 embodichain/learning/dynamics/trainer.py create mode 100644 embodichain/learning/dynamics/workflow.py create mode 100644 embodichain/learning/models/nerd/THIRD_PARTY_LICENSES create mode 100644 embodichain/learning/models/nerd/__init__.py create mode 100644 embodichain/learning/models/nerd/checkpoint.py create mode 100644 embodichain/learning/models/nerd/model.py create mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/ant/README.md create mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.full.yaml create mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.pretrained.yaml create mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.smoke.yaml create mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.full.yaml create mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.smoke.yaml create mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.warp-reference.yaml create mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.train.full.yaml create mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.train.smoke.yaml create mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/README.md create mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.full.yaml create mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.pretrained.yaml create mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.smoke.yaml create mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.full.yaml create mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.smoke.yaml create mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.warp-reference.yaml create mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.train.full.yaml create mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.train.smoke.yaml create mode 100644 scripts/benchmark/learning/dynamics_input.py create mode 100644 scripts/benchmark/learning/dynamics_sampling.py create mode 100644 scripts/benchmark/learning/dynamics_training.py create mode 100644 scripts/experiments/neural_dynamics/README.md create mode 100644 scripts/experiments/neural_dynamics/_historical_physics.py create mode 100644 scripts/experiments/neural_dynamics/_portable_reference.py create mode 100644 scripts/experiments/neural_dynamics/compare_official_model.py create mode 100644 scripts/experiments/neural_dynamics/matched_small_budget.py create mode 100644 scripts/experiments/neural_dynamics/official_evaluation_data.py create mode 100644 scripts/experiments/neural_dynamics/official_math_rollout.py create mode 100644 scripts/experiments/neural_dynamics/official_model_reference.py create mode 100644 scripts/experiments/neural_dynamics/official_small_budget.py create mode 100644 scripts/experiments/neural_dynamics/official_small_budget_evaluate.py create mode 100644 scripts/experiments/neural_dynamics/official_small_budget_rollout.py create mode 100644 tests/cli/test_dynamics.py create mode 100644 tests/compute/dynamics/data/ant-official-frame.json create mode 100644 tests/compute/dynamics/data/ant-official-frame.npz create mode 100644 tests/compute/dynamics/test_codec.py create mode 100644 tests/compute/dynamics/test_dynamics_metrics.py create mode 100644 tests/compute/dynamics/test_dynamics_task_reference.py create mode 100644 tests/data_pipeline/dynamics/test_generation.py create mode 100644 tests/data_pipeline/dynamics/test_shards.py create mode 100644 tests/learning/dynamics/test_comparison.py create mode 100644 tests/learning/dynamics/test_dynamics_runtime_identity.py create mode 100644 tests/learning/dynamics/test_reference.py create mode 100644 tests/learning/dynamics/test_reference_rollout.py create mode 100644 tests/learning/dynamics/test_reference_traces.py create mode 100644 tests/learning/dynamics/test_report.py create mode 100644 tests/learning/dynamics/test_sample_order.py create mode 100644 tests/learning/dynamics/test_small_budget_observation.py create mode 100644 tests/learning/dynamics/test_trainer.py create mode 100644 tests/learning/models/_nerd_reference_check.py create mode 100644 tests/learning/models/test_nerd.py diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index b4ac4f4c04..2f3a5ff947 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, 神经动力学, 动力学模型训练] + keywords: [NeRDStateCodec, NeRDModel, DynamicsGenerationCfg, + DynamicsShardWriter, DynamicsWindowDataset, DynamicsTrainer, + 官方训练精度对比] + paths: [topics/neural-dynamics/neural-dynamics.md] + source_of_truth: + - embodichain/compute/dynamics/ + - embodichain/learning/models/nerd/ + - embodichain/learning/dynamics/ + - embodichain/data_pipeline/dynamics/ + - embodichain/cli/dynamics.py + watch_paths: [embodichain/compute/dynamics/, + embodichain/learning/models/nerd/, embodichain/learning/dynamics/, + embodichain/data_pipeline/dynamics/, + tests/compute/dynamics/, tests/learning/dynamics/, tests/data_pipeline/dynamics/] + related_topics: [learning-layers, data-pipeline] + status: active - id: learning-layers title: Differentiable Learning Layers aliases: [learning layers, differentiable IK layer, 可微 IK, 训练层, 学习模块分层] diff --git a/agent_context/topics/data-pipeline/data-pipeline.md b/agent_context/topics/data-pipeline/data-pipeline.md index f795d24b6f..5db12b3faa 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` | +| Dynamics teacher generation and HDF5 windows | `embodichain/data_pipeline/dynamics/`; 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..7e609fb2f0 100644 --- a/agent_context/topics/learning-layers/learning-layers.md +++ b/agent_context/topics/learning-layers/learning-layers.md @@ -9,7 +9,10 @@ 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, +and configuration-driven MLP construction. `learning/models/nerd/` owns the NeRD causal model and portable weights; +offline dynamics training is in `learning/dynamics/`. See +[neural dynamics](../neural-dynamics/neural-dynamics.md) for dataset and checkpoint +contracts. These components do not import RL, TensorDict or simulation. RL policy interfaces, actor-critic composition and policy registration/builders belong to `learning/rl/policies/`; RL code consumes the shared networks and normalizer. See [RL learning](../rl-learning/rl-learning.md) 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..6403f7b44d --- /dev/null +++ b/agent_context/topics/neural-dynamics/neural-dynamics.md @@ -0,0 +1,64 @@ +# Neural dynamics models and training + +## Entry points and ownership + +| Request | Owner (relative to `embodichain/`) | +|---|---| +| Tensor state, contact and history codecs | `compute/dynamics/codec.py` | +| Model and portable checkpoints | `learning/models/nerd/` | +| Analytical teacher generation | `data_pipeline/dynamics/generation.py`; private `_newton.py` owns its Newton model | +| Reference assets and physical profiles | `data_pipeline/dynamics/_reference_assets.py`; packaged files in `data/nerd/` | +| Committed shards and sequence windows | `data_pipeline/dynamics/shards.py` | +| Training, resume, evaluation and export | `learning/dynamics/trainer.py`, `workflow.py` | +| Original-repository comparisons | `learning/dynamics/reference.py`, `comparison.py`, `report.py` | +| Command dispatch | `cli/dynamics.py`: `embodichain dynamics` | + +Model, codec, training and offline evaluation do not import simulation. Generation +uses Newton's built-in Featherstone/MuJoCo solver directly without DexSim, its +custom provider registry or Scene APIs. Native learned-solver deployment and +functional Gym integration are separate consumers, not dependencies of training. + +## Data and physical identity + +A trajectory contains T+1 states and T controls/contact inputs. Splits belong to +whole trajectories and windows cannot cross episodes or splits. Writers publish +only committed shards; resume advances candidate seeds past accepted and rejected +candidates. File handles and bounded caches belong to their reader/worker. + +The direct teacher stores its importer profile, asset/physical-array hashes, +solver settings and actual Newton/Warp versions. Its +`nerd-direct-newton-builder-v1` signature is distinct from Scene-built datasets. +Existing published Scene-generated data remains readable without simulation. +`warp_reference` reconstructs measured legacy asset parameters and the coupled +Ant Featherstone friction mapping; current Newton still owns integration. +Matching a profile name alone does not establish matching physics across +model-building paths, versions or solvers. + +Newton COM velocities and NeRD twists differ in ordering and reference point. +The codec owns transport, relative targets use the frame at t, and frozen train +normalizers are reused for evaluation. Ant uses 25 ordered plane-contact slots; +Cartpole is contact-free. Controls are generalized efforts with zero free-base +force. Physical records cannot be relabeled through training manifest overrides. + +The collector keeps one bounded batch of history on the model device, then writes +owned CPU arrays. CUDA teachers use graph capture by default where supported; +CPU MuJoCo execution does not capture. Snapshots retain independent ownership. +Mass-matrix update intervals divide the macro substep count so resets can replay. + +## Training and validation + +Training checkpoints preserve optimizer, scheduler, RNG and committed sample +progress. Portable exports contain tensors and JSON, including normalizers and +attention policy. Resume restores the same run; fine-tuning starts a new one. +Only explicitly trusted legacy conversion may deserialize Python model objects. + +Optional shared sample-order/statistics artifacts support controlled original +trainer comparisons; equal seeds alone do not prove identical batches or physics. +Original Warp references run in isolated supported environments. Keep virtualenv +executable paths intact instead of resolving their interpreter symlinks. + +Tests cover codecs, model conversion, bounded shards, reset/replay and partial +collection batches, training/resume and the actual offline CLI. Public usage is +in [Neural Dynamics Models](../../../docs/source/features/neural_dynamics.md). +Fixed-set transition error, live rollout error and paper protocol accuracy are +separate claims; this workflow does not establish their equivalence. diff --git a/docs/source/api_reference/embodichain/embodichain.data_pipeline.rst b/docs/source/api_reference/embodichain/embodichain.data_pipeline.rst index bfe2e2bf72..52476fc8ea 100644 --- a/docs/source/api_reference/embodichain/embodichain.data_pipeline.rst +++ b/docs/source/api_reference/embodichain/embodichain.data_pipeline.rst @@ -22,6 +22,7 @@ storage for LeRobot datasets on Python 3.10--3.12). depth_video engine recording + dynamics runtime Datasets diff --git a/docs/source/api_reference/embodichain/embodichain.learning.rst b/docs/source/api_reference/embodichain/embodichain.learning.rst index 10f9c879a4..f8c60d7ed6 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.rst +++ b/docs/source/api_reference/embodichain/embodichain.learning.rst @@ -10,6 +10,7 @@ embodichain.learning layers models rl + dynamics Reinforcement Learning ---------------------- diff --git a/docs/source/api_reference/public_api.rst b/docs/source/api_reference/public_api.rst index 095a621c3d..98eb579353 100644 --- a/docs/source/api_reference/public_api.rst +++ b/docs/source/api_reference/public_api.rst @@ -22,6 +22,56 @@ decoding and compiling the constrained Task Program schema surface. compile_mllm_task_program decode_mllm_task_program +embodichain.compute.dynamics +---------------------------- + +Simulation-independent NeRD coordinate, contact and history transformations. +State derivatives retain their graphs; discrete contact and orientation decisions +use branch-local derivatives. + +.. currentmodule:: embodichain.compute.dynamics + +.. autosummary:: + + reference_task_values + NeRDStateCodec + quat_multiply + quat_rotate + quat_inverse + plane_contact_features + append_history + reset_history + +embodichain.compute.dynamics.codec +---------------------------------- + +Convert public Newton COM velocities to NeRD spatial twists, encode relative +targets in the current body frame, and update per-world recurrent windows. + +.. currentmodule:: embodichain.compute.dynamics.codec + +.. autosummary:: + + NeRDStateCodec + quat_multiply + quat_rotate + quat_inverse + plane_contact_features + append_history + reset_history + +embodichain.compute.dynamics.reference +-------------------------------------- + +Shared Cartpole and Ant observation, reward and failure calculations used by +registered tasks and rollout comparison, with explicit quaternion validation. + +.. currentmodule:: embodichain.compute.dynamics.reference + +.. autosummary:: + + reference_task_values + embodichain.data.assets.demo_assets ----------------------------------- @@ -237,6 +287,37 @@ embodichain.data_pipeline.depth_video DEPTH_MILLIMETER_UNIT DEPTH_QMAX +embodichain.data_pipeline.dynamics.generation +--------------------------------------------- + +Generate bounded batches of Cartpole/Ant trajectories with Newton's built-in +analytical solvers. This offline collector does not initialize DexSim or register +a learned simulation solver. Committed shards support resumable generation. + +.. currentmodule:: embodichain.data_pipeline.dynamics.generation + +.. autosummary:: + + DynamicsGenerationCfg + generate_dataset + +embodichain.data_pipeline.dynamics.shards +----------------------------------------- + +Committed HDF5 trajectory shards and bounded sequence-window reads. State +sequences have T+1 entries, controls have T; windows respect trajectory and split +boundaries. Resume uses only manifest-committed shards. + +.. currentmodule:: embodichain.data_pipeline.dynamics.shards + +.. autosummary:: + + DynamicsShardWriter + DynamicsWindowDataset + trajectory_split + load_manifest + dataset_fingerprint + embodichain.gen_sim.env ----------------------- @@ -1655,6 +1736,162 @@ embodichain.lab.sim.motion.planners.trapezoidal_planner TrapezoidalPlanner TrapezoidalPlannerCfg +embodichain.learning.dynamics.comparison +---------------------------------------- + +Prepare and execute matched-data trainer comparisons with immutable source, +initialization, statistics and sample-order identities. Completion status reflects +which experiments actually ran. + +.. currentmodule:: embodichain.learning.dynamics.comparison + +.. autosummary:: + + prepare_comparison + run_comparison + run_candidate + +embodichain.learning.dynamics.qualification +------------------------------------------- + +Compare independent forward/gradient/update snapshots and evaluate paired +seed-level non-inferiority. Fewer than five seeds remains pending qualification. + +.. currentmodule:: embodichain.learning.dynamics.qualification + +.. autosummary:: + + capture_training_step + compare_training_steps + paired_noninferiority + +embodichain.learning.dynamics.reference +--------------------------------------- + +Explicit adapters for a trusted, pinned official source checkout. Reference +execution is isolated from the production runtime; record compatibility transport +and source identity rather than relabeling migrated code as an official baseline. + +.. currentmodule:: embodichain.learning.dynamics.reference + +.. autosummary:: + + OFFICIAL_COMMIT + reference_identity + run_reference + export_legacy_dataset + import_legacy_datasets + generate_official_dataset + repair_gravity_transport + +embodichain.learning.dynamics.reference_rollout +----------------------------------------------- + +Evaluate portable weights in the pinned original Warp environment and neural +integrator. This explicitly trusted, isolated reference process measures +historical-domain training quality; it does not deploy historical physics into +current Newton. Physical provenance, shared teacher traces and portable-network +equivalence are retained with failure and policy-return evidence. + +.. currentmodule:: embodichain.learning.dynamics.reference_rollout + +.. autosummary:: + + evaluate_reference_rollouts + +embodichain.learning.dynamics.report +------------------------------------ + +Write machine-readable and human-readable comparison evidence with explicit +complete, failed and pending qualification status. + +.. currentmodule:: embodichain.learning.dynamics.report + +.. autosummary:: + + write_comparison_report + annotate_rollout_evidence + +embodichain.learning.dynamics.trainer +------------------------------------- + +Single-device causal dynamics training with frozen train-only normalization, +committed sampling progress, optimizer/RNG resume and independent evaluation. + +.. currentmodule:: embodichain.learning.dynamics.trainer + +.. autosummary:: + + DynamicsTrainCfg + DynamicsTrainer + encode_batch + normalized_mse + compute_statistics + load_statistics + save_statistics + evaluate_model + +embodichain.learning.dynamics.workflow +-------------------------------------- + +Offline train, evaluate and export workflows validate dataset/model contracts +without importing or initializing the simulation runtime. + +.. currentmodule:: embodichain.learning.dynamics.workflow + +.. autosummary:: + + run_train + run_eval + run_export + +embodichain.learning.models.nerd +-------------------------------- + +Portable causal NeRD networks, registered normalization buffers and versioned +weights. Evaluation preserves gradients; legacy Python-object conversion requires +explicit trust and is separate from normal runtime loading. + +.. currentmodule:: embodichain.learning.models.nerd + +.. autosummary:: + + NeRDModelCfg + NeRDModel + FeatureNormalizer + export_model + load_model + convert_legacy_checkpoint + +embodichain.learning.models.nerd.checkpoint +------------------------------------------- + +Export tensor state and JSON metadata with checksums; load portable models or +convert an explicitly trusted official Transformer checkpoint. + +.. currentmodule:: embodichain.learning.models.nerd.checkpoint + +.. autosummary:: + + export_model + load_model + convert_legacy_checkpoint + +embodichain.learning.models.nerd.model +-------------------------------------- + +Causal Transformer architecture matching official NeRD parameter names. +Right-padded batches retain valid-prefix positional semantics and per-row output +selection; normalizer statistics move with the network. + +.. currentmodule:: embodichain.learning.models.nerd.model + +.. autosummary:: + + NeRDModelCfg + NeRDModel + FeatureNormalizer + embodichain.utils.warp.kinematics.trapezoidal_warp -------------------------------------------------- diff --git a/docs/source/features/neural_dynamics.md b/docs/source/features/neural_dynamics.md new file mode 100644 index 0000000000..c5b1c79d1f --- /dev/null +++ b/docs/source/features/neural_dynamics.md @@ -0,0 +1,170 @@ +# Neural Dynamics Models + +A neural dynamics model predicts how a robot's physical state changes under +applied joint efforts and contact. EmbodiChain provides a simulation-independent NeRD model, trajectory +storage, training, offline evaluation and portable export workflow. Data generation +uses Newton's built-in analytical solvers directly and does not require a custom +DexSim build. + +[NeRD (Neural Robot Dynamics)](https://neural-robot-dynamics.github.io/) is the +first implementation. Its causal Transformer consumes a short history of state, +efforts and contact features, then predicts the next state. The current reference +systems are contact-free Cartpole and free-base Ant on a static plane. + +## 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. + +## Published datasets and models + +Research datasets live in +[DexForceAI/embodichain_datasets](https://huggingface.co/datasets/DexForceAI/embodichain_datasets). +Deployable models live in +[DexForceAI/embodichain_model](https://huggingface.co/DexForceAI/embodichain_model). +Both repositories organize assets by model family, implementation and robot: + +| Reference system | Path in each repository | Deployment profile | +| --- | --- | --- | +| Cartpole | `neural_dynamics/nerd/cartpole/v1` | Contact-free, `warp_reference`, 1/60 s | +| Ant | `neural_dynamics/nerd/ant/v1` | Static plane, `warp_reference`, 1/60 s | + +Dataset directories contain a manifest and HDF5 shards with explicit +train/validation/test trajectory splits. Each trajectory stores T+1 states and T +efforts/contact inputs. Windows cannot cross trajectory or split boundaries. +Model directories contain `weights.pt` and `manifest.json`, including network +configuration, normalization buffers, codec, timestep and physical identity. +These are inference exports; resumable training checkpoints are local training +outputs, not part of the published model package. + +The dataset and model cards describe the generating solver, supported dependency +versions, data size and evaluation limits. In particular, long-horizon Ant +predictions can drift. Assess the model against its matching analytical solver +for the intended motion and prediction horizon before using it downstream. + +## Download and evaluate + +Run from the repository root. The Hugging Face client is an EmbodiChain +dependency. Download one matching dataset/model pair: + +```python +from huggingface_hub import snapshot_download + +robot = "cartpole" # or "ant" +prefix = f"neural_dynamics/nerd/{robot}/v1" +snapshot_download( + repo_id="DexForceAI/embodichain_datasets", + repo_type="dataset", + allow_patterns=[f"{prefix}/*"], + local_dir="outputs/datasets", +) +snapshot_download( + repo_id="DexForceAI/embodichain_model", + allow_patterns=[f"{prefix}/*"], + local_dir="outputs/models", +) +``` + +Evaluate the saved test transitions without starting a simulator: + +```bash +embodichain dynamics eval --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.pretrained.yaml +``` + +The pretrained configuration reads the downloaded paths above. Replace +`cartpole` with `ant` for Ant. Use `--device cpu` or `--device cuda:0` to choose +the tensor device. Offline evaluation measures held-out transition prediction; +it does not measure accumulated rollout error. + +The released datasets were collected through the matching Newton/Scene +integration. The independent collector below builds a new Newton model directly; +its manifests use a separate `nerd-direct-newton-builder-v1` physical identity. +Both datasets share the same state/contact schema and can be trained and evaluated +here. Do not treat different model-building paths or solver versions as identical +physical systems, or use offline transition error as a live rollout measurement. + +## Generate data, train and export + +The workflow configurations are under +`extensions/embodichain_tasks/configs/tasks/classic_control//`. +They configure `embodichain dynamics` commands and are separate from runnable +Gym deployment configs. + +For a short local workflow, use the smoke presets: + +```bash +embodichain dynamics generate --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.smoke.yaml +embodichain dynamics train --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.train.smoke.yaml +embodichain dynamics export --checkpoint outputs/dynamics/cartpole/smoke/train/best.pt --output outputs/dynamics/cartpole/smoke/export +embodichain dynamics eval --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.smoke.yaml +``` + +Smoke presets use a small network and short budget to exercise the workflow. +Full presets use the reference network scale and a substantially longer budget. +Neither preset name is an accuracy qualification. + +To train on a downloaded dataset, copy `dynamics.train.full.yaml`, then set +`dataset` to the downloaded robot's `v1` directory and `output_dir` to a new +training directory. Preserve its model feature dimensions and history length. +Run the copied configuration with `embodichain dynamics train --config PATH`, +then export its validation-selected `best.pt` with the `export` command. + +Training checkpoints retain model, optimizer, normalization, RNG and committed +sampling progress. Continue the same run with `train --resume PATH/last.pt`. +Initialize a new run from an inference export with `train --fine-tune PATH`. +Export destinations must be new directories. Paths written in configuration +files resolve relative to the file; explicit CLI path overrides resolve relative +to the working directory. + +## Use the exported model in PyTorch + +Load tensor weights and their recorded normalization buffers: + +```python +from embodichain.learning.models.nerd import load_model + +model, manifest = load_model( + "outputs/models/neural_dynamics/nerd/cartpole/v1", device="cpu" +) +``` + +`NeRDStateCodec` and `encode_batch` construct features from the dataset's state, +effort and contact fields. The model accepts feature windows and valid history +lengths; the training and evaluation workflows use this same preprocessing. + +Live Scene stepping, `solver_type: nerd`, and differentiable Gym integration are +provided separately by the [neural dynamics solver PR](https://github.com/DexForce/EmbodiChain/pull/822). +The comparison videos above were produced using that integration; they are +reference recordings, not a command available in this offline workflow. + +Keep state transforms in {mod}`embodichain.compute.dynamics`, model/checkpoint +code in {mod}`embodichain.learning.models.nerd`, training in +{mod}`embodichain.learning.dynamics.workflow`, and trajectory storage/generation in +{mod}`embodichain.data_pipeline.dynamics`. diff --git a/docs/source/guides/neural_dynamics.md b/docs/source/guides/neural_dynamics.md new file mode 100644 index 0000000000..74fcac64e4 --- /dev/null +++ b/docs/source/guides/neural_dynamics.md @@ -0,0 +1,8 @@ +--- +orphan: true +--- + +# Neural dynamics + +The feature overview, published datasets and models, comparison videos, and +workflow commands are now in {doc}`/features/neural_dynamics`. 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/dynamics.py b/embodichain/cli/dynamics.py new file mode 100644 index 0000000000..6f98492a14 --- /dev/null +++ b/embodichain/cli/dynamics.py @@ -0,0 +1,239 @@ +# ---------------------------------------------------------------------------- +# 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 +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: + """Dispatch one dynamics operation; offline paths never import simulation.""" + parser = argparse.ArgumentParser(prog="embodichain dynamics", description=__doc__) + subparsers = parser.add_subparsers(dest="operation", required=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, required=True) + 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: + 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 embodichain.learning.dynamics.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.learning.models.nerd import convert_legacy_checkpoint + + convert_legacy_checkpoint( + checkpoint, output, manifest, trust_pickle=args.trust_pickle + ) + result = {"output": str(output), "qualification": "unqualified"} + else: + path = args.config.expanduser().resolve() + config = _read_config(path) + 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 embodichain.data_pipeline.dynamics.generation import ( + generate_dataset, + ) + + result = generate_dataset(config) + elif args.operation == "train": + from embodichain.learning.dynamics.workflow import run_train + + result = run_train(config) + else: + mode = config.pop("mode", "offline") + if mode == "offline": + from embodichain.learning.dynamics.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/embodichain/cli/main.py b/embodichain/cli/main.py index 8a72100137..315ec493a1 100644 --- a/embodichain/cli/main.py +++ b/embodichain/cli/main.py @@ -112,6 +112,11 @@ class Command: target="embodichain.cli.dataset:main", help="Validate, split, and recover recorded expert datasets.", ), + Command( + name="dynamics", + target="embodichain.cli.dynamics:main", + help="Generate data, train and evaluate neural dynamics models.", + ), Command( name="preview_lerobot_data", target="embodichain.lab.scripts.preview_lerobot_data:cli", diff --git a/embodichain/compute/dynamics/__init__.py b/embodichain/compute/dynamics/__init__.py new file mode 100644 index 0000000000..ffc3c550f6 --- /dev/null +++ b/embodichain/compute/dynamics/__init__.py @@ -0,0 +1,41 @@ +# ---------------------------------------------------------------------------- +# 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 learned dynamics mathematics.""" + +from __future__ import annotations +from .codec import ( + NeRDStateCodec, + quat_multiply, + quat_rotate, + quat_inverse, + plane_contact_features, + append_history, + reset_history, +) + +from .reference import reference_task_values + +__all__ = [ + "reference_task_values", + "NeRDStateCodec", + "quat_multiply", + "quat_rotate", + "quat_inverse", + "plane_contact_features", + "append_history", + "reset_history", +] diff --git a/embodichain/compute/dynamics/_evaluation.py b/embodichain/compute/dynamics/_evaluation.py new file mode 100644 index 0000000000..e0aceb53b1 --- /dev/null +++ b/embodichain/compute/dynamics/_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/embodichain/compute/dynamics/codec.py b/embodichain/compute/dynamics/codec.py new file mode 100644 index 0000000000..57a7236f07 --- /dev/null +++ b/embodichain/compute/dynamics/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/embodichain/compute/dynamics/reference.py b/embodichain/compute/dynamics/reference.py new file mode 100644 index 0000000000..0b332e2b6b --- /dev/null +++ b/embodichain/compute/dynamics/reference.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 .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/embodichain/data/nerd/ATTRIBUTION.md b/embodichain/data/nerd/ATTRIBUTION.md new file mode 100644 index 0000000000..e198e27878 --- /dev/null +++ b/embodichain/data/nerd/ATTRIBUTION.md @@ -0,0 +1,17 @@ +# NeRD reference robot assets + +`cartpole_single.urdf` and `ant.xml` are unmodified copies from +NVlabs/neural-robot-dynamics, commit +`476dc67fccc6a99be293bf611f897b1c06f13e30`, under +`envs/newton_envs/assets/`. The upstream repository license is reproduced in +`LICENSE.md` (Apache License 2.0; NVIDIA Corporation and affiliates, 2025). + +Source: https://github.com/NVlabs/neural-robot-dynamics/tree/476dc67fccc6a99be293bf611f897b1c06f13e30/envs/newton_envs/assets + +EmbodiChain translates these assets using the installed Newton public importer, +then passes their physical descriptors through DexSim Spawn. This is a current +Newton teacher profile, not a claim to reproduce legacy Warp 1.5.1 physics. +Imported motor actuator declarations and XML solver options are not installed +into the Spawn scene: the host supplies actual generalized forces and declares +its own solver configuration. Masses, inertia, COMs, joint limits, passive joint +parameters and collision geometry/material parameters are retained. diff --git a/embodichain/data/nerd/LICENSE.md b/embodichain/data/nerd/LICENSE.md new file mode 100644 index 0000000000..2bb9ad240f --- /dev/null +++ b/embodichain/data/nerd/LICENSE.md @@ -0,0 +1,176 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS \ No newline at end of file diff --git a/embodichain/data/nerd/ant.xml b/embodichain/data/nerd/ant.xml new file mode 100644 index 0000000000..04c237b0d7 --- /dev/null +++ b/embodichain/data/nerd/ant.xml @@ -0,0 +1,92 @@ + + + + + + + + + + + + + \ No newline at end of file diff --git a/embodichain/data/nerd/cartpole_single.urdf b/embodichain/data/nerd/cartpole_single.urdf new file mode 100644 index 0000000000..3f5981a495 --- /dev/null +++ b/embodichain/data/nerd/cartpole_single.urdf @@ -0,0 +1,82 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/embodichain/data_pipeline/__init__.py b/embodichain/data_pipeline/__init__.py index 8e3ea1615f..9597ede6bd 100644 --- a/embodichain/data_pipeline/__init__.py +++ b/embodichain/data_pipeline/__init__.py @@ -27,9 +27,25 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from . import configuration, datasets, depth_video, engine, recording, runtime - -__all__ = ["configuration", "datasets", "depth_video", "engine", "recording", "runtime"] + from . import ( + configuration, + datasets, + depth_video, + dynamics, + engine, + recording, + runtime, + ) + +__all__ = [ + "configuration", + "datasets", + "depth_video", + "dynamics", + "engine", + "recording", + "runtime", +] def __getattr__(name: str) -> ModuleType: diff --git a/embodichain/data_pipeline/dynamics/__init__.py b/embodichain/data_pipeline/dynamics/__init__.py new file mode 100644 index 0000000000..b1f8e18a59 --- /dev/null +++ b/embodichain/data_pipeline/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. +# ---------------------------------------------------------------------------- + +"""Simulation-independent neural dynamics workflows.""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/embodichain/data_pipeline/dynamics/_newton.py b/embodichain/data_pipeline/dynamics/_newton.py new file mode 100644 index 0000000000..c09a147d27 --- /dev/null +++ b/embodichain/data_pipeline/dynamics/_newton.py @@ -0,0 +1,419 @@ +# ---------------------------------------------------------------------------- +# 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.compute.dynamics 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": { + name: hashlib.sha256( + Path(__file__).with_name(name).read_bytes() + ).hexdigest() + for name in ("_newton.py", "_reference_assets.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/embodichain/data_pipeline/dynamics/_reference_assets.py b/embodichain/data_pipeline/dynamics/_reference_assets.py new file mode 100644 index 0000000000..6928441396 --- /dev/null +++ b/embodichain/data_pipeline/dynamics/_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}") + import embodichain + + return Path(embodichain.__file__).parent / "data" / "nerd" / _ASSETS[robot] + + +def reference_asset_hash(robot: str) -> str: + """Return the SHA256 of the exact packaged 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/embodichain/data_pipeline/dynamics/_resident.py b/embodichain/data_pipeline/dynamics/_resident.py new file mode 100644 index 0000000000..9663efcaba --- /dev/null +++ b/embodichain/data_pipeline/dynamics/_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/embodichain/data_pipeline/dynamics/generation.py b/embodichain/data_pipeline/dynamics/generation.py new file mode 100644 index 0000000000..49e2f8b1b3 --- /dev/null +++ b/embodichain/data_pipeline/dynamics/generation.py @@ -0,0 +1,338 @@ +# ---------------------------------------------------------------------------- +# 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.compute.dynamics import quat_rotate +from embodichain.data_pipeline.dynamics.shards 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" + progress = ( + json.loads(progress_path.read_text()) + if cfg.resume and progress_path.exists() + else {"next_candidate": 0, "rejected": 0, "rejections": []} + ) + progress.setdefault("rejections", []) + 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: + 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"])) + 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 + 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 + 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() + progress.update(next_candidate=cursor, accepted=accepted) + _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/embodichain/data_pipeline/dynamics/shards.py b/embodichain/data_pipeline/dynamics/shards.py new file mode 100644 index 0000000000..0617787841 --- /dev/null +++ b/embodichain/data_pipeline/dynamics/shards.py @@ -0,0 +1,559 @@ +# ---------------------------------------------------------------------------- +# 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 metadata, schema, immutable shard hashes and split assignments.""" + return hashlib.sha256(_json_bytes(manifest)).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 the current shard atomically, including a content checksum.""" + if self._file is None: + 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/embodichain/learning/__init__.py b/embodichain/learning/__init__.py index 373314c615..b83c68b8a3 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", "dynamics"] def __getattr__(name: str) -> ModuleType: diff --git a/embodichain/learning/dynamics/__init__.py b/embodichain/learning/dynamics/__init__.py new file mode 100644 index 0000000000..b1f8e18a59 --- /dev/null +++ b/embodichain/learning/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. +# ---------------------------------------------------------------------------- + +"""Simulation-independent neural dynamics workflows.""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/embodichain/learning/dynamics/_numerics.py b/embodichain/learning/dynamics/_numerics.py new file mode 100644 index 0000000000..85d127da0b --- /dev/null +++ b/embodichain/learning/dynamics/_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/embodichain/learning/dynamics/_reference_traces.py b/embodichain/learning/dynamics/_reference_traces.py new file mode 100644 index 0000000000..def714dcfc --- /dev/null +++ b/embodichain/learning/dynamics/_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 embodichain.compute.dynamics._evaluation import _teacher_row_sha256 +from embodichain.data_pipeline.dynamics.shards 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/embodichain/learning/dynamics/_runtime.py b/embodichain/learning/dynamics/_runtime.py new file mode 100644 index 0000000000..ff0c3d1645 --- /dev/null +++ b/embodichain/learning/dynamics/_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/embodichain/learning/dynamics/_sample_order.py b/embodichain/learning/dynamics/_sample_order.py new file mode 100644 index 0000000000..593ae1b989 --- /dev/null +++ b/embodichain/learning/dynamics/_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/embodichain/learning/dynamics/comparison.py b/embodichain/learning/dynamics/comparison.py new file mode 100644 index 0000000000..74b5f5414d --- /dev/null +++ b/embodichain/learning/dynamics/comparison.py @@ -0,0 +1,525 @@ +# ---------------------------------------------------------------------------- +# 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.compute.dynamics import NeRDStateCodec + from embodichain.data_pipeline.dynamics.shards import DynamicsWindowDataset + from embodichain.learning.models.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) + module = ( + "embodichain.learning.dynamics.reference" + if trainer == "official" + else "embodichain.learning.dynamics.comparison" + ) + command = [python, "-m", module, "--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__)), + "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.data_pipeline.dynamics.shards 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"]) + pythonpaths.append(str(Path(__file__).resolve().parents[3])) + 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/embodichain/learning/dynamics/qualification.py b/embodichain/learning/dynamics/qualification.py new file mode 100644 index 0000000000..76ad6a6f1d --- /dev/null +++ b/embodichain/learning/dynamics/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/embodichain/learning/dynamics/reference.py b/embodichain/learning/dynamics/reference.py new file mode 100644 index 0000000000..6cfdff3dbb --- /dev/null +++ b/embodichain/learning/dynamics/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 ._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.compute.dynamics import NeRDStateCodec + from embodichain.data_pipeline.dynamics.shards import DynamicsWindowDataset + from embodichain.learning.models.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.compute.dynamics import NeRDStateCodec + from embodichain.data_pipeline.dynamics.shards 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.compute.dynamics import NeRDStateCodec + from embodichain.data_pipeline.dynamics.shards 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.compute.dynamics 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.compute.dynamics 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/embodichain/learning/dynamics/reference_rollout.py b/embodichain/learning/dynamics/reference_rollout.py new file mode 100644 index 0000000000..3226dd71f7 --- /dev/null +++ b/embodichain/learning/dynamics/reference_rollout.py @@ -0,0 +1,1116 @@ +# ---------------------------------------------------------------------------- +# 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.compute.dynamics 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 embodichain.compute.dynamics 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.compute.dynamics import NeRDStateCodec, reference_task_values + from embodichain.compute.dynamics._evaluation import ( + _error_values, + _safe_number, + _teacher_row_sha256, + ) + from embodichain.learning.models.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/embodichain/learning/dynamics/report.py b/embodichain/learning/dynamics/report.py new file mode 100644 index 0000000000..406fe33780 --- /dev/null +++ b/embodichain/learning/dynamics/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.data_pipeline.dynamics.shards 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/embodichain/learning/dynamics/trainer.py b/embodichain/learning/dynamics/trainer.py new file mode 100644 index 0000000000..e31c13ec59 --- /dev/null +++ b/embodichain/learning/dynamics/trainer.py @@ -0,0 +1,741 @@ +# ---------------------------------------------------------------------------- +# 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 +import os +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.data_pipeline.dynamics.shards import DynamicsWindowDataset +from ._sample_order import _CheckpointedOrder, _OfficialShuffleOrder, _open_order +from ._numerics import ( + _attention_context, + _checkpoint_numerical_profile, + _clip_grad_norm_, + _validate_numerical_profile, +) + +_PREPROCESSING_EXECUTION = "nerd-shared-frame-scripted-cuda-warmup-v2" + +__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 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.""" + 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 _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_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=torch.float32) + qd = batch["states"]["qd"].to(device=device, dtype=torch.float32) + 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, + } + + +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. + """ + + 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") + 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.manifest = manifest + self.global_step = 0 + self.best_validation = float("inf") + 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) + 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: + if source is None: + return 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") + self._sample_order_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 + + 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") + 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) + inputs, target, mask = encode_batch(batch, self.codec, self.cfg.device) + self.model.train() + self.optimizer.zero_grad(set_to_none=True) + lr = self._lr() + for group in self.optimizer.param_groups: + group["lr"] = lr + observed = None + with _attention_context(self.cfg.numerical_profile): + prediction = self.model(inputs) + loss = normalized_mse( + prediction, + target, + _model_variance(self.model), + mask if self.train_dataset.include_prefix else None, + ) + if self.cfg.numerical_profile == "default" and not torch.isfinite(loss): + raise FloatingPointError("Non-finite dynamics training loss") + loss.backward() + norm = _clip_grad_norm_( + self.model.parameters(), + self.cfg.max_grad_norm, + profile=self.cfg.numerical_profile, + error_if_nonfinite=self.cfg.numerical_profile == "default", + ) + if self.cfg.numerical_profile == "reference": + # One device-to-host observation covers loss, total norm and + # finite checks; no parameter norm is read back individually. + observed = torch.stack((loss.detach(), norm)).cpu().tolist() + if not all(math.isfinite(value) for value in observed): + raise FloatingPointError( + "Non-finite dynamics loss or gradient norm" + ) + self.optimizer.step() + self.sampler.set_state(candidate_sampler.get_state()) + self.global_step += 1 + return { + "step": self.global_step, + "loss": float(loss.detach()) if observed is None else observed[0], + "learning_rate": lr, + "grad_norm": float(norm) if observed is None else observed[1], + } + + def save(self, path: str | Path) -> None: + """Atomically save weights and exact committed training/RNG progress.""" + state = { + "format": "embodichain.nerd.training.v1", + "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, + "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.""" + 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("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", 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 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" + ) + 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): + self.sample_order = self._prepare_order(saved_order["path"]) + if self._sample_order_identity != saved_order: + raise ValueError("Resume sample order checksum mismatch") + else: + self.sample_order = saved_order + 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.sampler.set_state(state["sampler_state"]) + _restore_rng_state(state["rng"]) + + 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") + last: dict[str, Any] = {} + log_path = self.output_dir / "metrics.jsonl" + # Drop records newer than a restored checkpoint before appending. + if log_path.exists(): + temporary_log = log_path.with_suffix(".jsonl.pending") + with log_path.open() as source, temporary_log.open("w") as target: + for line in source: + try: + record = json.loads(line) + except json.JSONDecodeError: + break # Interrupted trailing write is not a committed record. + if record["step"] <= self.global_step: + target.write(line) + os.replace(temporary_log, log_path) + 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.save(self.output_dir / "best.pt") + 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/embodichain/learning/dynamics/workflow.py b/embodichain/learning/dynamics/workflow.py new file mode 100644 index 0000000000..4e152fe4be --- /dev/null +++ b/embodichain/learning/dynamics/workflow.py @@ -0,0 +1,257 @@ +# ---------------------------------------------------------------------------- +# 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 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.compute.dynamics import NeRDStateCodec + from embodichain.learning.models.nerd import NeRDModel, NeRDModelCfg, load_model + from embodichain.data_pipeline.dynamics.shards 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") + 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.compute.dynamics import NeRDStateCodec + from embodichain.learning.models.nerd import load_model + from embodichain.data_pipeline.dynamics.shards 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 a training checkpoint using the production model artifact writer.""" + import torch + from embodichain.learning.models.nerd import NeRDModel, NeRDModelCfg, export_model + from ._numerics import _checkpoint_numerical_profile + + state = torch.load(checkpoint, map_location="cpu", weights_only=True) + if state.get("format") != "embodichain.nerd.training.v1": + raise ValueError("Expected a dynamics training checkpoint") + profile = _checkpoint_numerical_profile(state) + model = NeRDModel(NeRDModelCfg(**state["model_cfg"])) + model.load_state_dict(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/embodichain/learning/models/nerd/THIRD_PARTY_LICENSES b/embodichain/learning/models/nerd/THIRD_PARTY_LICENSES new file mode 100644 index 0000000000..e30e79e719 --- /dev/null +++ b/embodichain/learning/models/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/embodichain/learning/models/nerd/__init__.py b/embodichain/learning/models/nerd/__init__.py new file mode 100644 index 0000000000..1d3ba033ea --- /dev/null +++ b/embodichain/learning/models/nerd/__init__.py @@ -0,0 +1,30 @@ +# ---------------------------------------------------------------------------- +# 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 networks, portable exports and trusted checkpoint conversion.""" + +from __future__ import annotations +from .model import NeRDModelCfg, NeRDModel, FeatureNormalizer +from .checkpoint import export_model, load_model, convert_legacy_checkpoint + +__all__ = [ + "NeRDModelCfg", + "NeRDModel", + "FeatureNormalizer", + "export_model", + "load_model", + "convert_legacy_checkpoint", +] diff --git a/embodichain/learning/models/nerd/checkpoint.py b/embodichain/learning/models/nerd/checkpoint.py new file mode 100644 index 0000000000..f161a55d18 --- /dev/null +++ b/embodichain/learning/models/nerd/checkpoint.py @@ -0,0 +1,225 @@ +# ---------------------------------------------------------------------------- +# 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 +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 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) + model = NeRDModel(cfg) + state = torch.load(path / "weights.pt", map_location="cpu", weights_only=True) + # Preserve exported floating dtype instead of silently downcasting to FP32. + dtype = next(v.dtype for v in state.values() if v.is_floating_point()) + model.to(dtype=dtype) + model.load_state_dict(state, strict=True) + return model.to(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/embodichain/learning/models/nerd/model.py b/embodichain/learning/models/nerd/model.py new file mode 100644 index 0000000000..3abd86c07b --- /dev/null +++ b/embodichain/learning/models/nerd/model.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. +# ---------------------------------------------------------------------------- + +"""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. + ``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 + 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] + 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_tasks/configs/tasks/classic_control/ant/README.md b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/README.md new file mode 100644 index 0000000000..d7d18c5508 --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/README.md @@ -0,0 +1,25 @@ +# Ant neural dynamics workflows + +These configuration files run the offline `embodichain dynamics` workflow; +they are not runnable Gym task deployments. + +```bash +embodichain dynamics generate --config extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.smoke.yaml +embodichain dynamics train --config extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.train.smoke.yaml +embodichain dynamics export --checkpoint outputs/dynamics/ant/smoke/train/best.pt --output outputs/dynamics/ant/smoke/export +embodichain dynamics eval --config extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.smoke.yaml +``` + +Generation calls Newton's built-in analytical solver directly. It does not +initialize DexSim or require its custom solver/Scene extension. The `full` and +`warp-reference` presets retain the recorded model/data scales and parameter +profiles; a profile does not promise identical dynamics across Newton versions. + +The pretrained evaluation preset reads the published data/model directories +shown in Features → Neural Dynamics Models. Downloaded datasets retain their +original teacher provenance. New direct-Newton datasets use a distinct physical +signature; replacing provenance fields does not align different teachers. + +See [Neural Dynamics Models](../../../../../../../docs/source/features/neural_dynamics.md) +for downloads, resume, fine-tuning and portable exports. Live neural simulation +and differentiable task integration are tracked separately in PR #822. diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.full.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.full.yaml new file mode 100644 index 0000000000..2e4a7e59ba --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/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/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.pretrained.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.pretrained.yaml new file mode 100644 index 0000000000..9a1c928796 --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/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/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.smoke.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.smoke.yaml new file mode 100644 index 0000000000..0caca1bfb6 --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/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/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.full.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.full.yaml new file mode 100644 index 0000000000..da72a9f087 --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/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/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.smoke.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.smoke.yaml new file mode 100644 index 0000000000..c134f5c0db --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/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/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.warp-reference.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.warp-reference.yaml new file mode 100644 index 0000000000..15ff72bec6 --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/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/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.train.full.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.train.full.yaml new file mode 100644 index 0000000000..212b3e9517 --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.train.full.yaml @@ -0,0 +1,38 @@ +# 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 + resident_cache_bytes: 4294967296 diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.train.smoke.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.train.smoke.yaml new file mode 100644 index 0000000000..8533943066 --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/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/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/README.md b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/README.md new file mode 100644 index 0000000000..5e0bac7331 --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/README.md @@ -0,0 +1,25 @@ +# Cartpole neural dynamics workflows + +These configuration files run the offline `embodichain dynamics` workflow; +they are not runnable Gym task deployments. + +```bash +embodichain dynamics generate --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.smoke.yaml +embodichain dynamics train --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.train.smoke.yaml +embodichain dynamics export --checkpoint outputs/dynamics/cartpole/smoke/train/best.pt --output outputs/dynamics/cartpole/smoke/export +embodichain dynamics eval --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.smoke.yaml +``` + +Generation calls Newton's built-in analytical solver directly. It does not +initialize DexSim or require its custom solver/Scene extension. The `full` and +`warp-reference` presets retain the recorded model/data scales and parameter +profiles; a profile does not promise identical dynamics across Newton versions. + +The pretrained evaluation preset reads the published data/model directories +shown in Features → Neural Dynamics Models. Downloaded datasets retain their +original teacher provenance. New direct-Newton datasets use a distinct physical +signature; replacing provenance fields does not align different teachers. + +See [Neural Dynamics Models](../../../../../../../docs/source/features/neural_dynamics.md) +for downloads, resume, fine-tuning and portable exports. Live neural simulation +and differentiable task integration are tracked separately in PR #822. diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.full.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.full.yaml new file mode 100644 index 0000000000..7cd4eef565 --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/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/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.pretrained.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.pretrained.yaml new file mode 100644 index 0000000000..73b9c4926b --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/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/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.smoke.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.smoke.yaml new file mode 100644 index 0000000000..6df395f4cb --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/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/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.full.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.full.yaml new file mode 100644 index 0000000000..8ee89e2b30 --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/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/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.smoke.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.smoke.yaml new file mode 100644 index 0000000000..c5e4d1431b --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/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/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.warp-reference.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.warp-reference.yaml new file mode 100644 index 0000000000..a7fd7e8b38 --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/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/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.train.full.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.train.full.yaml new file mode 100644 index 0000000000..cb287d7f1a --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/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/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.train.smoke.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.train.smoke.yaml new file mode 100644 index 0000000000..cf1815fde7 --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/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/scripts/benchmark/learning/dynamics_input.py b/scripts/benchmark/learning/dynamics_input.py new file mode 100644 index 0000000000..89772bd6ff --- /dev/null +++ b/scripts/benchmark/learning/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 scripts.benchmark.learning.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.data_pipeline.dynamics.shards 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.compute.dynamics import NeRDStateCodec + from embodichain.data_pipeline.dynamics.shards import DynamicsWindowDataset + from embodichain.learning.dynamics.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/scripts/benchmark/learning/dynamics_sampling.py b/scripts/benchmark/learning/dynamics_sampling.py new file mode 100644 index 0000000000..62360ee3a0 --- /dev/null +++ b/scripts/benchmark/learning/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 scripts.benchmark.learning.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 embodichain.learning.dynamics._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/scripts/benchmark/learning/dynamics_training.py b/scripts/benchmark/learning/dynamics_training.py new file mode 100644 index 0000000000..69dd5e9c65 --- /dev/null +++ b/scripts/benchmark/learning/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 scripts.benchmark.learning.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 embodichain.learning.dynamics.reference import reference_identity + from embodichain.learning.dynamics.trainer import normalized_mse + from embodichain.learning.models.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/scripts/experiments/neural_dynamics/README.md b/scripts/experiments/neural_dynamics/README.md new file mode 100644 index 0000000000..0cf9daa0ed --- /dev/null +++ b/scripts/experiments/neural_dynamics/README.md @@ -0,0 +1,51 @@ +# NeRD reference experiments + +These command-line tools compare an upstream NeRD implementation with the +EmbodiChain implementation. The normal dataset, training, export +workflow is documented in +[Neural Dynamics Models](../../../docs/source/features/neural_dynamics.md). +Downloadable datasets and inference models are linked there. Training logs, +checkpoints used for investigation, per-run reports and machine-specific plans +belong in an external output directory rather than this source tree. + +## Environments + +Run the unchanged upstream source in its supported Python/Torch environment. +Use the Warp release required by the selected upstream model. Run current +Newton collection and EmbodiChain training in a separate EmbodiChain environment. +The original-runtime helpers avoid importing the installed simulation SDK. +Use explicit interpreter and source paths; do not resolve a virtualenv executable +symlink to its base interpreter. + +External source/checkpoint execution requires the tools' explicit trust option. +Pin the source revision and keep any portability adapter separate from the +upstream numerical implementation. Use each tool's `--help` for its arguments. + +## Tool ownership + +`_portable_reference.py` and `_historical_physics.py` are private support modules +for the isolated original runtime. + +## Comparison contracts + +- Separate model/trainer changes, runtime kernels, coordinate transport and + analytical-solver versions instead of assigning every difference to migration. +- Preserve original sampling, input/output normalization, loss and model + selection when evaluating training-port parity. Equal random seeds alone do + not prove the same consumed sample sequence. +- Replay exactly the same realized open-loop efforts for predictive accuracy. + With feedback control, label comparisons that use each model's own state as + closed-loop behavior rather than same-action transition accuracy. +- Keep current Newton deployments and historical Warp models in their recorded + physical domains. Do not replace physics signatures or normalizers to make an + incompatible export load. +- Retain all requested trajectories and report nonfinite outputs and task + failures. Separate endpoint error from trajectory-time average error. +- Use the matching analytical solver as each model's accuracy reference. A + distance to another neural model is model similarity, not ground-truth error. +- State dataset size, update budget, evaluation population and horizon + independently. A README example or a complete optimizer budget does not prove + reproduction of a paper's training-data scale or final accuracy. + +These tools produce their outputs under caller-supplied directories. They do not +change the production numerical policy or provide a blanket accuracy acceptance. diff --git a/scripts/experiments/neural_dynamics/_historical_physics.py b/scripts/experiments/neural_dynamics/_historical_physics.py new file mode 100644 index 0000000000..f75718a722 --- /dev/null +++ b/scripts/experiments/neural_dynamics/_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/scripts/experiments/neural_dynamics/_portable_reference.py b/scripts/experiments/neural_dynamics/_portable_reference.py new file mode 100644 index 0000000000..c6cffc09ce --- /dev/null +++ b/scripts/experiments/neural_dynamics/_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 + +__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 / (name + ".py")) + 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/scripts/experiments/neural_dynamics/compare_official_model.py b/scripts/experiments/neural_dynamics/compare_official_model.py new file mode 100644 index 0000000000..81404b1979 --- /dev/null +++ b/scripts/experiments/neural_dynamics/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.compute.dynamics import NeRDStateCodec + from embodichain.learning.models.nerd import load_model + from embodichain.learning.dynamics.reference import ( + reference_identity, + import_legacy_datasets, + _official_batch, + ) + from embodichain.learning.dynamics.reference_rollout import ( + _cpu_stream_compatibility, + ) + from embodichain.learning.dynamics.trainer import encode_batch + from embodichain.data_pipeline.dynamics.shards 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/scripts/experiments/neural_dynamics/matched_small_budget.py b/scripts/experiments/neural_dynamics/matched_small_budget.py new file mode 100644 index 0000000000..ec4d2d8d35 --- /dev/null +++ b/scripts/experiments/neural_dynamics/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.compute.dynamics import NeRDStateCodec +from embodichain.data_pipeline.dynamics.shards import ( + DynamicsWindowDataset, + dataset_fingerprint, + load_manifest, +) +from embodichain.learning.dynamics.trainer import ( + DynamicsTrainCfg, + DynamicsTrainer, + encode_batch, + evaluate_model, + normalized_mse, +) +from embodichain.learning.dynamics.workflow import run_export +from embodichain.learning.models.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/scripts/experiments/neural_dynamics/official_evaluation_data.py b/scripts/experiments/neural_dynamics/official_evaluation_data.py new file mode 100644 index 0000000000..b409337763 --- /dev/null +++ b/scripts/experiments/neural_dynamics/official_evaluation_data.py @@ -0,0 +1,513 @@ +# ---------------------------------------------------------------------------- +# 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 _transport(directory: Path) -> tuple[Any, Any]: + package = types.ModuleType("_isolated_nerd_transport") + package.__path__ = [str(directory)] + sys.modules[package.__name__] = package + loaded = [] + for name in ("codec", "reference"): + spec = importlib.util.spec_from_file_location( + package.__name__ + "." + name, directory / (name + ".py") + ) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + 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.get( + "transport_root", + Path(__file__).resolve().parents[3] / "embodichain/compute/dynamics", + ) + ) + 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 / (name + ".py")), + "sha256": _hash(transport / (name + ".py")), + } + 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/scripts/experiments/neural_dynamics/official_math_rollout.py b/scripts/experiments/neural_dynamics/official_math_rollout.py new file mode 100644 index 0000000000..0f28e8e237 --- /dev/null +++ b/scripts/experiments/neural_dynamics/official_math_rollout.py @@ -0,0 +1,370 @@ +# ---------------------------------------------------------------------------- +# 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 +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", args.transport_root / "_evaluation.py" + ) + 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/scripts/experiments/neural_dynamics/official_model_reference.py b/scripts/experiments/neural_dynamics/official_model_reference.py new file mode 100644 index 0000000000..7d04d40807 --- /dev/null +++ b/scripts/experiments/neural_dynamics/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/scripts/experiments/neural_dynamics/official_small_budget.py b/scripts/experiments/neural_dynamics/official_small_budget.py new file mode 100644 index 0000000000..25b7af7b3e --- /dev/null +++ b/scripts/experiments/neural_dynamics/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/scripts/experiments/neural_dynamics/official_small_budget_evaluate.py b/scripts/experiments/neural_dynamics/official_small_budget_evaluate.py new file mode 100644 index 0000000000..8ab2ac9989 --- /dev/null +++ b/scripts/experiments/neural_dynamics/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/scripts/experiments/neural_dynamics/official_small_budget_rollout.py b/scripts/experiments/neural_dynamics/official_small_budget_rollout.py new file mode 100644 index 0000000000..f44fe077e9 --- /dev/null +++ b/scripts/experiments/neural_dynamics/official_small_budget_rollout.py @@ -0,0 +1,275 @@ +# ---------------------------------------------------------------------------- +# 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 + +__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", args.transport_root / "_evaluation.py" + ) + 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/setup.py b/setup.py index 367259147a..c0922dd293 100644 --- a/setup.py +++ b/setup.py @@ -177,6 +177,8 @@ def main(): package_dir=get_package_dir(), package_data={ "embodichain": ["VERSION"], + "embodichain.data.nerd": ["*.urdf", "*.xml", "*.md"], + "embodichain.learning.models.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/cli/test_dynamics.py b/tests/cli/test_dynamics.py new file mode 100644 index 0000000000..cf534b2da9 --- /dev/null +++ b/tests/cli/test_dynamics.py @@ -0,0 +1,298 @@ +# ---------------------------------------------------------------------------- +# 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 dynamics CLI routing, actual offline workflow and packaged configs.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +import subprocess +import sys +from types import ModuleType +from unittest.mock import Mock + +import numpy as np +import pytest +import yaml + +from embodichain.cli.dynamics import main + + +def test_help_and_dispatcher_are_simulation_free(): + code = "from embodichain.cli.main import main; import sys\ntry: main(['dynamics','--help'])\nexcept SystemExit as exc: assert exc.code == 0\nassert not any(k == 'dexsim' or k.startswith('embodichain.lab.sim') for k in sys.modules)" + subprocess.run( + [sys.executable, "-c", code], check=True, capture_output=True, text=True + ) + + +def test_generate_routing_and_path_resolution(tmp_path, monkeypatch): + host = ModuleType("embodichain.data_pipeline.dynamics.generation") + host.generate_dataset = Mock(return_value={"status": "ok"}) + monkeypatch.setitem(sys.modules, host.__name__, host) + config = tmp_path / "generate.yaml" + config.write_text("robot: cartpole\noutput_dir: data\ndevice: cpu\nresume: false\n") + main(["generate", "--config", str(config), "--seed", "7", "--resume"]) + received = host.generate_dataset.call_args.args[0] + assert received["output_dir"] == str(tmp_path / "data") + assert received["seed"] == 7 and received["resume"] is True + + +def test_invalid_paths_and_conversion_requires_explicit_trust(tmp_path): + config = tmp_path / "bad.yaml" + config.write_text("dataset: missing\n") + with pytest.raises(SystemExit) as error: + main(["train", "--config", str(config)]) + assert error.value.code == 2 + with pytest.raises(SystemExit) as error: + main( + [ + "convert-checkpoint", + "--checkpoint", + "model.pt", + "--output", + "export", + "--manifest", + "manifest.json", + ] + ) + assert error.value.code == 2 + + +def test_conversion_manifest_and_trust_reach_converter(tmp_path, monkeypatch): + import embodichain.learning.models.nerd as nerd + + checkpoint = tmp_path / "old.pt" + checkpoint.write_bytes(b"unused mock fixture") + manifest = tmp_path / "manifest.json" + manifest.write_text(json.dumps({"codec": {"robot": "cartpole"}, "dt": 1 / 60})) + convert = Mock() + monkeypatch.setattr(nerd, "convert_legacy_checkpoint", convert) + main( + [ + "convert-checkpoint", + "--checkpoint", + str(checkpoint), + "--output", + str(tmp_path / "export"), + "--manifest", + str(manifest), + "--trust-pickle", + ] + ) + assert convert.call_args.kwargs == {"trust_pickle": True} + assert convert.call_args.args[2]["codec"]["robot"] == "cartpole" + + +def test_real_offline_cli_train_export_eval(tmp_path): + from embodichain.data_pipeline.dynamics.shards import DynamicsShardWriter + + metadata = {"codec": {"robot": "cartpole"}, "dt": 1 / 60} + with DynamicsShardWriter(tmp_path / "data", metadata) as writer: + for split in ("train", "validation", "test"): + q = np.zeros((6, 2), np.float32) + q[:, 0] = np.arange(6) * 0.01 + qd = np.zeros_like(q) + qd[:, 0] = 0.6 + writer.add_trajectory( + split, + {"q": q, "qd": qd}, + {"joint_f": np.zeros((5, 2), np.float32)}, + split=split, + ) + training = { + "dataset": "data", + "output_dir": "train", + "model": { + "input_dims": {"states_embedding": 4, "joint_f": 2}, + "output_dim": 4, + "n_embd": 12, + "n_layer": 1, + "n_head": 3, + "block_size": 3, + "history_length": 3, + "head_hidden_dims": [12], + }, + "train": { + "total_steps": 2, + "batch_size": 2, + "device": "cpu", + "validation_interval": 1, + }, + } + train_cfg = tmp_path / "train.json" + train_cfg.write_text(json.dumps(training)) + evaluation = { + "mode": "offline", + "dataset": "data", + "model": "export", + "output": "reports/evaluation/result.json", + "device": "cpu", + } + eval_cfg = tmp_path / "eval.yaml" + eval_cfg.write_text(yaml.safe_dump(evaluation)) + code = """import sys +from embodichain.cli.main import main +main(['dynamics','train','--config',sys.argv[1]]) +main(['dynamics','export','--checkpoint',sys.argv[2],'--output',sys.argv[3]]) +main(['dynamics','eval','--config',sys.argv[4]]) +assert not any(k == 'dexsim' or k.startswith('embodichain.lab.sim') for k in sys.modules) +""" + env = {**os.environ, "OMP_NUM_THREADS": "1", "MKL_NUM_THREADS": "1"} + result = subprocess.run( + [ + sys.executable, + "-c", + code, + str(train_cfg), + str(tmp_path / "train" / "best.pt"), + str(tmp_path / "export"), + str(eval_cfg), + ], + capture_output=True, + text=True, + env=env, + ) + assert result.returncode == 0, result.stdout + result.stderr + report = json.loads((tmp_path / "reports/evaluation/result.json").read_text()) + assert report["split"] == "test" and report["windows"] == 3 + + +def test_packaged_workflow_configs_decode_through_cli(tmp_path, monkeypatch): + """Exercise the actual YAML loader/path resolver and model/train constructors.""" + import embodichain.learning.dynamics.workflow as workflow + from embodichain.learning.dynamics.trainer import DynamicsTrainCfg + from embodichain.learning.models.nerd import NeRDModelCfg + + host = ModuleType("embodichain.data_pipeline.dynamics.generation") + host.generate_dataset = Mock(return_value={"status": "not_run"}) + monkeypatch.setitem(sys.modules, host.__name__, host) + + def train(config): + model = NeRDModelCfg(**config["model"]) + model.validate() + cfg = DynamicsTrainCfg(**config["train"]) + cfg.validate() + assert "id" not in config + if cfg.total_steps > 100: + assert cfg.total_steps == 999 * 5000 + assert cfg.lr_update_interval == 5000 and cfg.lr_offset == 1 + return {"status": "not_run"} + + monkeypatch.setattr(workflow, "run_train", train) + monkeypatch.setattr(workflow, "run_eval", lambda config: {"status": "not_run"}) + repo = Path(__file__).resolve().parents[2] + (tmp_path / "data").mkdir() + (tmp_path / "export").mkdir() + files = [] + expected_configs = { + "dynamics.generate.smoke.yaml", + "dynamics.generate.full.yaml", + "dynamics.generate.warp-reference.yaml", + "dynamics.train.smoke.yaml", + "dynamics.train.full.yaml", + "dynamics.eval.smoke.yaml", + "dynamics.eval.full.yaml", + "dynamics.eval.pretrained.yaml", + } + for robot in ("cartpole", "ant"): + robot_configs = list( + ( + repo + / "extensions" + / "embodichain_tasks" + / "configs" + / "tasks" + / "classic_control" + / robot + ).glob("dynamics.*.yaml") + ) + assert {path.name for path in robot_configs} == expected_configs + files.extend(robot_configs) + for source in files: + config = yaml.safe_load(source.read_text()) + # Temporary destinations let path validation run without generating data. + if "dataset" in config: + config["dataset"] = "data" + if isinstance(config.get("model"), str): + config["model"] = "export" + if "output_dir" in config: + config["output_dir"] = "out" + if "output" in config: + config["output"] = "report.json" + copied = tmp_path / source.name + copied.write_text(yaml.safe_dump(config, sort_keys=False)) + operation = source.name.split(".")[1] + main([operation, "--config", str(copied)]) + + +def test_build_py_packages_reference_assets_and_workflow_resources(tmp_path): + """Use actual setup.py package-data rules, without compiling native modules.""" + import ast + import shutil + from setuptools import Distribution + from setuptools.command.build_py import build_py + + repo = Path(__file__).resolve().parents[2] + tree = ast.parse((repo / "setup.py").read_text()) + call = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "setup" + ) + package_data = ast.literal_eval( + next(item.value for item in call.keywords if item.arg == "package_data") + ) + source = tmp_path / "source" + assets = source / "nerd" + shutil.copytree(repo / "embodichain" / "data" / "nerd", assets) + configs = source / "configs" + configs.mkdir() + real_configs = repo / "extensions" / "embodichain_tasks" / "configs" + shutil.copyfile(real_configs / "__init__.py", configs / "__init__.py") + expected = [] + for robot in ("cartpole", "ant"): + relative = Path("tasks") / "classic_control" / robot + shutil.copytree(real_configs / relative, configs / relative) + expected.extend( + Path("embodichain_tasks/configs") / relative / file.name + for file in (configs / relative).iterdir() + if file.suffix in {".yaml", ".md"} + ) + expected.extend( + Path("embodichain/data/nerd") / name + for name in ("ant.xml", "cartpole_single.urdf", "ATTRIBUTION.md", "LICENSE.md") + ) + distribution = Distribution( + { + "packages": ["embodichain.data.nerd", "embodichain_tasks.configs"], + "package_dir": { + "embodichain.data.nerd": str(assets), + "embodichain_tasks.configs": str(configs), + }, + "package_data": package_data, + "include_package_data": False, + } + ) + distribution.script_name = str(repo / "setup.py") + command = build_py(distribution) + command.build_lib = str(tmp_path / "build") + command.ensure_finalized() + command.run() + assert all((Path(command.build_lib) / path).is_file() for path in expected) diff --git a/tests/compute/dynamics/data/ant-official-frame.json b/tests/compute/dynamics/data/ant-official-frame.json new file mode 100644 index 0000000000..8737a5f7c8 --- /dev/null +++ b/tests/compute/dynamics/data/ant-official-frame.json @@ -0,0 +1,64 @@ +{ + "source_commit": "370e89d2b8ce68aa3578aa7812859c681b7502bc", + "robot": "ant", + "python": "3.8.20", + "torch": "2.2.2+cpu", + "warp": "1.5.1", + "numpy": "1.24.4", + "device": "cpu", + "dtype": "float32", + "count": 32, + "horizon": 100, + "checkpoint_sha256": "be7e29f8834ab22d76efd3d101196c77b3b415c3ef91af48aed233c01eb4e950", + "seed": 500, + "runtime_imports_embodichain": false, + "cpu_stream_adapter": "CPU stream request mapped to None; numerical source unchanged", + "gravity_transport": { + "direction": [ + -0.0, + -1.0, + -0.0 + ], + "raw_preserved": true, + "reason": "Official sampler leaves non-up-axis gravity components uninitialized" + }, + "files": { + "official-windows.npz": "2ff996b31e0622e828254bda6703a1bb608dd8e451c20b0b9c4ad0e5fb242cd3", + "test.raw.h5": "2d86248861ace090f1d709f2f65ee96a1ce7a2c7e7a1f19eca5f30929ac940da", + "official-rollout.npz": "b315b2cd360729ddfd175be91bed5d137406c50c60b6e8937e48be98422a1431" + }, + "model_inputs": [ + "states_embedding", + "contact_normals", + "contact_points_1", + "contact_depths", + "joint_acts", + "gravity_dir" + ], + "state_frame": "body", + "body_com": [ + 0.0, + 0.0, + 0.0 + ], + "scope": "Four real source-world windows from official Torch2.2.2/Warp1.5.1; tests only codec preprocessing and targets", + "indices": [ + [ + 30, + 29 + ], + [ + 13, + 0 + ], + [ + 12, + 47 + ], + [ + 21, + 85 + ] + ], + "fixture_sha256": "4d4c4a1db07b35ba06eb94b810621c84d36d2799347864e443adaafb8383c8a2" +} diff --git a/tests/compute/dynamics/data/ant-official-frame.npz b/tests/compute/dynamics/data/ant-official-frame.npz new file mode 100644 index 0000000000000000000000000000000000000000..9c39fbd51a0a3f5ba9f8dc2193d74ff1d1d4c9fd GIT binary patch literal 34898 zcmb5UV~n6d)9*dlvF({1+qP}nwr$(CZSUB&ZQIt_J;{A??(^pP@cxs&K6F*m)!kQD zeyMIbN#I||0002se;Yo)`^NS+Ed&6-?%xIifDB+yV`JwI2mt@TJN|kA{4@SjcH#%; zqC4L32)=y8sQPrM6=9%J;#+6W0UP^tT%6I7oTh*b@e@@;JvEM*vzvlAadC@vD=G8w zLkDPwT8of+*hwjCp80(gm%bzi> zAKz34Ue}FYndh*>m5JfktvN!i&?XIqs}98dK#Wz>HNmDG1bA%@VdH@#EmwaaosA~N ztaA^%6r~h~`laQ(KFs(@yrlta71clZwiIE_2xti^P`eM8So>{>aGz`lOKdixGI2fg?77b4v;)ZjC5b=qZwzOHKAQO@Wn?F^IxYaCGv~egu{va^IIwDj~6_$4{2x z8*lg-$OX!|3pyYJa1!aRMC*ql#qo zp3DPsV8KaJccSJl1s}|;WFh`|CecTlPU?RZY4tm(1ufI-1ZAy@%Cx)jF^lk7Heq^g z8H3c@jF;W_$7*S?@v}x4l;<4+1c8sIz3pTOj%l_X=CzPq+w2H%#TX?%K5&J(;{lDs zS@ktZAt{kW&vx^sbsqW_OeA*~V#>75mg@{pE@lSaaJL!tl?cCbxU=q!_zKpv-TRTM zLp5XDlA)Fd+Kf6_(Xoq|Fe}^r9#6s4APQ{iObWMb^&19u7*p>v$f}tl_QZn>m!R+q zcLmt2+yTgKR@Rp?#^^pzOc-DAKrd6L>~3TL)wq(tXbw%uJEcff`Yey7eR>EV76Ej3 zGiQ#%HQ9X_W=pHb#Xq@#zDv>bo5ZmfNyKq~0);BrZM}+~Y_$*Nzx^c4u?A(5mpG}g zP!4OJMeRiGA>F$qabvRtBX|l`dt^#>aJd>NzOhM+2C0a{w=}Ath6`LY1Hx(rJV!rt zoPx<5bgyrbOI#Dm(T)$;*VYE_T#E5@yq+a*yjrcJN>$oY3USIL5qZ1BU#+(R!V+mi ze>o>ntqC@$#F6YXD2T3X`E(^qBPF&`_stW?!6i9<(+@ddxlZA(uw!luze~mmdT}L% z&)l3VJZ4`hg&o2;c9~5SJWW)c!EYqUFrvlnMY`y`-DF@nZ+3~O{mzGyKah@kRf8A> zC8@VVKx=KzJvkTkLCoFkL9V%#o$qP;va_hlP=kIVe1Qyf&Nx>anJaDwp_xskG@gI= z398ffHGBqky)2#E++m(NdnRn5vsx70*~SVotxeX8SHQ!awjbUwJa7_=7PWl;%-?QG ze|zBN<#Rk@GJ#S(7<}<<<2}qE4d)>(tlZ}OFd!SN58<^yn7oLe#9Oh{P8#<-RInK-nr7<_Pg;IlhqI zO+*dtH9;y`-gg`Geg1rLAp5!%!4aNubf!7_DF;wJsj+tU^3ywC#=cXe{JDgQ-TLJYb z-ESIOXqN~#J?9anNJU!}eHyID-p3|u_Uoj?$xY}uMqYL~e)`(z-X;dnP9lgXuyis@ zEH>y*zu_r`PIS$|m)hl7U*>fzvHh^_=h9KV#b);(0wVYDF`;uS=;aU6u^4k< zlVxIT8xNfRWg0l94`*;roIv}&CDJ^m>qrk&$z8>E8e6JQ%Pe%*-_PmgMT?`zph#MM zjM+n*<7j@djtDdvCxHxD>@~9p``9{^Ma~|E?$d~Ds-+W>V@~`vKH&8G*Z48oCVEIc zh%L0AVBHcVy)D$evE^P=9N(xTp34O?T2Rm>o_iQS#J5uNv8)8H{Xh%w{T7qNn$=W* zD(wOCt4vvs_Q8HrZ?Df$tf-mY1KHC)Y{uK*6PjG|DmKuvz)x;5&akW9%n@c<>!KN6 zRrse3>{3*j{=MGQ@U3m&C@f$HgQ2dct5?x&Bu{aScA@nhzBx`Ptm4oKY-ujm)LX)W zXQh!YdN_jv23sqGmQ(=H*Di7gNf>vAd+5n$Q~r_kRCPK822H*lVp_P&;!zwpbty

^WYWwH>-iX+nILYD@)gec~4 zUHMKVbk9-vPa7_*;+jME&4$SKRB!!_y7EF-pC98q;jX~j_^BRD!+vTBN}A3I+~p{e zCym)wK;V}+qRZh`Jn zC5Xpb+1V%t#HJtJzZO?Ll!iSeN+iH zDGRp~e~EJw-mwWD-kXHiMj*vNRh{}L!a8+N;_`4q2xiBN9@6;27ixvKH(CpA2>;A1215;N{i@D7GD2K< zmReSI2Av8xx7Yme3hzT?)>>iFZTM6t&RNBvVCNw-(QJ+&!jL|ZTV+g+n!W0Dxv3ol+di4CJoZ3+y%8FqeUBn zp{-eH9~EZWTcB!ETrg32TBiDHgIbz-!yY*UTE1eAhnZ#slT@~BV|rtXXW?{Q#)r>2 zhluseN*mE+^=c~z^BI4OS+|WGq*Wo(cYG&)m3BpW?Hm_St`fRbE)hszMw6*MGWXAh zw@{HyWwqQDf8z46KcGJ4&Wp4s9(MA*fzxsxaX;123S{Zw2#|{ss^E!)ilCH!;Fjfo zm=eFtIFwv;jb~CZC2sWeG(xYX=__+!#|NHjrKcE3TvuZA+nxh&Hw?iJr1cE2n-3FPfI}!0kfvJW z!r=90fC{XMnU9GjpL%>ps7{RtF)u%$YP zZp{Sdn0pOUb$TK?xIky}s+ooKM%a=ash9XMW4N2>=yrE%K%u;JLUngi(%}DM5ZEuG z|l_hxr(`LMlu8qGo z0D0A1$G8`f#XDu1Zu>Hd(j9xtA176#z`*a|u3E&U{TZjlbuG3wu;ZR?PQT(^4o2%q zC4V7;4&Nbuk9x-7lK%x~hC}7&xuQL{Z_W<&Wrfu}$~#hXJznZ-6YrPDh*gEsM)#Ho zi5=fXsBL+pcag<`8pltxwlaA6MQO>xp(%^V@q~lTD3@rqwIx$5sv0=DPg$lhpjTs! zf{M)p(y?D7unFm*yld4R}Kk_e%Z>Do71ks*?i z9TK99rI+!+(gBCJmOzguhh< zi;})G;46SrCy)VzQ@M@w)tFqReM7fV_rsj6K{miwn%Vg*Oi+ze0or%PRxcwLx}2c! z>D~VsP>JVB;u0TknMaZWdqnnk1A>nY|CHu zJJ1FOPW&-|o((F~ZJ&?%1b<;Q_1WjXNPyYp{@suuyIhJsFK(4)OWe4RoGIN%uy?>K z&Ez$?C`t!)tvHW$Jh_+K$riKU6R7gmKdY{n`zwK&w_UWz_CN; zTQ9otFkX;{e9;DfnTBMtoH8@RSH({Oj9cxr8OU+xQ0{D=?tYXyZhG@Mr-n8Xuu-(z zh^|&80Wy)Un>J#^1;@15g#sO^I zsP26O_S(7BFt-wlk#&{`O8L8B+M9r^@~gexiY^^|hBMR3c+H7OiC?k8HgUfF<1;f& zq<${nUJPCR*T;Kt-Qt-_VTdSX(L0R+J6;i45~RWL?!Xe1Pf_Y;?EFSFw^NrmOci!I zd2KXV{#x#|-+i*>-JeKcL{o$&Y`MY=M%ROiKE|S#PksU_`jE%K#It-tz-UfXlm(2l z;J8{keSBZ!x$QK(Q=}nkk5&Oa5{Q@HHQJzO4OhhSQu1@evq5K&oJa)U!X88Q>5v_>2XUh zd1KU7&1X9zy!%;2{x)dG4K>}cOglBQ!QiI#d4&@oDI-7IZppcHG@u>i&2_9asp2?> z`c8qI+(yWNFZ)85dm(~kYRPiV*5j`g#}MJ&xx;JXhHT_z{q>JWOjAd?Ie3rUlR4e6 zO^BaENa3K|lIY4VKQLYdt_6 zvwKD;v+Qnu*Q)ZRCDtGEJBOfQHqi;&$~9bWg#nV>z;7{^6dosCNI!#OkbPTAa=8!O ze|$axMMfD5rvd;5V7_%PqV>i-OW;@pEPRTh{LHCq2yQ-8a>0+A@}@ zPy-KLxzE>4`vON4>AhcRUuplFZt^BmlIj0)%)c5Ektsa;I}cbMf?}xZgZ9#U`fIP7XZu781s|cWz>HW0mU6~K(?CU_ zB)X$K+B6{&W@M3soZHdW#7jfVWWV$)9cm=$;<3YQIYqyg!Jf!hJkNYkzXK z!yhF_lmg4oyE+R)wk1gV08%_MgoY;R`Sx)PqQL4JWcXa(ZdC$8i5Ifh!5B??`%jLO z4t61Dzf;D85i^}nZPr@P7|^#zQQYL_Y`^QCq}b+-F;UF`R*ZIpUxfhk?fId9$cAMG z1r%w>l%Z4B8X>M-tsmn#%H#v0BioY_J_hbWKka+fSh?~}T<&ETz7{v)G8_e<=J~Ib zo6@tjJmiinJp~le)WXfqsWskxDi^foW`_Qthr?1`&SIDp0RE zQ;56kin3=SKP`0@X;#)DMR|3Ipm2(`p&z=yhE&Kg1a0a~Qb~NB5l8W5VV=P*Sl}g3 z7CnW05R02tGuF-+E|l)lX9`_;=~Ay#ZH2x`=1l+H8i59|1DZvH##nEHVH@ zc?a&=JdnJWRl?dXsvJDvlb7>Y`t@lz?+-KE&+2PVn{BoN7_$Z!ou}jTFjEm_WC|aC zqrjzc4fl>rLC~2p^V2NX4r&%4G}!fHM7yFe)v93v<9Ws~4veS3N7WLMaFnFks0*n{ z>3iz|Tt?>|?zUq7N#d5pY#UQs#F_D2N(SgSqMkI4H2iHWYdGP;*M$nHf2K`TO+~$*242^Ct=UOQ-s5?9qFcOnv|JU8GoWZS)SV$*~L0-m@|<} z%FQRSbe6yVsAi_ut39z){|6ufRXp5t+E=c;7a$Q5$d1*Olh(M0Lrc=q_}4w*?fWy6 zK_#ij63_A5QlVRGIyLkm!u#!;;b=_+aodFGZ3hgGn?5rb42Q?i%6B0gy<&CLp9bhd z!kap^%+L*EH209Xo1xO4Z|E9G*#)>xeYS^V26{X!k!@wri{=!)u0ozAT8H3#FY-k8 zN7=B|4?xEo_kFtPBkoy-LeAKYUN+CL?oLj~*l#&N>tuYcl#?hJ+mh2X(A;HM5P?Q< ztcS~;9X`z`s^buk-N$@hZamoc1(7!9KQZgOMtf|7-&1~=z5va?O3Jcc5>ggX<=js} z%X$`ui#pH*ujAP}<9blxe7Xl0?v`Qy)`T!EVuJVY+R*G|SbxES8@sV-U`4(>Ddum$ zaM;m<5(r;OS<542>C%V#Guxc*kXxbVQ0Gq(NRq+A3{O)J zHuz|4A;Db1ZD7dD0ABjiM`m(r&C>BHW}SHZr#PqcpH2Tp#=-NG!@(o7Zj0vr)@HId z%US<^P9fn(UxoBFRNy%V{FqQrBDSliC|CRBTNc{^MHhU5^mQt`2y{x-gEv|-WbH`U z-YkhyMUzrkfa$9vrvm~zfblae9kJ~h)O#LhGiYl|yG&aOnwYsThgZ!G&jt^SS+kXo z!-_1<)wVw zQDr7ZRwKk(57sUXV_4n_ph|rWtfTTfg}Cbdu$CuR`Qm#lSxzA?Nu+4HviJPL)o7)P zf$;_tg4H|5Kr|OE zn`wQTW-!66BTv+Nw)vkQ%7SMB@4slKI z>inZ{CUjjMumNx!f&%yYN)kGu3%0uWtvhHB(V6&EOJG=7tGBzlpkA$@urXHB#HCnp zh=ht*b-E9Jr;sLl0~)HkX1GHz-}-Kq)`)NzZh4oNw0g#igQwr`-1b)%$#tGMb<^|- z5~jh6y5H<|=Rs-m?6QX(`0#Nap%u|IMKUpgPP(&L_ZNoBzj>sics?M*R}v)pJ8*96 zt=V974kT?4t3w0}ejBBKOJ^Mi^pW;t$6%VzpC=)#ceUmo5|{Z&h#)b#|1i))B}zDz zV$Zf|@Oc-2YkOw{JE*hCs7NE~C&P${i$UNXKU?xU68A+WPncuX?{U{~|O~ zm@aNxO|npLoiIeL<@7ZbPk~We{A>U_d%=I|Q=LmB48nICp+C;pQm|yXJ&y#-%_Ty| z&jicFIAn1$HkQE>PPGV88Df1_a3Pb_F$)`!8oZyNgprPe#+Pye&#g6_-#nzqv}=)M z>O5(gKxS@2q;nD`Fh_z=#=e&{nvjy^k{5a8a{I0S>s;C?enY@_l$}WcN%E>J?hEa! z{LHdmcpiMM65qrHHix)k`YHi#T}G|nc!BrAn*2l6bGCXjz)6!LX#z>ze-GP+@J`oq5DsfTz-*;Q?;1{jAiD%oe$571#!EiSQA%i$w|r>!PeGK zzjvq34PML(+=%SR_l(NxJ4hvXX0Xvk*l1Am-7*x>GOqkX^%e25eEnVK{)!cJbaZUk zXRdB12?_pvwyat!3C+EBYq~A1y^F_)m!_w9#NkhPNu4DJMlpG(tDFq%SKe5S&wWA& zvGs6(^YEhP{#edpw_hY&GyUz@4HNGZyzehvy$^N4Vz((>Of=oWHQAHmX4%@gb9CDF zqc3cy#aUw(W*4f5a8Bm+7krsnkg6?L2CNKQkmYk13Hr8(Fx5>qO9XJ)7k@DB8ml`> zS+ACSFxbnF2oCGmzSvaKz2r%#+t6=`Ti=|(b2r3#)nrI<7p*dl2}S`8{LO|CA0J_; z8MEWLhYMT&8wBr7B_k2IXP=3so)47u6qUTM8=BwPA_NxGdjcoL@3CIf{mGe%tIBHx zOb)3_$tOgOHQEkRi%I%c;h?f~KP%YK z*64-Q)WTHMqPZ9&+I}-#&c+{woAU2#=$=nhzZ@-D{fY|1$fl#?dh_h*MlpyhJQdwE zbuo{*s*E*%?)x{58C=-(xD~bKqDnP+^qY_O+*XGHgNo1;_ENyUG!TODKvX1bD{cxV zdT{-pfl!)8*+SA2kEn#y`f5;nhRUV|RJ0$G{-i?v4vV$S#%@s@i$wIgey5HeKhJ+b zww$v}=~g%rhdD1ESAXFWG&zbZ^Qy_7&ZBbpklzYxtma*&1lqymm zYrP>S#jz~SC-kGeOx4@ZOK)J;OPCznCysgQDw6R(Dq?L6iex6)ZN#-)G>-w; zXA@o{^BMN(xWhBUr0}b;#ASQYhk(j$!^u7V>Bm%fIJ=r+fPb{=t6@ZI+Myhhf3uu# zd4^^2mYWJ_46~Y1Cii^UjEP>u9PKad2|*ghWXSFySK{i1ByvpUTeARtz#gyLT~%Lw z8|P!!at?1{Umx2~tvMOL=(ZZ(j~%;QnV_3L2R1QO5gK`F;5Rj-h^nb7EA(WTEuz9S zal26+FK+xz*(bJqOl8U}T_c6dA~cd{Teu^_i7Es;8*%{&@9=QZBbx!4p0c!>{sr@i zbsbe~fZ0^#B^WiGZNv&4Q-2yrd-^EA9`a}VFj=;BetgdXo=X;?gF0hYNhylA=KL3T zB)Qf0r{3mgVOTN^K?gxlV_=JGpOYptERl(z@gnB=$u$3I*<+pNCqmYfh49RU#pyu@ zs?ua<{e}}F7If~w=JlpgB8mf7c@Llguugz(6)(zUJ6{zUY0t~@z*`Prr*{LuCX9@t zVq&q=+?FVEw2yK0+!WRPgQUf1X1SFQ7#a+)*y_2r(AlXh+SO&zio!Xc;hqHBYhFAE z8Yh2SC5)s5uYd_y%siA2>UheWlLvt>W(56}P<>&>XO<>k<_N2NXq;x;@6mj1%X%uC zKDMkt7_SRG_7h8!I$;_A@ge~jm*aI{f@zgvbS*LLl&ie^I;IS%XTIAmuQdGzp_1P) zzg02feBK=4s!rk8RM%j~^OOv=a$m9cedm6B5YpgxXmL`Favtjh|ZyYAQ+hQXo2u8hU;5j@$u-au1;8I?IN9A;isc=g4 z41B3%%{Yf5aF*zJWbj5{VfvH!z^t~Z9gJFrD;COhfv5 zy_j{lyOUMajs*K#u+dAlIvvS?4nma6ZuK&ngeU_;jJ1k@OL`P^t8oUGj3`&RT!Lzz zr7mCWOjU*z#{`;r%Qbs`At&c9lZPaP-xf7*^LK~Fgt*Q(-NPzqZB6xr)h1@(X8TbW z>eTu21sT`0o;RVE@_YE9ZM>)m<*UUpnEKYJ*QM%!bI5R>dz^ASK&sJmXsTl<3%zvm zLP@Ap4Jlqvi)|9kx(Nl!WOpI06=!o<)uQFDdikj^R~587VwwrWQ$j;(Ovo~?c^vX$6nv|Fw z4X;0}(rlEcJt`zGrHJmb9m*A!h8J1AGdb7_`Q1?<-Af9PyEET1{^<4RN>)AmFhaAI+x*m;ryPNRQm|PRki&zfhU?qe?VroS2_3UQq5o2%Mbb{xtJcKIZGj1* zok?`5^C*vO6;U?qqU@z}hZcS@y?Ysjk%h3hCv<$mH1LI&gZB=zis=Kxp&3w=7>>#< zimv?h4g6{7lBcPTntb>r9S3rUo(@9v{v*wm@pv{Ri6Hc>1|XOwxwx=&D%OB`>g^kZ ztF%yrx&l`a@6IS@6CJ;ujLT)^O^W6%@apP0ewxi%^U1A<(7!&BZ?eO>duIctCDY#@ z|M8Ul$i?E}Iji-Wx*H_M;WCWy?R+XoauvEwDYIywYq)fqaw(OD^`i zjmwGvhf(y~>qXw0;#*HS7Rh_FhBAXIszVy-M{UG*!zU`k1{GdI3FPh`wxCg{moxo= zjl3%g7|dt?2jssI2XZFJ2MYuM!0&JShd36t<~B~cCjYrI``7(X|6KklJN|=r$rXDf z+*v;TC@5ojZipXQm`4A^YukZ5YF1_<#32(oMk|HO!3T`Sa@O^X$ou3(7mt>+{*%g2=MSzMl{g4w6mAD#o#h zdL&^gJPHXLB?t#194Scj^JHtl0EzI}+rUZ9Yw(k$vDs)tO|t=Kl6vjz@+TPfhYMJj z=y^ecx$&3L~wq0HDYz{APq7bE{ zOS6ue`j^_uB);2_Y}}=g{C;xmzOPo^pXsnL(<;BOQ;I!Pm{~}FWNeY{VZ7#pIPakk$^$ z*^VZV&ZA(ZxWgz{l9k=vWJ0n^kLic!ZqfDi&{4Z7g}M|Hqy}H59C`&n(=b+aiNk5z==y?lWbQY~&m#Silq||y(`;Md8=qF#~Sta=i5XMs2 zz$Y3r{$?9wPg?1bm<)_=L+xRBJShe71MpqGc5%p4)o<8<1v}JyI`^VfoeUkc-*ZMy;t>4EP~+k**j*_+LbyN2I-4 zaSIER!U_3`rI%krPz9!5_=J3+7!ZE~Ogggj8c0EJz$oipw@k(v#(}Fv@zolQAjEW< zkxH^;eXK16n#ZIGJXM1mA+qhfs>#`CYK6Ne#F2Yi^QWDG;;z^MbB=N8?o-1P?8p4; zi>r<9+XC9xRs#HMab&n>Nhb{6#gCswyf)l0wJ&J(oJL42dh&sTBcbYAonq$Y4XreF zC9WYo$y8lq(EgMMD6Lfw%H^u*RJ@~RnD>xcw2I?u_;pN^SD8PS*wMw`{X7u-R;5R$~0OH@U3~g#Zg~EF)E1qQFwOp@L>fdiCFxMyqvGETf8P;_nfYslg8!m zt0_&~PO96lxk5ydD#dj^A@lmLzpMxfPhZTB<#tEPn?TBBmfxeW%K}4X#(Bgf~{w2G^G3C1o zSBOjrG+?i_h=xtYNZ4#|HJv;%nQ|*a&*7&V(w+5J**HKuC?)U>nJ504$D_T`GsJo( z36$s$2MSEgs#&{G2a~~O^;>smR^4Kcwys+TO>6Y{S@Ab^jsq@zXgs&uPHlFpRrmFk%#lfBvo~kV7w)*y#lSyAX{-=5IBHM z16H0S?1KBSp&>Jv`csC?RI0PVU8i07IB(gaWf8sgzK~=xeY{~jDc(zyv>_g%+ zZ5$9xzK>yWp&WoczcImVq+9^-ygSjO6$3Cjhh)Q9o=XvBy1ne;7DzZt)8lis*>E+} zRnZCYzQ={l*v=p*Qp)=LB9VG6b6oU>H@eMWsd4rX=X}dV*Xv=jxSAhON~K5I@i5$j zN`GyfsFu*(KdBz|L!@HtK5d%+-3P4OxnG=g-H!MfMU$CXn|I};>nOrgSzdjqBCBOM z9O!MFsd^h+^a=4jYbEg4JP!ac;tPDuE(GY-yaJ7h4|AJyTDymm30=V&^r{QK?*MAFQ@|r-k-g0kviu|0v2%;2+sB8fvXQgf zn#YiEDN6c;ILW_$7U*F%lt)hbcFv}ebux?nG=c7Ws|qpYT8MG73b$iS3#gJQzG_0*kV%X$>RT}ttOeb^4>t1}j+CW7fM z<~_}m@7-Wt)Sc0e224mR)O+*F!o>D?i`}yB0DpT>K~Zt214vXY^pbYt6{P2P4UGBv z$ObM?$^zO(R}T&!wtBe@wga!`*knWhd@b_#k=vvFk@z&-r>Y&$(=vs0JJv(a0Gr3k z`UK;JS?_*>blY>)usjlkh;+3sshq!gX2g2ghEPUHB^zqBr9;;ymM5ooaXSOy9q!4MLHxq(gPgjN^ zEdxEEWUe_pal+)J(0fA=_Jv0~Mmo}*{?lw#$I5rBBnXV;5}`$B9U!Du67Z(j*xwV| zC13@PMzGPYu1RI%x5f|c$bY+j@2U+g zda6Y4k%!~#r7p$oJ`{-C$GoIW^oA=&)aExQRMUWedeo;^8WzjU%lFu*=K4;*3)^Uj zF5*|_8LYs7n&YiP49IGm2#AlnFrbX3+f=!_9cOU-4qlxAr-dfM1(J8G5hTlc_zU{m zw~eip=W?2(gLquwt>luMol7L~pW9?gw}&3c+ou?MGm)Vay8Dk8>mYT54dx=#!(ia+ z)yK8JqYfDB2(cUuMVT#^6ba5cpZ`BV%%_Zw|+#*aWlbTZ_I%2rLjJ2A5 zj&AE6{ee_AKw6m}8)Z}|G2B4>Oo&IjbV5AtsRK?`Uk!*2J=&$d^Z88yI+A}DPHrMG zw*%fHva#fwk{-rWy$q6h96QZ;kAF8^WPLhL8r9jqBYw-nh#>XS3?gki){8`$wA(8p z$@JWUDZ=wwlWZquH9(4Ed`MzGy=oS-On*O{TLP&YwqX^ygW>M#uOye_0MM${2m*ze z-RCFAe!5~3ShCd<65qP4mVcF*IWg)S9J9)_e$03;xl>u^-Cdp>-_fe_pg6}Nqr_bW zAh`vzyV%Y#J8Dz}(NSW0xaBx82iiGpa(Exwyc%@1$Z%hg2(!wPL#BVILi+gWmIv55 zs|;b&xDXhxF4{|tD>{2QUTt*flHkzfeFPKrx%y;b#Wa^3#kJjqun_~pA!2i(HgWT^ z9I!I2;vX3Rtv6Z&Ivcm+XR?8UL-;eH9{Fryi9`7ERv6MlZPTxVYjSAY!fbHiSdxN5>% zX?WPC?BOHEMQY=t1cKVU;#Vqo-YX@}oiTq*i_ykZf@rNKhd|gGSvf57V+}^ADsA4| za@065sqExz0h4Q~1mlxe1C&h+YIpk#Zi-he(&#Zx&m#HgXN9vq8aiZ^e!4MbSy1MAlUE+G!fDo*^@zYx#p1UO_o(OP36hbCS{(H`$V{#A!Rr*&Nm+!8 zl|CRZ;$x_Y)51BJ!}a_j@x9oCO(m!7_ky>{CtY;UQwS?3rz!}OGOKuDWmGYeBzY%7>dZ0(WD@P`&a;d^^&X} z^*OE0&4LG=O7Acs&Ackpp*1<0iKcuzT7Y=#Dvd#Qwc>C5rD12{MvvL!?_TQWqPg}7c`Io%(XaN9_{x43mv30Q4xB5Sz z_RsuJ|6Kp+YLp(bpAT7huHG4?=L1ra2p&PGu2E9VP#*=9u}?rM4T1%K7q>ybFc_+? zuDzAxMtX>xb_VK!E^wJgC*WG>MY=?#$y^X`v&%y8YUgKet04Es%{e>wg`zQmM5QIO zIa%Jv+?=vRXr(9OfK`u1! zigwe^mv`piMXd&{@-JFwu;K)b(JUhsmwKvhNZb^@B7T%wBZKL0n|{+n6qADeV2^j7 ztXr_vEx9eWE@^+=J5`TAusF_)Qnbjb%gi-OxQ>4(qYwr{Y7D7J1&RN$aqU=oQ`DkKXk2~NOFms%h&XC#siWJ#9Gu!T^c z2{e=?(-L8uI@QamuX~l+-A6PtzEN?<$CGx^1AY~?%YuuI{EU|=Uf_k`Ey_*#gM^a zoDhFBd6yXF!}9*~iDiN1?9rWdnY&ZOB2!M2WxfDSKJp^j6w8RiTm$?GKSEa`In5g2 zyYoBZOj%K*xAh8A-(z*1C++9}XVOU~*8rB+aKSaqJkbR4Y zCv_b$8{$BHhQlZmh_3$o(NBh-dyE&?f@ssmZ$rEEWeoh4;bvK`fu<)F@>fUCn7Qs@ z7eCudaV}qBrhuJ}v||_EL_gje?x){lXTj<(FF3{=mstkWS~@1qJWsoEN5Vpa6=4diHerzFL)j<{oPw#con%Q{5%&94duzgGp$j|j>@FD1^ zpw^}z>nfw;*n1+}6nIvEuuKJ3>bmSfjIO8tNTlzR=%H$E8np{m!Z1t~H>JWYRd0U! zd>^hU+_o|!nfBMa+|H~nY^L)>@7hWiwLL^@_TESo;eJMW=%>7O5I%DFsEqEJ>R9FB zT8N7K;XfR*ltgA=;WK{F+#xTo0A*$oJ13XR1EvFZl!qQJlh{DTt}Wrmn_K^=j0#cC zah#lMwn3}z$8L~;=+679kky!h{haWc+0$rYN@)Q#aJnFyVO(M zrkmar;Ojv6v|F(N)fl6xpvp`F$^JQ6z`i&Td}yKve*f-&Ch6N1AXU#kh3Vday87sB zma?<#FlrWF7V5q{)yq{H)$YKW#CM#|@O6o3Px<16Sgr0X1nHAOZ~u+C1?nRWrG0V+ zjMaCOF#q(aJ`dELP51;~z6sSh`K+LtXahICXR5W0WAv^r!$?O&Vr{ZM`@(EA{5D%v zo@wD%&8`R@vAB^AU9d+#A3mp zE%8h^>(MW2$x3cDp(y3ZPxv$(@_qV_n|k_Lo9Ji*DJ)P9&Dkt zr6*)EPfTNB>+kGJT7~GUX*}@`O*^T|z?yq9{MkQ$IX~_D{bDfwGTq{nT*K&&Sr(uF zyLVk-Yy;FAy+4JCKC0?%?SP*=C>a-qb3osL^os*g!d-L|G^ymZ@*lU#H>?-q%^5FcseiWAmnk-;Ar=XGsO7GBlld5(5 zCl2|vRdo`>QlH$cF|JD62c7N2e68`$#kFHV5ears3qvO;2*hWJg!+Wu5FfW<3 z#O_1lq6xWRJ|)P`GthsNm{v2d5)@ToOL2Cup!Wb<`e16dYK_24==*`N5YdwimJhW^ZNN=0zNw%a~i$-AB)e3Nd$Fis(iFlBg-GOf8#ZY%)^EV^$Z7G?Z6xelj zp!2ckEXUGx^X?9^GQ-ysD<-yYiU*TLP{r7~3dA zbcYyii!YFsw1@qoqgQ-b?%mI0+}VJ$doT9tSJoK_7D2ZWka4<@?fS$9b~)Kd^8LqS z5?l98iMV5}Md8l}&pq+ZQXfY)NYnJq^QrG+(F&PMzTbzfil0@AJJMOx+&DJx0-|^9 znWeA=cJ3aQ^}4P3X&i03nmS2O(&sW%b1#eL1II!;8gArkdiCG*Q;Xx^_g3)DSQ4!z z_`C5NZ(VY$GPVF>Czbw7_FgaMCe!km?n`hG{#G^^Q8w6;y9((u?wDm{F`$;QIKw+6~4kma2?0-aIiZcPygncNZ^knl#gy#yTy;cS3o*F@q=yy|&X^BH3{C z^ipH?-~THYGkxSymS&*OS4)!X`dPYeI>C~O^<0knb$^&t3U3B`@XnN*vY@76Wz6x% zSZ0)gx{@*Nd6Rng6|UDQsaiGJALJ+HaF}YB`WkcgC}Bu~=JwXYATpEypK~ThA0; zV~^9e89GXR#`>iT*xgr*(oLxnR%peOCO7CtTX|IjTjrQtj4VNJE7_aZ{uF#esOi$0 z_ktn<9}@a-okC5K7eDulcEl1*?7(iL?>OV-8Zwjj1ew@VPVaVB(p-$eyRQ?))5&sk zdB7d&I5!riLhtGu?qZy0)G$@6ecTwVLB_6*OOH{=$e(h*7<({qw)#_li(aOG5E z5t-S1NG>QjI(`}7_+NDrXJ1Kx@r1u(4B-!|8zuuy(!TG+TD=4H80Abb48@w!0@m_| zcjCeOeth|3@7RVFFp`dX|9UN%2@`ca=2Sabrqk#V-otqPC|(?pSb3qdd@w&hAIy{) z)p`rRF1(9h|9h^Ip|r8y>pA|OwF%{Oe2o{vx~-|ankq%!$X5(Q(sm!*AW~pQAq$!X z;8bk?sH1_U0$=ArPv^ZLxQTQfUk&Bam#yazx~~#3oKIojD%18&ez4CVB9sy|bPKA&^=7*SwS%RXL!(rw zp7adzX`>G7zmhyc&ctg)`H#)HMm*3{tEs)l4OK9NwB6rM`ADkmQ6`$+EWn8gtxd$0b?Y`BvT!v z*Wh#A>O(AaKU?gEisTx-UrWK<7G|7rxH2t+^pLucEv&O->{e3DceQk>tGmUqQZ}td znDbWvFeryC3TBqG(EWLlR~_Lwn?k>UIU4uIW4el<5JG`X(_U)YW`rSY8 z@7)XrFj#X^wa51rwbzV&mn$2~#|y?p!9HG}KXI9oMu=J<2f7wucP znp6z5z)?qu?^W)P6j-{n}p7hbn{TBNgq#$?pPAvfVRaz9PW;yT14c4Hm%|(wAIxeo-&ameROZ za}Xg?znXn0j!t2D+|pB3`=sUW=#ehA>K3{I#?M?@pd+70OX>*-2p7Oj=a`@;mU3$# zv7da8bdPQq&>*B+OrQ7pfghnuVi@s1F9Ll$>TpJiGVgG@UX0Ir8yA3gCzRvPabkR- zudqdTW|@=$N*{KtVdS|fWu<4fk~4eutK&6J$_PBl`j`Y-)(&Z7w63xRN}n-l(Lkh? zbM}AgoBVo^8@3cUQEe7^Owd+Bsr)sk|Mubu1mJA+^-X|0W9vgh_(5$DjR(Lse^+~V zNH+(wP-LUB)#b$Ej1kzFm%5n1#_TI-e{_y?aToI)@i%nrk;Zc_Q2o`X zf%K_Q`M}2av^{<7vE#?UH+|p+(iO=^0teA_*|-a!O+i)U13mAhKwgn@Vrf3a9juS< zo9-^sbeplKdmjsfi*)dFAinJ=cSt&=APpei$DIxYbbiul+;{pjMKq0j3~3S#M`sDn zgVUWju4%_FRB-<1&%&IC{b%q*4Gv;(yKB6nh-qvi_@*8`!g8$vsSI?wBR)r{IDd6y(wJ@k68udn;Gjk4chCRSF3E(7}-z`ld% z4!y)uNl)7?07F#lz^&Arb?%t0u;Lkd$dwY;a#-a%Zn!l88+&9FV42lN6VOu(xHDFD z^6M>6j-PpoOfy@^V%WV8-fHj~2p|tXkWkDW21IV`Z+%evh92k34Qz2`mrOU{u+{h4 zUKJlUerwt9t8`KFXill+z-Rn6y=nUVAfwVdlOwQ+Sr}CjIPoTT`>M3031^j|a+r4E zkpqjJ+pAq{J*MQ>)*Qxrd3G!1H-QD}QOwX#;8$@;{$4#UvgKP#CObRn0BBi3ij~X7 zdR?$)ei(Y823+Em5RjAa^;Htv9y94LHaU%)b&2x4+SS5f!Ntm}g$6sv&1Q-C!phGl=a_`{*F} z+_3JjG$^~~tHPh{VGDGd@?ZNVT zDH`BV>H&E@UkWoD8}I^L<`lK{2R8JvO@1S2zC`Y)-W{A@xMRVNYr_)cNmEq?7Bciz zL#RKE$VJx8BOrhX_n-77c2;5JL1g4AD1%%FF6&*h3&M&J+$A>36!xU#(HUW}B6;ua zHUw-^6U*wX!T>IvP-AEQV1$iedqVub(C76hI55;(b_-a5x z%56(4uogH+w7sNwq&uh~Sj~OaXen-pn4Z}0=mA17^gTbtrwtg}cZqKD67CJP)JMF1 z6CP`%seyePG3`U}DuI0CF}zEjHfRW6y)l^8%v+&eDaVgiY}nHf8(_{Y(`YTnf+oVD zPWm+JaIo%cimNY57X--dz`POJ>*V3AdbA6x5g>l;qA%(lQh6NA#FU}O!J;E80;5ar zSGp<~k*rspot_IFkZv6;`-`b+ppO>t0j4)E-UFVm)g{?bUXgmgwj{!|Pq3&G=0B;!l7Ka|FL1EHc+xk=jGjhOs-)KcxO%n`YzG5PW3M6x$kiel=jJb9}Jjy$Ki-3vfr{9~;X_Y=VOz6UeairPwJ|I;eTFXJ+Jcf!-1G-Dapf5F9T*`(_+RI z48y}d4*QzveE%%@(?utPcW#cA8jO!DnZ2yGCiHs_#lY@TGUEG~{Vo2LA20G_ioqRg+-eXCz5@IGf`uEDD?D>VLTdNUL@lwbx+=VrL0-+JFlP~8wG{zVioBHhyx^PyQ zcSpARC*F?4jx5<8%l6dhX-Wa%RnHU7Bd>&i&4AS$ig399%4q$rMPvSvF?9NBZe(d= z;^=7d-&xRqm8*Y;|6c#Ug>=#4e^k=uoZ|*c-&8aut<%_Ic`pMZr5s^A26blR# z+!Z2gJqnHnWHbXKG7dNmv@XF}hueL^mLLCyYqQQ;Sc` zgnVnlt-6mN92NnWu}7!Lrj0sKc11`Gtx>nnk^J|gZ#^f5Mgiib5u$#SJz76X@%fv8Z0V%Z||(!l_Ls?DvUW6O;6Igdg*tb zgfWc+hl5l)T!U>ewy)Oi66DnmK#(6+#cgKn5)4lh)$d9AT{Op^_Hb>B8?@W!%$RP& zpNG9PW;(u6OU{Ly=afJ;eRBCIX9@WXQX>qVSDf_R<35~5-RCF zaSCHl2HIEk#QJghtH{wGh=x~T3QJ*m7=%F(Jbny&{Q8M$h2}F@gA=TAyu!1UrX(D0 z)sIAk4rA$Wn^A$?t76xxwFQGyI_7s+Ld++IYvDhqWK#LiSJJL1m54dC!F{NagxJO( ztbDf7ez28fEX=6XGgg2>NznvGMWieD-EquM;tMoWWSMqYMlY{|%h@m8r|Y-d5EJs5 z(DWdCBiG!ehg@9lxMBKaeb70|hP?hNU!Mj`4#}#%tM%g3iw8NBHO}b58Q1B5CaT z*+b&5Ry|}m=^2n@xi@>`A}aTTnGsP+YvYUqMW~f)gbf_gSvBZ%d3RF(s`-%=j{BmW zg-_N#4+?zYf{rjf`1j#-;n`s1l%}PQq)@DIBjCVc1qS5@KzHdG746g}Cz?>6(!j;erL)jAT@SP*Q1B$PJ-*nR|;QM!D9D=3L6Pm8D z>Si7?Xo20<<)UyOkJI(^^s8jb5vv+4EsE9*CR9z&vhR@SMqax1fq1-mkyF2%KT@ep z`>)HL=4*}P7KBoxoS=uj7Q_`(+SAcgf^J$ftlFQG;bUDBX1WM&$_f^MeadIzO>rOJ zB4uf2`FKEWcPSgEu0u`xH65rc2zBi$KIvGBl=e_li`9g(HNNijRq>JwrQ_5-8=YQF zWnEBrN@Pl>ulfykpg-o726DGslHOrR7quXGZ3~b|+yZd|e1msU_T;p7a4@N)Vo9lZ?3CGj9Jg%LV>S%2gK!NEVPG~7~V zV{%04-8gix%iXRYj@NLTRW}Y>*z@&?>%nK?Lg-`CD)MxamE3wMy?uw%JRq|ypKP5`xXKq;l{;0r!=$WQo@ z&CKN+bUKKGjT`xdG=%b98Fc5FaC6!?)+Tej#S5W89GHQ1(e+^QCfosweN<(rxxqYo zF>RY>es0_BFU+(JfhMbx+f}D@w?^i?_S$0NN42A8edQHxm8KNy*OcTf1K%RlCSoR3 zWK|Qip|#R*X?!*~nytAFR-MP=9waqB8zAOj_D2{K))%{?Uu2sGibcuDv-JqypQ7Tc zzs!cQ%Si+PwuoPwsz7g5IR&j29}JD19@eY ze6ZyX|Javqn_>?2KC`taJNvK;fGe*Sw#ONyrg5HNTrl8_H|l%scfMZR?PS+Y8oOgg z3>?Q5t3~9Rig^9lY7x`U)2$P(ze`DL?qn*K>31ge>Ehj0;6WdKklrZ1iu<|LDVbmF z3-0NChtBPl55|-Hg`QVA2$OHV;lu_dMe_>vGul_7{N=H4tBKrGeKfB?A9?By|C)~F67gP`=Rg*AmE)Cs!n z)&eAl%0iAuBp1!gt*A=8eBhf0W3z6qYqPH?2=@;DcL>2)y9=n_mQolY$DLGNCU8Zy zH4#NCfo`{`FB4OMmGux)^qL2I-X`24nDJmsn&P^#hmUo#18)xp9%iU06Fjxn%s%M3 zY&;-!rlMS2hXOc~^O^)eVl8ZTeuRT0omcCFwxrbe$954Nj%||aH$jl)CbYIolFzU~ z8y~d+%@fRLXT)CZP|1luX~ZXOAU)0wY@3w!9DsZ{4r09KP+m&CTc-@tI0&}jn`~8E zgW4}STwN7V};bmizF$6Y7&`H^@ZM zS-%Iqd?60fW2e8!J7_gZDKmlWQIMw>D>Vt~5*fNNf6$ZGhEz5Lbz0KXBMM1oe{q#i z-7A>Qn%>2OGS}m#z4WOa-zhAeVys@?G1-8ib`B`!HdgFT?hKmX(MnIQ%s|l zxF5=L3M1EP$aASA3{~$yp$4=#&6FXkI(k{m2VfNxX%pxGi%YU~BrCH~7P8kMJ++0t}D$p$P?H!f=!KgJo=${t`dc9H~hL4FGI9 zh3DwPEL5lQa@WisJ}2iL&DgW1uiDE*9kmIB<)%wraR&iS$a1Bg*VH6XlCrdLCZi!4BWilC=*o5O*t?miEMb6q9zg0od1Xi8G3g%5}3I6SwSmMsI8FhABmbV^$OrC+T8Q!`SkX1y%aXwrr^x z$6H{Bl-Y(<(?zN;yoJlfb9QqS^_3Fw%~5wjo7T-n?}U5>^xtr#Cba8@YY_8pi9yew z!J-(~!!aJq|eq_ zjyZS3GzNtegLn|H!Ji2;EGtbSjE8pjF?g@lMUoQM-_Z#rM^bD=t-z(^hA<2#(YVE(TIAzC(rwbDuo*tSLVn_F?8pdlj|O$%j=e6LtNc>YRF7juxj zxLcZ^1CgGm4U>{A$nA7rlI1&U_7cBDqpmLaQ6^J6J%y5W5VCt3*o?&npwlLE7?_4w z$|Jg`ohOy;9XS(bbi3M|_{j)|YqnB9+Q(t|_EyzeJ_~{Ro>Y)H*pKJjVHX9c&0p;i z{~&9=>0YQBrv&{7 z+yLnR#$si$i`Zrsm|$Sp0_S$6pCzEVl=tEidU zJSSB-Kl3w2CJT(o=-)hnvpbIh%P+|jag?xZOO8S%G%NR-Au<$#BJ${a@JyIb*ui%< z%f_3QfCEgUP6y^K=8mqcm!cdT1^ofA47M&1xW_d2Cy+8OZo23lJ?%+B&r|}l9mni&P7xY(KnTlh?-zj20G^4Xkv27zOEIeiRbM!gvAtG|s z5;s+U2J9V^G^?3GoEEu7Q_ATHNRNZd(-URS&tuom6}jfK(}fMNMu)oJc+}E(HtTDf z1fLE2t=7kK8mrARw|t-w@mL5*uKXie)r#q`Ff-!IYzZ7#gE}Hzxx&)7s$vDUadcHMZXkvAECc>`mN^4P-L3<%NPaG&O-`Ev$0n6qvo>-Ud zG!_fc4wZFTgs$?TArqLx@6)JB2)~2>X%NZTY? z#ToHHd%&O-i$U4!-2{Vm`<=n%!@$-)tEo5L(6iNq1lBEWTA@1vkU=}L(n9J^)(&%L zqc^gq9z4-!&KOjja$~Z%!5gb8u4C@iAudbu2jaG=AxNb1>V&qOObdGxl^$M}g1RmsxEB(K?d3bJ5UyRYLTSw>p4204A>Sw=X#c0&SD38$R; zIvh!p=FwZT@I{Y*JLNY}M)#ln(!u}M{owt1M*f#D{r@3&`i=R$&eQb7mF&@m52YfK z)8Y_sBp=SJu{ce>?5$6OClI4~iqEK5ZdX%N>tv>DRY98=e|XVnf;!J8=V->B7Hj{=x`ogtV&87hA1I<@ThP?F3 zya|gZOw-}nW+_O8$2z%xM)1|ce)Sv*pW*K%h%~0E`F8)U*r4(_cBkVXMN(>Ywe+e}yzh$@Sy5o9`ch+^(J3 z!#jtfC%7&%q!Io}KHq{S?ZdG(K5;McUh4 zncM}}$r~T&G5&>e-x0fz4`Z3so*uU?t_lR{sPI4~a)Ste?Pczm?$^7uYz)i~#?S$! zFuR+Z7t$upSM?m%%tQyrbmq4&Etu-N2c;}MNCYUXrF;v#Ij>FCVbx7*y86+X-xgm5 z-Ncvk-j$D~uRXH`b~X;~Ejr2DU^5X%CvFk!2LX!Grj7*BOQf#o4llNU&5xY~rXR8X z>d*Ln=lmtR^shVTzod2lmsIRG=pVOEhP!-{7|sw65*j*W8ypT+n+n)I5_ZT#qR<`= z7PbwtFt7kLQSHg6;jjmP;-XSNcE5-S?E<8vd|X5{OyTelVKp+~-SR7bFk#>ZtA%%l z?DC%{>zWbYcb;rFo)=c-_%@e#nDl%h*E`+3dg_c=V}2G8qTsZUtLFAE9)_eJa z+No~R+h!mU-goY(H8}RDommP0`WV~>XEWkB6gd@H_V`L+ZZ~qvtBBu`HDcUhPx`#w z+nrIWMM6O}SHg95)$_Br=~}zIZL`NP=mGL^r zfEAThz1p-IKB3FrBDjgOxUhL@SLfh?>PZ+a*r6+V=qJ~Z$fP-2dwk-8 zZv`$UIs?V80~uqDqKV3cv=Q&A)>eAU3v@TNjF2{NdAJR4UWX`&LXPE2-1_s~GXJ?) zc_rF~_v8_Fy3t~JEEwJVLxW}xD{xt*zHgnz6(WxSa(q7fYIWe3ilNizsK+dZ6kq!L z$Vp@6Z#WC3qAcYK(H6?zo)rgk6$Z1L$I+b5O5j&7pI)Y^dp9=`@y7PbxUsld`qpv$ zx3(&T2gj?L`d}&}aDF<6D0@9=@O~W$mv@W9?z^<%9AWXO0rg}=3#F-`bj238x(9ue zQ7Ywb?i|%3%Tr6mB!apk|b0H^5uuEsMEE zhX}rJRL?kY0xdS6mvVbyWNztTH|ISienH=C6aenfnc+tTcY09yDXYisQx0X_#7O*z zkm{(7y()oZEx~pnFYXb`Ztz*V*v|PWUdI;~HZ9rrs*@F1mY)E{dyZb%B_IL`4Te*A| zo|i~Tqn#zJfSaH721hfg+fO@C7G^Dni)N9o7;f|Dz=Vb_Qb8xMDf#R&>)W^_`8}M4O1)Ud zPp3n=Z4iQ1H<^JhSw|?JK68VnRRjbzWha#2M)vRAW_s*nB-3WFbb#4YjqbQ*(gKfW zIg-)3m5jtRyiL0>oXd1)`3uElVS`fygF8tz;8S@pIIMK3MQ#82&rSjSm_siDMfH{mSVM zp#r>P8um^^Xel>9855+0M_RKJ6)QBjIG!t*1mllqOFY2$8k8si7xFr`cc?S-x8s=z z6buz)FO{k+_F{1>XO;V?dERk#-BRZoNNH+z22u?bPI7uLbkF5HAO9oY^tn*k_cNGK zio;J5h)p#??5r?EBpj%}sr#|%vYsc9Qf&%vdI~uc(zllU8&%^plGiKCF>EJlBU(K? z#2xY!T9iKR=2=3?1So#*52(wPm$`9}S0lfKlu1E)@MsNG6-7uG5-(RDB~*LY_cN2@ z6k!=I=38X6o!bMPN?)%Osd?#nA4Fs?=)L#j;91~TOB)y8Ebwky~rB1K8SmjCSb5eWvKdmF9-E%gd6 z-kYd2Gqylk?2UCa3C*1wKnb_10@<|NNk0t*WvN5MkWW z$2qn2p+3GVUZ2xB3NAu(6XO2mQvq#NY@^5EH^f@uNzmeyJtU*6SEaODB(AhswN-~kjb}K*`OS*HUW_kHtHsh5YkgtZcH?{FdpWXQ7!E5pj+o8Sq{cn)rxDM4R2Mwt&&YPS14cyC_%Ck=#cgQ<7p zN*&Hth$>+)5)L{ETb4O*z^y8e7m3vKC6&aqs@Qp2Q3Lkt?5mu{+ll_*3VDAr6FRYK zBp0>M3ABS}L%ea1Gt^-Y*(fU46u=s$~; zxc}++@}K)^e?$K97A(b6T1y3Q7#pR4T1wy0zq-G@4t?ZKgj2}wusyA9q8CayOeK>) ze^6jqKbRgxU7dL>M;`niX?<-h z)_CvW-Td+G=EeK+E$8OgeQnN>oW)UZMe7t~4CD)%rjA6lZak^l*@V&;thP@z;(Fyv zcoScKSUTT65y%*XPd#)R_Y|ui&echZzVEfjz@6o%7 zksDJKdAGDIn|rJS!?F(Jqeh(wrO!;_%C_mq3Q%}J{N~wfd4%`Q@EPJZh+L4Kgr<<- zt}LQja3!r4nSpU5)t~nw9Bm@(?U?z*k)aAAcmMLuf3t)SJSuWp3Y^Jh{ShelTak?T zYJm+g&(joyeo&z;icjjO<>2}+T^bfxW2Oqwp)_2wW3J}Q3Xv4*b?)-fL zo)csn!f1-7+WE}uSpS)fhf*r5w4Wv*Z^o(Pmu<`zzNCSoCSCX|!at1Txups2=>+&BCF~r_@{hgU(X4Le~5+UzwS{V@j+t7bdw3uCPahmhjY)G%qByWL~&<#zH(7cmvpNj5g|P2IxCyqc@e-4ONP;(-4(VEJDImj5+i z`CkK;|21Iw|1)6W>byJP_?=n)^@;W83+(^k#PU1-_sSkQmr}CF8oH8*GbNW2m!!yc zHCxC~N>zBt#swdxFYlU}skLe&(UK}z&9dQ$RhK+9P*g3YtR}UCCdNTNKXVl$Hdd6U zCC5lsq(VvH#^B{={s^vC%ROCnFWjr`rj|WO$~KDb?D{cuVWzHFWTFD6Wzxi96_No{ zF*#Cm2O84g!n=4*E<{84B3TtN`lB>03;)Yd|LYRU-0rZs!*<%TqD`OecM?uEH0UzW zh0-qjR2S}=}1kloZ&edrv2N_jhG*PA*MCgPG(v>_tnBenHQc% z;^C+=*e3aBQ?T=Q3*pl7Qg|1RN+PBt2hm_o+Y}yS_Z{{w#DqiVxLdMWuxh&55X0T1 zVVetnDMnR*G`+11f<>dcRmhWdWSg0@-VtlUEas1G5WaJ_^CRapTe*N^Xxy01vyg6U zt-i8~{vv0eEj>m*1FGhCoS?QrJoRQ(BS$y1PCKYQI|wNf7@L}TaC2+~V~o^M<6V=- zw0&XUYGDg89st*hA7s<$S`}R~Qd28^jbSpL8k#@vj_3y>DnrjGSX()?$v5Y|fS=0C z-|q7GsR`gF3^;j;Jr{DjJPn8d5Q%n0Dt_=K2|9~WX>D;p8~+&P2B(eKi=ZEq~ z>?FmDji{Bxna`T;b)D_UorFf5E(H=ej8DinZ=Nu_Ph0NZ_uY9;-#T+28k&P;aWGj2 z!Ki_$vI`EWv$FRJRXO^w{j5@hI z+si^aI?A}cwAYj?ls|HfINnb=kMoC#kGE}6$^1I5OKuvruQ0{LQDsCvNo`N>DL_sA z1k|xamR)}R-f?42FJEz3<(o0~tU(ndsuFjW1?lz8!W9}BHDttAdG?jUw&mKO2XAP4 zl01daD;hBE-?xw06Iw@B7rg&&$3Q}!JFb0hSSUAYjsj@hq1fbwYZ{M{fL2pN-nha) zG3x;KoK|GKGQ|fk~B}v%jR&M@fg--diO3>F?Qq8WEre%Fh2I^s1phIj! z0RG}|Z)8K3ee)vJi}aMbKCFJwEK%FeL!&z3ZAQ@Bd7gJvzJ=e=pJICrTaIYsBi++( zY2+pjewL6MScS%g*)#uF_u_#Gk+cryyD|HlsI6eT08htMpp_AA+@6JTQWuK z^c`A(@{;-I!)WhIMC0cOrLR^!(%Q@}Y&bfbI&~RM0tUdG+@IcwTlQkIlkC4Jq9k70 z4P4Tky*E&4TZxNoJqpR*J$N4ZkVLR9I_GVzVMr*K7g8m{e(rY5Y`n|%pNg{Qm+mI< zf+M6}ZR95T@&vW%>X>543BaTE(5hK`Mj97PUEd$24{?w*QW*8BLcy9_80NMK3xYiz zZu>gQ=DoGW=ZniwAt@oTM$!G_g~-`Qwv6~02(lT=K3Hv_{RoPt6E)z+`Gmq9 zC8hw@53$(;v@7RbCGf2IVDN70p?1qtG@WsT6owtwBunJH`Sk3e6kOI?0@|*T@6nSd zo6{d95OtxgY$UpZIu^bKo`wU`UVYodHSVe{X<zr$!d=T)mwE;C=*I*j^eH~`6Wqsdpq9Gth^(+UL0qSa z;qpy^9ZQ}xkhBt=KtxU_D)QWs6dGZK&wb_OfxO=V7xl~lx{T#F7hXY0Z!(8wO4zJm$;D-0U(n|H` zKsEPjlAA%ghq9c+MOCNEBO3)?gl4V0=_nEDlTF0eFS0 zeJB*hW9%K&BZ{^X7&M2>P`44#L#fLrz;WscJF2^>eJz(#YQ#qCsZhojTe<>i_TU1_4x=og1Q(;mNastf6xe3QrpF|0% z`{NA9dF*r?u4(uFg5$KNHM*Q_N4a1S?rCiEjbuZFa_`ZJf;!Ij>y4 zSU7kKv1lAV-D#b=xg5fPJskACZu-j~or!}&`W~pv%0nqKCf&IR{haB^47xST4NfUn zsDEGZbgD81Lb!>&y-Uv@wEY%g=d6vB!%~Z$-PaEwHK9w|oDPMY-bn1F)2#k6LpUVt zCf}45p%~+mbj4@jwJAs%h5>8f?8yF6YK7{{;6Xp>t~BBqRE_@SvRkC!qeieQCG~Eu zASiN0@9Cs801TOPeZeCBZlumzyyHQ1NMiZ3c^twQY~ISH2kStG)lof26| z;R(HTCJBcBkw!3aqRmUw%meG$u2VN_s&AW#(q_0va|;-jCa2|vgM}ZY?{uOy_U?jy zTpUa1%>}Yn`}u=aO;KUgU!cDA21wyaZp-=j#l4O-r^Xug(N-6Io_?2oz(}UlxgZ@# z+xE-UXdsv>PU%2(Wa1&b-gXD&d9fGQdi&an6-jxfc}#!7;&n?oo_+55x1Vr1g0z~` z9L9r4`aOYMSC!bD;!Rl0?%FGq5>UOU|mzA0zc3 zNKtkTx=_&s!AP<%-!I80=kn9UXYM}tZQ!`y2sN}6E9JnLtjNHM_01IN$MGr~I!3kk zsh3Ra?x5n-Rx84)*D>z94w>DoFuyh9OYxDW<9;ab)uSFu>U8X0wzVU~SYpClow?R7CK?)~l{!o8a) zkjT?Jxv|a16{&z{zt|chrVBoy>GtM3Hy~@o+{41)j^G>b zN2%Y=7-*sMJyc_1!8behPdXYYF+oGf7+yC^l%4C%l9kt8F^sf1OuVrM>&@82Uo4wZKGIz*$mK#X)nX5n%c1w2l)iYZ_{IV8;y|6&;skW;lDd|m z>Sc1vnwl1Kk34@4Ui(mM?mUU=J#!_U+4R?i>!#Qpu8@$6bvXNYj&r#EhN*R7UPGEo zgPEn|W8$KV8}~rY#yu!OXN~NVSb(l-!SEG^n~gYCYI!X^&4=ZliCq}TdlLyrnV0(V zvprqLAln6O{Lc`O?=gGG-sbSReC27Xyr)s`O$aDXS7whgRL{utMA5^S9~kq`Ao)dw z;$`bEK<7VALw75ksFxccsWxqCWqdm#-Sk{NJvcW!G$xyyUW0Ml%%Ew3HVQOI8yCHkG-L5n@3Gx&RGXm+ zMu|ZVT*x8RXr^JJ-4Cyf%WIwm;P#Em5GImPpEbVZnr8eGa#uq+Yh>>4rh3jHV(NIM zYxrWAG0MlIaNC(*)>Yn5U!$d&x}-Su_V3 z`zm0$FQ$IEV}RZm1JNdi?D37{Q_eGy|HAHPyYoJtofep~scgW(kgjB!CPgKmeI(fZ z2n`o>Pwa*buqf0sW-wQbRHZR5BpgPn;+Kp zJReb)g7F5{Gj7T#s=(exkOe+;hQ}SurXDIT`$cTG_8L+#Zd0{>Co+~aMDX#v0Q!m^ z%&~7XI<2dp!+wVggm)2nh8;BG8#gh%Ok%lf`!$Xy>^TPXY<2Z(=U+T7wPLp4^T7a`4q;eYT6*jFiQ({T%zTBH>2(y zm$X9_=aMTq@AzG}k>|7P)ngKHHyGKC)2ljFDwgt^48xV2n_T^R2_#DwUm?^W_Ln`? z|5hLBSvQKW%Ve4;Mz0DneL{+4FtNBvq1d^essC^jRJQIW#(oN-sHq5b=pp6`Iy*X?ehyuC=#(Sg zd&jOtPI1wE4E!}Ug87GK0pr8rAQFDdSpXGX#sl$YrN>P#hQhru?7V1K{2+P(H`hMB zhx!rK9^uJn$YRcs$cXRgUbI=UT)sMGGo7NNJgH3fHf%l#+f9ZZ;v`3rO+`xcJX`Z_ z6aF)tO`*|&8Oj)rv0Bj|_5Su?V9$<3*ItyWqR7%GWJ1Gjq;@r!=nqMfGCo;9n%CazRY1WXv&1Qt299cUdceW=_cV+t_8Y%Z}vj`pcYSxK@kNl1&ib6BOZ$Fl_whr zxUz32>z((gW-9BtD5;mJvLze4l|wpTs6GtVOajFK3-QBY%f*nDIdv&mtrF)z(${I^;cYcVu9pFALD5_bxzTCQ$uGp2ce!8z@D><1iJ&hJ zf{`{h%(ca+=x%dUU-8&ao8jbe%1d2BdnC{eW5?X>rCYi!KVBUxP*$ouFwcV=foLu-EzTQ-gvT|xU z-JXF&3Xtrd09#B#s-C24-ko_AFl3P1L^|lWfx>#Lph*q|HWnU(VAY!e)7vB?a z^^r6k>}x@P%`WO|PBP4t(mt(55`oky64PWE+nNZzrKjUq++|&B=~20-(vK+iE;DqA z5UX=kIcqTsOcG07qO1b-+(v~rYF(rCv6KCnV2nE+D{^kwzV@zmNQ~xI0p}2YjZN)s z0FnLXK)huyiM`xYHF0nNY?0=bv3K2Fx?YB}J@NQc5^&fskG2fm^5q2z8dG=d!TWl) z_G&u5(<|I?`3ksdnJH#5eO(x{^0`hR=CsYvgFsgskcXzMtI~7rx|3iX+GqNWA`AOM zbdUrCOgW0YGA{yka9!BRDFvbxsho0X;Qj>t$ssthQ`Q%#N=6ZyB~^Mq%T3HsL!+|s zwN~X*hh~FFYJL2BCJB?Y0xIKFRIn_2;5(wc6fg)HDB7<)%d&C|ALtQBKhYy>c3nk{5$Ti z6N>)C)#3gN_m6o+e+T|qe}ert4#*EUi4^G1DFA;5{?*v}-$*~;GfJSp+g$%0{8#(w zKf!@tfc}FS_22P-wMG0BPsRLi{2!X(zZ3k`hwx8=mp=&pbS3;d!(W~G{$!wH{r5xo zQ+)Y%hQD?;|H(ij@o$E|^*R5Y;M1& literal 0 HcmV?d00001 diff --git a/tests/compute/dynamics/test_codec.py b/tests/compute/dynamics/test_codec.py new file mode 100644 index 0000000000..f2dbb6cdbe --- /dev/null +++ b/tests/compute/dynamics/test_codec.py @@ -0,0 +1,331 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +from __future__ import annotations +from pathlib import Path +import numpy as np +import pytest +import torch +from torch.autograd import gradcheck +from embodichain.compute.dynamics import ( + NeRDStateCodec, + quat_rotate, + plane_contact_features, + append_history, + reset_history, +) + + +def ant_state(batch: int = 2) -> tuple[torch.Tensor, torch.Tensor]: + torch.manual_seed(7) + q = torch.randn(batch, 15, dtype=torch.double) * 0.1 + rotation = q[:, 3:7] + q.new_tensor((0.1, 0.2, 0.3, 1.0)) + q[:, 3:7] = rotation / rotation.norm(dim=-1, keepdim=True) + return q, torch.randn(batch, 14, dtype=torch.double) + + +def test_ant_real_official_frame_features_match_reference_rounding() -> None: + """Frozen actual Warp1.5.1/Torch2.2.2 windows catch FP32 formula drift.""" + with np.load(Path(__file__).parent / "data/ant-official-frame.npz") as data: + values = {key: torch.from_numpy(data[key].copy()) for key in data.files} + contacts = { + key: value for key, value in values.items() if key.startswith("contact_") + } + thickness = contacts.pop("contact_thicknesses") + contacts["contact_thicknesses_0"] = thickness + contacts["contact_thicknesses_1"] = torch.zeros_like(thickness) + codec = NeRDStateCodec("ant") + actual = codec.encode( + values["q"], values["qd"], values["joint_f"], contacts, values["gravity_dir"] + ) + for name, value in actual.items(): + torch.testing.assert_close(value, values["expected_" + name], atol=0, rtol=0) + target = codec.targets( + values["q"], values["qd"], values["next_q"], values["next_qd"] + ) + torch.testing.assert_close(target, values["expected_target"], atol=1e-5, rtol=1e-4) + + +def test_com_reference_shift_and_inverse() -> None: + codec = NeRDStateCodec("ant", base_com=(0.1, -0.2, 0.3)) + q, qd = ant_state() + legacy = codec.native_to_legacy(q, qd) + center = q[:, :3] + quat_rotate(q[:, 3:7], q.new_tensor(codec.base_com)) + # A spatial twist evaluated at the body's COM reproduces its linear velocity. + torch.testing.assert_close( + legacy[:, 3:6] + torch.linalg.cross(legacy[:, :3], center), qd[:, :3] + ) + torch.testing.assert_close(codec.legacy_to_native(q, legacy), qd) + + +def test_original_velocity_override_retains_values_and_input_gradients( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Imported twists need not pass through the lossy COM representation.""" + codec = NeRDStateCodec("ant") + q = torch.zeros(2, 3, 15, dtype=torch.double) + q[..., 6] = 1.0 + qd = torch.zeros(2, 3, 14, dtype=torch.double, requires_grad=True) + original = torch.arange(2 * 6 * 14, dtype=torch.double).reshape(2, 6, 14) / 100 + original = original[:, ::2].requires_grad_() + assert not original.is_contiguous() + following = (original.detach() + 0.25).requires_grad_() + before = original.detach().clone() + contacts = { + "contact_depths": torch.zeros(2, 3, 2, dtype=torch.double), + "contact_normals": torch.ones(2, 3, 6, dtype=torch.double), + "contact_points_1": torch.ones(2, 3, 6, dtype=torch.double), + "contact_masks": torch.tensor([True, False]).expand(2, 3, 2), + } + + def reject_native_conversion(*args: object) -> None: + raise AssertionError("Original velocity must bypass COM conversion") + + monkeypatch.setattr(codec, "native_to_legacy", reject_native_conversion) + inputs = codec.encode( + q, + qd, + torch.zeros(2, 3, 8, dtype=torch.double), + contacts, + legacy_velocity=original, + ) + target = codec.targets( + q, + qd, + q, + qd, + current_embedding=inputs["states_embedding"], + legacy_velocity=original, + next_legacy_velocity=following, + ) + torch.testing.assert_close( + inputs["states_embedding"][..., 15:], original, atol=0, rtol=0 + ) + torch.testing.assert_close(target[..., 15:], following - original, atol=0, rtol=0) + assert torch.count_nonzero(inputs["contact_points_1"][..., 3:]) == 0 + torch.testing.assert_close(original, before, atol=0, rtol=0) + (3 * inputs["states_embedding"][..., 15:].sum() + target[..., 15:].sum()).backward() + torch.testing.assert_close(original.grad, torch.full_like(original, 2)) + torch.testing.assert_close(following.grad, torch.ones_like(following)) + assert qd.grad is None + + +@pytest.mark.parametrize("which", ["encode", "current_target", "next_target"]) +@pytest.mark.parametrize("invalid", ["shape", "dtype", "device"]) +def test_legacy_velocity_override_rejects_mismatched_metadata( + which: str, invalid: str +) -> None: + codec = NeRDStateCodec("cartpole") + q = torch.zeros(2, 3, 2) + qd = torch.zeros_like(q) + value = { + "shape": torch.zeros(2, 3, 1), + "dtype": torch.zeros(2, 3, 2, dtype=torch.float64), + "device": torch.empty(2, 3, 2, device="meta"), + }[invalid] + with pytest.raises(ValueError, match="Legacy velocity override"): + if which == "encode": + codec.encode(q, qd, qd, legacy_velocity=value) + else: + kwargs = { + ( + "legacy_velocity" + if which == "current_target" + else "next_legacy_velocity" + ): value + } + codec.targets(q, qd, q, qd, **kwargs) + + +@pytest.mark.parametrize("robot", ["ant", "cartpole"]) +def test_explicit_absent_legacy_velocity_preserves_native_contract(robot: str) -> None: + codec = NeRDStateCodec(robot) + q = torch.zeros(2, codec.q_dim) + qd = torch.ones(2, codec.qd_dim) + if robot == "ant": + q[:, 6] = 1.0 + contacts = { + "contact_depths": torch.zeros(2, 1), + "contact_normals": torch.zeros(2, 3), + "contact_points_1": torch.zeros(2, 3), + } + default = codec.encode(q, qd, torch.zeros_like(qd), contacts) + optional = codec.encode(q, qd, torch.zeros_like(qd), contacts, legacy_velocity=None) + for name in default: + torch.testing.assert_close(default[name], optional[name], atol=0, rtol=0) + torch.testing.assert_close( + codec.targets(q, qd, q, qd + 0.1), + codec.targets( + q, qd, q, qd + 0.1, legacy_velocity=None, next_legacy_velocity=None + ), + atol=0, + rtol=0, + ) + + +def test_ant_targets_decode_in_frame_at_t() -> None: + codec = NeRDStateCodec("ant", base_com=(0.1, -0.2, 0.3)) + q, qd = ant_state() + next_q, next_qd = ant_state() + next_q = next_q.clone() + next_q[:, :3] += q.new_tensor((0.03, -0.07, 0.02)) + next_q[:, 3:7] = ( + q.new_tensor((0.2, -0.1, 0.3, 0.9)) / q.new_tensor((0.2, -0.1, 0.3, 0.9)).norm() + ) + next_q[:, 7:] += 0.01 + next_qd = next_qd + 0.13 + prediction = codec.targets(q, qd, next_q, next_qd) + actual_q, actual_qd = codec.decode(q, qd, prediction) + torch.testing.assert_close(actual_q, next_q) + torch.testing.assert_close(actual_qd, next_qd) + # Rotation target is relative to orientation at t, not identity at t+1. + assert prediction[:, 3:6].abs().sum() > 0.1 + + +def test_ant_targets_reuse_the_current_encoded_frame() -> None: + codec = NeRDStateCodec("ant", base_com=(0.1, -0.2, 0.3)) + q, qd = ant_state() + next_q, next_qd = q.clone(), qd + 0.13 + next_q[:, :3] += 0.02 + contacts = { + "contact_depths": q.new_full((len(q), 1), 1e4), + "contact_normals": q.new_zeros(len(q), 3), + "contact_points_1": q.new_zeros(len(q), 3), + } + inputs = codec.encode(q, qd, torch.zeros_like(qd), contacts) + target = codec.targets( + q, qd, next_q, next_qd, current_embedding=inputs["states_embedding"] + ) + actual_q, actual_qd = codec.decode( + q, qd, target, current_embedding=inputs["states_embedding"] + ) + torch.testing.assert_close(actual_q, next_q) + torch.testing.assert_close(actual_qd, next_qd) + with pytest.raises(ValueError, match="embedding shape"): + codec.targets(q, qd, next_q, next_qd, current_embedding=q) + + +def test_cartpole_large_angle_wrapping_matches_official_preprocessing() -> None: + codec = NeRDStateCodec("cartpole") + q = torch.tensor([[0.25, 123.4], [0.5, -321.7]]) + next_q = q + torch.tensor([[0.01, 0.007], [-0.02, 0.009]]) + velocity = torch.zeros_like(q) + + def official_wrap(value): + return value - torch.floor((value + torch.pi) / (2 * torch.pi)) * (2 * torch.pi) + + target = codec.targets(q, velocity, next_q, velocity) + expected_angle = official_wrap(official_wrap(next_q[:, 1]) - official_wrap(q[:, 1])) + torch.testing.assert_close(target[:, 1], expected_angle, atol=0, rtol=0) + actual_q, _ = codec.decode(q, velocity, target) + expected_decoded = official_wrap(official_wrap(q[:, 1]) + expected_angle) + torch.testing.assert_close(actual_q[:, 1], expected_decoded, atol=0, rtol=0) + + +def test_cartpole_wrap_and_decode() -> None: + codec = NeRDStateCodec() + q = torch.tensor([[0.4, torch.pi - 0.01]]) + next_q = torch.tensor([[0.5, -torch.pi + 0.02]]) + qd = torch.ones_like(q) + target = codec.targets(q, qd, next_q, qd * 2) + torch.testing.assert_close(target[0, 1], torch.tensor(0.03), atol=1e-6, rtol=0) + actual = codec.decode(q, qd, target) + torch.testing.assert_close(actual[0], next_q) + torch.testing.assert_close(actual[1], qd * 2) + + +def test_static_plane_mask_and_gradient() -> None: + points = torch.tensor( + [[[0.2, 0.1, 0.05], [0.4, -0.1, 0.5]]], dtype=torch.double, requires_grad=True + ) + normal = points.new_tensor((0.0, 0.0, 1.0)) + result = plane_contact_features(points, points.new_tensor([0.02, 0.02]), normal) + torch.testing.assert_close( + result["contact_points_1"], points.new_tensor([[0.2, 0.1, 0.0, 0.4, -0.1, 0.0]]) + ) + assert result["contact_masks"].tolist() == [[True, False]] + assert gradcheck( + lambda x: tuple( + plane_contact_features(x, x.new_tensor([0.02, 0.02]), normal)[k] + for k in ("contact_points_1", "contact_depths") + ), + (points,), + ) + + +def test_ant_encode_mask_and_gradients() -> None: + codec = NeRDStateCodec("ant", base_com=(0.1, -0.2, 0.3)) + q, qd = ant_state(1) + effort = torch.randn(1, 8, dtype=torch.double, requires_grad=True) + point = torch.tensor([[[0.2, 0.1, 0.05], [0.4, -0.1, 0.5]]], dtype=torch.double) + contacts = plane_contact_features( + point, point.new_tensor([0.02, 0.02]), point.new_tensor((0.0, 0.0, 1.0)) + ) + + def function( + q: torch.Tensor, qd: torch.Tensor, effort: torch.Tensor + ) -> torch.Tensor: + # Quaternion constraint stays on the manifold during finite differences. + q = torch.cat( + (q[:, :3], q[:, 3:7] / q[:, 3:7].norm(dim=-1, keepdim=True), q[:, 7:]), -1 + ) + return torch.cat(tuple(codec.encode(q, qd, effort, contacts).values()), -1) + + assert gradcheck(function, (q.requires_grad_(), qd.requires_grad_(), effort)) + output = codec.encode(q, qd, effort, contacts) + assert torch.count_nonzero(output["contact_points_1"][:, 3:]) == 0 + + +def test_decode_gradients_nonzero_com_and_orientation() -> None: + codec = NeRDStateCodec("ant", base_com=(0.1, -0.2, 0.3)) + q, qd = ant_state(1) + pred = torch.randn(1, 29, dtype=torch.double) * 0.01 + pred[:, 6] = 1.0 + + def function( + q: torch.Tensor, qd: torch.Tensor, pred: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + q = torch.cat( + (q[:, :3], q[:, 3:7] / q[:, 3:7].norm(dim=-1, keepdim=True), q[:, 7:]), -1 + ) + return codec.decode(q, qd, pred) + + assert gradcheck( + function, (q.requires_grad_(), qd.requires_grad_(), pred.requires_grad_()) + ) + + +def test_history_partial_reset_retains_other_graphs() -> None: + history = {"x": torch.zeros(2, 3, 1, dtype=torch.double)} + lengths = torch.zeros(2, dtype=torch.long) + tokens = [ + torch.tensor( + [[float(i)], [float(i + 10)]], dtype=torch.double, requires_grad=True + ) + for i in range(4) + ] + for token in tokens[:3]: + history, lengths = append_history(history, lengths, {"x": token}) + saved = history["x"] + history, lengths = reset_history(history, lengths, torch.tensor([True, False])) + history, lengths = append_history(history, lengths, {"x": tokens[3]}) + assert lengths.tolist() == [1, 3] + assert history["x"].squeeze(-1).tolist() == [[3.0, 0.0, 0.0], [11.0, 12.0, 13.0]] + assert saved.squeeze(-1).tolist() == [[0.0, 1.0, 2.0], [10.0, 11.0, 12.0]] + history["x"].sum().backward() + torch.testing.assert_close( + tokens[1].grad, torch.tensor([[0.0], [1.0]], dtype=torch.double) + ) diff --git a/tests/compute/dynamics/test_dynamics_metrics.py b/tests/compute/dynamics/test_dynamics_metrics.py new file mode 100644 index 0000000000..ba93266d37 --- /dev/null +++ b/tests/compute/dynamics/test_dynamics_metrics.py @@ -0,0 +1,99 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Independent geometry and simulator-free contracts for rollout metrics.""" + +from __future__ import annotations + +import math +import subprocess +import sys + +import numpy as np +import torch + +from embodichain.compute.dynamics._evaluation import ( + _error_values, + _safe_number, + _teacher_row_sha256, +) + + +def test_metrics_import_without_any_simulation_runtime(): + subprocess.run( + [ + sys.executable, + "-c", + "import sys; from embodichain.compute.dynamics._evaluation import _error_values; assert not any(k == 'newton' or k == 'dexsim' or k.startswith('embodichain.lab.sim') for k in sys.modules)", + ], + check=True, + capture_output=True, + text=True, + ) + + +def test_continuous_angle_and_bounded_joint_errors_have_distinct_meaning(): + cart = torch.tensor([[0.0, 2 * math.pi]], dtype=torch.float64) + error = _error_values( + "cartpole", + cart, + torch.zeros_like(cart), + torch.zeros_like(cart), + torch.zeros_like(cart), + ) + torch.testing.assert_close(error, torch.zeros_like(error), atol=1e-12, rtol=0) + q = torch.zeros(1, 15, dtype=torch.float64) + q[:, 6] = 1 + target = q.clone() + q[:, 7] = 2 * math.pi + v = torch.zeros(1, 14, dtype=torch.float64) + error = _error_values("ant", q, v, target, v) + assert error[0, 0].item() == math.pi / 4 + + +def test_orientation_metric_is_geodesic_and_quaternion_sign_invariant(): + q = torch.zeros(2, 15, dtype=torch.float64) + q[:, 6] = 1 + target = q.clone() + q[0, 5:7] = math.sqrt(0.5) + q[1, 6] = -1 + v = torch.zeros(2, 14, dtype=torch.float64) + error = _error_values("ant", q, v, target, v) + torch.testing.assert_close( + error[:, 3], + torch.tensor([math.pi / 2, 0], dtype=torch.float64), + atol=1e-12, + rtol=0, + ) + assert _safe_number(float("nan")) is None + assert _safe_number(float("inf")) is None + assert _safe_number(0.0) == 0.0 + + +def test_teacher_content_hash_changes_with_values_and_identity(): + row = {"id": "seed:1", "metadata": {"sampling_seed": 12}} + values = [ + np.zeros((3, 2), np.float32), + np.zeros((3, 2), np.float32), + np.zeros((2, 2), np.float32), + ] + expected = _teacher_row_sha256(row, values) + assert expected == _teacher_row_sha256( + {**row, "file": "different.h5"}, [v.copy() for v in values] + ) + assert expected != _teacher_row_sha256({**row, "id": "seed:2"}, values) + values[-1][0, 0] = 1 + assert expected != _teacher_row_sha256(row, values) diff --git a/tests/compute/dynamics/test_dynamics_task_reference.py b/tests/compute/dynamics/test_dynamics_task_reference.py new file mode 100644 index 0000000000..6f65154d54 --- /dev/null +++ b/tests/compute/dynamics/test_dynamics_task_reference.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. +# ---------------------------------------------------------------------------- + +from __future__ import annotations +import pytest +import torch +from embodichain.compute.dynamics import reference_task_values + + +def test_cartpole_observation_and_official_cost_definition() -> None: + q = torch.tensor([[1.0, 0.2], [-0.7, 2 * torch.pi + 0.2]], dtype=torch.double) + qd = torch.tensor([[0.3, -0.4], [0.1, 0.2]], dtype=torch.double) + obs, reward, failed = reference_task_values("cartpole", q, qd) + torch.testing.assert_close(obs[:, 1], torch.full((2,), 0.2, dtype=torch.double)) + expected = -( + torch.full((2,), 0.2**2, dtype=torch.double) + + 0.05 * q[:, 0] ** 2 + + 0.1 * (qd[:, 0] ** 2 + qd[:, 1] ** 2) + ) + torch.testing.assert_close(reward, expected) + assert obs.shape == (2, 4) and not failed.any() + assert torch.autograd.gradcheck( + lambda a, b: reference_task_values("cartpole", a, b)[:2], + (q.requires_grad_(), qd.requires_grad_()), + ) + + +def test_ant_com_velocity_forward_reward_and_translation_invariance() -> None: + q = torch.zeros(2, 15, dtype=torch.double) + q[:, 2] = 1.0 + angle = torch.tensor([0.2, -0.4], dtype=torch.double) + q[:, 5] = torch.sin(angle / 2) + q[:, 6] = torch.cos(angle / 2) + qd = torch.zeros(2, 14, dtype=torch.double) + qd[:, 1] = torch.tensor([0.3, 0.7], dtype=torch.double) + obs, reward, failed = reference_task_values("ant", q, qd) + torch.testing.assert_close(reward, qd[:, 1] + 0.1 + torch.cos(angle)) + moved = q.clone() + moved[:, :2] += torch.tensor([100.0, -100.0], dtype=torch.double) + changed = reference_task_values("ant", moved, qd) + torch.testing.assert_close(changed[0], obs) + torch.testing.assert_close(changed[1], reward) + assert obs.shape == (2, 27) and not failed.any() + + +def test_reference_failures_are_explicit_and_shapes_checked() -> None: + q = torch.zeros(2, 15) + q[:, 2] = torch.tensor([0.2, 1.0]) + q[0, 6] = 1.0 + _, _, failed = reference_task_values("ant", q, torch.zeros(2, 14)) + assert failed.tolist() == [True, True] # low torso / invalid quaternion + _, _, failed = reference_task_values( + "cartpole", torch.zeros(1, 2), torch.tensor([[float("nan"), 0.0]]) + ) + assert failed.item() + with pytest.raises(ValueError, match="dimensions"): + reference_task_values("ant", torch.zeros(2, 2), torch.zeros(2, 2)) diff --git a/tests/data_pipeline/dynamics/test_generation.py b/tests/data_pipeline/dynamics/test_generation.py new file mode 100644 index 0000000000..c04d60fd1d --- /dev/null +++ b/tests/data_pipeline/dynamics/test_generation.py @@ -0,0 +1,206 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Real analytical collection and an SDK-free generate/train/export/eval path.""" + +from __future__ import annotations + +import os +from pathlib import Path +import subprocess +import sys + +import h5py +import pytest +import torch + +from embodichain.data_pipeline.dynamics.generation import ( + DynamicsGenerationCfg, + _sample, + generate_dataset, +) +from embodichain.data_pipeline.dynamics._newton import _NewtonTeacher +from embodichain.data_pipeline.dynamics.shards import load_manifest, trajectory_split + + +@pytest.mark.parametrize( + "field,value", + [ + ("num_envs", True), + ("horizon", 2.5), + ("resume", "false"), + ("seed", False), + ("use_cuda_graph", 1), + ("effort_scale", float("nan")), + ("physics_profile", "unregistered"), + ], +) +def test_generation_rejects_invalid_config(field, value): + with pytest.raises((ValueError, TypeError)): + DynamicsGenerationCfg(**{field: value}) + + +@pytest.mark.parametrize("robot", ["cartpole", "ant"]) +def test_real_generation_partial_batch_resume_and_teacher_replay(tmp_path, robot): + config = { + "robot": robot, + "output_dir": str(tmp_path / "data"), + "num_envs": 2, + "trajectories": 3, + "horizon": 3, + "num_substeps": 2, + "seed": 7, + "split_seed": 11, + "trajectories_per_shard": 2, + "initial_states": "default", + "effort_scale": 0.01, + } + assert generate_dataset(config)["transitions"] == 9 + first = load_manifest(tmp_path / "data", verify=True) + original = [row for shard in first["shards"] for row in shard["trajectories"]] + assert ( + generate_dataset({**config, "trajectories": 5, "resume": True})["transitions"] + == 15 + ) + manifest = load_manifest(tmp_path / "data", verify=True) + rows = [row for shard in manifest["shards"] for row in shard["trajectories"]] + assert [row["id"] for row in rows] == [f"7:{i}" for i in range(5)] + assert rows[:3] == original + assert [row["split"] for row in rows] == [ + trajectory_split(row["id"], 11) for row in rows + ] + assert manifest["metadata"]["teacher_origin"]["collector"] == "direct_newton" + shard = manifest["shards"][0] + with h5py.File(tmp_path / "data" / shard["file"]) as file: + row = file[shard["trajectories"][0]["group"]] + q, qd, force = [ + torch.tensor(row[key][()]) + for key in ("states/q", "states/qd", "transitions/joint_f") + ] + contacts = { + key: torch.tensor(value[:1]) + for key, value in row["transitions"].items() + if key.startswith("contact_") + } + assert q.shape[0] == force.shape[0] + 1 == 4 + with _NewtonTeacher(robot, num_substeps=2) as teacher: + assert ( + teacher.manifest["physics_signature"] + == manifest["metadata"]["physics_signature"] + ) + teacher.reset(q[:1], qd[:1]) + previous, actual_contacts = teacher.snapshot() + for key, value in contacts.items(): + torch.testing.assert_close(actual_contacts[key], value) + teacher.step(force[:1]) + actual, _ = teacher.snapshot() + torch.testing.assert_close(actual["q"], q[1:2], atol=1e-5, rtol=1e-4) + torch.testing.assert_close(actual["qd"], qd[1:2], atol=1e-5, rtol=1e-4) + torch.testing.assert_close(previous["q"], q[:1], atol=0, rtol=0) + teacher.reset(q[:1], qd[:1]) + teacher.step(force[:1]) + repeated, _ = teacher.snapshot() + torch.testing.assert_close(repeated["q"], actual["q"], atol=1e-6, rtol=1e-5) + + +@pytest.mark.parametrize("robot", ["cartpole", "ant"]) +def test_actual_cli_pipeline_with_dexsim_imports_forbidden(tmp_path, robot): + """A native SDK or its custom solver extension cannot mask a dependency.""" + script = """ +import importlib.abc +import json +from pathlib import Path +import sys + +class NoSDK(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname == 'dexsim' or fullname.startswith(('dexsim.', 'embodichain.lab.')): + raise AssertionError('SDK-free workflow imported ' + fullname) +sys.meta_path.insert(0, NoSDK()) +from embodichain.cli.main import main +from embodichain.data_pipeline.dynamics.shards import load_manifest +root, robot = Path(sys.argv[1]), sys.argv[2] +def run(operation, config): + path = root / (operation + '.json') + path.write_text(json.dumps(config)) + main(['dynamics', operation, '--config', str(path)]) +run('generate', dict(robot=robot, output_dir='data', num_envs=2, trajectories=64, + horizon=4, num_substeps=2, initial_states='default', effort_scale=0.01)) +import yaml +model_cfg = yaml.safe_load((Path.cwd() / 'extensions/embodichain_tasks/configs/tasks/classic_control' / robot / 'dynamics.train.smoke.yaml').read_text())['model'] +model_cfg.update(n_embd=12, history_length=2, head_hidden_dims=[12]) +run('train', dict(dataset='data', output_dir='train', + model=model_cfg, + train=dict(total_steps=2, batch_size=2, validation_interval=1, checkpoint_interval=1, + device='cpu'))) +main(['dynamics', 'export', '--checkpoint', str(root/'train/best.pt'), + '--output', str(root/'export')]) +run('eval', dict(dataset='data', model='export', split='test', device='cpu', + output='reports/evaluation.json')) +report=json.loads((root/'reports/evaluation.json').read_text()) +assert report['windows'] > 0 +assert not any(name == 'dexsim' or name.startswith('embodichain.lab.') for name in sys.modules) +""" + result = subprocess.run( + [sys.executable, "-c", script, str(tmp_path), robot], + capture_output=True, + text=True, + env={**os.environ, "CUDA_VISIBLE_DEVICES": "", "OMP_NUM_THREADS": "2"}, + ) + assert result.returncode == 0, result.stdout + result.stderr + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "robot,substeps,profile", + [ + ("cartpole", 5, "warp_reference"), + ("ant", 16, "warp_reference"), + ("cartpole", 10, "native"), + ], +) +def test_cuda_graph_and_direct_teacher_replay(robot, substeps, profile): + if not torch.cuda.is_available(): + pytest.skip("CUDA teacher required") + q, qd, effort = _sample(robot, 5, 3, "random_effort", "default", 0.01) + q, qd, effort = q[None].cuda(), qd[None].cuda(), effort.cuda() + trajectories = [] + for capture in (False, True): + with _NewtonTeacher( + robot, + device="cuda:0", + num_substeps=substeps, + physics_profile=profile, + use_cuda_graph=capture, + ) as teacher: + assert (teacher.graph is not None) is capture + teacher.reset(q, qd) + rows = [] + for force in effort: + teacher.step(force[None]) + state, _ = teacher.snapshot() + rows.append(torch.cat((state["q"], state["qd"]), -1)) + trajectories.append(torch.stack(rows)) + teacher.reset(q, qd) + teacher.step(effort[:1]) + repeated, _ = teacher.snapshot() + torch.testing.assert_close( + torch.cat((repeated["q"], repeated["qd"]), -1), + rows[0], + atol=1e-6, + rtol=1e-5, + ) + torch.testing.assert_close(trajectories[0], trajectories[1], atol=1e-5, rtol=1e-4) diff --git a/tests/data_pipeline/dynamics/test_shards.py b/tests/data_pipeline/dynamics/test_shards.py new file mode 100644 index 0000000000..380a3401b0 --- /dev/null +++ b/tests/data_pipeline/dynamics/test_shards.py @@ -0,0 +1,343 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Persistence, split and bounded-window contracts.""" + +from __future__ import annotations + +import pickle +import numpy as np +import pytest +import torch +from torch.utils.data import DataLoader + +from embodichain.data_pipeline.dynamics.shards import ( + DynamicsShardWriter, + DynamicsWindowDataset, + load_manifest, + trajectory_split, +) + + +def trajectory(offset=0, length=5): + q = np.arange(length + 1, dtype=np.float32)[:, None] + offset + return {"q": q, "qd": q * 2}, { + "joint_f": q[:-1], + "contact_depths": np.zeros((length, 0), dtype=np.float32), + } + + +def test_commit_resume_and_final_partial_shard(tmp_path): + writer = DynamicsShardWriter(tmp_path, {"dt": 1 / 60}, trajectories_per_shard=2) + writer.add_trajectory("a", *trajectory(), split="train") + assert not load_manifest(tmp_path)["shards"] + writer.add_trajectory("b", *trajectory(20), split="train") + writer.add_trajectory("c", *trajectory(40), split="train") + writer.abort() + resumed = DynamicsShardWriter( + tmp_path, {"dt": 1 / 60}, trajectories_per_shard=2, resume=True + ) + assert resumed.committed_ids == {"a", "b"} + resumed.add_trajectory("c", *trajectory(40), split="test") + resumed.close() + manifest = load_manifest(tmp_path, verify=True) + assert [len(shard["trajectories"]) for shard in manifest["shards"]] == [2, 1] + assert len(DynamicsWindowDataset(tmp_path, "train", 3)) == 6 + assert len(DynamicsWindowDataset(tmp_path, "test", 3)) == 3 + + +def test_windows_do_not_cross_trajectory_and_prefix_masks(tmp_path): + with DynamicsShardWriter(tmp_path, {}) as writer: + writer.add_trajectory("a", *trajectory(), split="train") + writer.add_trajectory("b", *trajectory(100), split="validation") + dataset = DynamicsWindowDataset(tmp_path, "train", 3, include_prefix=True) + assert len(dataset) == 5 + assert dataset[0]["valid_mask"].tolist() == [True, False, False] + assert dataset[0]["states"]["q"].flatten().tolist() == [0.0, 1.0, 1.0, 1.0] + assert dataset[-1]["states"]["q"].flatten().tolist() == [2.0, 3.0, 4.0, 5.0] + assert all(dataset[i]["trajectory_id"] == "a" for i in range(len(dataset))) + with pytest.raises(IndexError): + dataset[5] + + +def test_independent_worker_handles_bounded_lru(tmp_path): + with DynamicsShardWriter(tmp_path, {}, trajectories_per_shard=1) as writer: + for index in range(4): + writer.add_trajectory(str(index), *trajectory(index * 10), split="train") + dataset = DynamicsWindowDataset(tmp_path, "train", 2, max_open_shards=1) + for index in range(len(dataset)): + dataset[index] + assert len(dataset._handles) == 1 + copy = pickle.loads(pickle.dumps(dataset)) + assert not copy._handles + assert copy[0]["states"]["q"].shape == (3, 1) + batches = list(DataLoader(dataset, batch_size=3, num_workers=2)) + assert sum(len(batch["length"]) for batch in batches) == len(dataset) + dataset.close() + copy.close() + + +def test_validation_and_integrity(tmp_path): + with DynamicsShardWriter(tmp_path, {}) as writer: + writer.add_trajectory("a", *trajectory(), split="train") + with pytest.raises(ValueError, match="duplicate"): + writer.add_trajectory("a", *trajectory()) + states, transitions = trajectory() + states["q"][0] = np.nan + with pytest.raises(ValueError, match="Non-finite"): + writer.add_trajectory("invalid", states, transitions) + with pytest.raises(ValueError, match="metadata"): + DynamicsShardWriter(tmp_path, {"changed": True}, resume=True) + shard = tmp_path / load_manifest(tmp_path)["shards"][0]["file"] + with shard.open("ab") as stream: + stream.write(b"changed") + with pytest.raises(ValueError, match="checksum"): + load_manifest(tmp_path, verify=True) + + +def test_stable_split_not_ingestion_order(): + forward = {str(i): trajectory_split(str(i), 7) for i in range(100)} + backward = {str(i): trajectory_split(str(i), 7) for i in reversed(range(100))} + assert forward == backward + assert set(forward.values()) == {"train", "validation", "test"} + + +@pytest.mark.parametrize("prefix", [False, True]) +def test_resident_windows_preserve_order_padding_ownership_and_bounded_staging( + tmp_path, monkeypatch, prefix +): + with DynamicsShardWriter(tmp_path, {}, trajectories_per_shard=2) as writer: + for index, length in enumerate([3, 5, 11]): + writer.add_trajectory( + str(index), *trajectory(index * 100, length), split="train" + ) + dataset = DynamicsWindowDataset( + tmp_path, "train", 3, include_prefix=prefix, read_page_steps=2 + ) + requests = [-1, 0, 3, 0, len(dataset) // 2] + expected = dataset.get_batch(requests) + original = dataset._read_block + spans = [] + + def tracked(source, name, start, end, shape, dtype): + spans.append(end - start) + return original(source, name, start, end, shape, dtype) + + monkeypatch.setattr(dataset, "_read_block", tracked) + used = dataset.enable_device_cache("cpu", 1024 * 1024) + assert used > 0 and max(spans) <= 2 + + def forbidden(*args, **kwargs): + raise AssertionError("Resident batches must not reopen trajectory data") + + monkeypatch.setattr(dataset, "_read_block", forbidden) + actual = dataset.get_batch(requests) + _assert_same_batch(actual, expected) + actual["states"]["q"][0].zero_() + _assert_same_batch(dataset.get_batch(requests), expected) + monkeypatch.delattr(dataset, "_read_block") + copied = pickle.loads(pickle.dumps(dataset)) + assert copied._resident is None + result = dataset.get_batch(requests) + dataset.close() + _assert_same_batch(result, expected) + + +def test_resident_budget_failure_keeps_lazy_dataset_usable(tmp_path): + with DynamicsShardWriter(tmp_path, {}) as writer: + writer.add_trajectory("a", *trajectory(), split="train") + dataset = DynamicsWindowDataset(tmp_path, "train", 3) + with pytest.raises(ValueError, match="exceeding cache budget"): + dataset.enable_device_cache("cpu", 1) + assert dataset._resident is None + assert dataset[0]["length"] == 3 + + +@pytest.mark.gpu +def test_resident_cuda_windows_match_cpu_and_survive_close(tmp_path): + with DynamicsShardWriter(tmp_path, {}) as writer: + writer.add_trajectory("a", *trajectory(), split="train") + dataset = DynamicsWindowDataset(tmp_path, "train", 3, include_prefix=True) + expected = dataset.get_batch([0, -1, 0]) + dataset.enable_device_cache("cuda:0", 1024 * 1024) + actual = dataset.get_batch([0, -1, 0]) + dataset.close() + for kind in ["states", "transitions"]: + for name, value in actual[kind].items(): + assert value.device.type == "cuda" + torch.testing.assert_close( + value.cpu(), expected[kind][name], atol=0, rtol=0 + ) + + +def test_finalizer_is_best_effort_but_explicit_close_surfaces_errors(): + from collections import OrderedDict + from types import SimpleNamespace + from unittest.mock import Mock + + handle = SimpleNamespace(close=Mock(side_effect=RuntimeError("h5py finalized"))) + dataset = object.__new__(DynamicsWindowDataset) + dataset._handles = OrderedDict([("shard.h5", handle)]) + with pytest.raises(RuntimeError, match="h5py finalized"): + dataset.close() + dataset.__del__() + assert handle.close.call_count == 2 + + +def _assert_same_batch(actual, expected): + import torch + + if isinstance(actual, dict): + assert actual.keys() == expected.keys() + for key in actual: + _assert_same_batch(actual[key], expected[key]) + elif isinstance(actual, torch.Tensor): + assert actual.dtype == expected.dtype + assert torch.equal(actual, expected) + else: + assert actual == expected + + +@pytest.mark.parametrize("history", [1, 3, 10]) +@pytest.mark.parametrize("include_prefix", [False, True]) +def test_grouped_batch_matches_individual_windows_across_shards( + tmp_path, history, include_prefix +): + import torch + from torch.utils.data import default_collate + + with DynamicsShardWriter(tmp_path, {}, trajectories_per_shard=1) as writer: + for index, length in enumerate([1, 3, 9, 25, 30, 11, 7, 41]): + writer.add_trajectory( + str(index), + *trajectory(index * 100, length), + split="train" if index % 3 else "validation", + ) + reader = DynamicsWindowDataset( + tmp_path, + "train", + history, + include_prefix=include_prefix, + max_open_shards=2, + read_page_steps=8, + ) + baseline = DynamicsWindowDataset( + tmp_path, "train", history, include_prefix=include_prefix, max_open_shards=2 + ) + generator = torch.Generator().manual_seed(7) + for size in (1, 16, 512): + indices = torch.randint(len(reader), (size,), generator=generator) + indices[0] = -1 + _assert_same_batch( + reader.get_batch(indices), + default_collate([baseline[int(index)] for index in indices]), + ) + assert len(reader._handles) <= 2 + reader.close() + baseline.close() + + +def test_grouped_batch_preserves_duplicate_and_cross_call_ownership(tmp_path): + import torch + + with DynamicsShardWriter(tmp_path, {}) as writer: + writer.add_trajectory("a", *trajectory(), split="train") + dataset = DynamicsWindowDataset(tmp_path, "train", 3, include_prefix=True) + first = dataset.get_batch([0, 0, 4]) + untouched = first["states"]["q"][1].clone() + first["states"]["q"][0].fill_(-999) + assert torch.equal(first["states"]["q"][1], untouched) + second = dataset.get_batch([0, 4]) + assert torch.equal(second["states"]["q"][0], untouched) + assert torch.all(first["states"]["q"][0] == -999) + dataset.close() + assert torch.equal(second["states"]["q"][0], untouched) + + +def test_grouped_reads_are_bounded_for_sparse_long_trajectory(tmp_path, monkeypatch): + import h5py + from torch.utils.data import default_collate + + with DynamicsShardWriter(tmp_path, {}) as writer: + writer.add_trajectory("long", *trajectory(length=4097), split="train") + dataset = DynamicsWindowDataset(tmp_path, "train", 7, read_page_steps=32) + selections = [] + original = dataset._read_block + + def observed(group, name, start, stop, shape, dtype): + block = original(group, name, start, stop, shape, dtype) + assert block.shape == (stop - start, *shape) + selections.append(slice(start, stop)) + return block + + monkeypatch.setattr(dataset, "_read_block", observed) + indices = [0, 31, 32, 63, 2048, len(dataset) - 1, -1, 1] + batch = dataset.get_batch(indices) + assert selections + assert max(selection.stop - selection.start for selection in selections) <= 39 + _assert_same_batch(batch, default_collate([dataset[index] for index in indices])) + dataset.close() + + +def test_grouped_batch_validates_indices_before_io(tmp_path): + with DynamicsShardWriter(tmp_path, {}) as writer: + writer.add_trajectory("one", *trajectory(), split="train") + dataset = DynamicsWindowDataset(tmp_path, "train", 3) + with pytest.raises(ValueError, match="at least one"): + dataset.get_batch([]) + with pytest.raises(TypeError): + dataset.get_batch([1.5]) + with pytest.raises(IndexError): + dataset.get_batch([len(dataset)]) + assert not dataset._handles + + +def test_grouped_low_level_read_checks_shape_before_buffer_write(tmp_path): + import h5py + + with DynamicsShardWriter(tmp_path, {}) as writer: + writer.add_trajectory("one", *trajectory(), split="train") + manifest = load_manifest(tmp_path) + with h5py.File(tmp_path / manifest["shards"][0]["file"], "r+") as file: + del file["0/states/q"] + file["0/states"].create_dataset("q", data=np.zeros((6, 2), np.float32)) + reader = DynamicsWindowDataset(tmp_path, "train", 3, verify=False) + with pytest.raises(ValueError, match="shape/range mismatch"): + reader.get_batch([0]) + reader.close() + + +def test_grouped_hyperslabs_preserve_dtype_multiaxis_and_utf8_fields(tmp_path): + from torch.utils.data import default_collate + + states = { + "q": np.arange(12, dtype=np.float64).reshape(6, 2), + "qd": np.arange(6, dtype=np.int16).reshape(6, 1), + "额外值": np.arange(36, dtype=np.float32).reshape(6, 2, 3), + } + transitions = { + "joint_f": np.arange(5, dtype=np.uint8).reshape(5, 1), + "contact_masks": np.array([[True], [False], [True], [True], [False]]), + "empty": np.zeros((5, 0), np.float32), + } + with DynamicsShardWriter(tmp_path, {}) as writer: + writer.add_trajectory("typed", states, transitions, split="train") + reader = DynamicsWindowDataset(tmp_path, "train", 3) + indices = [2, 0, 1, 1] + _assert_same_batch( + reader.get_batch(indices), default_collate([reader[index] for index in indices]) + ) + reader.close() diff --git a/tests/learning/dynamics/test_comparison.py b/tests/learning/dynamics/test_comparison.py new file mode 100644 index 0000000000..c7da565317 --- /dev/null +++ b/tests/learning/dynamics/test_comparison.py @@ -0,0 +1,190 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Immutable experiment preparation and truthful pending qualification reports.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pytest +import torch + +from embodichain.data_pipeline.dynamics.shards import DynamicsShardWriter +from embodichain.learning.dynamics.comparison import ( + prepare_comparison, + run_comparison, + run_candidate, + _read_plan, +) +from embodichain.learning.dynamics.report import write_comparison_report + + +def configuration(tmp_path, monkeypatch): + import embodichain.learning.dynamics.comparison as comparison + from embodichain.learning.dynamics.reference import OFFICIAL_COMMIT + + monkeypatch.setattr( + comparison, + "reference_identity", + lambda path: { + "root": str(tmp_path / "source"), + "commit": OFFICIAL_COMMIT, + "files": {}, + "source_patches": [], + }, + ) + with DynamicsShardWriter( + tmp_path / "data", {"codec": {"robot": "cartpole"}, "dt": 1 / 60} + ) as writer: + for split in ("train", "validation", "test"): + q = np.zeros((6, 2), np.float32) + qd = np.zeros_like(q) + writer.add_trajectory( + split, + {"q": q, "qd": qd}, + {"joint_f": np.ones((5, 2), np.float32)}, + split=split, + ) + return { + "output_dir": str(tmp_path / "experiment"), + "source": str(tmp_path / "source"), + "datasets": {"official_warp": None, "current_newton": str(tmp_path / "data")}, + "scope": "smoke", + "seeds": [0], + "model": { + "input_dims": {"states_embedding": 4, "joint_f": 2}, + "output_dim": 4, + "n_embd": 12, + "n_head": 3, + "n_layer": 1, + "history_length": 3, + "block_size": 3, + "head_hidden_dims": [12], + }, + "train": { + "total_steps": 2, + "batch_size": 2, + "validation_interval": 1, + "device": "cpu", + }, + } + + +def test_preparation_shared_artifacts_and_missing_row_are_explicit( + tmp_path, monkeypatch +): + torch.set_num_threads(1) + path = prepare_comparison(configuration(tmp_path, monkeypatch)) + root, plan = _read_plan(path) + assert len(plan["cells"]) == 4 + assert sum(cell["status"] == "pending" for cell in plan["cells"]) == 2 + specs = [ + json.loads(Path(cell["spec"]).read_text()) + for cell in plan["cells"] + if cell["status"] == "prepared" + ] + assert specs[0]["initial_model_sha256"] == specs[1]["initial_model_sha256"] + assert specs[0]["sample_order_sha256"] == specs[1]["sample_order_sha256"] + with pytest.raises(ValueError, match="trust_external_code"): + run_comparison(path) + report = write_comparison_report(path) + assert report["status"] == "pending" + assert any("Five" in reason for reason in report["pending"]) + assert (root / "report" / "training-curves.png").is_file() + assert (root / "report" / "REPORT.md").is_file() + + +def test_changed_common_artifact_fails_before_execution(tmp_path, monkeypatch): + path = prepare_comparison(configuration(tmp_path, monkeypatch)) + _, plan = _read_plan(path) + name = next(iter(plan["artifacts"])) + with (path.parent / name).open("ab") as stream: + stream.write(b"tampered") + with pytest.raises(ValueError, match="Immutable comparison artifact"): + _read_plan(path) + + +def test_candidate_worker_uses_frozen_artifacts(tmp_path, monkeypatch): + path = prepare_comparison(configuration(tmp_path, monkeypatch)) + _, plan = _read_plan(path) + cell = next( + cell + for cell in plan["cells"] + if cell["status"] == "prepared" and cell["trainer"] == "embodichain" + ) + result = run_candidate(json.loads(Path(cell["spec"]).read_text())) + assert result["completed_updates"] == 2 + assert result["test"]["split"] == "test" + assert result["qualification"] == "matched_training_run_only" + + +def test_rollout_attachment_preserves_measurements_and_requires_identity(tmp_path): + from embodichain.learning.dynamics.report import annotate_rollout_evidence + + source = tmp_path / "source.json" + source.write_text( + json.dumps({"trajectories": 2048, "horizons": {"1": {"q_mae": 0.1}}}) + ) + with pytest.raises(ValueError, match="model_sha256"): + annotate_rollout_evidence(source, "run", tmp_path / "bad.json") + raw = { + "model_sha256": "abc", + "motion_mode": "passive", + "trajectories": 2048, + "horizons": { + "1": { + "q_mae": 0.1, + "qd_mae": 0.2, + "failure_rate": 0.0, + "nonfinite_rate": 0.0, + }, + "10": { + "q_mae": 0.2, + "qd_mae": 0.3, + "failure_rate": 0.1, + "nonfinite_rate": 0.05, + }, + }, + } + source.write_text(json.dumps(raw)) + result = annotate_rollout_evidence(source, "run", tmp_path / "attached.json") + assert result["horizons"] == [1, 10] + assert result["errors"] == raw["horizons"] + assert result["failure_rate"] == 0.1 + assert "policy_return" not in result + + +def test_preparation_keeps_selected_virtualenv_interpreter_symlink( + tmp_path, monkeypatch +): + import sys + + chosen = tmp_path / "selected-env/bin/python" + chosen.parent.mkdir(parents=True) + chosen.symlink_to(sys.executable) + config = configuration(tmp_path, monkeypatch) + config["python"] = str(chosen) + path = prepare_comparison(config) + _, plan = _read_plan(path) + assert plan["python"] == str(chosen) + assert all( + cell["command"][0] == str(chosen) + for cell in plan["cells"] + if cell["status"] == "prepared" + ) diff --git a/tests/learning/dynamics/test_dynamics_runtime_identity.py b/tests/learning/dynamics/test_dynamics_runtime_identity.py new file mode 100644 index 0000000000..15042ae47e --- /dev/null +++ b/tests/learning/dynamics/test_dynamics_runtime_identity.py @@ -0,0 +1,183 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Execution provenance is read-only, simulation-free and best effort.""" + +from __future__ import annotations + +import json +import subprocess +import sys +from types import SimpleNamespace + +import pytest +import torch + +from embodichain.learning.dynamics import _runtime + + +def test_cpu_identity_reflects_actual_parameters_without_side_effects(monkeypatch): + model = torch.nn.Linear(2, 1, dtype=torch.float64) + rng = torch.get_rng_state().clone() + modules = set(sys.modules) + monkeypatch.delitem(sys.modules, "warp", raising=False) + + def unexpected(*args, **kwargs): + pytest.fail("CPU identity must not query CUDA or nvidia-smi") + + monkeypatch.setattr(_runtime, "_driver_inventory", unexpected) + monkeypatch.setattr(torch.cuda, "get_device_name", unexpected) + result = _runtime._training_runtime_identity(model) + assert result["parameter_placements"] == [ + {"device": "cpu", "dtype": "torch.float64"} + ] + assert result["cuda"]["used"] is False + assert result["cuda"]["driver_inventory"] == {"status": "not_used"} + assert result["torch_version"] == str(torch.__version__) + assert result["python"] and result["platform"] + assert result["warp"] == {"status": "not_loaded"} + assert not any( + name.split(".")[0] in {"warp", "dexsim", "newton"} + for name in set(sys.modules) - modules + ) + assert torch.equal(rng, torch.get_rng_state()) + assert json.loads(json.dumps(result)) == result + + +def test_empty_and_mixed_parameter_metadata_and_loaded_warp(monkeypatch): + assert ( + _runtime._training_runtime_identity(torch.nn.Identity())[ + "parameter_placement_status" + ] + == "no_parameters" + ) + model = torch.nn.Module() + model.register_parameter( + "a", torch.nn.Parameter(torch.zeros(1, dtype=torch.float32)) + ) + model.register_parameter( + "b", torch.nn.Parameter(torch.zeros(1, dtype=torch.float64)) + ) + monkeypatch.setitem( + sys.modules, "warp", SimpleNamespace(__version__="test-version") + ) + result = _runtime._training_runtime_identity(model) + assert result["parameter_placements"] == [ + {"device": "cpu", "dtype": "torch.float32"}, + {"device": "cpu", "dtype": "torch.float64"}, + ] + assert result["warp"] == {"status": "already_loaded", "version": "test-version"} + + +def _mock_cuda_model(): + # Metadata only: no CUDA allocations or initialization in these tests. + return SimpleNamespace( + parameters=lambda: iter( + [ + SimpleNamespace(device=torch.device("cuda:1"), dtype=torch.float32), + SimpleNamespace(device=torch.device("cuda:1"), dtype=torch.float32), + ] + ) + ) + + +def test_cuda_inventory_is_bounded_and_keeps_physical_ids_separate(monkeypatch): + monkeypatch.setattr(_runtime.shutil, "which", lambda name: "/usr/bin/nvidia-smi") + calls = [] + + def run(command, **kwargs): + calls.append((command, kwargs)) + return SimpleNamespace( + returncode=0, stdout='0, GPU-physical, "Example, GPU", 595.1\n', stderr="" + ) + + monkeypatch.setattr(_runtime.subprocess, "run", run) + monkeypatch.setattr(torch.cuda, "get_device_name", lambda device: "Logical GPU") + result = _runtime._training_runtime_identity(_mock_cuda_model()) + assert result["cuda"]["devices"] == [ + {"device": "cuda:1", "status": "available", "name": "Logical GPU"} + ] + inventory = result["cuda"]["driver_inventory"] + assert inventory == { + "status": "available", + "physical_gpus": [ + { + "index": "0", + "uuid": "GPU-physical", + "name": "Example, GPU", + "driver_version": "595.1", + } + ], + } + assert calls[0][0][1:] == [ + "--query-gpu=index,uuid,name,driver_version", + "--format=csv,noheader,nounits", + ] + assert calls[0][1]["timeout"] == 2.0 + assert not calls[0][1].get("shell", False) + + +@pytest.mark.parametrize( + "failure, expected", + [ + ("missing", "unavailable"), + ("timeout", "error"), + ("exit", "error"), + ("oserror", "error"), + ("malformed", "error"), + ("oversized", "error"), + ], +) +def test_cuda_inventory_failure_does_not_abort_receipt(monkeypatch, failure, expected): + monkeypatch.setattr( + _runtime.shutil, + "which", + lambda name: None if failure == "missing" else "/mock/nvidia-smi", + ) + + def run(command, **kwargs): + if failure == "timeout": + raise subprocess.TimeoutExpired(command, kwargs["timeout"]) + if failure == "oserror": + raise OSError("driver unavailable") + return SimpleNamespace( + returncode=1 if failure == "exit" else 0, + stdout="x" * 65537 if failure == "oversized" else "bad,row", + stderr="driver error", + ) + + monkeypatch.setattr(_runtime.subprocess, "run", run) + + def device_name(device): + raise RuntimeError("device unavailable") + + monkeypatch.setattr(torch.cuda, "get_device_name", device_name) + result = _runtime._training_runtime_identity(_mock_cuda_model()) + assert result["cuda"]["devices"][0]["status"] == "error" + assert result["cuda"]["driver_inventory"]["status"] == expected + json.dumps(result, allow_nan=False) + + +def test_fresh_process_identity_does_not_import_simulation(): + code = """ +import sys +import torch +from embodichain.learning.dynamics._runtime import _training_runtime_identity +_training_runtime_identity(torch.nn.Linear(1, 1)) +assert not any(name.split('.')[0] in {'warp', 'dexsim', 'newton'} for name in sys.modules) +assert not any(name.startswith('embodichain.lab.sim') for name in sys.modules) +""" + subprocess.run([sys.executable, "-c", code], check=True, timeout=30) diff --git a/tests/learning/dynamics/test_reference.py b/tests/learning/dynamics/test_reference.py new file mode 100644 index 0000000000..470e14e1bc --- /dev/null +++ b/tests/learning/dynamics/test_reference.py @@ -0,0 +1,935 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Legacy transport and opt-in independent official-trainer execution.""" + +from __future__ import annotations + +import json +import hashlib +import os +from pathlib import Path +import subprocess +import sys +from types import ModuleType, SimpleNamespace + +import h5py +import numpy as np +import pytest +import torch + +from embodichain.data_pipeline.dynamics.shards import ( + DynamicsShardWriter, + DynamicsWindowDataset, +) +from embodichain.learning.dynamics.reference import ( + export_legacy_dataset, + run_reference, + reference_identity, + OFFICIAL_COMMIT, +) + + +def test_external_execution_requires_explicit_trust(): + with pytest.raises(ValueError, match="trust_external_code"): + run_reference({}) + + +@pytest.mark.parametrize("profile", ["reference", "unknown"]) +def test_original_method_adapter_rejects_unsupported_profile_before_source(profile): + with pytest.raises(ValueError, match="numerical_profile"): + run_reference( + {"train": {"numerical_profile": profile}}, trust_external_code=True + ) + + +def test_original_method_adapter_rejects_math_model_before_data(tmp_path, monkeypatch): + import embodichain.learning.dynamics.reference as reference + import embodichain.learning.models.nerd as models + + for name in ("algorithms", "utils"): + package = ModuleType(name) + package.__path__ = [] + monkeypatch.setitem(sys.modules, name, package) + vanilla = ModuleType("algorithms.vanilla_trainer") + vanilla.VanillaTrainer = object + timers = ModuleType("utils.time_report") + timers.TimeReport = object + monkeypatch.setitem(sys.modules, vanilla.__name__, vanilla) + monkeypatch.setitem(sys.modules, timers.__name__, timers) + monkeypatch.setattr( + reference, "reference_identity", lambda source: {"root": str(tmp_path)} + ) + monkeypatch.setattr( + models, + "load_model", + lambda *args, **kwargs: ( + SimpleNamespace(cfg=SimpleNamespace(attention_backend="math")), + {}, + ), + ) + monkeypatch.setattr(sys, "path", list(sys.path)) + threads, precision = torch.get_num_threads(), torch.get_float32_matmul_precision() + try: + with pytest.raises(ValueError, match="math attention initial_model"): + run_reference( + { + "source": "unused", + "initial_model": "unused", + "train": {"device": "cpu"}, + }, + trust_external_code=True, + ) + finally: + torch.set_num_threads(threads) + torch.set_float32_matmul_precision(precision) + + +def test_official_transport_uses_the_exact_native_codec_velocity_order(): + from embodichain.compute.dynamics import NeRDStateCodec + from embodichain.learning.dynamics.reference import _legacy_states + + with np.load( + Path(__file__).parents[2] / "compute/dynamics/data/ant-official-frame.npz" + ) as fixture: + q = torch.from_numpy(fixture["q"].copy()) + qd = torch.from_numpy(fixture["qd"].copy()) + codec = NeRDStateCodec("ant") + actual = _legacy_states(q, qd, codec) + torch.testing.assert_close(actual[..., :15], q, atol=0, rtol=0) + torch.testing.assert_close( + actual[..., 15:], codec.native_to_legacy(q, qd), atol=0, rtol=0 + ) + + +@pytest.mark.parametrize( + "device", ["cpu", pytest.param("cuda:0", marks=pytest.mark.gpu)] +) +def test_training_transport_device_is_independent_of_raw_residency(device): + from embodichain.compute.dynamics import NeRDStateCodec + from embodichain.learning.dynamics.reference import _official_batch + + if device.startswith("cuda") and not torch.cuda.is_available(): + pytest.skip("CUDA is unavailable") + with np.load( + Path(__file__).parents[2] / "compute/dynamics/data/ant-official-frame.npz" + ) as fixture: + q = torch.from_numpy(fixture["q"].copy()) + qd = torch.from_numpy(fixture["qd"].copy()) + shape = (q.shape[0], q.shape[1] - 1) + batch = { + "states": {"q": q, "qd": qd}, + "transitions": { + "joint_f": torch.zeros(*shape, 14), + "contact_depths": torch.zeros(*shape, 25), + }, + } + resident = { + key: {name: tensor.to(device) for name, tensor in fields.items()} + for key, fields in batch.items() + } + codec = NeRDStateCodec("ant") + lazy = _official_batch(batch, codec, device=device) + cached = _official_batch(resident, codec, device=device) + for name, actual in lazy.items(): + assert actual.device == torch.device(device) + torch.testing.assert_close(actual, cached[name], atol=0, rtol=0) + + +def test_legacy_export_preserves_lengths_splits_and_native_to_twist(tmp_path): + metadata = {"codec": {"robot": "ant", "base_com": [0.0, 0.0, 0.0]}, "dt": 1 / 60} + with DynamicsShardWriter(tmp_path / "data", metadata) as writer: + for index, length in enumerate((3, 5)): + q = np.zeros((length + 1, 15), np.float32) + q[:, 0] = 1.0 + q[:, 6] = 1.0 + qd = np.zeros((length + 1, 14), np.float32) + qd[:, 1] = 2.0 + qd[:, 5] = 1.0 + transitions = { + "joint_f": np.zeros((length, 14), np.float32), + "contact_depths": np.zeros((length, 1), np.float32), + "contact_thicknesses_0": np.full((length, 1), 0.02, np.float32), + "contact_thicknesses_1": np.full((length, 1), 0.03, np.float32), + "contact_normals": np.tile([0.0, 0.0, 1.0], (length, 1)).astype( + np.float32 + ), + "contact_points_1": np.zeros((length, 3), np.float32), + } + writer.add_trajectory( + str(index), {"q": q, "qd": qd}, transitions, split="train" + ) + result = export_legacy_dataset(tmp_path / "data", tmp_path / "legacy.h5", "train") + with h5py.File(tmp_path / "legacy.h5") as file: + data = file["data"] + assert data["states"].shape == (5, 2, 29) + assert data["traj_lengths"][()].tolist() == [3, 5] + np.testing.assert_allclose( + data["states"][0, 0, 15:21], [0.0, 0.0, 1.0, 0.0, 1.0, 0.0] + ) + np.testing.assert_allclose(data["contact_thicknesses"][0], 0.05) + assert data["joint_acts"].shape == (5, 2, 8) + assert result["native_to_legacy_qd"] is True + + +@pytest.mark.skipif( + not os.getenv("NERD_REFERENCE_SOURCE"), + reason="Explicit trusted official checkout/dependency environment required", +) +def test_actual_official_main_two_updates_match_migrated(tmp_path): + """External reference is opt-in; synthetic data tests migration, not accuracy.""" + from embodichain.compute.dynamics import NeRDStateCodec + from embodichain.learning.models.nerd import ( + NeRDModel, + NeRDModelCfg, + export_model, + load_model, + ) + from embodichain.learning.dynamics.trainer import DynamicsTrainer, DynamicsTrainCfg + from embodichain.learning.dynamics.reference import _hash + from embodichain.learning.dynamics.workflow import run_train, run_export + + metadata = {"codec": {"robot": "cartpole"}, "dt": 1 / 60} + with DynamicsShardWriter(tmp_path / "data", metadata) as writer: + for split in ("train", "validation", "test"): + q = np.zeros((7, 2), np.float32) + q[:, 0] = np.arange(7) * 0.02 + qd = np.zeros_like(q) + qd[:, 0] = 1.2 + writer.add_trajectory( + split, + {"q": q, "qd": qd}, + {"joint_f": np.zeros((6, 2), np.float32)}, + split=split, + ) + source = os.environ["NERD_REFERENCE_SOURCE"] + assert reference_identity(source)["commit"] == OFFICIAL_COMMIT + torch.set_num_threads(1) + torch.manual_seed(4) + model = NeRDModel( + NeRDModelCfg( + input_dims={"states_embedding": 4, "joint_f": 2}, + output_dim=4, + n_embd=12, + n_head=3, + n_layer=1, + history_length=3, + block_size=3, + head_hidden_dims=(12,), + ) + ) + train = DynamicsWindowDataset(tmp_path / "data", "train", 3) + validation = DynamicsWindowDataset(tmp_path / "data", "validation", 3) + cfg = DynamicsTrainCfg(total_steps=2, batch_size=2, validation_interval=1, seed=4) + trainer = DynamicsTrainer( + model, + NeRDStateCodec("cartpole"), + train, + validation, + cfg, + tmp_path / "scratch", + manifest=metadata, + ) + export_model(tmp_path / "initial", trainer.model, metadata) + np.save(tmp_path / "order.npy", np.array([[0, 1], [1, 2]], np.int64)) + spec = { + "source": source, + "dataset": str(tmp_path / "data"), + "dataset_fingerprint": train.fingerprint, + "initial_model": str(tmp_path / "initial"), + "initial_model_sha256": _hash(tmp_path / "initial" / "weights.pt"), + "sample_order": str(tmp_path / "order.npy"), + "sample_order_sha256": _hash(tmp_path / "order.npy"), + "train": cfg.to_dict(), + "output_dir": str(tmp_path / "official"), + } + path = tmp_path / "spec.json" + path.write_text(json.dumps(spec)) + env = { + **os.environ, + "OMP_NUM_THREADS": "1", + "MKL_NUM_THREADS": "1", + "CUDA_VISIBLE_DEVICES": "", + "PYTHONPATH": os.pathsep.join( + [ + os.environ.get("NERD_REFERENCE_DEPS", ""), + source, + str(Path(__file__).resolve().parents[3]), + ] + ), + } + process = subprocess.run( + [ + sys.executable, + "-m", + "embodichain.learning.dynamics.reference", + "--spec", + str(path), + "--trust-external-code", + ], + env=env, + capture_output=True, + text=True, + ) + assert process.returncode == 0, process.stdout + process.stderr + receipt = json.loads((tmp_path / "official" / "result.json").read_text()) + assert receipt["runtime_identity"]["parameter_placements"] == [ + {"device": "cpu", "dtype": "torch.float32"} + ] + assert ( + receipt["runtime_identity"]["cuda"]["driver_inventory"]["status"] == "not_used" + ) + run_train( + { + "dataset": spec["dataset"], + "output_dir": str(tmp_path / "candidate"), + "fine_tune": spec["initial_model"], + "sample_order": spec["sample_order"], + "train": cfg.to_dict(), + } + ) + run_export(tmp_path / "candidate" / "last.pt", tmp_path / "candidate-export") + reference, _ = load_model(tmp_path / "official" / "last-export") + candidate, _ = load_model(tmp_path / "candidate-export") + for name, value in candidate.state_dict().items(): + torch.testing.assert_close( + value, reference.state_dict()[name], atol=1e-5, rtol=1e-4 + ) + report = json.loads((tmp_path / "official" / "result.json").read_text()) + assert report["completed_updates"] == 2 + assert report["reference"]["source_patches"] == [] + train.close() + validation.close() + + +def test_legacy_roundtrip_preserves_native_qd_and_declared_split(tmp_path): + from embodichain.learning.dynamics.reference import import_legacy_datasets + + metadata = {"codec": {"robot": "cartpole"}, "dt": 1 / 60} + with DynamicsShardWriter(tmp_path / "original", metadata) as writer: + q = np.arange(12, dtype=np.float32).reshape(6, 2) * 0.01 + writer.add_trajectory( + "sample", + {"q": q, "qd": q * 0.5}, + {"joint_f": np.zeros((5, 2), np.float32)}, + split="validation", + ) + export_legacy_dataset(tmp_path / "original", tmp_path / "legacy.h5", "validation") + imported = import_legacy_datasets( + {"validation": tmp_path / "legacy.h5"}, tmp_path / "imported", metadata + ) + assert imported["trajectories"] == 1 + data = DynamicsWindowDataset(tmp_path / "imported", "validation", 3) + np.testing.assert_allclose(data[0]["states"]["q"], q[:4]) + np.testing.assert_allclose(data[0]["states"]["qd"], q[:4] * 0.5) + assert len(DynamicsWindowDataset(tmp_path / "imported", "train", 3)) == 0 + data.close() + + +def _write_source_ant_velocities(path: Path) -> np.ndarray: + """Make continuous source states exposing cancellation and signed-zero loss.""" + horizon, trajectories, state_dim = 5, 2, 29 + states = np.zeros((horizon + 1, trajectories, state_dim), np.float32) + states[..., :3] = [10.0, 20.0, 30.0] + states[..., 6] = 1.0 + states[..., 15:18] = [5.0, 6.0, 7.0] + states[..., 18:21] = [0.3, 0.1, 0.2] + states[:, :, 18] += np.arange(horizon + 1, dtype=np.float32)[:, None] * 0.01 + states[0, 1, 19] = -0.0 + with h5py.File(path, "w") as file: + data = file.create_group("data") + data.attrs["mode"] = "trajectory" + data.create_dataset("states", data=states[:-1]) + data.create_dataset("next_states", data=states[1:]) + data.create_dataset("traj_lengths", data=np.array([3, 5], np.int32)) + data.create_dataset("root_body_q", data=states[:-1, :, :7]) + data.create_dataset( + "joint_acts", data=np.zeros((horizon, trajectories, 8), np.float32) + ) + data.create_dataset( + "contact_depths", + data=np.broadcast_to( + np.array([0.01, 100.0], np.float32), (horizon, trajectories, 2) + ), + ) + data.create_dataset( + "contact_thicknesses", + data=np.full((horizon, trajectories, 2), 0.02, np.float32), + ) + data.create_dataset( + "contact_points_1", data=np.zeros((horizon, trajectories, 6), np.float32) + ) + data.create_dataset( + "contact_normals", data=np.zeros((horizon, trajectories, 6), np.float32) + ) + data.create_dataset( + "gravity_dir", + data=np.broadcast_to( + np.array([0.0, -1.0, 0.0], np.float32), (horizon, trajectories, 3) + ), + ) + return states + + +def test_original_velocity_import_roundtrip_preserves_fp32_bits_and_defaults(tmp_path): + from embodichain.compute.dynamics import NeRDStateCodec + from embodichain.data_pipeline.dynamics.shards import load_manifest + from embodichain.learning.dynamics.reference import ( + import_legacy_datasets, + _official_batch, + ) + + path = tmp_path / "source.h5" + original = _write_source_ant_velocities(path) + source_sha = hashlib.sha256(path.read_bytes()).hexdigest() + metadata = {"codec": {"robot": "ant"}, "dt": 1 / 60} + imported = import_legacy_datasets( + {"validation": path}, + tmp_path / "preserved", + metadata, + preserve_legacy_velocity=True, + trajectories_per_shard=1, + ) + import_legacy_datasets({"validation": path}, tmp_path / "default", metadata) + manifest = load_manifest(tmp_path / "preserved", verify=True) + baseline = load_manifest(tmp_path / "default", verify=True) + assert imported["trajectories"] == 2 + assert manifest["fields"]["states"]["legacy_velocity"]["shape"] == [14] + assert "legacy_velocity" not in baseline["fields"]["states"] + assert "legacy_velocity" not in baseline["metadata"] + without_contract = dict(manifest["metadata"]) + contract = without_contract.pop("legacy_velocity") + assert without_contract == baseline["metadata"] + assert contract["world_frame"] == "source" + assert contract["field"] == "states/legacy_velocity" + assert manifest["metadata"]["legacy_sources"]["validation"]["sha256"] == source_sha + assert metadata == {"codec": {"robot": "ant"}, "dt": 1 / 60} + codec = NeRDStateCodec("ant") + for shard in manifest["shards"]: + with h5py.File(tmp_path / "preserved" / shard["file"], "r") as file: + for record in shard["trajectories"]: + index, length = record["metadata"]["legacy_index"], record["length"] + group = file[record["group"]] + stored = group["states/legacy_velocity"][()] + expected = original[: length + 1, index, 15:] + np.testing.assert_array_equal( + stored.view(np.uint32), expected.view(np.uint32) + ) + batch = { + kind: { + name: torch.from_numpy(field[()]).unsqueeze(0) + for name, field in group[kind].items() + } + for kind in ("states", "transitions") + } + raw = _official_batch(batch, codec) + np.testing.assert_array_equal( + raw["states"][0].numpy().view(np.uint32), + original[:length, index].view(np.uint32), + ) + np.testing.assert_array_equal( + raw["next_states"][0].numpy().view(np.uint32), + original[1 : length + 1, index].view(np.uint32), + ) + exported = export_legacy_dataset( + tmp_path / "preserved", tmp_path / "export.h5", "validation" + ) + assert exported["native_to_legacy_qd"] is False + with h5py.File(tmp_path / "export.h5", "r") as file: + for index, length in enumerate((3, 5)): + np.testing.assert_array_equal( + file["data/states"][:length, index].view(np.uint32), + original[:length, index].view(np.uint32), + ) + np.testing.assert_array_equal( + file["data/next_states"][:length, index].view(np.uint32), + original[1 : length + 1, index].view(np.uint32), + ) + assert hashlib.sha256(path.read_bytes()).hexdigest() == source_sha + + +def test_preserved_velocity_follows_prefix_masks_and_cpu_residency(tmp_path): + from torch.utils.data._utils.collate import default_collate + from embodichain.compute.dynamics import NeRDStateCodec + from embodichain.learning.dynamics.reference import ( + import_legacy_datasets, + _official_batch, + ) + + source = tmp_path / "source.h5" + original = _write_source_ant_velocities(source) + import_legacy_datasets( + {"train": source}, + tmp_path / "data", + {"codec": {"robot": "ant"}}, + preserve_legacy_velocity=True, + trajectories_per_shard=1, + ) + dataset = DynamicsWindowDataset(tmp_path / "data", "train", 3, include_prefix=True) + indices = [0, 2, 3, 7, 0] + individually = default_collate([dataset[index] for index in indices]) + grouped = dataset.get_batch(indices) + for name, value in grouped["states"].items(): + torch.testing.assert_close(value, individually["states"][name], atol=0, rtol=0) + expected = np.concatenate( + (original[:2, 0, 15:], np.repeat(original[1:2, 0, 15:], 2, axis=0)) + ) + np.testing.assert_array_equal( + grouped["states"]["legacy_velocity"][0].numpy().view(np.uint32), + expected.view(np.uint32), + ) + assert grouped["valid_mask"][0].tolist() == [True, False, False] + codec = NeRDStateCodec("ant") + transported = _official_batch(grouped, codec) + dataset.enable_device_cache("cpu", max_bytes=1024 * 1024) + cached = dataset.get_batch(indices) + for name, value in cached["states"].items(): + torch.testing.assert_close(value, grouped["states"][name], atol=0, rtol=0) + for name, value in _official_batch(cached, codec).items(): + torch.testing.assert_close(value, transported[name], atol=0, rtol=0) + cached["states"]["legacy_velocity"][0, 0, 0] = 999 + assert dataset.get_batch(indices)["states"]["legacy_velocity"][0, 0, 0] != 999 + dataset.close() + + +def test_disabling_velocity_preservation_clears_only_copied_contract(tmp_path): + from embodichain.data_pipeline.dynamics.shards import load_manifest + from embodichain.learning.dynamics.reference import import_legacy_datasets + + path = tmp_path / "source.h5" + _write_source_ant_velocities(path) + import_legacy_datasets( + {"train": path}, + tmp_path / "preserved", + {"codec": {"robot": "ant"}}, + preserve_legacy_velocity=True, + ) + metadata = load_manifest(tmp_path / "preserved")["metadata"] + original = json.loads(json.dumps(metadata)) + import_legacy_datasets( + {"train": path}, + tmp_path / "native", + metadata, + preserve_legacy_velocity=False, + ) + native = load_manifest(tmp_path / "native") + assert "legacy_velocity" not in native["metadata"] + assert "legacy_velocity" not in native["fields"]["states"] + assert metadata == original + + +@pytest.mark.parametrize("mismatch", ["small_delta", "signed_zero"]) +def test_lossless_import_rejects_nonidentical_adjacent_fp32_states(tmp_path, mismatch): + from embodichain.learning.dynamics.reference import import_legacy_datasets + + path = tmp_path / "source.h5" + _write_source_ant_velocities(path) + with h5py.File(path, "a") as file: + if mismatch == "small_delta": + file["data/next_states"][0, 0, 18] += np.float32(1e-7) + else: + file["data/next_states"][0, 0, 21] = -0.0 + # The old tolerance remains available to existing callers. + import_legacy_datasets( + {"train": path}, tmp_path / "default", {"codec": {"robot": "ant"}} + ) + with pytest.raises(ValueError, match="bitwise continuous"): + import_legacy_datasets( + {"train": path}, + tmp_path / "preserved", + {"codec": {"robot": "ant"}}, + preserve_legacy_velocity=True, + ) + + +def test_lossless_import_rejects_even_explicit_identity_world_rotation(tmp_path): + from embodichain.learning.dynamics.reference import import_legacy_datasets + + with pytest.raises(ValueError, match="source world frame"): + import_legacy_datasets( + {"train": tmp_path / "unused.h5"}, + tmp_path / "preserved", + {"codec": {"robot": "ant"}}, + preserve_legacy_velocity=True, + world_rotation_xyzw=(0.0, 0.0, 0.0, 1.0), + ) + assert not (tmp_path / "preserved").exists() + + +def test_official_transport_rejects_malformed_preserved_velocity_shape(): + from embodichain.compute.dynamics import NeRDStateCodec + from embodichain.learning.dynamics.reference import _official_batch + + q = torch.zeros(1, 2, 15) + q[..., 6] = 1 + batch = { + "states": { + "q": q, + "qd": torch.zeros(1, 2, 14), + "legacy_velocity": torch.zeros(1, 2, 13), + }, + "transitions": {"joint_f": torch.zeros(1, 1, 14)}, + } + with pytest.raises(ValueError, match="Preserved legacy velocity shape"): + _official_batch(batch, NeRDStateCodec("ant")) + + +def test_gravity_repair_is_explicit_preserves_raw_and_other_fields(): + from embodichain.learning.dynamics.reference import repair_gravity_transport + + raw = np.array([[[np.nan, -1.0, 8.49e5], [1e-45, -1.0, np.nan]]], np.float32) + before = raw.copy() + with pytest.raises(ValueError, match="uninitialized gravity"): + repair_gravity_transport(raw, (0.0, -1.0, 0.0)) + repaired, evidence = repair_gravity_transport(raw, (0.0, -1.0, 0.0), enabled=True) + np.testing.assert_array_equal(raw, before) + np.testing.assert_allclose( + repaired, np.array([[[0.0, -1.0, 0.0], [0.0, -1.0, 0.0]]]) + ) + assert evidence["nonfinite_components"] == 2 + assert evidence["enabled"] is True + + +def test_import_explicit_world_rotation_and_gravity_transport(tmp_path): + from embodichain.learning.dynamics.reference import import_legacy_datasets + import math + + metadata = {"codec": {"robot": "ant", "base_com": [0.0, 0.0, 0.0]}, "dt": 1 / 60} + q = np.zeros((4, 15), np.float32) + q[:, 1] = 1.0 + q[:, 3] = -math.sqrt(0.5) + q[:, 6] = math.sqrt(0.5) + qd = np.zeros((4, 14), np.float32) + qd[:, 0] = 1.0 + qd[:, 4] = 1.0 + with DynamicsShardWriter(tmp_path / "source", metadata) as writer: + writer.add_trajectory( + "one", + {"q": q, "qd": qd}, + { + "joint_f": np.zeros((3, 14), np.float32), + "contact_depths": np.zeros((3, 1), np.float32), + "gravity_dir": np.tile([0.0, -1.0, 0.0], (3, 1)).astype(np.float32), + }, + split="train", + ) + export_legacy_dataset(tmp_path / "source", tmp_path / "raw.h5", "train") + with h5py.File(tmp_path / "raw.h5", "a") as file: + file["data/gravity_dir"][:, :, 0] = np.nan + import_legacy_datasets( + {"train": tmp_path / "raw.h5"}, + tmp_path / "canonical", + metadata, + world_rotation_xyzw=(math.sqrt(0.5), 0.0, 0.0, math.sqrt(0.5)), + gravity_direction=(0.0, -1.0, 0.0), + ) + dataset = DynamicsWindowDataset(tmp_path / "canonical", "train", 3) + item = dataset[0] + np.testing.assert_allclose( + item["states"]["q"][0, :7], [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0], atol=1e-6 + ) + np.testing.assert_allclose( + item["states"]["qd"][0, :6], [1.0, 0.0, 0.0, 0.0, 0.0, 1.0], atol=1e-6 + ) + np.testing.assert_allclose( + item["transitions"]["gravity_dir"][0], [0.0, 0.0, -1.0], atol=1e-6 + ) + assert dataset.manifest["metadata"]["gravity_transport_override"] == [ + 0.0, + -1.0, + 0.0, + ] + dataset.close() + + +@pytest.mark.skipif( + not os.getenv("NERD_REFERENCE_SOURCE"), + reason="Explicit trusted official checkout/dependency environment required", +) +def test_real_official_collection_streams_batches_and_preserves_raw(tmp_path): + spec = { + "source": os.environ["NERD_REFERENCE_SOURCE"], + "output_dir": str(tmp_path / "collection"), + "robot": "cartpole", + "warp_version": "1.8.0", + "device": "cpu", + "num_envs": 2, + "horizon": 4, + "splits": {"train": 5, "validation": 2, "test": 2}, + "repair_uninitialized_gravity": True, + "canonical_world": "z_up", + } + spec_path = tmp_path / "generate.json" + spec_path.write_text(json.dumps(spec)) + env = { + **os.environ, + "OMP_NUM_THREADS": "1", + "MKL_NUM_THREADS": "1", + "CUDA_VISIBLE_DEVICES": "", + "PYTHONPATH": os.pathsep.join( + [ + os.environ.get("NERD_REFERENCE_DEPS", ""), + spec["source"], + str(Path(__file__).resolve().parents[3]), + ] + ), + } + process = subprocess.run( + [ + sys.executable, + "-m", + "embodichain.learning.dynamics.reference", + "--spec", + str(spec_path), + "--trust-external-code", + "--generate", + ], + env=env, + capture_output=True, + text=True, + ) + assert process.returncode == 0, process.stdout + process.stderr + root = tmp_path / "collection" + result = json.loads((root / "generation-result.json").read_text()) + assert result["source_trajectory_counts"] == { + "train": 5, + "validation": 2, + "test": 2, + } + assert result["gravity_repair"]["train"]["sampler_calls"] == 3 + assert result["source_bug_corrected"] is True + assert result["teacher_origin"]["integrator"] == "featherstone" + physics = json.loads((root / "source-physics.json").read_text()) + assert physics["physics_contract_version"] == "official-warp-physical-v2" + assert ( + result["teacher_origin"]["physics_contract_version"] + == physics["physics_contract_version"] + ) + assert result["teacher_origin"]["up_axis"] == 1 + np.testing.assert_allclose( + physics["source_physical_arrays"]["gravity"], [0.0, -9.81, 0.0], atol=1e-6 + ) + assert ( + physics["source_physical_arrays"]["gravity"] + == result["teacher_origin"]["gravity"] + ) + assert result["teacher_origin"]["gravity_magnitude"] == pytest.approx( + 9.81, abs=1e-6 + ) + from embodichain.learning.dynamics.reference import _hash + import hashlib + + assert physics["source_generator_files"]["generate/trajectory_sampler.py"] == _hash( + Path(spec["source"]) / "generate/trajectory_sampler.py" + ) + expected_generator_hash = hashlib.sha256( + json.dumps( + physics["source_generator_files"], sort_keys=True, separators=(",", ":") + ).encode() + ).hexdigest() + assert physics["source_generator_sha256"] == expected_generator_hash + assert physics["generation_adapter_sha256"] == _hash( + Path(__file__).resolve().parents[3] + / "embodichain/learning/dynamics/reference.py" + ) + with h5py.File(root / "train.raw.h5") as file: + assert file["data/states"].shape == (4, 5, 4) + dataset = DynamicsWindowDataset(root / "canonical", "train", 3) + assert len(dataset) == 10 + assert ( + dataset.manifest["metadata"]["source_physical_arrays"]["gravity"] + == physics["source_physical_arrays"]["gravity"] + ) + assert ( + dataset.manifest["metadata"]["generation_adapter_sha256"] + == physics["generation_adapter_sha256"] + ) + np.testing.assert_allclose( + dataset[0]["transitions"]["gravity_dir"], + np.tile([0.0, 0.0, -1.0], (3, 1)), + atol=1e-6, + ) + dataset.close() + + +@pytest.mark.skipif( + not os.getenv("NERD_REFERENCE_SOURCE"), + reason="Explicit trusted official checkout required", +) +def test_portable_to_official_restoration_preserves_state_dtype_and_ownership(): + code = r""" +import sys,torch +from embodichain.learning.dynamics.reference import reference_identity,_restore_official_model +source=sys.argv[1] +reference_identity(source) +sys.path.insert(0,source) +from embodichain.learning.models.nerd import NeRDModel,NeRDModelCfg +for dtype in (torch.float32,torch.float64): + torch.manual_seed(5) + common=NeRDModel(NeRDModelCfg(input_dims={"states_embedding":4,"joint_f":2},output_dim=4,n_embd=12,n_head=3,n_layer=1,block_size=4,history_length=3,head_hidden_dims=(12,),dropout=.25,output_tanh=True)).to(dtype=dtype).eval() + common.set_normalization({"states_embedding":(torch.arange(4,dtype=dtype),torch.arange(4,dtype=dtype)+.5),"joint_f":(torch.tensor([.1,-.3],dtype=dtype),torch.tensor([.9,1.2],dtype=dtype))},torch.arange(4,dtype=dtype)*.2,torch.arange(4,dtype=dtype)+.3) + old=_restore_official_model(common,device="cpu") + assert not old.training + for name,value in old.named_parameters(): + assert value.dtype==dtype + torch.testing.assert_close(value,dict(common.named_parameters())[name],atol=0,rtol=0) + inputs={"states_embedding":torch.randn(2,3,4,dtype=dtype),"joint_f":torch.randn(2,3,2,dtype=dtype)} + legacy=lambda:{("joint_acts" if name=="joint_f" else name):value.clone() for name,value in inputs.items()} + torch.testing.assert_close(old(legacy()),common(inputs),atol=1e-6,rtol=1e-5) + torch.testing.assert_close(old.evaluate(legacy()),common(inputs)[:,-1:,:],atol=1e-6,rtol=1e-5) + before=common.input_rms["joint_f"].mean.clone() + old.input_rms["joint_acts"].mean.add_(1.) + torch.testing.assert_close(common.input_rms["joint_f"].mean,before,atol=0,rtol=0) +assert "dexsim" not in sys.modules +""" + env = {**os.environ, "OMP_NUM_THREADS": "1", "MKL_NUM_THREADS": "1"} + result = subprocess.run( + [sys.executable, "-c", code, os.environ["NERD_REFERENCE_SOURCE"]], + env=env, + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stdout + result.stderr + + +@pytest.mark.skipif( + not os.getenv("NERD_REFERENCE_SOURCE"), reason="Explicit trusted reference required" +) +def test_actual_official_trainer_resume_preserves_optimizer_rng_and_best_model( + tmp_path, +): + """Synthetic CPU data verify lifecycle, not learned-dynamics accuracy.""" + from embodichain.compute.dynamics import NeRDStateCodec + from embodichain.learning.models.nerd import ( + NeRDModel, + NeRDModelCfg, + export_model, + load_model, + ) + from embodichain.learning.dynamics.trainer import compute_statistics + from embodichain.learning.dynamics._sample_order import _write_checkpointed_order + from embodichain.learning.dynamics.reference import _hash + + metadata = {"codec": {"robot": "cartpole"}, "dt": 1 / 60} + generator = np.random.default_rng(3) + with DynamicsShardWriter(tmp_path / "data", metadata) as writer: + for split in ["train", "validation", "test"]: + for index in range(2): + q = generator.normal(size=(13, 2)).astype(np.float32) + qd = generator.normal(size=(13, 2)).astype(np.float32) + writer.add_trajectory( + f"{split}:{index}", + {"q": q, "qd": qd}, + {"joint_f": np.zeros((12, 2), np.float32)}, + split=split, + ) + dataset = DynamicsWindowDataset(tmp_path / "data", "train", 3) + codec = NeRDStateCodec("cartpole") + stats = compute_statistics(dataset, codec) + torch.manual_seed(4) + model = NeRDModel( + NeRDModelCfg( + input_dims={"states_embedding": 4, "joint_f": 2}, + output_dim=4, + n_embd=12, + n_head=3, + n_layer=1, + block_size=4, + history_length=3, + dropout=0.1, + ) + ) + model.set_normalization( + {k: (v["mean"], v["var"]) for k, v in stats["inputs"].items()}, + stats["output"]["mean"], + stats["output"]["var"], + ) + export_model(tmp_path / "initial", model, metadata) + order, _ = _write_checkpointed_order( + tmp_path / "order.json", + dataset_size=len(dataset), + steps=6, + batch_size=4, + seed=7, + block_steps=2, + ) + cfg = { + "total_steps": 6, + "batch_size": 4, + "seed": 7, + "validation_interval": 2, + "checkpoint_interval": 2, + "log_interval": 2, + "resident_cache_bytes": 1024 * 1024, + "device": "cpu", + } + source = os.environ["NERD_REFERENCE_SOURCE"] + spec = { + "source": source, + "dataset": str(tmp_path / "data"), + "dataset_fingerprint": dataset.fingerprint, + "initial_model": str(tmp_path / "initial"), + "initial_model_sha256": _hash(tmp_path / "initial/weights.pt"), + "sample_order": str(order), + "sample_order_sha256": _hash(order), + "train": cfg, + } + environment = { + **os.environ, + "CUDA_VISIBLE_DEVICES": "", + "OMP_NUM_THREADS": "1", + "MKL_NUM_THREADS": "1", + "PYTHONPATH": os.pathsep.join( + [ + os.environ.get("NERD_REFERENCE_DEPS", ""), + source, + str(Path(__file__).resolve().parents[3]), + ] + ), + } + + def execute(name, extra): + path = tmp_path / (name + ".json") + path.write_text(json.dumps({**spec, "output_dir": str(tmp_path / name)})) + result = subprocess.run( + [ + sys.executable, + "-m", + "embodichain.learning.dynamics.reference", + "--spec", + str(path), + "--trust-external-code", + *extra, + ], + env=environment, + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stdout + result.stderr + + execute("full", []) + execute("resumed", ["--stop-after", "3"]) + interrupted = json.loads((tmp_path / "resumed/result.json").read_text()) + assert ( + interrupted["status"] == "interrupted" and interrupted["completed_updates"] == 3 + ) + execute("resumed", ["--resume", str(tmp_path / "resumed/reference-last.pt")]) + for artifact in ["best-export", "last-export"]: + first, _ = load_model(tmp_path / "full" / artifact) + second, _ = load_model(tmp_path / "resumed" / artifact) + for key, value in first.state_dict().items(): + torch.testing.assert_close(second.state_dict()[key], value, atol=0, rtol=0) + dataset.close() diff --git a/tests/learning/dynamics/test_reference_rollout.py b/tests/learning/dynamics/test_reference_rollout.py new file mode 100644 index 0000000000..98a62860cd --- /dev/null +++ b/tests/learning/dynamics/test_reference_rollout.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. +# ---------------------------------------------------------------------------- + +"""Historical runtime qualification through explicitly trusted isolated Warp.""" + +from __future__ import annotations +import hashlib +import importlib.util +import json +import os +from pathlib import Path +import subprocess +import sys +import pytest +from embodichain.learning.dynamics.reference_rollout import evaluate_reference_rollouts + +_ROOT = Path(__file__).resolve().parents[3] + + +def test_external_reference_requires_explicit_trust() -> None: + with pytest.raises(ValueError, match="trust_external_code"): + evaluate_reference_rollouts({}) + + +@pytest.mark.parametrize("robot,version", [("cartpole", "1.5.1"), ("ant", "1.8.0")]) +def test_wrong_historical_version_rejected_before_source_import( + robot: str, version: str +) -> None: + with pytest.raises(ValueError, match="pinned Warp"): + evaluate_reference_rollouts( + {"robot": robot, "warp_version": version}, trust_external_code=True + ) + + +def _check_original_runtime(tmp_path: Path, robot: str, source: str) -> None: + import numpy as np + import h5py + import torch + import warp as wp + from embodichain.learning.dynamics.reference import generate_official_dataset + from embodichain.learning.models.nerd import NeRDModel, NeRDModelCfg, export_model + from embodichain.data_pipeline.dynamics.shards import load_manifest + + # Generate new physical provenance from the actual original teacher. No + # current-Newton manifest is rewritten to pretend it describes old physics. + with wp.ScopedDevice("cpu"): + generation = generate_official_dataset( + { + "source": source, + "robot": robot, + "output_dir": str(tmp_path / "source-data"), + "device": "cpu", + "num_envs": 2, + "horizon": 4, + "splits": {"train": 2, "validation": 2, "test": 2}, + "seed": 3, + "initial_states_source": "env", + "canonical_world": "z_up", + "repair_uninitialized_gravity": True, + "cpu_threads": 2, + }, + trust_external_code=True, + ) + metadata = load_manifest(tmp_path / "source-data/canonical")["metadata"] + assert "gravity" in metadata["source_physical_arrays"] + 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, + } + ) + torch.manual_seed(7) + model = NeRDModel( + NeRDModelCfg( + input_dims=dims, + output_dim=4 if robot == "cartpole" else 29, + n_embd=24, + n_head=3, + n_layer=1, + history_length=3, + block_size=4, + ) + ) + with torch.no_grad(): + model.model.output_net.weight.mul_(0.02) + model.model.output_net.bias.zero_() + if robot == "ant": + model.model.output_net.bias[6] = 1.0 + export = tmp_path / "model" + export_model(export, model, metadata) + reports = {} + for mode in ["passive", "random_effort", "fixed_policy"]: + spec = { + "source": source, + "robot": robot, + "model": str(export), + "output": str(tmp_path / f"{mode}.json"), + "device": "cpu", + "num_envs": 2, + "trajectories": 3, + "horizons": [1, 4], + "seed": 17, + "initial_states": "default", + "motion_mode": mode, + "cpu_threads": 2, + } + if mode == "fixed_policy": + spec["policy"] = {"kind": "pd"} + reports[mode] = evaluate_reference_rollouts(spec, trust_external_code=True) + report = reports[mode] + assert report["teacher_origin"]["kind"] == "official_warp" + assert report["runtime_identity"]["engine"] == "official_warp" + assert report["physical_verification"]["complete"] + assert report["physical_verification"]["missing_recorded_fields"] == [] + assert report["portable_runtime_equivalence"]["status"] == "passed" + assert report["portable_runtime_equivalence"]["max_abs_error"] <= 1e-5 + assert report["trajectories"] == 3 and report["num_envs"] == 2 + assert len(Path(report["trajectory_evidence"]).read_text().splitlines()) == 3 + with h5py.File(report["trace_evidence"], "r") as evidence: + assert evidence["0/q"].shape[0] == 2 and evidence["1/q"].shape[0] == 1 + if mode == "fixed_policy": + assert not np.allclose( + evidence["0/teacher_joint_f"][:, 1:], + evidence["0/joint_f"][:, 1:], + atol=1e-8, + ) + assert report["selected_teacher_sha256"] + assert ( + report["trace_evidence_sha256"] + == hashlib.sha256(Path(report["trace_evidence"]).read_bytes()).hexdigest() + ) + assert ( + report["horizons"]["4"]["valid_trajectories"] + + report["horizons"]["4"]["numerically_invalid_trajectories"] + == 3 + ) + # Run an independently restored equal export: teacher identity and content + # are independent of the neural artifact's path, and results remain exact. + second = tmp_path / "equal-model" + export_model(second, model, metadata) + again = evaluate_reference_rollouts( + { + "source": source, + "robot": robot, + "model": str(second), + "output": str(tmp_path / "equal.json"), + "device": "cpu", + "num_envs": 2, + "trajectories": 3, + "horizons": [1, 4], + "seed": 17, + "initial_states": "default", + "motion_mode": "passive", + "cpu_threads": 2, + }, + trust_external_code=True, + ) + assert ( + again["selected_teacher_sha256"] + == reports["passive"]["selected_teacher_sha256"] + ) + assert again["horizons"] == reports["passive"]["horizons"] + # Original per-trajectory sampling is deterministic across replay batching. + sample_reports = [] + for size in [1, 2]: + sample_reports.append( + evaluate_reference_rollouts( + { + "source": source, + "robot": robot, + "model": str(export), + "output": str(tmp_path / f"sample-{size}.json"), + "device": "cpu", + "num_envs": size, + "trajectories": 3, + "horizons": [1], + "seed": 31, + "initial_states": "sample", + "motion_mode": "passive", + "cpu_threads": 2, + }, + trust_external_code=True, + ) + ) + assert ( + sample_reports[0]["selected_teacher_sha256"] + == sample_reports[1]["selected_teacher_sha256"] + ) + if robot == "cartpole": + with torch.no_grad(): + model.model.output_net.bias[2:] = torch.finfo(torch.float32).max + divergent = tmp_path / "divergent-model" + export_model(divergent, model, metadata) + failure = evaluate_reference_rollouts( + { + "source": source, + "robot": robot, + "model": str(divergent), + "output": str(tmp_path / "nonfinite.json"), + "num_envs": 2, + "trajectories": 3, + "horizons": [1, 3], + "initial_states": "default", + "motion_mode": "passive", + "cpu_threads": 2, + }, + trust_external_code=True, + ) + assert failure["nonfinite_rate"] == 1.0 and failure["failure_rate"] == 1.0 + assert failure["horizons"]["3"]["valid_trajectories"] == 0 + assert len(Path(failure["trajectory_evidence"]).read_text().splitlines()) == 3 + assert failure["portable_runtime_equivalence"]["status"] == "matched_nonfinite" + with torch.no_grad(): + model.model.output_net.bias[2:] = 0.0 + # Missing gravity is a diagnostic run, never complete physical qualification. + legacy = json.loads(json.dumps(metadata)) + legacy["source_physical_arrays"].pop("gravity") + legacy["teacher_origin"].pop("gravity", None) + old = tmp_path / "legacy-model" + export_model(old, model, legacy) + old_report = evaluate_reference_rollouts( + { + "source": source, + "robot": robot, + "model": str(old), + "output": str(tmp_path / "legacy.json"), + "num_envs": 2, + "trajectories": 3, + "horizons": [1], + "initial_states": "default", + "motion_mode": "passive", + "cpu_threads": 2, + }, + trust_external_code=True, + ) + assert old_report["qualification_status"] == "pending_missing_recorded_physics" + assert old_report["physical_verification"]["missing_recorded_fields"] == ["gravity"] + # Independent official-version generation records the physical replica + # count in its source split receipt rather than the old sampling field. + split_metadata = json.loads(json.dumps(metadata)) + batch = split_metadata["sampling"].pop("batch_size") + split_metadata["source_split_identities"] = {"train": {"num_envs": batch}} + split_export = tmp_path / "split-source-model" + export_model(split_export, model, split_metadata) + split_spec = { + "source": source, + "robot": robot, + "model": str(split_export), + "output": str(tmp_path / "split-source.json"), + "num_envs": 2, + "trajectories": 3, + "horizons": [1], + "initial_states": "default", + "motion_mode": "passive", + "cpu_threads": 2, + } + split_report = evaluate_reference_rollouts(split_spec, trust_external_code=True) + assert split_report["physical_verification"]["complete"] + assert ( + split_report["physical_verification"]["sha256"] + == reports["passive"]["physical_verification"]["sha256"] + ) + split_metadata["sampling"]["batch_size"] = batch + 1 + conflicting_export = tmp_path / "conflicting-source-model" + export_model(conflicting_export, model, split_metadata) + with pytest.raises(ValueError, match="Conflicting source sampling"): + evaluate_reference_rollouts( + { + **split_spec, + "model": str(conflicting_export), + "output": str(tmp_path / "conflicting-source.json"), + }, + trust_external_code=True, + ) + # A declared change to gravity must reject, not silently reuse old dynamics. + changed = json.loads(json.dumps(metadata)) + changed["source_physical_arrays"]["gravity"] = [0.0, -1.0, 0.0] + bad = tmp_path / "bad-model" + export_model(bad, model, changed) + with pytest.raises(ValueError, match="gravity"): + evaluate_reference_rollouts( + { + "source": source, + "robot": robot, + "model": str(bad), + "output": str(tmp_path / "bad.json"), + "num_envs": 2, + "trajectories": 3, + "horizons": [1], + }, + trust_external_code=True, + ) + + +@pytest.mark.parametrize("robot", ["cartpole", "ant"]) +def test_real_historical_teacher_neural_runtime_and_source_contract( + tmp_path: Path, robot: str +) -> None: + source = os.environ.get("NERD_MAIN_REFERENCE_ROOT") + dependencies = os.environ.get( + "NERD_WARP18_ROOT" if robot == "cartpole" else "NERD_WARP151_ROOT" + ) + if not source or not dependencies: + pytest.skip("Set explicit trusted reference and legacy Warp dependency paths") + code = "import importlib.util,sys;from pathlib import Path;s=importlib.util.spec_from_file_location('reference_rollout_test',sys.argv[1]);m=importlib.util.module_from_spec(s);s.loader.exec_module(m);m._check_original_runtime(Path(sys.argv[2]),sys.argv[3],sys.argv[4])" + subprocess.run( + [ + sys.executable, + "-c", + code, + str(Path(__file__).resolve()), + str(tmp_path), + robot, + source, + ], + check=True, + env={ + **os.environ, + "PYTHONPATH": dependencies + os.pathsep + str(_ROOT), + "OMP_NUM_THREADS": "2", + "MKL_NUM_THREADS": "2", + "CUDA_VISIBLE_DEVICES": "", + }, + ) + + +def _check_supported_replay( + tmp_path: Path, robot: str, source: str, data: Path +) -> None: + import shutil + import h5py + import numpy as np + import torch + from embodichain.learning.models.nerd import NeRDModel, NeRDModelCfg, export_model + + metadata = json.loads((data / f"{robot}-passive/manifest.json").read_text())[ + "metadata" + ] + 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, + } + ) + torch.manual_seed(7) + model = NeRDModel( + NeRDModelCfg( + input_dims=dims, + output_dim=4 if robot == "cartpole" else 29, + n_embd=24, + n_layer=1, + n_head=3, + history_length=3, + block_size=4, + ) + ) + with torch.no_grad(): + model.model.output_net.weight.mul_(0.02) + model.model.output_net.bias.zero_() + if robot == "ant": + model.model.output_net.bias[6] = 1.0 + export = tmp_path / "export" + export_model(export, model, metadata) + reports = {} + specs = {} + for mode in ["passive", "random_effort", "fixed_policy"]: + spec = { + "source": source, + "robot": robot, + "model": str(export), + "output": str(tmp_path / (mode + ".json")), + "device": "cpu", + "num_envs": 2, + "trajectories": 3, + "horizons": [1, 4, 12], + "seed": 1000000, + "initial_states": "sample", + "motion_mode": mode, + "cpu_threads": 2, + "teacher_dataset": str(data / (robot + "-" + mode)), + } + if mode == "fixed_policy": + spec["policy"] = {"kind": "pd"} + result = evaluate_reference_rollouts(spec, trust_external_code=True) + reports[mode], specs[mode] = result, spec + assert result["teacher_generation_environment"]["python"] == "3.8.20" + assert result["teacher_generation_environment"]["torch"].startswith("2.2.2") + assert ( + result["teacher_generation_environment"]["runtime_imports_embodichain"] + is False + ) + assert result["physical_verification"]["complete"] + assert result["teacher_physical_verification"]["complete"] + assert result["teacher_transport_max_abs_error"] <= 1e-5 + assert result["trajectories"] == 3 and set(result["horizons"]) == { + "1", + "4", + "12", + } + assert result["portable_runtime_equivalence"]["status"] == "passed" + if mode == "fixed_policy": + assert ( + result["teacher_feedback_verification"] + == "all_selected_actions_verified_against_teacher_states" + ) + assert result["teacher_feedback_max_abs_error"] <= 1e-5 + with h5py.File(result["trace_evidence"], "r") as file: + assert np.any( + file["0/joint_f"][:, 1:] != file["0/teacher_joint_f"][:, 1:] + ) + # Different evaluation batching cannot select different teacher contents. + again = evaluate_reference_rollouts( + {**specs["passive"], "num_envs": 1, "output": str(tmp_path / "batch1.json")}, + trust_external_code=True, + ) + assert ( + again["selected_teacher_sha256"] + == reports["passive"]["selected_teacher_sha256"] + ) + # A rehashed but inconsistent raw/canonical frame must not become a new oracle. + changed = tmp_path / "changed-transport" + shutil.copytree(data / (robot + "-passive"), changed) + manifest_path = changed / "manifest.json" + manifest = json.loads(manifest_path.read_text()) + shard = changed / manifest["shards"][0]["file"] + with h5py.File(shard, "a") as file: + file["0/states/q"][1, 0] += 1.0 + manifest["shards"][0]["sha256"] = hashlib.sha256(shard.read_bytes()).hexdigest() + manifest_path.write_text(json.dumps(manifest)) + with pytest.raises(ValueError, match="state transport differs"): + evaluate_reference_rollouts( + { + **specs["passive"], + "teacher_dataset": str(changed), + "output": str(tmp_path / "bad-transport.json"), + }, + trust_external_code=True, + ) + # PD action verification also rejects consistent raw/canonical effort tampering. + changed = tmp_path / "changed-feedback" + shutil.copytree(data / (robot + "-fixed_policy"), changed) + manifest_path = changed / "manifest.json" + manifest = json.loads(manifest_path.read_text()) + shard = changed / manifest["shards"][0]["file"] + with h5py.File(shard, "a") as file: + file["0/transitions/raw_joint_f"][0, 0] += 1.0 + file["0/transitions/joint_f"][0, 6 if robot == "ant" else 0] += 1.0 + manifest["shards"][0]["sha256"] = hashlib.sha256(shard.read_bytes()).hexdigest() + manifest_path.write_text(json.dumps(manifest)) + with pytest.raises(ValueError, match="actions differ from declared feedback"): + evaluate_reference_rollouts( + { + **specs["fixed_policy"], + "teacher_dataset": str(changed), + "output": str(tmp_path / "bad-feedback.json"), + }, + trust_external_code=True, + ) + missing = tmp_path / "missing-teacher-physics" + shutil.copytree(data / (robot + "-passive"), missing) + manifest_path = missing / "manifest.json" + manifest = json.loads(manifest_path.read_text()) + manifest["metadata"]["source_physical_arrays"].pop("gravity") + manifest["metadata"]["teacher_origin"].pop("gravity", None) + manifest_path.write_text(json.dumps(manifest)) + with pytest.raises(ValueError, match="missing measured physical fields"): + evaluate_reference_rollouts( + { + **specs["passive"], + "teacher_dataset": str(missing), + "output": str(tmp_path / "missing-physics.json"), + }, + trust_external_code=True, + ) + + +@pytest.mark.parametrize("robot", ["cartpole", "ant"]) +def test_real_supported_environment_teacher_replay(tmp_path: Path, robot: str) -> None: + source = os.environ.get("NERD_MAIN_REFERENCE_ROOT") + dependencies = os.environ.get( + "NERD_WARP18_ROOT" if robot == "cartpole" else "NERD_WARP151_ROOT" + ) + data = os.environ.get("NERD_SUPPORTED_EVAL_DATA") + if not source or not dependencies or not data: + pytest.skip( + "Set explicit trusted original source/dependencies and supported-environment evaluation corpus" + ) + code = "import importlib.util,sys;from pathlib import Path;s=importlib.util.spec_from_file_location('reference_replay_test',sys.argv[1]);m=importlib.util.module_from_spec(s);s.loader.exec_module(m);m._check_supported_replay(Path(sys.argv[2]),sys.argv[3],sys.argv[4],Path(sys.argv[5]))" + subprocess.run( + [ + sys.executable, + "-c", + code, + str(Path(__file__).resolve()), + str(tmp_path), + robot, + source, + data, + ], + check=True, + env={ + **os.environ, + "PYTHONPATH": dependencies + os.pathsep + str(_ROOT), + "CUDA_VISIBLE_DEVICES": "", + "OMP_NUM_THREADS": "2", + "MKL_NUM_THREADS": "2", + }, + ) diff --git a/tests/learning/dynamics/test_reference_traces.py b/tests/learning/dynamics/test_reference_traces.py new file mode 100644 index 0000000000..7257b0c7bc --- /dev/null +++ b/tests/learning/dynamics/test_reference_traces.py @@ -0,0 +1,247 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Original trace provenance, ordering, bounded batching and failure retention.""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path +from typing import Any + +import h5py +import numpy as np +import pytest +import torch + +from embodichain.learning.dynamics._reference_traces import ( + _load_reference_trace, + _reference_trace_batches, +) +from embodichain.learning.dynamics.reference import OFFICIAL_COMMIT + + +def _fixture( + root: Path, robot: str +) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: + root.mkdir() + ant = robot == "ant" + q_dim, qd_dim, force_dim = (15, 14, 8) if ant else (2, 2, 2) + codec = { + "robot": robot, + "base_com": [0.0, 0.0, 0.0], + "continuous_q_indices": [] if ant else [1], + } + shared = { + "robot": robot, + "asset_sha256": "fixture-asset", + "dt": 1 / 60, + "codec": codec, + "world_rotation_xyzw": [1.0, 0.0, 0.0, 0.0], + } + source = {"runtime_files": {"fixture-source.py": "fixture-source-hash"}} + version = "1.5.1" if ant else "1.8.0" + meta = { + **shared, + "source_runtime_files": source["runtime_files"], + "teacher_origin": { + "kind": "official_warp", + "source_commit": OFFICIAL_COMMIT, + "warp_version": version, + }, + "generation_environment": { + "python": "3.8.20", + "torch": "2.2.2+cpu", + "numpy": "1.24.4", + "warp": version, + "runtime_imports_embodichain": False, + "runtime_imports_dexsim": False, + "runtime_imports_newton": False, + }, + } + rows = [] + file = root / "teacher.h5" + with h5py.File(file, "w") as stream: + for index in range(3): + group = stream.create_group(str(index)) + q = np.full((5, q_dim), index + 0.25, dtype=np.float32) + qd = np.zeros((5, qd_dim), dtype=np.float32) + if index == 2: + q[-1, 0] = np.nan + raw = np.concatenate((q, qd), -1) + force = np.full((4, force_dim), index, dtype=np.float32) + canonical = ( + np.concatenate((np.zeros((4, 6), dtype=np.float32), force), -1) + if ant + else force + ) + for key, value in { + "states/raw": raw, + "states/q": q, + "states/qd": qd, + "transitions/raw_joint_f": force, + "transitions/joint_f": canonical, + }.items(): + group.create_dataset(key, data=value) + rows.append( + { + "id": f"71:{index}", + "group": str(index), + "length": 4, + "split": "test", + "metadata": {"sampling_seed": 71 + index}, + } + ) + manifest = { + "schema_version": 1, + "trace_type": "evaluation_including_failures", + "metadata": meta, + "evaluation": { + "seed": 71, + "trajectories": 3, + "horizon": 4, + "initial_states": "sample", + "motion_mode": "random_effort", + "policy": None, + }, + "shards": [ + { + "file": file.name, + "sha256": hashlib.sha256(file.read_bytes()).hexdigest(), + "trajectories": rows, + } + ], + } + (root / "manifest.json").write_text(json.dumps(manifest)) + return {"robot": robot, "seed": 71, "motion_mode": "random_effort"}, source, shared + + +@pytest.mark.parametrize("robot", ["cartpole", "ant"]) +def test_original_trace_preserves_order_padding_failure_and_ownership( + tmp_path: Path, robot: str +) -> None: + directory = tmp_path / "teacher" + spec, source, metadata = _fixture(directory, robot) + _, records = _load_reference_trace(directory, spec, source, metadata, None, 3, 4) + batches = list(_reference_trace_batches(directory, records, 2, 4, robot, "cpu")) + assert [[row["id"] for row in selected] for selected, _ in batches] == [ + ["71:0", "71:1"], + ["71:2"], + ] + assert batches[0][1]["q"][0, 0, 0] == 0.25 + assert batches[1][1]["q"][0, 0, 0] == 2.25 + assert torch.isnan(batches[1][1]["q"][:, -1, 0]).all() + assert all(row["reference_sha256"] for row in records) + batches[1][1]["q"][0, 0, 0] = 99 + assert batches[1][1]["q"][1, 0, 0] == 2.25 + assert batches[0][1]["q"][0, 0, 0] == 0.25 + + +@pytest.mark.parametrize( + "field,value", + [ + ("python", "3.11.14"), + ("torch", "2.7.0"), + ("warp", "1.17.0"), + ("runtime_imports_embodichain", True), + ("runtime_imports_newton", True), + ], +) +def test_original_trace_rejects_wrong_generation_environment( + tmp_path: Path, field: str, value: Any +) -> None: + directory = tmp_path / "teacher" + spec, source, metadata = _fixture(directory, "ant") + path = directory / "manifest.json" + manifest = json.loads(path.read_text()) + manifest["metadata"]["generation_environment"][field] = value + path.write_text(json.dumps(manifest)) + with pytest.raises(ValueError, match="supported isolated environment"): + _load_reference_trace(directory, spec, source, metadata, None, 3, 4) + + +@pytest.mark.parametrize("change", ["seed", "order", "count", "asset", "policy"]) +def test_original_trace_rejects_mispaired_evaluation( + tmp_path: Path, change: str +) -> None: + directory = tmp_path / "teacher" + spec, source, metadata = _fixture(directory, "cartpole") + path = directory / "manifest.json" + manifest = json.loads(path.read_text()) + if change == "seed": + manifest["evaluation"]["seed"] += 1 + elif change == "order": + manifest["shards"][0]["trajectories"].reverse() + elif change == "count": + manifest["evaluation"]["trajectories"] += 1 + elif change == "asset": + manifest["metadata"]["asset_sha256"] = "another-asset" + else: + manifest["evaluation"]["policy"] = {"kind": "pd"} + path.write_text(json.dumps(manifest)) + with pytest.raises(ValueError): + _load_reference_trace(directory, spec, source, metadata, None, 3, 4) + + +def test_original_trace_rejects_rehashed_control_mismatch(tmp_path: Path) -> None: + directory = tmp_path / "teacher" + spec, source, metadata = _fixture(directory, "ant") + file = directory / "teacher.h5" + with h5py.File(file, "a") as stream: + stream["0/transitions/joint_f"][1, 0] = 1 + path = directory / "manifest.json" + manifest = json.loads(path.read_text()) + manifest["shards"][0]["sha256"] = hashlib.sha256(file.read_bytes()).hexdigest() + path.write_text(json.dumps(manifest)) + _, records = _load_reference_trace(directory, spec, source, metadata, None, 3, 4) + with pytest.raises(ValueError, match="controls differ"): + list(_reference_trace_batches(directory, records, 2, 4, "ant", "cpu")) + + +def test_original_trace_rejects_raw_shard_tampering(tmp_path: Path) -> None: + directory = tmp_path / "teacher" + spec, source, metadata = _fixture(directory, "cartpole") + with (directory / "teacher.h5").open("ab") as stream: + stream.write(b"tamper") + with pytest.raises(ValueError, match="checksum mismatch"): + _load_reference_trace(directory, spec, source, metadata, None, 3, 4) + + +@pytest.mark.parametrize("change", ["missing", "dtype", "length"]) +def test_original_trace_rejects_rehashed_invalid_state_schema( + tmp_path: Path, change: str +) -> None: + directory = tmp_path / "teacher" + spec, source, metadata = _fixture(directory, "cartpole") + file = directory / "teacher.h5" + with h5py.File(file, "a") as stream: + value = stream["0/states/raw"][()] + del stream["0/states/raw"] + if change != "missing": + stream.create_dataset( + "0/states/raw", + data=value.astype(np.float64) if change == "dtype" else value[:-1], + ) + path = directory / "manifest.json" + manifest = json.loads(path.read_text()) + manifest["shards"][0]["sha256"] = hashlib.sha256(file.read_bytes()).hexdigest() + path.write_text(json.dumps(manifest)) + _, records = _load_reference_trace(directory, spec, source, metadata, None, 3, 4) + with pytest.raises( + ValueError, match="missing raw/canonical fields|dtype, shape or length" + ): + list(_reference_trace_batches(directory, records, 2, 4, "cartpole", "cpu")) diff --git a/tests/learning/dynamics/test_report.py b/tests/learning/dynamics/test_report.py new file mode 100644 index 0000000000..788bca3ab3 --- /dev/null +++ b/tests/learning/dynamics/test_report.py @@ -0,0 +1,414 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Qualification cannot substitute pass flags or mismatched rollout populations.""" + +from __future__ import annotations + +import copy +import json +from pathlib import Path + +import pytest + +from embodichain.learning.dynamics.report import ( + _curve, + _dataset_evidence, + _matched_rollouts, + _numerical_evidence, + _rollout_evidence, + annotate_rollout_evidence, +) +from embodichain.learning.dynamics.reference import OFFICIAL_COMMIT, _hash + + +def test_learning_curve_retains_sparse_full_budget_observations(tmp_path: Path) -> None: + path = tmp_path / "metrics.jsonl" + budget = 4995000 + with path.open("w") as stream: + for step in range(500, budget + 1, 500): + stream.write(json.dumps({"step": step, "loss": 1.0 / step}) + "\n") + steps, losses = _curve(path, budget) + assert steps[0] == 500 and steps[-1] == budget + assert 900 <= len(steps) <= 1000 + assert losses[0] == 1.0 / 500 and losses[-1] == 1.0 / budget + assert all(left < right for left, right in zip(steps, steps[1:])) + + +def test_learning_curve_keeps_irregular_final_observation(tmp_path: Path) -> None: + path = tmp_path / "metrics.jsonl" + observed = [300, 650, 3000, 4500, 12000] + path.write_text( + "".join( + json.dumps({"step": step, "loss": float(step)}) + "\n" for step in observed + ) + ) + steps, losses = _curve(path, 12000, max_points=4) + assert steps[0] == 300 and steps[-1] == 12000 + assert len(steps) <= 4 and all(step in observed for step in steps) + assert losses == [float(step) for step in steps] + + +def test_isolated_supported_source_import_has_verified_dataset_lineage(tmp_path): + import h5py + import numpy as np + from embodichain.learning.dynamics.reference import import_legacy_datasets + + sources = {} + identities = {} + for split in ("train", "validation", "test"): + raw = tmp_path / f"{split}.h5" + with h5py.File(raw, "w") as file: + data = file.create_group("data") + data.attrs["mode"] = "trajectory" + data.create_dataset("states", data=np.zeros((2, 1, 4), np.float32)) + data.create_dataset("next_states", data=np.zeros((2, 1, 4), np.float32)) + data.create_dataset("joint_acts", data=np.zeros((2, 1, 2), np.float32)) + sources[split] = raw + identities[split] = { + "source_commit": OFFICIAL_COMMIT, + "robot": "cartpole", + "python": "3.8.20", + "numpy": "1.24.4", + "torch": "2.2.2+cpu", + "warp": "1.8.0", + "runtime_imports_embodichain": False, + "raw_sha256": _hash(raw), + } + metadata = { + "codec": {"robot": "cartpole"}, + "source_split_identities": identities, + "source_physical_arrays": {"body_mass": [50.0, 2.5]}, + "teacher_origin": {"kind": "official_warp", "source_commit": OFFICIAL_COMMIT}, + } + output = tmp_path / "canonical" + result = import_legacy_datasets(sources, output, metadata) + declaration = { + "path": str(output), + "status": "available", + "fingerprint": result["fingerprint"], + } + evidence = _dataset_evidence("official_warp", declaration) + assert evidence["status"] == "available" + assert evidence["transitions"] == {"train": 2, "validation": 2, "test": 2} + with h5py.File(sources["test"], "r+") as file: + file["data/states"][0, 0, 0] = 1.0 + assert _dataset_evidence("official_warp", declaration)["status"] == "pending" + + +@pytest.mark.parametrize( + "field,value", + [ + ("source_commit", "unrelated"), + ("python", "3.11.14"), + ("torch", "2.7.0+cu128"), + ("warp", "1.17.0"), + ("raw_sha256", "0" * 64), + ("runtime_imports_embodichain", True), + ], +) +def test_source_split_labels_do_not_replace_supported_runtime_receipts( + tmp_path, field, value +): + from embodichain.learning.dynamics.report import _supported_split_sources + + raw = tmp_path / "source.h5" + raw.write_bytes(b"synthetic raw-file identity fixture") + identity = { + "source_commit": OFFICIAL_COMMIT, + "robot": "cartpole", + "python": "3.8.20", + "numpy": "1.24.4", + "torch": "2.2.2+cpu", + "warp": "1.8.0", + "runtime_imports_embodichain": False, + "raw_sha256": _hash(raw), + } + splits = ("train", "validation", "test") + metadata = { + "codec": {"robot": "cartpole"}, + "legacy_sources": { + split: {"path": str(raw), "sha256": _hash(raw)} for split in splits + }, + "source_split_identities": {split: dict(identity) for split in splits}, + } + assert _supported_split_sources(metadata) + metadata["source_split_identities"]["test"][field] = value + assert not _supported_split_sources(metadata) + + +def numerical_fixture() -> dict: + return { + "robot": "Cartpole", + "reference_commit": "476dc67fccc6a99be293bf611f897b1c06f13e30", + "checkpoint_sha256": "a" * 64, + "codec_fixtures_sha256": "b" * 64, + "history_lengths": [1, 9, 10], + "adam_updates": 3, + "precision": "float32", + "atol": 1e-5, + "rtol": 1e-4, + **{ + key: 1e-7 + for key in ( + "max_codec_feature_error", + "max_codec_target_error", + "max_codec_decode_error", + "max_forward_error", + "max_input_gradient_error", + "max_parameter_gradient_error", + "max_adam_parameter_error", + ) + }, + } + + +def test_numerical_receipt_requires_measurements_and_matching_identity(tmp_path): + path = tmp_path / "numerical.json" + path.write_text(json.dumps(numerical_fixture())) + assert _numerical_evidence(path, "cartpole")["status"] == "passed" + assert _numerical_evidence(path, "ant")["status"] == "failed" + path.write_text( + json.dumps( + { + "checks": { + key: {"passed": True} + for key in ( + "preprocessing", + "forward", + "input_gradients", + "parameter_gradients", + "adam_update", + ) + } + } + ) + ) + assert _numerical_evidence(path, "cartpole")["status"] == "pending" + + +@pytest.mark.parametrize( + "commit,reference_kind", + [ + (OFFICIAL_COMMIT, "official_warp_main"), + ("476dc67fccc6a99be293bf611f897b1c06f13e30", "newton_preview"), + ], +) +def test_pinned_main_and_preview_numerical_sources_remain_identifiable( + tmp_path, commit, reference_kind +): + raw = numerical_fixture() + raw["reference_commit"] = commit + path = tmp_path / "numerical.json" + path.write_text(json.dumps(raw)) + result = _numerical_evidence(path, "cartpole") + assert result["status"] == "passed" + assert result["reference_commit"] == commit + assert result["reference_kind"] == reference_kind + + +@pytest.mark.parametrize( + "field,value", + [ + ("max_forward_error", float("nan")), + ("max_codec_decode_error", -1.0), + ("atol", 1.0), + ("reference_commit", "unrelated"), + ("reference_commit", []), + ("reference_commit", {}), + ], +) +def test_invalid_numerical_measurements_cannot_pass(tmp_path, field, value): + raw = numerical_fixture() + raw[field] = value + path = tmp_path / "numerical.json" + path.write_text(json.dumps(raw)) + assert _numerical_evidence(path, "cartpole")["status"] == "failed" + + +def rollout_fixture(tmp_path: Path, *, motion: str = "passive"): + model = tmp_path / "run" / "best-export" + model.mkdir(parents=True) + (model / "weights.pt").write_bytes(b"synthetic test artifact") + (model / "manifest.json").write_text("{}") + trace = tmp_path / "trace.h5" + trace.write_bytes(b"synthetic test trace") + trajectories = tmp_path / "rows.jsonl" + trajectories.write_bytes(b"synthetic test row index") + raw = { + "robot": "cartpole", + "motion_mode": motion, + "state_dtype": "float32", + "teacher_origin": {"kind": "current_newton"}, + "teacher_split": "test", + "selected_teacher_sha256": "a" * 64, + "selected_teacher_hash_schema": "ordered-teacher-content-v1", + "runtime_identity": { + "device": "cpu", + "engine": "current_newton", + "state_dtype": "float32", + "python": "3.11", + "versions": {"torch": "2.7.1", "newton": "1.6.0", "warp-lang": "1.17.0"}, + }, + "trajectories": 2048, + "horizons": {str(h): {"q_mae": 0.1, "qd_mae": 0.2} for h in (1, 10, 100, 1000)}, + "model_sha256": _hash(model / "weights.pt"), + "model_manifest_sha256": _hash(model / "manifest.json"), + "trace_evidence": str(trace), + "trace_evidence_sha256": _hash(trace), + "trajectory_evidence": str(trajectories), + "trajectory_evidence_sha256": _hash(trajectories), + "failure_rate": 0.0, + "nonfinite_rate": 0.0, + "policy_return": 1.0 if motion == "fixed_policy" else None, + "policy_sha256": "b" * 64 if motion == "fixed_policy" else None, + } + source = tmp_path / "source.json" + source.write_text(json.dumps(raw)) + record = annotate_rollout_evidence( + source, "current_newton-official-0", tmp_path / "attached.json" + ) + cell = {"teacher": "current_newton", "output_dir": str(model.parent)} + return record, cell + + +def test_rollout_source_model_and_trace_identity_are_verified(tmp_path): + record, cell = rollout_fixture(tmp_path) + assert _rollout_evidence(record, cell, "cartpole")[0] == "available" + record["errors"]["1"]["q_mae"] = 0.0 + assert "Annotated" in _rollout_evidence(record, cell, "cartpole")[1] + record["errors"]["1"]["q_mae"] = 0.1 + Path(record["trace_evidence"]).write_bytes(b"changed") + assert "changed" in _rollout_evidence(record, cell, "cartpole")[1] + + +@pytest.mark.parametrize( + "field,value", + [ + ("teacher_origin", {"kind": "official_warp"}), + ("teacher_split", "train"), + ("failure_rate", -1.0), + ("nonfinite_rate", 1.5), + ("robot", "ant"), + ], +) +def test_wrong_domain_split_or_invalid_rates_fail(tmp_path, field, value): + record, cell = rollout_fixture(tmp_path) + record[field] = value + assert _rollout_evidence(record, cell, "cartpole")[0] == "failed" + + +def test_undefined_error_is_failed_instead_of_crashing_statistics(tmp_path): + record, cell = rollout_fixture(tmp_path) + record["errors"]["1000"]["q_mae"] = None + assert _rollout_evidence(record, cell, "cartpole")[0] == "failed" + + +def test_unmatched_teacher_traces_or_fixed_policy_never_form_seed_pair(tmp_path): + record, cell = rollout_fixture(tmp_path, motion="fixed_policy") + assert _rollout_evidence(record, cell, "cartpole")[0] == "available" + other = copy.deepcopy(record) + assert _matched_rollouts(record, other) + other["selected_teacher_sha256"] = "c" * 64 + assert not _matched_rollouts(record, other) + other = copy.deepcopy(record) + other["policy_sha256"] = "d" * 64 + assert not _matched_rollouts(record, other) + + +def test_historical_rollout_uses_pinned_warp_identity_without_newton(tmp_path): + from embodichain.learning.dynamics.reference import OFFICIAL_COMMIT + + record, cell = rollout_fixture(tmp_path) + source = Path(record["source_report"]) + raw = json.loads(source.read_text()) + raw["teacher_origin"] = { + "kind": "official_warp", + "source_commit": OFFICIAL_COMMIT, + "warp_version": "1.8.0", + } + raw["runtime_identity"] = { + "device": "cpu", + "engine": "official_warp", + "state_dtype": "float32", + "python": "3.11", + "reference_commit": OFFICIAL_COMMIT, + "versions": {"torch": "2.7.1", "warp-lang": "1.8.0"}, + } + raw["physical_verification"] = { + "schema": "official-warp-physical-arrays-v2", + "complete": True, + "matched_recorded_fields": True, + "missing_recorded_fields": [], + "matched_fields": ["gravity"], + "sha256": "f" * 64, + } + source.write_text(json.dumps(raw)) + record = annotate_rollout_evidence( + source, "official_warp-official-0", tmp_path / "historical.json" + ) + cell["teacher"] = "official_warp" + assert _rollout_evidence(record, cell, "cartpole")[0] == "available" + record["physical_verification"]["missing_recorded_fields"] = ["gravity"] + assert _rollout_evidence(record, cell, "cartpole")[0] == "pending" + record["physical_verification"]["missing_recorded_fields"] = [] + record["runtime_identity"]["versions"]["warp-lang"] = "1.17.0" + assert "pinned original Warp" in _rollout_evidence(record, cell, "cartpole")[1] + + +def test_full_profile_checks_actual_feature_order_precision_and_training_controls( + tmp_path, +): + import torch + from embodichain.learning.dynamics.report import _full_profile_issues + from embodichain.learning.dynamics.trainer import DynamicsTrainCfg + from embodichain.learning.models.nerd import NeRDModelCfg + + model = NeRDModelCfg(input_dims={"states_embedding": 4, "joint_f": 2}).to_dict() + root = tmp_path / "initial" + root.mkdir() + manifest = root / "manifest.json" + manifest.write_text(json.dumps({"model_config": model})) + torch.save({"weight": torch.ones(1)}, root / "weights.pt") + spec = tmp_path / "run.json" + spec.write_text(json.dumps({"initial_model": str(root)})) + train = DynamicsTrainCfg( + total_steps=4995000, + batch_size=512, + validation_interval=5000, + validation_batches=100, + lr_update_interval=5000, + lr_offset=1, + ).to_dict() + # Plan JSON sorts dict keys. This must not reorder the actual network inputs. + plan = { + "model": json.loads(json.dumps(model, sort_keys=True)), + "train": train, + "cells": [{"id": "fixture", "status": "prepared", "spec": str(spec)}], + } + assert _full_profile_issues(plan) == [] + model["input_dims"] = {"joint_f": 2, "states_embedding": 4} + manifest.write_text(json.dumps({"model_config": model})) + assert any("feature order" in issue for issue in _full_profile_issues(plan)) + model["input_dims"] = {"states_embedding": 4, "joint_f": 2} + model["head_hidden_dims"] = [128] + manifest.write_text(json.dumps({"model_config": model})) + torch.save({"weight": torch.ones(1, dtype=torch.float64)}, root / "weights.pt") + plan["train"]["batch_size"] = 16 + issues = _full_profile_issues(plan) + assert any("head_hidden_dims" in issue for issue in issues) + assert any("FP32" in issue for issue in issues) + assert any("batch_size" in issue for issue in issues) diff --git a/tests/learning/dynamics/test_sample_order.py b/tests/learning/dynamics/test_sample_order.py new file mode 100644 index 0000000000..7886d24b0b --- /dev/null +++ b/tests/learning/dynamics/test_sample_order.py @@ -0,0 +1,328 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Exact, bounded shared minibatch order and tamper detection.""" + +from __future__ import annotations + +import json +from pathlib import Path +import numpy as np +import pytest +import torch + +from embodichain.learning.dynamics._sample_order import ( + _CheckpointedOrder, + _OfficialShuffleOrder, + _open_order, + _permutation_sha256, + _write_checkpointed_order, + _write_official_shuffle_order, +) + + +def test_checkpointed_stream_matches_the_original_matrix_and_random_seeks(tmp_path): + steps, batch, block_steps = 2071, 16, 1024 + generator = torch.Generator().manual_seed(19) + expected = torch.cat( + [ + torch.randint( + 1037, (min(block_steps, steps - start), batch), generator=generator + ) + for start in range(0, steps, block_steps) + ] + ).numpy() + path, rng = _write_checkpointed_order( + tmp_path / "order.json", + dataset_size=1037, + steps=steps, + batch_size=batch, + seed=19, + block_steps=block_steps, + ) + actual = _open_order(path) + assert isinstance(actual, _CheckpointedOrder) + for index in [0, 1023, 1024, 2070, 2, 2048, 1024]: + np.testing.assert_array_equal(actual[index], expected[index]) + row = actual[0] + row[:] = -1 + np.testing.assert_array_equal(actual[0], expected[0]) + assert actual._block.nbytes <= block_steps * batch * 8 + assert rng.stat().st_size < expected.nbytes + with pytest.raises(IndexError): + actual[steps] + with pytest.raises(IndexError): + actual[-1] + + +def test_checkpoint_stream_first_block_is_independent_of_declared_budget(tmp_path): + for name, steps in [("a", 8), ("b", 3000)]: + _write_checkpointed_order( + tmp_path / (name + ".json"), + dataset_size=91, + steps=steps, + batch_size=8, + seed=7, + ) + left, right = [_open_order(tmp_path / (name + ".json")) for name in ["a", "b"]] + for step in range(8): + np.testing.assert_array_equal(left[step], right[step]) + + +def test_changed_rng_artifact_or_torch_version_fails_before_training(tmp_path): + path, rng = _write_checkpointed_order( + tmp_path / "order.json", + dataset_size=97, + steps=40, + batch_size=4, + seed=3, + block_steps=8, + ) + spec = json.loads(path.read_text()) + spec["torch_version"] = "other" + path.write_text(json.dumps(spec)) + with pytest.raises(ValueError, match="Torch version"): + _open_order(path) + spec["torch_version"] = str(torch.__version__) + path.write_text(json.dumps(spec)) + values = np.load(rng).copy() + values[-1, 20] ^= 1 + np.save(rng, values) + with pytest.raises(ValueError, match="checksum"): + _open_order(path) + + +def test_legacy_matrix_still_opens_read_only(tmp_path): + values = np.arange(24).reshape(6, 4) + np.save(tmp_path / "order.npy", values) + actual = _open_order(tmp_path / "order.npy") + assert isinstance(actual, np.memmap) + np.testing.assert_array_equal(actual, values) + + +def _official_order(tmp_path: Path) -> tuple[Path, Path, Path, np.ndarray]: + # Five complete batches/pass; each twelve-update epoch discards the unused + # tail of its third pass. The final three-update prefix uses one more pass. + size, batch, epoch_steps, steps = 23, 4, 12, 27 + seeds = np.arange(101, 108, dtype=np.int64) + hashes = [ + _permutation_sha256( + torch.randperm( + size, generator=torch.Generator().manual_seed(int(seed)) + ).numpy() + ) + for seed in seeds + ] + paths = _write_official_shuffle_order( + tmp_path / "official.json", + dataset_size=size, + steps=steps, + batch_size=batch, + epoch_steps=epoch_steps, + pass_seeds=seeds, + permutation_sha256=hashes, + provenance={ + "source_rng_evidence": "Synthetic conditional RandomSampler fixture", + "producer_torch_version": "2.2.2+cpu", + }, + ) + # Independent oracle: actual Torch RandomSampler/DataLoader, dropping + # incomplete minibatches and abandoning the final iterator each epoch. + from torch.utils.data import DataLoader, RandomSampler + + result = [] + seed_iterator = iter(seeds) + while len(result) < steps: + remaining = min(epoch_steps, steps - len(result)) + while remaining: + sampler = RandomSampler( + range(size), + replacement=False, + generator=torch.Generator().manual_seed(int(next(seed_iterator))), + ) + loader = DataLoader( + range(size), + batch_size=batch, + sampler=sampler, + drop_last=True, + num_workers=0, + generator=torch.Generator().manual_seed(8), + ) + for indices in loader: + result.append(indices.numpy()) + remaining -= 1 + if not remaining: + break + return *paths, np.stack(result) + + +def test_official_order_matches_real_random_sampler_and_epoch_resets( + tmp_path: Path, +) -> None: + path, _, _, expected = _official_order(tmp_path) + actual = _open_order(path) + assert isinstance(actual, _OfficialShuffleOrder) + assert actual.shape == expected.shape + for index in range(len(expected)): + np.testing.assert_array_equal(actual[index], expected[index]) + for index in (26, 0, 5, 11, 12, 9, 25, 1): + np.testing.assert_array_equal(actual[index], expected[index]) + assert actual._permutation.nbytes == actual.dataset_size * 8 + assert not actual.seeds.flags.writeable + for index in (-1, len(expected)): + with pytest.raises(IndexError): + actual[index] + + +def test_official_rows_remain_owned_across_seeks_and_reopen(tmp_path: Path) -> None: + path, _, _, expected = _official_order(tmp_path) + actual = _open_order(path) + row = actual[0] + actual[26] + np.testing.assert_array_equal(row, expected[0]) + row[:] = -1 + np.testing.assert_array_equal(actual[0], expected[0]) + resumed = _open_order(path) + for index in (12, 25, 26): + np.testing.assert_array_equal(resumed[index], expected[index]) + + +def test_trainer_admits_official_order_and_rejects_a_different_dataset_size( + tmp_path: Path, +) -> None: + from types import SimpleNamespace + from embodichain.learning.dynamics.trainer import DynamicsTrainer + + path, _, _, expected = _official_order(tmp_path) + # Exercise the production order admission boundary without creating a + # model, optimizer, simulation or allocating CUDA memory. + owner = DynamicsTrainer.__new__(DynamicsTrainer) + owner.cfg = SimpleNamespace(total_steps=27, batch_size=4) + owner.train_dataset = range(23) + order = owner._prepare_order(path) + assert isinstance(order, _OfficialShuffleOrder) + np.testing.assert_array_equal(order[12], expected[12]) + assert owner._sample_order_identity["path"] == str(path.resolve()) + owner.train_dataset = range(24) + with pytest.raises(ValueError, match="dataset size mismatch"): + owner._prepare_order(path) + + +@pytest.mark.parametrize("payload", ["seeds", "hashes"]) +def test_official_payload_tampering_is_rejected(tmp_path: Path, payload: str) -> None: + path, seeds, hashes, _ = _official_order(tmp_path) + target = seeds if payload == "seeds" else hashes + values = np.load(target, allow_pickle=False).copy() + values[0] = int(values[0]) ^ 1 if payload == "seeds" else b"0" * 64 + np.save(target, values, allow_pickle=False) + with pytest.raises(ValueError, match="payload checksum"): + _open_order(path) + + +def test_official_rehashed_wrong_permutation_still_fails_before_yield( + tmp_path: Path, +) -> None: + import hashlib + + path, _, hashes, _ = _official_order(tmp_path) + values = np.load(hashes, allow_pickle=False).copy() + values[0] = b"0" * 64 + np.save(hashes, values, allow_pickle=False) + spec = json.loads(path.read_text()) + spec["permutation_hashes_sha256"] = hashlib.sha256(hashes.read_bytes()).hexdigest() + path.write_text(json.dumps(spec)) + reader = _open_order(path) + with pytest.raises(ValueError, match="permutation checksum"): + reader[0] + + +@pytest.mark.parametrize( + "change", + [ + {"drop_last": False}, + {"discard_epoch_tail": False}, + {"permutation_algorithm": "other"}, + {"seeds_file": "../outside.npy"}, + ], +) +def test_official_unknown_or_escaping_contract_is_rejected( + tmp_path: Path, change: dict +) -> None: + path, _, _, _ = _official_order(tmp_path) + spec = json.loads(path.read_text()) + spec.update(change) + path.write_text(json.dumps(spec)) + with pytest.raises(ValueError, match="contract|sibling file"): + _open_order(path) + + +def test_official_full_budget_storage_is_not_a_dense_minibatch_matrix( + tmp_path: Path, +) -> None: + size, batch, steps, epoch_steps = 931840, 512, 4995000, 5000 + passes = 2997 + # This is a scale fixture, not claimed original RNG capture: repeated seeds + # let the real reader exercise the full formal shape with one hash oracle. + seeds = np.full(passes, 107, dtype=np.int64) + expected = torch.randperm( + size, generator=torch.Generator().manual_seed(107) + ).numpy() + digest = _permutation_sha256(expected) + path, seed_path, hash_path = _write_official_shuffle_order( + tmp_path / "full.json", + dataset_size=size, + steps=steps, + batch_size=batch, + epoch_steps=epoch_steps, + pass_seeds=seeds, + permutation_sha256=[digest] * passes, + provenance={"source_rng_evidence": "Synthetic storage/shape fixture"}, + ) + reader = _open_order(path) + np.testing.assert_array_equal(reader[0], expected[:batch]) + # The final epoch's third pass consumes only 1360 minibatches. + start = (1360 - 1) * batch + np.testing.assert_array_equal(reader[steps - 1], expected[start : start + batch]) + assert reader._permutation.nbytes == size * 8 + assert sum(p.stat().st_size for p in (path, seed_path, hash_path)) < 256 * 1024 + + +@pytest.mark.parametrize( + "change", + [ + {"dataset_size": 0}, + {"dataset_size": 4194305}, + {"epoch_steps": 1}, + {"pass_seeds": [1]}, + {"permutation_sha256": ["wrong"]}, + {"provenance": {}}, + ], +) +def test_official_writer_rejects_invalid_or_unbounded_inputs( + tmp_path: Path, change: dict +) -> None: + options = dict( + dataset_size=23, + steps=27, + batch_size=4, + epoch_steps=12, + pass_seeds=np.arange(7, dtype=np.int64), + permutation_sha256=["0" * 64] * 7, + provenance={"source_rng_evidence": "Synthetic negative fixture"}, + ) + options.update(change) + with pytest.raises(ValueError): + _write_official_shuffle_order(tmp_path / "invalid.json", **options) diff --git a/tests/learning/dynamics/test_small_budget_observation.py b/tests/learning/dynamics/test_small_budget_observation.py new file mode 100644 index 0000000000..9f711f3292 --- /dev/null +++ b/tests/learning/dynamics/test_small_budget_observation.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. +# ---------------------------------------------------------------------------- + +"""Regression coverage for original-model in-place input normalization.""" + +from __future__ import annotations + +import importlib.util +import json +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + + +@pytest.fixture(scope="module") +def workers() -> tuple[Any, Any]: + directory = ( + Path(__file__).resolve().parents[3] / "scripts/experiments/neural_dynamics" + ) + modules = [] + for filename in ("official_small_budget", "matched_small_budget"): + spec = importlib.util.spec_from_file_location( + "_test_" + filename, directory / (filename + ".py") + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + modules.append(module) + return tuple(modules) + + +def _inputs() -> tuple[Any, dict[str, torch.Tensor], torch.Tensor]: + # Nonzero mean and nonunit variance expose accidental double normalization. + rms = SimpleNamespace(mean=torch.tensor([3.0]), var=torch.tensor([4.0])) + model = SimpleNamespace(input_rms={"joint_f": rms}) + return model, {"joint_f": torch.tensor([[[7.0], [11.0]]])}, torch.tensor([1.5]) + + +def test_snapshot_retains_physical_features_after_model_mutation( + workers: tuple[Any, Any], tmp_path: Path +) -> None: + source, _ = workers + _, inputs, targets = _inputs() + data = {"joint_acts": inputs["joint_f"].clone(), "target": targets.clone()} + source._save_first_batch(tmp_path, data) + data["joint_acts"].sub_(3.0).div_(torch.sqrt(torch.tensor(4.0 + 1e-5))) + stored = torch.load(tmp_path / "first-preprocessed-batch.pt", weights_only=True) + assert torch.equal(stored["joint_acts"], inputs["joint_f"]) + + +def test_loss_observer_saves_before_original_forward_normalizes_in_place( + workers: tuple[Any, Any], tmp_path: Path +) -> None: + source, _ = workers + _, inputs, targets = _inputs() + data = {"joint_acts": inputs["joint_f"].clone(), "target": targets} + state = {"steps": 0, "last_loss": None} + calls = [] + + def original_loss(trainer: Any, batch: dict[str, Any], train: bool) -> Any: + calls.append(train) + batch["joint_acts"] = (batch["joint_acts"] - 3.0) / (4.0 + 1e-5) ** 0.5 + return torch.tensor(2.0), {} + + source._observed_loss(None, data, True, original_loss, state, tmp_path) + stored = torch.load(tmp_path / "first-preprocessed-batch.pt", weights_only=True) + assert torch.equal(stored["joint_acts"], inputs["joint_f"]) + assert not torch.equal(data["joint_acts"], inputs["joint_f"]) + assert state["last_loss"] == 2.0 and calls == [True] + + +@pytest.mark.parametrize("train,step", [(False, 0), (True, 1)]) +def test_loss_observer_only_captures_the_first_training_batch( + workers: tuple[Any, Any], tmp_path: Path, train: bool, step: int +) -> None: + source, _ = workers + _, inputs, targets = _inputs() + state = {"steps": step, "last_loss": None} + source._observed_loss( + None, + {"joint_acts": inputs["joint_f"], "target": targets}, + train, + lambda *_: (torch.tensor(2.0), {}), + state, + tmp_path, + ) + assert not (tmp_path / "first-preprocessed-batch.pt").exists() + + +@pytest.mark.parametrize("legacy", [False, True]) +def test_feature_comparison_respects_observation_normalization_stage( + workers: tuple[Any, Any], tmp_path: Path, legacy: bool +) -> None: + source, candidate = workers + model, inputs, targets = _inputs() + if legacy: + rms = model.input_rms["joint_f"] + normalized = (inputs["joint_f"] - rms.mean) / torch.sqrt(rms.var + 1e-5) + torch.save( + {"joint_acts": normalized, "target": targets}, + tmp_path / "first-preprocessed-batch.pt", + ) + else: + source._save_first_batch( + tmp_path, {"joint_acts": inputs["joint_f"], "target": targets} + ) + differences, stage = candidate._first_batch_differences( + model, inputs, targets, tmp_path + ) + assert differences == {"joint_f": 0.0, "target": 0.0} + assert stage == ( + "legacy_unlabeled_normalized_after_model_forward" + if legacy + else "physical_before_model_forward" + ) + + +def test_feature_comparison_still_detects_changed_inputs_and_targets( + workers: tuple[Any, Any], tmp_path: Path +) -> None: + source, candidate = workers + model, inputs, targets = _inputs() + source._save_first_batch( + tmp_path, {"joint_acts": inputs["joint_f"], "target": targets} + ) + differences, _ = candidate._first_batch_differences( + model, {"joint_f": inputs["joint_f"] + 2.0}, targets + 0.25, tmp_path + ) + assert differences["joint_f"] == pytest.approx(2.0 / (4.0 + 1e-5) ** 0.5) + assert differences["target"] == 0.25 + + +@pytest.mark.parametrize( + "change", + [ + {"feature_stage": "unknown"}, + {"format_version": 2}, + {"tensor_sha256": "0" * 64}, + ], +) +def test_feature_comparison_rejects_unknown_or_unbound_observations( + workers: tuple[Any, Any], tmp_path: Path, change: dict[str, Any] +) -> None: + source, candidate = workers + model, inputs, targets = _inputs() + source._save_first_batch( + tmp_path, {"joint_acts": inputs["joint_f"], "target": targets} + ) + metadata = tmp_path / "first-preprocessed-batch.json" + value = json.loads(metadata.read_text()) + value.update(change) + metadata.write_text(json.dumps(value)) + with pytest.raises(ValueError, match="observation metadata"): + candidate._first_batch_differences(model, inputs, targets, tmp_path) diff --git a/tests/learning/dynamics/test_trainer.py b/tests/learning/dynamics/test_trainer.py new file mode 100644 index 0000000000..195fb75f18 --- /dev/null +++ b/tests/learning/dynamics/test_trainer.py @@ -0,0 +1,901 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""CPU training, exact resume, train-only statistics and standard exports.""" + +from __future__ import annotations + +import json +import math +import subprocess +import sys + +import numpy as np +import pytest +import torch +from torch.utils.data import default_collate + +from embodichain.compute.dynamics import NeRDStateCodec +from embodichain.data_pipeline.dynamics.shards import ( + DynamicsShardWriter, + DynamicsWindowDataset, +) +from embodichain.learning.models.nerd import NeRDModel, NeRDModelCfg, load_model +from embodichain.learning.dynamics.trainer import ( + DynamicsTrainCfg, + DynamicsTrainer, + compute_statistics, + encode_batch, + evaluate_model, + normalized_mse, +) +from embodichain.learning.dynamics.workflow import run_train, run_export, run_eval +from embodichain.learning.dynamics.qualification import ( + paired_noninferiority, + compare_training_steps, + capture_training_step, +) + + +@pytest.fixture(autouse=True) +def bounded_threads(): + previous = torch.get_num_threads() + torch.set_num_threads(1) + yield + torch.set_num_threads(previous) + + +def make_data(root, validation_offset=0.0): + metadata = {"codec": {"robot": "cartpole"}, "dt": 1 / 60} + rng = np.random.default_rng(5) + with DynamicsShardWriter(root, metadata, trajectories_per_shard=3) as writer: + for split in ("train", "validation", "test"): + for index in range(4): + n = 12 + effort = rng.normal(size=(n, 2)).astype(np.float32) + q, qd = np.zeros((n + 1, 2), np.float32), np.zeros( + (n + 1, 2), np.float32 + ) + q[0] = validation_offset if split != "train" else 0 + for step in range(n): + qd[step + 1] = qd[step] + effort[step] * 0.1 + q[step + 1] = q[step] + qd[step] * 0.1 + writer.add_trajectory( + f"{split}-{index}", + {"q": q, "qd": qd}, + {"joint_f": effort}, + split=split, + ) + return metadata + + +def architecture(dropout=0.0): + return NeRDModelCfg( + input_dims={"states_embedding": 4, "joint_f": 2}, + output_dim=4, + n_embd=12, + n_head=3, + n_layer=1, + block_size=4, + history_length=3, + head_hidden_dims=(12,), + dropout=dropout, + ) + + +def build( + root, + out, + *, + steps=30, + dropout=0.0, + log_interval=1, + resident_cache_bytes=0, + numerical_profile="default", + device="cpu", +): + torch.manual_seed(2) + model = NeRDModel(architecture(dropout)) + codec = NeRDStateCodec("cartpole") + train = DynamicsWindowDataset(root, "train", 3) + validation = DynamicsWindowDataset(root, "validation", 3) + cfg = DynamicsTrainCfg( + total_steps=steps, + batch_size=8, + learning_rate=0.005, + final_learning_rate=0.002, + validation_interval=10, + checkpoint_interval=10, + seed=9, + log_interval=log_interval, + resident_cache_bytes=resident_cache_bytes, + numerical_profile=numerical_profile, + device=device, + ) + return DynamicsTrainer( + model, codec, train, validation, cfg, out, manifest=train.manifest["metadata"] + ) + + +def test_statistics_only_training_windows(tmp_path): + make_data(tmp_path / "a", 0.0) + make_data(tmp_path / "b", 1000.0) + codec = NeRDStateCodec("cartpole") + first = compute_statistics(DynamicsWindowDataset(tmp_path / "a", "train", 3), codec) + second = compute_statistics( + DynamicsWindowDataset(tmp_path / "b", "train", 3), codec + ) + for key in first["inputs"]: + torch.testing.assert_close( + first["inputs"][key]["mean"], second["inputs"][key]["mean"] + ) + with pytest.raises(ValueError, match="training split"): + compute_statistics( + DynamicsWindowDataset(tmp_path / "a", "validation", 3), codec + ) + + +@pytest.mark.parametrize("numerical_profile", ["default", "reference"]) +def test_exact_resume_with_dropout_optimizer_sampler_rng(tmp_path, numerical_profile): + make_data(tmp_path / "data") + full = build( + tmp_path / "data", + tmp_path / "full", + steps=20, + dropout=0.1, + numerical_profile=numerical_profile, + ) + full.train() + expected = {key: value.clone() for key, value in full.model.state_dict().items()} + interrupted = build( + tmp_path / "data", + tmp_path / "resume", + steps=20, + dropout=0.1, + numerical_profile=numerical_profile, + ) + assert interrupted.train(stop_after=7)["status"] == "interrupted" + restored = build( + tmp_path / "data", + tmp_path / "resume", + steps=20, + dropout=0.1, + numerical_profile=numerical_profile, + ) + restored.resume(tmp_path / "resume" / "last.pt") + restored.train() + for key, value in expected.items(): + torch.testing.assert_close( + value, restored.model.state_dict()[key], rtol=0, atol=0 + ) + assert torch.equal(full.sampler.get_state(), restored.sampler.get_state()) + full_lines = [ + json.loads(line) + for line in (tmp_path / "full" / "metrics.jsonl").read_text().splitlines() + ] + resumed_lines = [ + json.loads(line) + for line in (tmp_path / "resume" / "metrics.jsonl").read_text().splitlines() + ] + assert full_lines == resumed_lines + + +def test_training_improves_validation_and_export_matches(tmp_path): + make_data(tmp_path / "data") + trainer = build(tmp_path / "data", tmp_path / "run", steps=100) + initial = evaluate_model(trainer.model, trainer.codec, trainer.validation_dataset) + trainer.train() + final = evaluate_model(trainer.model, trainer.codec, trainer.validation_dataset) + assert final["normalized_mse"] < initial["normalized_mse"] * 0.6 + run_export(tmp_path / "run" / "last.pt", tmp_path / "export") + exported, _ = load_model(tmp_path / "export") + inputs, _, _ = encode_batch( + default_collate([trainer.validation_dataset[0]]), trainer.codec, "cpu" + ) + trainer.model.eval() + torch.testing.assert_close(exported(inputs), trainer.model(inputs), rtol=0, atol=0) + report = run_eval( + {"dataset": str(tmp_path / "data"), "model": str(tmp_path / "export")} + ) + assert report["split"] == "test" + assert report["qualification"] == "offline_only" + + +def test_workflow_fine_tune_and_resume_contract(tmp_path): + metadata = make_data(tmp_path / "data") + cfg = { + "dataset": str(tmp_path / "data"), + "output_dir": str(tmp_path / "first"), + "manifest": metadata, + "model": architecture().to_dict(), + "train": { + "total_steps": 2, + "batch_size": 4, + "validation_interval": 1, + "checkpoint_interval": 1, + }, + } + run_train(cfg) + run_export(tmp_path / "first" / "last.pt", tmp_path / "export") + fine = { + **cfg, + "output_dir": str(tmp_path / "fine"), + "fine_tune": str(tmp_path / "export"), + } + assert run_train(fine)["global_step"] == 2 + mismatch = { + **cfg, + "output_dir": str(tmp_path / "bad"), + "resume": str(tmp_path / "first" / "last.pt"), + "train": {**cfg["train"], "total_steps": 3}, + } + with pytest.raises(ValueError, match="train_cfg mismatch"): + run_train(mismatch) + + +def test_resume_rejects_unknown_preprocessing_execution(tmp_path): + make_data(tmp_path / "data") + trainer = build(tmp_path / "data", tmp_path / "run", steps=2) + trainer.train(stop_after=1) + state = torch.load(tmp_path / "run/last.pt", map_location="cpu", weights_only=True) + state.pop("preprocessing_execution") + path = tmp_path / "old-execution.pt" + torch.save(state, path) + with pytest.raises(ValueError, match="preprocessing execution mismatch"): + trainer.resume(path) + trainer.train_dataset.close() + trainer.validation_dataset.close() + + +def test_official_loss_order_and_mask(): + prediction = torch.tensor([[[1.0, 2.0], [500.0, 500.0]]], requires_grad=True) + target = torch.tensor([[[2.0, 5.0], [0.0, 0.0]]]) + variance = torch.tensor([4.0, 9.0]) + weights = 1.0 / torch.sqrt(variance + 1e-5) + expected = torch.nn.functional.mse_loss( + prediction[:, :1] * weights, target[:, :1] * weights + ) + actual = normalized_mse(prediction, target, variance, torch.tensor([[True, False]])) + torch.testing.assert_close(actual, expected) + actual.backward() + assert torch.count_nonzero(prediction.grad[:, 1]) == 0 + + +@pytest.mark.parametrize( + "device", ["cpu", pytest.param("cuda:0", marks=pytest.mark.gpu)] +) +def test_heterogeneous_target_statistics_preserve_official_loss_and_gradient(device): + if device.startswith("cuda") and not torch.cuda.is_available(): + pytest.skip("CUDA is unavailable") + generator = torch.Generator().manual_seed(42) + prediction = torch.randn(4, 10, 29, generator=generator).to(device).requires_grad_() + target = torch.randn(4, 10, 29, generator=generator).to(device) + variance = torch.logspace(-6, 3, 29, device=device) + weights = 1.0 / torch.sqrt(variance + 1e-5) + expected = torch.nn.functional.mse_loss(prediction * weights, target * weights) + actual = normalized_mse(prediction, target, variance) + expected_gradient = torch.autograd.grad(expected, prediction, retain_graph=True)[0] + actual_gradient = torch.autograd.grad(actual, prediction)[0] + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + torch.testing.assert_close(actual_gradient, expected_gradient, atol=0, rtol=0) + + +def test_qualification_gate_requires_five_seeds_and_detects_regression(): + assert paired_noninferiority([1.0, 1.0], [1.0, 1.0])["status"] == "pending" + assert paired_noninferiority([1.0] * 5, [1.01] * 5)["status"] == "passed" + assert paired_noninferiority([1.0] * 5, [1.1] * 5)["status"] == "failed" + assert paired_noninferiority([0.0] * 5, [1e-6] * 5)["difference_mode"] == "absolute" + + +def test_training_capture_detects_altered_update(): + model = NeRDModel(architecture()) + inputs = {"states_embedding": torch.randn(2, 3, 4), "joint_f": torch.randn(2, 3, 2)} + reference = capture_training_step(model, inputs, torch.randn(2, 3, 4)) + # Self-comparison only tests the checker, not official parity. + assert compare_training_steps(reference, reference)["passed"] + candidate = dict(reference, loss=reference["loss"] + 1) + assert not compare_training_steps(reference, candidate)["passed"] + + +def test_offline_import_does_not_load_simulation(): + code = "import sys; import embodichain.learning.dynamics.workflow; import embodichain.learning.dynamics.trainer; assert not any(k == 'dexsim' or k.startswith('embodichain.lab.sim') for k in sys.modules)" + subprocess.run([sys.executable, "-c", code], check=True) + + +def test_interrupted_workflow_before_validation_is_json_serializable(tmp_path): + metadata = make_data(tmp_path / "data") + result = run_train( + { + "dataset": str(tmp_path / "data"), + "output_dir": str(tmp_path / "run"), + "manifest": metadata, + "model": architecture().to_dict(), + "train": {"total_steps": 10, "batch_size": 4, "validation_interval": 10}, + "stop_after": 1, + } + ) + assert result["best_validation"] is None + assert result["status"] == "interrupted" + assert result["runtime_identity"]["parameter_placements"] == [ + {"device": "cpu", "dtype": "torch.float32"} + ] + assert ( + json.loads((tmp_path / "run" / "result.json").read_text())["runtime_identity"] + == result["runtime_identity"] + ) + + +def test_ant_offline_training_and_metrics(tmp_path): + """A floating-base model trains through native-to-legacy targets without sim.""" + metadata = {"codec": {"robot": "ant"}, "dt": 1 / 60} + with DynamicsShardWriter(tmp_path / "data", metadata) as writer: + for split in ("train", "validation", "test"): + q = np.zeros((7, 15), np.float32) + q[:, 2] = 0.5 + q[:, 6] = 1.0 + q[:, 0] = np.arange(7) * 0.01 + qd = np.zeros((7, 14), np.float32) + qd[:, 0] = 0.6 + normal = np.tile(np.array([[0.0, 0.0, 1.0]], np.float32), (6, 1)) + transitions = { + "joint_f": np.zeros((6, 14), np.float32), + "contact_normals": normal, + "contact_points_1": np.zeros((6, 3), np.float32), + "contact_depths": np.full((6, 1), 0.1, np.float32), + "contact_thicknesses_0": np.zeros((6, 1), np.float32), + "contact_thicknesses_1": np.zeros((6, 1), np.float32), + "gravity_dir": -normal, + } + writer.add_trajectory(split, {"q": q, "qd": qd}, transitions, split=split) + codec = NeRDStateCodec("ant") + train = DynamicsWindowDataset(tmp_path / "data", "train", 3) + validation = DynamicsWindowDataset(tmp_path / "data", "validation", 3) + inputs, _, _ = encode_batch(default_collate([train[0]]), codec, "cpu") + cfg = architecture() + cfg.input_dims = {key: value.shape[-1] for key, value in inputs.items()} + cfg.output_dim = 29 + model = NeRDModel(cfg) + trainer = DynamicsTrainer( + model, + codec, + train, + validation, + DynamicsTrainCfg(total_steps=2, batch_size=2, validation_interval=1), + tmp_path / "run", + manifest=metadata, + ) + trainer.train() + report = evaluate_model(model, codec, validation) + assert np.isfinite(report["normalized_mse"]) + assert report["base_orientation_error_rad"] >= 0.0 + assert report["base_position_error"] >= 0.0 + + +def test_memory_mapped_order_resume_and_integrity(tmp_path): + make_data(tmp_path / "data") + order = np.zeros((10, 8), np.int64) + path = tmp_path / "order.npy" + np.save(path, order) + trainer = build(tmp_path / "data", tmp_path / "run", steps=10) + trainer.sample_order = trainer._prepare_order(path) + assert isinstance(trainer.sample_order, np.memmap) + trainer.train(stop_after=3) + snapshot = torch.load(tmp_path / "run" / "last.pt", weights_only=True) + assert isinstance(snapshot["sample_order"], dict) + assert "sha256" in snapshot["sample_order"] + restored = build(tmp_path / "data", tmp_path / "restored", steps=10) + restored.resume(tmp_path / "run" / "last.pt") + restored.train(stop_after=4) + order[0, 0] = 1 + np.save(path, order) + invalid = build(tmp_path / "data", tmp_path / "invalid", steps=10) + with pytest.raises(ValueError, match="checksum"): + invalid.resume(tmp_path / "run" / "last.pt") + + +def test_checkpointed_order_resume_is_identical_to_uninterrupted_training(tmp_path): + from embodichain.learning.dynamics._sample_order import _write_checkpointed_order + + make_data(tmp_path / "data") + first = build(tmp_path / "data", tmp_path / "full", steps=10) + order, _ = _write_checkpointed_order( + tmp_path / "order.json", + dataset_size=len(first.train_dataset), + steps=10, + batch_size=first.cfg.batch_size, + seed=4, + block_steps=3, + ) + first.sample_order = first._prepare_order(order) + first.train() + interrupted = build(tmp_path / "data", tmp_path / "partial", steps=10) + interrupted.sample_order = interrupted._prepare_order(order) + interrupted.train(stop_after=4) + resumed = build(tmp_path / "data", tmp_path / "resumed", steps=10) + resumed.resume(tmp_path / "partial/last.pt") + resumed.train() + for name, value in first.model.state_dict().items(): + torch.testing.assert_close( + resumed.model.state_dict()[name], value, atol=0, rtol=0 + ) + + +def test_sparse_logging_retains_validation_and_final_committed_step(tmp_path): + make_data(tmp_path / "data") + trainer = build(tmp_path / "data", tmp_path / "run", steps=13, log_interval=4) + result = trainer.train() + rows = [ + json.loads(line) + for line in (tmp_path / "run/metrics.jsonl").read_text().splitlines() + ] + assert [row["step"] for row in rows] == [4, 8, 10, 12, 13] + assert "validation" in rows[2] and "validation" in rows[-1] + assert result["global_step"] == 13 + + +def test_resident_raw_cache_preserves_training_updates_and_validation(tmp_path): + make_data(tmp_path / "data") + lazy = build(tmp_path / "data", tmp_path / "lazy", steps=12) + lazy_result = lazy.train() + cached = build( + tmp_path / "data", + tmp_path / "cached", + steps=12, + resident_cache_bytes=1024 * 1024, + ) + cached_result = cached.train() + assert lazy_result["last"] == cached_result["last"] + for name, value in lazy.model.state_dict().items(): + torch.testing.assert_close( + cached.model.state_dict()[name], value, atol=0, rtol=0 + ) + + +def test_validation_batch_cap_reports_actual_windows(tmp_path): + make_data(tmp_path / "data") + trainer = build(tmp_path / "data", tmp_path / "run") + report = evaluate_model( + trainer.model, + trainer.codec, + trainer.validation_dataset, + batch_size=2, + max_batches=1, + ) + assert report["windows"] == 2 + assert report["valid_positions"] == 6 + + +def test_grouped_input_batches_preserve_exact_legacy_training_updates(tmp_path): + make_data(tmp_path / "data") + old = build(tmp_path / "data", tmp_path / "old", steps=20, dropout=0.1) + old.train_dataset.get_batch = lambda indices: default_collate( + [old.train_dataset[int(index)] for index in indices] + ) + old.validation_dataset.get_batch = lambda indices: default_collate( + [old.validation_dataset[int(index)] for index in indices] + ) + old.train() + expected = {key: value.clone() for key, value in old.model.state_dict().items()} + grouped = build(tmp_path / "data", tmp_path / "grouped", steps=20, dropout=0.1) + grouped.train() + for key, value in expected.items(): + torch.testing.assert_close( + value, grouped.model.state_dict()[key], rtol=0, atol=0 + ) + assert (tmp_path / "old" / "metrics.jsonl").read_text() == ( + tmp_path / "grouped" / "metrics.jsonl" + ).read_text() + old.train_dataset.close() + old.validation_dataset.close() + grouped.train_dataset.close() + grouped.validation_dataset.close() + + +def _add_physical_metadata(root, signature): + path = root / "manifest.json" + data = json.loads(path.read_text()) + data["metadata"].update( + physics_signature=signature, + physics_signature_schema=2, + gravity=[0.0, 0.0, -9.81], + joint_names=["slider", "pole"], + asset_sha256="asset-digest", + contact_schema="fixed-slots", + ) + path.write_text(json.dumps(data)) + return data["metadata"] + + +@pytest.mark.parametrize( + "mutation", + [ + "changed_signature", + "dropped_signature", + "changed_schema", + "dropped_gravity", + "changed_asset", + "empty_override", + ], +) +def test_training_manifest_cannot_relabel_dataset_physics(tmp_path, mutation): + make_data(tmp_path / "data") + metadata = _add_physical_metadata(tmp_path / "data", "old-physics") + override = dict(metadata) + if mutation == "changed_signature": + override["physics_signature"] = "current-physics" + elif mutation == "dropped_signature": + override.pop("physics_signature") + elif mutation == "changed_schema": + override["physics_signature_schema"] = 3 + elif mutation == "dropped_gravity": + override.pop("gravity") + elif mutation == "changed_asset": + override["asset_sha256"] = "other-asset" + else: + override = {} + with pytest.raises(ValueError, match="manifest/data .* mismatch"): + run_train( + { + "dataset": str(tmp_path / "data"), + "output_dir": str(tmp_path / "run"), + "manifest": override, + "model": architecture().to_dict(), + "train": {"total_steps": 1, "batch_size": 2}, + } + ) + assert not (tmp_path / "run").exists() + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("physics_signature", "invented"), + ("legacy_velocity", {"field": "states/legacy_velocity", "version": 1}), + ], +) +def test_training_annotations_cannot_invent_missing_physics_contract( + tmp_path, field, value +): + metadata = make_data(tmp_path / "data") + with pytest.raises(ValueError, match=f"{field} mismatch"): + run_train( + { + "dataset": str(tmp_path / "data"), + "output_dir": str(tmp_path / "run"), + "manifest": {**metadata, field: value}, + "model": architecture().to_dict(), + "train": {"total_steps": 1, "batch_size": 2}, + } + ) + + +def test_unrelated_annotations_and_fine_tune_to_new_measured_physics(tmp_path): + make_data(tmp_path / "old-data") + make_data(tmp_path / "new-data") + _add_physical_metadata(tmp_path / "old-data", "old-physics") + new_metadata = _add_physical_metadata(tmp_path / "new-data", "new-physics") + config = { + "dataset": str(tmp_path / "old-data"), + "output_dir": str(tmp_path / "old-run"), + "model": architecture().to_dict(), + "train": {"total_steps": 1, "batch_size": 2}, + } + run_train(config) + run_export(tmp_path / "old-run" / "last.pt", tmp_path / "old-export") + run_train( + { + **config, + "dataset": str(tmp_path / "new-data"), + "output_dir": str(tmp_path / "new-run"), + "fine_tune": str(tmp_path / "old-export"), + "manifest": {**new_metadata, "notes": "adapt to new measured teacher"}, + } + ) + state = torch.load(tmp_path / "new-run" / "last.pt", weights_only=True) + assert state["manifest"]["physics_signature"] == "new-physics" + assert state["manifest"]["notes"] == "adapt to new measured teacher" + + +def _gradient_parameters(): + generator = torch.Generator().manual_seed(14) + parameters = [torch.nn.Parameter(torch.zeros(width)) for width in (17, 65, 513)] + for parameter, scale in zip(parameters, (1e-3, 1.0, 10.0), strict=True): + parameter.grad = torch.randn(parameter.shape, generator=generator) * scale + return parameters + + +@pytest.mark.parametrize("profile", ["default", "reference"]) +def test_numerical_clipping_matches_its_independent_cpu_oracle(profile): + from embodichain.learning.dynamics._numerics import _clip_grad_norm_ + + parameters = _gradient_parameters() + expected = [parameter.grad.clone() for parameter in parameters] + if profile == "reference": + # This is the scalar policy validated with the historical CPU trainer; + # production must preserve its left-to-right FP64 rounding order. + norms = [float(torch.linalg.vector_norm(gradient)) for gradient in expected] + total = torch.tensor(math.sqrt(sum(value * value for value in norms))) + coefficient = (1.0 / (total + 1e-6)).clamp(max=1.0) + for gradient in expected: + gradient.mul_(coefficient) + else: + originals = [ + torch.nn.Parameter(torch.zeros_like(gradient)) for gradient in expected + ] + for parameter, gradient in zip(originals, expected, strict=True): + parameter.grad = gradient + total = torch.nn.utils.clip_grad_norm_(originals, 1.0, error_if_nonfinite=True) + actual = _clip_grad_norm_(parameters, 1.0, profile=profile) + assert torch.equal(actual, total) + for parameter, gradient in zip(parameters, expected, strict=True): + assert torch.equal(parameter.grad, gradient) + + +def test_reference_attention_context_covers_backward_eval_and_restores_flags( + tmp_path, monkeypatch +): + make_data(tmp_path / "data") + trainer = build(tmp_path / "data", tmp_path / "run", numerical_profile="reference") + original = torch.nn.functional.scaled_dot_product_attention + before = ( + torch.backends.cuda.flash_sdp_enabled(), + torch.backends.cuda.math_sdp_enabled(), + torch.backends.cuda.mem_efficient_sdp_enabled(), + ) + forward_flags, backward_flags = [], [] + + def record(*args, **kwargs): + forward_flags.append( + ( + torch.backends.cuda.flash_sdp_enabled(), + torch.backends.cuda.math_sdp_enabled(), + torch.backends.cuda.mem_efficient_sdp_enabled(), + ) + ) + return original(*args, **kwargs) + + def backward(gradient): + backward_flags.append( + ( + torch.backends.cuda.flash_sdp_enabled(), + torch.backends.cuda.math_sdp_enabled(), + torch.backends.cuda.mem_efficient_sdp_enabled(), + ) + ) + return gradient + + monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", record) + handle = next(trainer.model.parameters()).register_hook(backward) + try: + trainer.step() + evaluate_model(trainer.model, trainer.codec, trainer.validation_dataset) + finally: + handle.remove() + trainer.train_dataset.close() + trainer.validation_dataset.close() + assert forward_flags and backward_flags + assert all( + flags == (False, True, False) for flags in forward_flags + backward_flags + ) + assert before == ( + torch.backends.cuda.flash_sdp_enabled(), + torch.backends.cuda.math_sdp_enabled(), + torch.backends.cuda.mem_efficient_sdp_enabled(), + ) + + +def test_attention_context_restores_flags_on_failure(): + from embodichain.learning.dynamics._numerics import _attention_context + + before = ( + torch.backends.cuda.flash_sdp_enabled(), + torch.backends.cuda.math_sdp_enabled(), + ) + with pytest.raises(RuntimeError, match="injected"): + with _attention_context("reference"): + raise RuntimeError("injected") + assert before == ( + torch.backends.cuda.flash_sdp_enabled(), + torch.backends.cuda.math_sdp_enabled(), + ) + + +def test_reference_nonfinite_gradient_does_not_commit_optimizer_or_cursor(tmp_path): + make_data(tmp_path / "data") + trainer = build(tmp_path / "data", tmp_path / "run", numerical_profile="reference") + before = {key: value.clone() for key, value in trainer.model.state_dict().items()} + sampler = trainer.sampler.get_state().clone() + handle = next(trainer.model.parameters()).register_hook( + lambda gradient: torch.full_like(gradient, float("nan")) + ) + try: + with pytest.raises(FloatingPointError, match="gradient norm"): + trainer.step() + finally: + handle.remove() + trainer.train_dataset.close() + trainer.validation_dataset.close() + assert trainer.global_step == 0 + assert not trainer.optimizer.state + assert torch.equal(trainer.sampler.get_state(), sampler) + assert all( + torch.equal(value, trainer.model.state_dict()[key]) + for key, value in before.items() + ) + + +def test_unknown_numerical_profile_is_rejected_before_output(tmp_path): + make_data(tmp_path / "data") + with pytest.raises(ValueError, match="numerical_profile"): + build(tmp_path / "data", tmp_path / "run", numerical_profile="unknown") + assert not (tmp_path / "run").exists() + + +def test_reference_clipping_rejects_unsupported_gradient_precision(): + from embodichain.learning.dynamics._numerics import _clip_grad_norm_ + + parameter = torch.nn.Parameter(torch.ones(2, dtype=torch.float64)) + parameter.grad = torch.ones_like(parameter) + with pytest.raises(ValueError, match="FP32 gradients"): + _clip_grad_norm_([parameter], 1.0, profile="reference") + + +@pytest.mark.parametrize( + "old_profile,new_profile", [("default", "reference"), ("reference", "default")] +) +def test_resume_rejects_numerical_profile_changes(tmp_path, old_profile, new_profile): + make_data(tmp_path / "data") + original = build(tmp_path / "data", tmp_path / "old", numerical_profile=old_profile) + original.step() + original.save(tmp_path / "saved.pt") + candidate = build( + tmp_path / "data", tmp_path / "new", numerical_profile=new_profile + ) + try: + with pytest.raises(ValueError, match="numerical_profile mismatch"): + candidate.resume(tmp_path / "saved.pt") + finally: + for trainer in (original, candidate): + trainer.train_dataset.close() + trainer.validation_dataset.close() + + +def test_legacy_default_checkpoint_resume_and_export(tmp_path): + make_data(tmp_path / "data") + trainer = build(tmp_path / "data", tmp_path / "run") + trainer.step() + trainer.save(tmp_path / "new.pt") + state = torch.load(tmp_path / "new.pt", weights_only=True) + state.pop("numerical_profile") + state.pop("attention_backend") + state["train_cfg"].pop("numerical_profile") + state["model_cfg"].pop("attention_backend") + torch.save(state, tmp_path / "legacy.pt") + restored = build(tmp_path / "data", tmp_path / "restored") + try: + restored.resume(tmp_path / "legacy.pt") + assert restored.global_step == 1 + assert restored.step() == trainer.step() + run_export(tmp_path / "legacy.pt", tmp_path / "legacy-export") + model, manifest = load_model(tmp_path / "legacy-export") + assert model.cfg.attention_backend == manifest["attention_backend"] == "default" + assert manifest["numerical_profile"] == "default" + finally: + for instance in (trainer, restored): + instance.train_dataset.close() + instance.validation_dataset.close() + + +def test_reference_profile_export_and_offline_evaluation_preserve_execution(tmp_path): + make_data(tmp_path / "data") + trainer = build( + tmp_path / "data", tmp_path / "run", steps=2, numerical_profile="reference" + ) + try: + trainer.train() + run_export(tmp_path / "run/last.pt", tmp_path / "export") + model, manifest = load_model(tmp_path / "export") + assert model.cfg.attention_backend == manifest["attention_backend"] == "math" + assert ( + manifest["numerical_profile"] + == manifest["training"]["numerical_profile"] + == "reference" + ) + report = run_eval( + {"dataset": str(tmp_path / "data"), "model": str(tmp_path / "export")} + ) + assert ( + report["attention_backend"] == "math" + and report["numerical_profile"] == "reference" + ) + finally: + trainer.train_dataset.close() + trainer.validation_dataset.close() + + +@pytest.mark.gpu +def test_reference_cuda_resume_preserves_model_adam_and_rng(tmp_path): + if not torch.cuda.is_available(): + pytest.skip("CUDA is unavailable") + make_data(tmp_path / "data") + options = { + "steps": 8, + "dropout": 0.1, + "numerical_profile": "reference", + "device": "cuda:0", + } + full = build(tmp_path / "data", tmp_path / "full", **options) + full.train() + partial = build(tmp_path / "data", tmp_path / "partial", **options) + partial.train(stop_after=3) + resumed = build(tmp_path / "data", tmp_path / "resume", **options) + try: + resumed.resume(tmp_path / "partial/last.pt") + resumed.train() + for key, value in full.model.state_dict().items(): + assert torch.equal(value, resumed.model.state_dict()[key]) + left, right = full.optimizer.state_dict(), resumed.optimizer.state_dict() + assert left["param_groups"] == right["param_groups"] + for index, state in left["state"].items(): + for field, tensor in state.items(): + assert torch.equal(tensor, right["state"][index][field]) + assert torch.equal(full.sampler.get_state(), resumed.sampler.get_state()) + finally: + for trainer in (full, partial, resumed): + trainer.train_dataset.close() + trainer.validation_dataset.close() + + +@pytest.mark.gpu +def test_reference_cuda_clipping_stays_on_device_and_preserves_ordered_norm(): + from embodichain.learning.dynamics._numerics import _clip_grad_norm_ + + if not torch.cuda.is_available(): + pytest.skip("CUDA is unavailable") + # The official six-layer networks have 45 parameter tensors. Exercise both + # cold and optimized TorchScript execution at that reduction length. + parameters = [ + torch.nn.Parameter(torch.zeros(17 + index, device="cuda:0")) + for index in range(45) + ] + for index, parameter in enumerate(parameters): + parameter.grad = torch.full_like(parameter, (index + 1) * 1e-3) + original = [parameter.grad.clone() for parameter in parameters] + norms = torch.stack(torch._foreach_norm(original, 2.0)).cpu().tolist() + expected = torch.tensor( + math.sqrt(sum(value * value for value in norms)), device="cuda:0" + ) + for _ in range(10): + for parameter, gradient in zip(parameters, original, strict=True): + parameter.grad.copy_(gradient) + actual = _clip_grad_norm_( + parameters, 0.25, profile="reference", error_if_nonfinite=False + ) + assert torch.equal(actual, expected) + for parameter, gradient in zip(parameters, original, strict=True): + parameter.grad.copy_(gradient) + with torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ] + ) as profile: + actual = _clip_grad_norm_( + parameters, 0.25, profile="reference", error_if_nonfinite=False + ) + keys = {event.key for event in profile.key_averages()} + assert actual.device == parameters[0].device + assert not {"aten::item", "aten::_local_scalar_dense"} & keys diff --git a/tests/learning/models/_nerd_reference_check.py b/tests/learning/models/_nerd_reference_check.py new file mode 100644 index 0000000000..5c31b58f27 --- /dev/null +++ b/tests/learning/models/_nerd_reference_check.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. +# ---------------------------------------------------------------------------- + +"""Explicit trusted reference qualification helper; never imported by production.""" + +from __future__ import annotations +import json +import hashlib +import platform +import subprocess +import sys +from pathlib import Path +import torch +from embodichain.learning.models.nerd import convert_legacy_checkpoint, load_model + + +def compare_codec( + robot: str, destination: str, *, legacy_warp: bool = False +) -> dict[str, float]: + """Compare the actual upstream solver's pure encoding/target/decode methods. + + Construct only its math metadata: no reference simulation/model is initialized. + Native qd is independently converted to the legacy world-origin twist at the + comparison boundary, matching the upstream solver's expected input contract. + """ + if legacy_warp: + from types import SimpleNamespace + import warp as wp + import warp.sim + from integrators.integrator_neural import NeuralIntegrator as NeuralSolver + + JointType = SimpleNamespace( + FREE=wp.sim.JOINT_FREE, + REVOLUTE=wp.sim.JOINT_REVOLUTE, + PRISMATIC=wp.sim.JOINT_PRISMATIC, + ) + else: + from newton import JointType + from solvers.solver_neural import NeuralSolver + from embodichain.compute.dynamics import NeRDStateCodec, quat_rotate + + is_ant = robot == "Ant" + codec = NeRDStateCodec( + robot.lower(), base_com=(0.1, -0.2, 0.3) if is_ant else (0.0, 0.0, 0.0) + ) + reference = NeuralSolver.__new__(NeuralSolver) + reference.torch_device = "cpu" + reference.dof_q_per_env, reference.dof_qd_per_env = codec.q_dim, codec.qd_dim + reference.state_dim = codec.q_dim + codec.qd_dim + reference.prediction_dim = reference.state_dim + reference.num_envs = 2 + reference.num_contacts_per_env = 25 if is_ant else (16 if legacy_warp else 0) + reference.joint_types = ( + [JointType.FREE] + [JointType.REVOLUTE] * 8 + if is_ant + else [JointType.PRISMATIC, JointType.REVOLUTE] + ) + reference.num_joints_per_env = len(reference.joint_types) + reference.joint_q_start = [0] + list(range(7, 15)) if is_ant else [0, 1] + reference.joint_q_end = [7] + list(range(8, 16)) if is_ant else [1, 2] + reference.is_continuous_dof = torch.zeros(reference.state_dim, dtype=torch.bool) + if not is_ant: + reference.is_continuous_dof[1] = True + reference.states_frame = "body" if is_ant else "world" + reference.anchor_frame_step = "every" + reference.states_embedding_type = "identical" + reference.orientation_prediction_parameterization = "quaternion" + reference.prediction_type = "relative" + result = { + "max_codec_feature_error": 0.0, + "max_codec_target_error": 0.0, + "max_codec_decode_error": 0.0, + } + fixtures = {} + for length in [1, 9, 10]: + q = torch.randn(2, length, codec.q_dim) * 0.2 + qd = torch.randn(2, length, codec.qd_dim) * 0.2 + next_q = q + torch.randn_like(q) * 0.03 + next_qd = qd + torch.randn_like(qd) * 0.03 + if is_ant: + q[..., 3:7] += torch.tensor([0.2, -0.1, 0.3, 1.0]) + q[..., 3:7] = q[..., 3:7] / q[..., 3:7].norm(dim=-1, keepdim=True) + next_q[..., 3:7] = q[..., 3:7] + torch.randn_like(q[..., 3:7]) * 0.03 + next_q[..., 3:7] = next_q[..., 3:7] / next_q[..., 3:7].norm( + dim=-1, keepdim=True + ) + effort = torch.randn(2, length, codec.qd_dim) + gravity = torch.tensor([0.0, 0.0, -1.0]).expand(2, length, 3) + contacts = None + if is_ant: + depth = torch.randn(2, length, 25) * 0.1 + contacts = { + "contact_depths": depth, + "contact_normals": torch.tensor([0.0, 0.0, 1.0]) + .expand(2, length, 25, 3) + .flatten(-2), + "contact_points_1": torch.cat( + (torch.randn(2, length, 25, 2), torch.zeros(2, length, 25, 1)), -1 + ).flatten(-2), + "contact_thicknesses_0": torch.full((2, length, 25), 0.05), + "contact_thicknesses_1": torch.zeros(2, length, 25), + } + contacts["contact_masks"] = ( + reference.get_contact_masks( + depth, + contacts["contact_thicknesses_0"] + + contacts["contact_thicknesses_1"], + ) + if legacy_warp + else reference.get_contact_masks( + depth, + contacts["contact_thicknesses_0"], + contacts["contact_thicknesses_1"], + ) + ) + + def legacy(q: torch.Tensor, qd: torch.Tensor) -> torch.Tensor: + if not is_ant: + return torch.cat((q, qd), -1) + # Definition v(x)=omega cross x + nu, evaluated at root COM. + com = q[..., :3] + quat_rotate(q[..., 3:7], q.new_tensor(codec.base_com)) + omega = qd[..., 3:6] + nu = qd[..., :3] - torch.linalg.cross(omega, com) + return torch.cat((q, omega, nu, qd[..., 6:]), -1) + + state, next_state = legacy(q, qd), legacy(next_q, next_qd) + root_q = ( + q[..., :7] + if is_ant + else torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]).expand(2, length, 7) + ) + raw = { + "root_body_q": root_q.clone(), + "states": state.clone(), + "next_states": next_state.clone(), + "joint_f": effort[..., -8:].clone() if is_ant else effort.clone(), + "gravity_dir": gravity.clone(), + } + if contacts: + raw.update({k: v.clone() for k, v in contacts.items()}) + elif legacy_warp: + # The original contact-free Cartpole still authors the two boxes' + # 16 inactive slots; its preprocessing expects their shape metadata. + raw.update( + { + "contact_points_1": torch.zeros(2, length, 48), + "contact_normals": torch.zeros(2, length, 48), + "contact_depths": torch.full((2, length, 16), 1000.0), + "contact_masks": torch.zeros(2, length, 16, dtype=torch.bool), + } + ) + processed = reference.process_neural_model_inputs(raw) + encoded = codec.encode(q, qd, effort, contacts, gravity) + for key, value in encoded.items(): + torch.testing.assert_close(value, processed[key], atol=1e-5, rtol=1e-4) + result["max_codec_feature_error"] = max( + result["max_codec_feature_error"], + (value - processed[key]).abs().max().item(), + ) + target = codec.targets(q, qd, next_q, next_qd) + expected = reference.convert_next_states_to_prediction( + processed["states"], processed["next_states"] + ) + torch.testing.assert_close(target, expected, atol=1e-5, rtol=1e-4) + result["max_codec_target_error"] = max( + result["max_codec_target_error"], (target - expected).abs().max().item() + ) + # Arbitrary predictions also exercise decode independently of target inverse. + prediction = torch.randn_like(target) * 0.03 + if is_ant: + prediction[..., 6] += 1.0 + actual_q, actual_qd = codec.decode(q, qd, prediction) + predicted_body = reference.convert_prediction_to_next_states( + processed["states"], prediction + ) + # The upstream method converts one last step; flatten B*T into batch and T=1. + world = reference.convert_states_back_to_world( + root_q.reshape(-1, 1, 7), predicted_body.reshape(-1, reference.state_dim) + ).reshape_as(state) + reference.wrap2PI(world) + actual = legacy(actual_q, actual_qd) + torch.testing.assert_close(actual, world, atol=1e-5, rtol=1e-4) + result["max_codec_decode_error"] = max( + result["max_codec_decode_error"], (actual - world).abs().max().item() + ) + fixtures[str(length)] = { + "q": q, + "qd": qd, + "next_q": next_q, + "next_qd": next_qd, + "effort": effort, + "contacts": contacts, + "expected_features": {k: processed[k] for k in encoded}, + "expected_target": expected, + "prediction": prediction, + "expected_decoded_legacy": world, + } + torch.save(fixtures, Path(destination) / "codec-fixtures.pt") + return result + + +def main() -> None: + root, robot, destination = sys.argv[1:] + reference_commit = subprocess.check_output( + ["git", "-C", root, "rev-parse", "HEAD"], text=True + ).strip() + if reference_commit not in { + "370e89d2b8ce68aa3578aa7812859c681b7502bc", + "476dc67fccc6a99be293bf611f897b1c06f13e30", + }: + raise ValueError("Use one of the explicitly pinned official references") + sys.path.insert(0, root) + torch.set_num_threads(2) + path = ( + Path(root) + / "pretrained_models" + / "NeRD_models" + / robot + / "model" + / "nn" + / "model.pt" + ) + convert_legacy_checkpoint( + path, + destination, + {"codec": {"robot": robot.lower()}, "dt": 1 / 60}, + trust_pickle=True, + ) + actual, metadata = load_model(destination) + official = torch.load(path, map_location="cpu", weights_only=False)[0].eval() + torch.manual_seed(7) + result = { + "robot": robot, + "checkpoint_sha256": metadata["legacy_checkpoint_sha256"], + "max_forward_error": 0.0, + "max_input_gradient_error": 0.0, + "max_parameter_gradient_error": 0.0, + "max_adam_parameter_error": 0.0, + } + + def reference_inputs(values: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + return {"joint_acts" if k == "joint_f" else k: v for k, v in values.items()} + + for length in [1, 9, 10]: + features = { + k: torch.randn(2, length, d, requires_grad=True) + for k, d in actual.cfg.input_dims.items() + } + reference = { + k: v.detach().clone().requires_grad_() for k, v in features.items() + } + predicted = actual(features) + expected = official(reference_inputs(reference)) + torch.testing.assert_close(predicted, expected, atol=1e-5, rtol=1e-4) + result["max_forward_error"] = max( + result["max_forward_error"], (predicted - expected).abs().max().item() + ) + predicted.square().mean().backward() + expected.square().mean().backward() + for k in features: + torch.testing.assert_close( + features[k].grad, reference[k].grad, atol=1e-5, rtol=1e-4 + ) + result["max_input_gradient_error"] = max( + result["max_input_gradient_error"], + (features[k].grad - reference[k].grad).abs().max().item(), + ) + for (name, p), (ref_name, rp) in zip( + actual.named_parameters(), official.named_parameters() + ): + assert name == ref_name + torch.testing.assert_close(p.grad, rp.grad, atol=1e-5, rtol=1e-4) + result["max_parameter_gradient_error"] = max( + result["max_parameter_gradient_error"], + (p.grad - rp.grad).abs().max().item(), + ) + actual.zero_grad() + official.zero_grad() + # Deterministic three-step normalized MSE/Adam equivalence with shared stats. + opt = torch.optim.Adam(actual.parameters(), lr=1e-3) + ref_opt = torch.optim.Adam(official.parameters(), lr=1e-3) + for _ in range(3): + features = {k: torch.randn(2, 10, d) for k, d in actual.cfg.input_dims.items()} + target = torch.randn(2, 10, actual.cfg.output_dim) + out = actual(features) + ref = official(reference_inputs({k: v.clone() for k, v in features.items()})) + # Same variance normalizer in both runs, isolating trainer update math. + loss = ((out - target).square() / (actual.output_rms.var + 1e-5)).mean() + ref_loss = ((ref - target).square() / (official.output_rms.var + 1e-5)).mean() + torch.testing.assert_close(loss, ref_loss, atol=1e-5, rtol=1e-4) + opt.zero_grad() + ref_opt.zero_grad() + loss.backward() + ref_loss.backward() + norm = torch.nn.utils.clip_grad_norm_(actual.parameters(), 1.0) + ref_norm = torch.nn.utils.clip_grad_norm_(official.parameters(), 1.0) + torch.testing.assert_close(norm, ref_norm, atol=1e-5, rtol=1e-4) + opt.step() + ref_opt.step() + for (name, p), (ref_name, rp) in zip( + actual.named_parameters(), official.named_parameters() + ): + assert name == ref_name + torch.testing.assert_close(p, rp, atol=1e-5, rtol=1e-4) + result["max_adam_parameter_error"] = max( + result["max_adam_parameter_error"], (p - rp).abs().max().item() + ) + result.update( + compare_codec( + robot, destination, legacy_warp=reference_commit.startswith("370e89d") + ) + ) + result["reference_commit"] = subprocess.check_output( + ["git", "-C", root, "rev-parse", "HEAD"], text=True + ).strip() + result["precision"] = "float32" + result["device"] = "cpu" + result["torch_version"] = str(torch.__version__) + result["python_version"] = platform.python_version() + result["history_lengths"] = [1, 9, 10] + result["adam_updates"] = 3 + result["atol"] = 1e-5 + result["rtol"] = 1e-4 + result["qualification_scope"] = ( + "network, gradient, short optimizer and codec parity; not full-budget training quality" + ) + result["codec_fixtures_sha256"] = hashlib.sha256( + (Path(destination) / "codec-fixtures.pt").read_bytes() + ).hexdigest() + (Path(destination) / "parity-result.json").write_text(json.dumps(result, indent=2)) + print(json.dumps(result)) + + +if __name__ == "__main__": + main() diff --git a/tests/learning/models/test_nerd.py b/tests/learning/models/test_nerd.py new file mode 100644 index 0000000000..77ce6d093c --- /dev/null +++ b/tests/learning/models/test_nerd.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. +# ---------------------------------------------------------------------------- + +from __future__ import annotations +import json +import os +from pathlib import Path +import subprocess +import sys +import pytest +import torch +from embodichain.compute.dynamics import append_history +from embodichain.learning.models.nerd import ( + NeRDModel, + NeRDModelCfg, + export_model, + load_model, + convert_legacy_checkpoint, +) + + +def tiny() -> NeRDModel: + return NeRDModel( + NeRDModelCfg( + input_dims={"states_embedding": 4, "joint_f": 2}, + output_dim=4, + n_embd=12, + n_head=3, + n_layer=2, + block_size=10, + history_length=10, + ) + ) + + +def test_normalizers_and_parameters_follow_dtype() -> None: + model = tiny().double() + assert all(v.dtype == torch.double for v in model.state_dict().values()) + model.set_normalization( + { + k: (torch.ones(d), torch.full((d,), 2.0)) + for k, d in model.cfg.input_dims.items() + }, + torch.ones(4), + torch.full((4,), 3.0), + ) + value = torch.randn(2, 3, 4, dtype=torch.double) + torch.testing.assert_close( + model.input_rms["states_embedding"](value), (value - 1) / (2 + 1e-5) ** 0.5 + ) + + +def test_variable_prefix_matches_individual_unpadded_calls() -> None: + model = tiny().eval() + inputs = {k: torch.randn(3, 10, d) for k, d in model.cfg.input_dims.items()} + copies = {k: v.clone() for k, v in inputs.items()} + lengths = torch.tensor([1, 9, 10]) + outputs = model.evaluate(inputs, lengths) + for row, length in enumerate(lengths): + expected = model.evaluate( + {k: v[row : row + 1, :length] for k, v in inputs.items()} + ) + torch.testing.assert_close( + outputs[row : row + 1], expected, atol=1e-6, rtol=1e-5 + ) + for k in inputs: + torch.testing.assert_close(inputs[k], copies[k]) + + +def test_multistep_history_control_and_parameter_gradients() -> None: + model = tiny().double().eval() + history = { + k: torch.zeros(1, 10, d, dtype=torch.double) + for k, d in model.cfg.input_dims.items() + } + lengths = torch.zeros(1, dtype=torch.long) + initial = torch.randn(1, 4, dtype=torch.double, requires_grad=True) + state = initial + controls = [ + torch.randn(1, 2, dtype=torch.double, requires_grad=True) for _ in range(12) + ] + for control in controls: + history, lengths = append_history( + history, lengths, {"states_embedding": state, "joint_f": control} + ) + state = state + model.evaluate(history, lengths).squeeze(1) + state.square().mean().backward() + assert initial.grad is not None and torch.isfinite(initial.grad).all() + assert all(c.grad is not None and c.grad.abs().sum() > 0 for c in controls) + assert all( + p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters() + ) + + +def test_export_roundtrip_and_checksum(tmp_path: Path) -> None: + model = tiny().double().eval() + inputs = { + k: torch.randn(2, 3, d, dtype=torch.double) + for k, d in model.cfg.input_dims.items() + } + export_model(tmp_path, model, {"codec": {"robot": "cartpole"}, "dt": 1 / 60}) + loaded, manifest = load_model(tmp_path) + assert manifest["codec"]["robot"] == "cartpole" + torch.testing.assert_close(model(inputs), loaded(inputs), atol=0, rtol=0) + with pytest.raises(FileExistsError): + export_model(tmp_path, model, {}) + with (tmp_path / "weights.pt").open("ab") as stream: + stream.write(b"changed") + with pytest.raises(ValueError, match="checksum"): + load_model(tmp_path) + + +def test_legacy_pickle_requires_explicit_trust(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="trust_pickle"): + convert_legacy_checkpoint(tmp_path / "missing", tmp_path, {}) + + +def test_attention_backend_preserves_initial_weights_and_rng() -> None: + torch.manual_seed(19) + default = tiny() + expected_rng = torch.get_rng_state().clone() + cfg = default.cfg.copy() + cfg.attention_backend = "math" + torch.manual_seed(19) + math_model = NeRDModel(cfg) + assert torch.equal(expected_rng, torch.get_rng_state()) + assert all( + torch.equal(value, math_model.state_dict()[key]) + for key, value in default.state_dict().items() + ) + + +def test_math_attention_export_load_and_provider_prediction_restore_flags( + tmp_path: Path, monkeypatch +) -> None: + model = tiny().eval() + model.cfg.attention_backend = "math" + export_model( + tmp_path, + model, + { + "codec": {"robot": "cartpole"}, + "dt": 1 / 60, + "numerical_profile": "reference", + }, + ) + loaded, manifest = load_model(tmp_path) + assert loaded.cfg.attention_backend == manifest["attention_backend"] == "math" + assert manifest["numerical_profile"] == "reference" + original = torch.nn.functional.scaled_dot_product_attention + flags = [] + before = ( + torch.backends.cuda.flash_sdp_enabled(), + torch.backends.cuda.math_sdp_enabled(), + torch.backends.cuda.mem_efficient_sdp_enabled(), + ) + + def record(*args, **kwargs): + flags.append( + ( + torch.backends.cuda.flash_sdp_enabled(), + torch.backends.cuda.math_sdp_enabled(), + torch.backends.cuda.mem_efficient_sdp_enabled(), + ) + ) + return original(*args, **kwargs) + + monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", record) + inputs = { + key: torch.randn(2, 3, width) for key, width in loaded.cfg.input_dims.items() + } + assert torch.equal(loaded.evaluate(inputs), model.evaluate(inputs)) + assert flags and all(row == (False, True, False) for row in flags) + assert before == ( + torch.backends.cuda.flash_sdp_enabled(), + torch.backends.cuda.math_sdp_enabled(), + torch.backends.cuda.mem_efficient_sdp_enabled(), + ) + + +def test_legacy_portable_export_without_execution_fields_remains_default( + tmp_path: Path, +) -> None: + model = tiny().eval() + export_model(tmp_path, model, {"codec": {"robot": "cartpole"}, "dt": 1 / 60}) + manifest = json.loads((tmp_path / "manifest.json").read_text()) + manifest.pop("attention_backend") + manifest.pop("numerical_profile") + manifest["model_config"].pop("attention_backend") + (tmp_path / "manifest.json").write_text(json.dumps(manifest)) + loaded, _ = load_model(tmp_path) + assert loaded.cfg.attention_backend == "default" + assert all( + torch.equal(value, loaded.state_dict()[key]) + for key, value in model.state_dict().items() + ) + + +@pytest.mark.parametrize( + "mutation", ["attention_backend", "numerical_profile", "reference_requires_math"] +) +def test_export_rejects_conflicting_execution_metadata( + tmp_path: Path, mutation: str +) -> None: + model = tiny() + model.cfg.attention_backend = "math" + export_model(tmp_path, model, {"numerical_profile": "reference"}) + manifest = json.loads((tmp_path / "manifest.json").read_text()) + if mutation == "attention_backend": + manifest["attention_backend"] = "default" + elif mutation == "numerical_profile": + manifest["numerical_profile"] = "unknown" + else: + manifest["attention_backend"] = "default" + manifest["model_config"]["attention_backend"] = "default" + (tmp_path / "manifest.json").write_text(json.dumps(manifest)) + with pytest.raises(ValueError, match="attention_backend|numerical_profile"): + load_model(tmp_path) + + +@pytest.mark.parametrize("robot", ["Cartpole", "Ant"]) +def test_explicit_official_checkpoint_forward_gradient_and_adam( + robot: str, tmp_path: Path +) -> None: + # Reference code is only loaded when the operator explicitly opts into a local + # trusted checkout. Run it in a subprocess to avoid polluting production imports. + reference = os.environ.get("NERD_REFERENCE_ROOT") + if reference is None: + pytest.skip("Set NERD_REFERENCE_ROOT to the pinned trusted upstream checkout") + script = Path(__file__).with_name("_nerd_reference_check.py") + subprocess.run( + [sys.executable, str(script), reference, robot, str(tmp_path)], + check=True, + env={ + **os.environ, + "OMP_NUM_THREADS": "2", + "MKL_NUM_THREADS": "2", + "PYTHONPATH": os.pathsep.join( + filter( + None, + ( + os.environ.get("NERD_REFERENCE_DEPS"), + str(Path(__file__).resolve().parents[3]), + ), + ) + ), + }, + ) diff --git a/tests/test_agent_context_map.py b/tests/test_agent_context_map.py index 6e51206086..1b9ac04d67 100644 --- a/tests/test_agent_context_map.py +++ b/tests/test_agent_context_map.py @@ -62,6 +62,7 @@ def test_map_registers_the_supported_context_domains() -> None: "motion-planning", "rl-learning", "learning-layers", + "neural-dynamics", "configclass-pattern", "randomization", "atomic-actions", From a58de92a82b4e1c1c260083b5b1930173322aceb Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 24 Oct 2026 14:50:39 +0000 Subject: [PATCH 02/21] docs(dynamics): register the offline workflow in architecture navigation --- docs/architecture/curated.json | 37 +++++++++++++++++++++++++--- embodichain/data/nerd/ATTRIBUTION.md | 18 ++++++++------ 2 files changed, 45 insertions(+), 10 deletions(-) diff --git a/docs/architecture/curated.json b/docs/architecture/curated.json index 00aa7240a4..0d05ee8810 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,35 @@ "label": "Shared numerical geometry and layout" } ] + }, + { + "id": "neural-dynamics", + "label": "NeRD offline workflow", + "kind": "package", + "topic_ids": [ + "neural-dynamics" + ], + "summary": "Generates analytical trajectories and trains, evaluates and exports causal neural dynamics models.", + "details": "State/contact codecs live in compute/dynamics, bounded trajectory collection and HDF5 storage in data_pipeline/dynamics, NeRD networks in learning/models/nerd, and offline learning in learning/dynamics. The independent workflow uses built-in Newton teachers without initializing DexSim; live Scene solver integration is a separate consumer.", + "boundaries": [ + "Model training and offline evaluation do not own a simulation World.", + "A collector profile records its physical domain; exported weights do not establish equivalence across teachers." + ], + "evidence": [ + { + "path": "embodichain/learning/models/nerd/model.py", + "symbol": "NeRDModel", + "start_line": 201, + "end_line": 201, + "excerpt": "class NeRDModel(nn.Module):" + } + ], + "documentation": [ + { + "docname": "features/neural_dynamics", + "label": "NeRD data and training workflow" + } + ] } ], "edges": [ @@ -5774,7 +5803,8 @@ "grasp-tools", "urdf-assembly", "convex-decomposition", - "scene-services" + "scene-services", + "neural-dynamics" ], "edge_ids": [ "adapter-calls-compiler", @@ -5989,7 +6019,8 @@ "differentiable-trainer", "differentiable-collector", "observation-normalizer", - "policy-evaluation" + "policy-evaluation", + "neural-dynamics" ] }, { diff --git a/embodichain/data/nerd/ATTRIBUTION.md b/embodichain/data/nerd/ATTRIBUTION.md index e198e27878..b791343085 100644 --- a/embodichain/data/nerd/ATTRIBUTION.md +++ b/embodichain/data/nerd/ATTRIBUTION.md @@ -8,10 +8,14 @@ NVlabs/neural-robot-dynamics, commit Source: https://github.com/NVlabs/neural-robot-dynamics/tree/476dc67fccc6a99be293bf611f897b1c06f13e30/envs/newton_envs/assets -EmbodiChain translates these assets using the installed Newton public importer, -then passes their physical descriptors through DexSim Spawn. This is a current -Newton teacher profile, not a claim to reproduce legacy Warp 1.5.1 physics. -Imported motor actuator declarations and XML solver options are not installed -into the Spawn scene: the host supplies actual generalized forces and declares -its own solver configuration. Masses, inertia, COMs, joint limits, passive joint -parameters and collision geometry/material parameters are retained. +EmbodiChain interprets these assets using the installed Newton public importer. +The independent data collector finalizes an analytical Newton model directly; +Scene-based neural solver integration is a separate consumer of the same asset +interpretation. Neither path claims to reproduce legacy Warp 1.5.1 physics. + +The host supplies generalized forces and its own solver configuration rather +than installing the source motor actuators. Dataset manifests identify the +importer profile, asset hash, actual parameter hashes and solver versions. +`warp_reference` reconstructs measured legacy mass, inertia and contact-parameter +semantics; it does not select the historical simulator or guarantee identical +trajectories across solver versions. From cc1a67e71902b4a4ff32c4dd818ff2e8566dc714 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 24 Oct 2026 16:54:57 +0000 Subject: [PATCH 03/21] feat(sl): add shared supervised batch and optimizer update contracts --- embodichain/learning/sl/__init__.py | 29 +++++ embodichain/learning/sl/learner.py | 170 ++++++++++++++++++++++++++++ embodichain/learning/sl/types.py | 109 ++++++++++++++++++ tests/learning/sl/test_learner.py | 107 +++++++++++++++++ 4 files changed, 415 insertions(+) create mode 100644 embodichain/learning/sl/__init__.py create mode 100644 embodichain/learning/sl/learner.py create mode 100644 embodichain/learning/sl/types.py create mode 100644 tests/learning/sl/test_learner.py diff --git a/embodichain/learning/sl/__init__.py b/embodichain/learning/sl/__init__.py new file mode 100644 index 0000000000..fd9479cf52 --- /dev/null +++ b/embodichain/learning/sl/__init__.py @@ -0,0 +1,29 @@ +# ---------------------------------------------------------------------------- +# 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 + +__all__ = [ + "LossResult", + "SupervisedBatch", + "SupervisedObjective", + "SupervisedLearner", +] diff --git a/embodichain/learning/sl/learner.py b/embodichain/learning/sl/learner.py new file mode 100644 index 0000000000..b7d0636ce8 --- /dev/null +++ b/embodichain/learning/sl/learner.py @@ -0,0 +1,170 @@ +# ---------------------------------------------------------------------------- +# 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().float(), + torch.as_tensor(norm).detach().float(), + ) + ) + .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 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(), + "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}") + count = state["num_updates"] + if type(count) is not int or count < 0: + raise ValueError("Invalid committed update count") + self.model.load_state_dict(saved) + self.optimizer.load_state_dict(state["optimizer"]) + self.num_updates = count 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/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 From dd526602dc7dc800ee04c03ee45f6fb7d6da1984 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 24 Oct 2026 16:54:57 +0000 Subject: [PATCH 04/21] feat(sl): add shared supervised batch and optimizer update contracts --- embodichain/learning/sl/__init__.py | 29 +++++ embodichain/learning/sl/learner.py | 170 ++++++++++++++++++++++++++++ embodichain/learning/sl/types.py | 109 ++++++++++++++++++ tests/learning/sl/test_learner.py | 107 +++++++++++++++++ 4 files changed, 415 insertions(+) create mode 100644 embodichain/learning/sl/__init__.py create mode 100644 embodichain/learning/sl/learner.py create mode 100644 embodichain/learning/sl/types.py create mode 100644 tests/learning/sl/test_learner.py diff --git a/embodichain/learning/sl/__init__.py b/embodichain/learning/sl/__init__.py new file mode 100644 index 0000000000..fd9479cf52 --- /dev/null +++ b/embodichain/learning/sl/__init__.py @@ -0,0 +1,29 @@ +# ---------------------------------------------------------------------------- +# 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 + +__all__ = [ + "LossResult", + "SupervisedBatch", + "SupervisedObjective", + "SupervisedLearner", +] diff --git a/embodichain/learning/sl/learner.py b/embodichain/learning/sl/learner.py new file mode 100644 index 0000000000..b7d0636ce8 --- /dev/null +++ b/embodichain/learning/sl/learner.py @@ -0,0 +1,170 @@ +# ---------------------------------------------------------------------------- +# 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().float(), + torch.as_tensor(norm).detach().float(), + ) + ) + .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 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(), + "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}") + count = state["num_updates"] + if type(count) is not int or count < 0: + raise ValueError("Invalid committed update count") + self.model.load_state_dict(saved) + self.optimizer.load_state_dict(state["optimizer"]) + self.num_updates = count 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/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 From 8408d9f9775e2e349353073edc8e66e2ce40a5ea Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 24 Oct 2026 17:00:51 +0000 Subject: [PATCH 05/21] feat(sl): add offline run recovery and recipe-based train-sl CLI --- agent_context/MAP.yaml | 12 + .../topics/learning-layers/learning-layers.md | 3 + .../supervised-learning.md | 33 +++ .../embodichain/embodichain.learning.rst | 12 + .../embodichain/embodichain.learning.sl.rst | 125 ++++++++++ embodichain/cli/main.py | 5 + embodichain/learning/__init__.py | 2 +- embodichain/learning/sl/__init__.py | 6 + embodichain/learning/sl/learner.py | 6 +- embodichain/learning/sl/sources.py | 149 +++++++++++ embodichain/learning/sl/trainer.py | 231 ++++++++++++++++++ embodichain/learning/sl/workflow.py | 73 ++++++ tests/learning/sl/test_trainer.py | 115 +++++++++ 13 files changed, 769 insertions(+), 3 deletions(-) create mode 100644 agent_context/topics/supervised-learning/supervised-learning.md create mode 100644 docs/source/api_reference/embodichain/embodichain.learning.sl.rst create mode 100644 embodichain/learning/sl/sources.py create mode 100644 embodichain/learning/sl/trainer.py create mode 100644 embodichain/learning/sl/workflow.py create mode 100644 tests/learning/sl/test_trainer.py diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index b4ac4f4c04..c80ee880eb 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -523,3 +523,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] + status: active diff --git a/agent_context/topics/learning-layers/learning-layers.md b/agent_context/topics/learning-layers/learning-layers.md index fd46bc6486..bb76d70034 100644 --- a/agent_context/topics/learning-layers/learning-layers.md +++ b/agent_context/topics/learning-layers/learning-layers.md @@ -21,6 +21,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/supervised-learning/supervised-learning.md b/agent_context/topics/supervised-learning/supervised-learning.md new file mode 100644 index 0000000000..8e7cfc0874 --- /dev/null +++ b/agent_context/topics/supervised-learning/supervised-learning.md @@ -0,0 +1,33 @@ +# Supervised learning (SL) + +`embodichain.learning.sl` owns common supervised optimizer updates and offline +training orchestration. Models remain in `learning.models`; 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`. + +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. +Validation callbacks use frozen models and restored RNG; 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/source/api_reference/embodichain/embodichain.learning.rst b/docs/source/api_reference/embodichain/embodichain.learning.rst index 10f9c879a4..da7b1b942a 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.rst +++ b/docs/source/api_reference/embodichain/embodichain.learning.rst @@ -10,6 +10,7 @@ embodichain.learning layers models rl + sl Reinforcement Learning ---------------------- @@ -64,3 +65,14 @@ standalone IK layers do not import RL, simulation or TensorDict. :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..f016e712bb --- /dev/null +++ b/docs/source/api_reference/embodichain/embodichain.learning.sl.rst @@ -0,0 +1,125 @@ +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 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 index fd9479cf52..d63ae13c3a 100644 --- a/embodichain/learning/sl/__init__.py +++ b/embodichain/learning/sl/__init__.py @@ -20,10 +20,16 @@ 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/learner.py b/embodichain/learning/sl/learner.py index b7d0636ce8..1ade37c9b5 100644 --- a/embodichain/learning/sl/learner.py +++ b/embodichain/learning/sl/learner.py @@ -113,8 +113,8 @@ def update(self, batch: SupervisedBatch) -> dict[str, float | int]: observed = ( torch.stack( ( - result.loss.detach().float(), - torch.as_tensor(norm).detach().float(), + result.loss.detach(), + torch.as_tensor(norm, device=result.loss.device).detach(), ) ) .cpu() @@ -162,6 +162,8 @@ def load_state_dict(self, state: Mapping[str, Any]) -> None: 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") diff --git a/embodichain/learning/sl/sources.py b/embodichain/learning/sl/sources.py new file mode 100644 index 0000000000..bf4e7f9443 --- /dev/null +++ b/embodichain/learning/sl/sources.py @@ -0,0 +1,149 @@ +# ---------------------------------------------------------------------------- +# 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. + + 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 structure(value: Any) -> Any: + if isinstance(value, torch.Tensor): + return {"shape": list(value.shape), "dtype": str(value.dtype)} + if isinstance(value, Mapping): + return {str(key): structure(item) for key, item in value.items()} + if isinstance(value, (tuple, list)): + return [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), + { + **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..d12eabfc43 --- /dev/null +++ b/embodichain/learning/sl/trainer.py @@ -0,0 +1,231 @@ +# ---------------------------------------------------------------------------- +# 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 .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.learner.model.to(self.cfg.device) + + 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") + 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 + ): + rng = self._rng_state() + was_training = self.learner.model.training + try: + self.learner.model.eval() + with torch.no_grad(): + metrics.update( + { + f"validation/{key}": float(value) + for key, value in self.evaluate(self.learner).items() + } + ) + finally: + self._restore_rng(rng) + self.learner.model.train(was_training) + 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 _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) + path.parent.mkdir(parents=True, exist_ok=True) + state = { + "version": 1, + "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": ( + None if self.scheduler is None else self.scheduler.state_dict() + ), + } + 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") + 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"]) diff --git a/embodichain/learning/sl/workflow.py b/embodichain/learning/sl/workflow.py new file mode 100644 index 0000000000..467bf1781c --- /dev/null +++ b/embodichain/learning/sl/workflow.py @@ -0,0 +1,73 @@ +# ---------------------------------------------------------------------------- +# 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 +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() -> None: + """Train a supervised recipe from a YAML/JSON file and optional checkpoint.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", required=True, type=Path) + parser.add_argument("--resume") + args = parser.parse_args() + 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/tests/learning/sl/test_trainer.py b/tests/learning/sl/test_trainer.py new file mode 100644 index 0000000000..184f5a1be2 --- /dev/null +++ b/tests/learning/sl/test_trainer.py @@ -0,0 +1,115 @@ +# ---------------------------------------------------------------------------- +# 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 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) From 11142d50e86e2b723f077e1329d45b56f911b46f Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 24 Oct 2026 17:03:03 +0000 Subject: [PATCH 06/21] test(sl): register the supervised learning context domain --- tests/test_agent_context_map.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_agent_context_map.py b/tests/test_agent_context_map.py index 6e51206086..ae9cf4b51d 100644 --- a/tests/test_agent_context_map.py +++ b/tests/test_agent_context_map.py @@ -62,6 +62,7 @@ def test_map_registers_the_supported_context_domains() -> None: "motion-planning", "rl-learning", "learning-layers", + "supervised-learning", "configclass-pattern", "randomization", "atomic-actions", From 77869ee3ba13f3b1d4ae50ddba85ef105e665f49 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 24 Oct 2026 17:15:16 +0000 Subject: [PATCH 07/21] fix(sl): accept argument lists from the unified CLI dispatcher --- embodichain/learning/sl/workflow.py | 11 ++++++++--- tests/learning/sl/test_trainer.py | 11 +++++++++++ 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/embodichain/learning/sl/workflow.py b/embodichain/learning/sl/workflow.py index 467bf1781c..9d6c68275f 100644 --- a/embodichain/learning/sl/workflow.py +++ b/embodichain/learning/sl/workflow.py @@ -19,6 +19,7 @@ from __future__ import annotations import argparse +from collections.abc import Sequence import importlib import json from pathlib import Path @@ -55,12 +56,16 @@ def train_from_config( return trainer.fit() -def cli() -> None: - """Train a supervised recipe from a YAML/JSON file and optional checkpoint.""" +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() + args = parser.parse_args(argv) import yaml config = yaml.safe_load(args.config.read_text()) diff --git a/tests/learning/sl/test_trainer.py b/tests/learning/sl/test_trainer.py index 184f5a1be2..421d777561 100644 --- a/tests/learning/sl/test_trainer.py +++ b/tests/learning/sl/test_trainer.py @@ -113,3 +113,14 @@ def test_source_snapshots_and_sampled_batches_own_storage(): 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 From 736a127a45529229710a08e5d2a1b176353a0d25 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 24 Oct 2026 17:19:01 +0000 Subject: [PATCH 08/21] test(sl): isolate supervised test modules from dynamics tests --- tests/learning/sl/__init__.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) create mode 100644 tests/learning/sl/__init__.py 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 From 1c1441f47505146a263ada51e34b4e1d59b06160 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 24 Oct 2026 17:24:17 +0000 Subject: [PATCH 09/21] feat(sl): integrate NeRD fitting and reference acceptance --- .../topics/neural-dynamics/neural-dynamics.md | 11 + .../embodichain.learning.dynamics.sl.rst | 14 + embodichain/learning/dynamics/sl_recipe.py | 222 ++++++++ embodichain/learning/dynamics/supervised.py | 102 ++++ embodichain/learning/dynamics/trainer.py | 52 +- .../tasks/classic_control/ant/sl.train.yaml | 45 ++ .../classic_control/cartpole/sl.train.yaml | 41 ++ scripts/experiments/supervised/nerd.md | 67 +++ .../experiments/supervised/nerd_acceptance.py | 501 ++++++++++++++++++ tests/learning/dynamics/test_trainer.py | 55 ++ 10 files changed, 1078 insertions(+), 32 deletions(-) create mode 100644 docs/source/api_reference/embodichain/embodichain.learning.dynamics.sl.rst create mode 100644 embodichain/learning/dynamics/sl_recipe.py create mode 100644 embodichain/learning/dynamics/supervised.py create mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/ant/sl.train.yaml create mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/sl.train.yaml create mode 100644 scripts/experiments/supervised/nerd.md create mode 100644 scripts/experiments/supervised/nerd_acceptance.py diff --git a/agent_context/topics/neural-dynamics/neural-dynamics.md b/agent_context/topics/neural-dynamics/neural-dynamics.md index 6403f7b44d..8db8d85b98 100644 --- a/agent_context/topics/neural-dynamics/neural-dynamics.md +++ b/agent_context/topics/neural-dynamics/neural-dynamics.md @@ -10,6 +10,8 @@ | Reference assets and physical profiles | `data_pipeline/dynamics/_reference_assets.py`; packaged files in `data/nerd/` | | Committed shards and sequence windows | `data_pipeline/dynamics/shards.py` | | Training, resume, evaluation and export | `learning/dynamics/trainer.py`, `workflow.py` | +| Shared SL objective and encoded batches | `learning/dynamics/supervised.py` | +| Common offline trainer recipe | `learning/dynamics/sl_recipe.py`; `embodichain train-sl` | | Original-repository comparisons | `learning/dynamics/reference.py`, `comparison.py`, `report.py` | | Command dispatch | `cli/dynamics.py`: `embodichain dynamics` | @@ -52,6 +54,15 @@ progress. Portable exports contain tensors and JSON, including normalizers and attention policy. Resume restores the same run; fine-tuning starts a new one. Only explicitly trusted legacy conversion may deserialize Python model objects. +`DynamicsTrainer` delegates optimizer updates to `learning.sl.SupervisedLearner` +while retaining its historical sampling, numerical profiles, LR schedule and +training-checkpoint format. The standalone `sl_recipe.build_supervised_trainer` +instead composes the common trainer and uses SL checkpoints. Both reuse the +same codec and normalized physical-target objective. Models and datasets do not +depend on either trainer. The paired acceptance runner under +`scripts/experiments/supervised/nerd_acceptance.py` compares independent reference +and candidate processes; fixed-budget parity is separate from full convergence. + Optional shared sample-order/statistics artifacts support controlled original trainer comparisons; equal seeds alone do not prove identical batches or physics. Original Warp references run in isolated supported environments. Keep virtualenv diff --git a/docs/source/api_reference/embodichain/embodichain.learning.dynamics.sl.rst b/docs/source/api_reference/embodichain/embodichain.learning.dynamics.sl.rst new file mode 100644 index 0000000000..e619e89692 --- /dev/null +++ b/docs/source/api_reference/embodichain/embodichain.learning.dynamics.sl.rst @@ -0,0 +1,14 @@ +NeRD supervised-learning adapters +================================ + +NeRD uses the shared supervised optimizer update while retaining its physical +codec, normalized loss and numerical execution policy. The standalone recipe +connects immutable NeRD windows to the common trainer; its checkpoints use the +SL format. Historical NeRD training checkpoints remain owned by the existing +``dynamics`` workflow. + +.. automodule:: embodichain.learning.dynamics.supervised + :members: NeRDObjective, make_nerd_batch + +.. automodule:: embodichain.learning.dynamics.sl_recipe + :members: NeRDWindowSource, build_supervised_trainer diff --git a/embodichain/learning/dynamics/sl_recipe.py b/embodichain/learning/dynamics/sl_recipe.py new file mode 100644 index 0000000000..cec2ec19ed --- /dev/null +++ b/embodichain/learning/dynamics/sl_recipe.py @@ -0,0 +1,222 @@ +# ---------------------------------------------------------------------------- +# 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 embodichain.compute.dynamics import NeRDStateCodec +from embodichain.data_pipeline.dynamics.shards import DynamicsWindowDataset +from embodichain.learning.models.nerd 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 .trainer 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 existing ``dynamics`` CLI + continues reading and writing historical 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(), + "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/embodichain/learning/dynamics/supervised.py b/embodichain/learning/dynamics/supervised.py new file mode 100644 index 0000000000..8a4583ad56 --- /dev/null +++ b/embodichain/learning/dynamics/supervised.py @@ -0,0 +1,102 @@ +# ---------------------------------------------------------------------------- +# 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. + """ + # Keep the historical public helper and its arithmetic as the owner; + # importing lazily also permits DynamicsTrainer to use this adapter. + from .trainer 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 .trainer 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/embodichain/learning/dynamics/trainer.py b/embodichain/learning/dynamics/trainer.py index e31c13ec59..87b3bec97e 100644 --- a/embodichain/learning/dynamics/trainer.py +++ b/embodichain/learning/dynamics/trainer.py @@ -32,6 +32,8 @@ from embodichain.utils import configclass from embodichain.data_pipeline.dynamics.shards import DynamicsWindowDataset +from embodichain.learning.sl import SupervisedLearner +from .supervised import NeRDObjective, make_nerd_batch from ._sample_order import _CheckpointedOrder, _OfficialShuffleOrder, _open_order from ._numerics import ( _attention_context, @@ -475,6 +477,19 @@ def __init__( 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) @@ -541,46 +556,18 @@ def step(self) -> dict[str, float | int]: 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) - inputs, target, mask = encode_batch(batch, self.codec, self.cfg.device) - self.model.train() - self.optimizer.zero_grad(set_to_none=True) + supervised_batch = make_nerd_batch(batch, self.codec, self.cfg.device) lr = self._lr() for group in self.optimizer.param_groups: group["lr"] = lr - observed = None - with _attention_context(self.cfg.numerical_profile): - prediction = self.model(inputs) - loss = normalized_mse( - prediction, - target, - _model_variance(self.model), - mask if self.train_dataset.include_prefix else None, - ) - if self.cfg.numerical_profile == "default" and not torch.isfinite(loss): - raise FloatingPointError("Non-finite dynamics training loss") - loss.backward() - norm = _clip_grad_norm_( - self.model.parameters(), - self.cfg.max_grad_norm, - profile=self.cfg.numerical_profile, - error_if_nonfinite=self.cfg.numerical_profile == "default", - ) - if self.cfg.numerical_profile == "reference": - # One device-to-host observation covers loss, total norm and - # finite checks; no parameter norm is read back individually. - observed = torch.stack((loss.detach(), norm)).cpu().tolist() - if not all(math.isfinite(value) for value in observed): - raise FloatingPointError( - "Non-finite dynamics loss or gradient norm" - ) - self.optimizer.step() + metrics = self.learner.update(supervised_batch) self.sampler.set_state(candidate_sampler.get_state()) self.global_step += 1 return { "step": self.global_step, - "loss": float(loss.detach()) if observed is None else observed[0], + "loss": metrics["loss"], "learning_rate": lr, - "grad_norm": float(norm) if observed is None else observed[1], + "grad_norm": metrics["grad_norm"], } def save(self, path: str | Path) -> None: @@ -670,6 +657,7 @@ def resume(self, path: str | Path) -> None: state["best_validation"], ) self.statistics = state["statistics"] + self.learner.num_updates = self.global_step self.sampler.set_state(state["sampler_state"]) _restore_rng_state(state["rng"]) diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/sl.train.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/sl.train.yaml new file mode 100644 index 0000000000..57355107ac --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/sl.train.yaml @@ -0,0 +1,45 @@ +# Standalone SL recipe; run from the repository root after downloading data. +factory: embodichain.learning.dynamics.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 + 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/ant + checkpoint_interval: 500 + validation_interval: 500 +evaluation: + max_batches: 100 diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/sl.train.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/sl.train.yaml new file mode 100644 index 0000000000..8b6b323d88 --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/sl.train.yaml @@ -0,0 +1,41 @@ +# Standalone SL recipe; run from the repository root after downloading data. +factory: embodichain.learning.dynamics.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/scripts/experiments/supervised/nerd.md b/scripts/experiments/supervised/nerd.md new file mode 100644 index 0000000000..16e162965f --- /dev/null +++ b/scripts/experiments/supervised/nerd.md @@ -0,0 +1,67 @@ +# NeRD supervised-learning acceptance + +The production NeRD trainer delegates optimizer updates to +`SupervisedLearner`. Its existing data, codec, numerical execution, learning-rate +schedule, export and checkpoint contracts remain authoritative. The separate +`train-sl` recipe builds the common offline trainer and uses the common SL +checkpoint format. + +Run a bounded standalone SL recipe from the repository root after downloading +the published dataset into `outputs/datasets`: + +```bash +embodichain train-sl --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/sl.train.yaml +embodichain train-sl --config extensions/embodichain_tasks/configs/tasks/classic_control/ant/sl.train.yaml +``` + +These 1,000-update recipes demonstrate the public entry point; they do not claim +full-budget convergence. Configure the output directory explicitly for each run. +The existing `embodichain dynamics` CLI retains its historical training/export +formats. New common-format checkpoints resume through `train-sl` with the same +recipe and data identity. + +## Frozen comparison + +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 python scripts/experiments/supervised/nerd_acceptance.py \ + --reference-source /path/to/pr855-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. + +`summary.json` is the machine-readable gate; the command exits nonzero if a gate +fails or fewer than five paired seeds are supplied. `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/scripts/experiments/supervised/nerd_acceptance.py b/scripts/experiments/supervised/nerd_acceptance.py new file mode 100644 index 0000000000..0decba9c17 --- /dev/null +++ b/scripts/experiments/supervised/nerd_acceptance.py @@ -0,0 +1,501 @@ +# ---------------------------------------------------------------------------- +# 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 scripts/experiments/supervised/nerd_acceptance.py \ + --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 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_identity(source: Path) -> dict[str, Any]: + def git(*args: str) -> bytes: + return subprocess.check_output(["git", "-C", str(source), *args]) + + 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") + sys.path.insert(0, request["source"]) + import torch + + from embodichain.compute.dynamics import NeRDStateCodec + from embodichain.data_pipeline.dynamics.shards import ( + DynamicsWindowDataset, + dataset_fingerprint, + load_manifest, + ) + from embodichain.learning.dynamics.trainer import ( + DynamicsTrainCfg, + DynamicsTrainer, + evaluate_model, + ) + from embodichain.learning.models.nerd import NeRDModel, 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") + 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() + 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"]) + 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=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 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[3] + ) + 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", 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 + + sys.path.insert(0, str(args.candidate_source.resolve())) + from embodichain.data_pipeline.dynamics.shards import load_manifest + from embodichain.learning.dynamics.qualification import ( + compare_training_steps, + paired_noninferiority, + ) + + seeds = [int(value) for value in args.seeds.split(",")] + 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/tests/learning/dynamics/test_trainer.py b/tests/learning/dynamics/test_trainer.py index 195fb75f18..e0def16c6c 100644 --- a/tests/learning/dynamics/test_trainer.py +++ b/tests/learning/dynamics/test_trainer.py @@ -148,6 +148,61 @@ def test_statistics_only_training_windows(tmp_path): ) +@pytest.mark.parametrize("numerical_profile", ["default", "reference"]) +def test_shared_recipe_matches_legacy_updates_and_restores_sampler( + tmp_path, numerical_profile +): + """The new CLI recipe retains the established NeRD schedule and updates.""" + from embodichain.learning.dynamics.sl_recipe import build_supervised_trainer + + make_data(tmp_path / "data") + legacy = build( + tmp_path / "data", + tmp_path / "legacy", + steps=4, + numerical_profile=numerical_profile, + ) + config = { + "model": architecture().to_dict(), + "data": {"path": str(tmp_path / "data")}, + "objective": {"numerical_profile": numerical_profile}, + "optimizer": { + "learning_rate": 0.005, + "final_learning_rate": 0.002, + "max_grad_norm": 1.0, + }, + "trainer": { + "total_updates": 4, + "batch_size": 8, + "seed": 9, + "output_dir": str(tmp_path / "shared"), + }, + } + shared = build_supervised_trainer(config) + shared.learner.model.load_state_dict(legacy.model.state_dict()) + for _ in range(4): + legacy.step() + shared.fit(stop_after=2) + restored = build_supervised_trainer(config) + restored.resume(tmp_path / "shared/last.pt") + restored.fit() + for name, value in legacy.model.state_dict().items(): + torch.testing.assert_close( + restored.learner.model.state_dict()[name], value, atol=0, rtol=0 + ) + assert torch.equal(restored.sampler.get_state(), legacy.sampler.get_state()) + from embodichain.learning.sl.workflow import train_from_config + + report = train_from_config( + { + "factory": "embodichain.learning.dynamics.sl_recipe:build_supervised_trainer", + **config, + }, + resume=str(tmp_path / "shared/last.pt"), + ) + assert report["num_updates"] == legacy.global_step + + @pytest.mark.parametrize("numerical_profile", ["default", "reference"]) def test_exact_resume_with_dropout_optimizer_sampler_rng(tmp_path, numerical_profile): make_data(tmp_path / "data") From 1395871e8627143bd9ed7ce197743b3cea48a13a Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 24 Oct 2026 17:28:43 +0000 Subject: [PATCH 10/21] fix(sl): isolate metadata and validate optimizer recovery before mutation --- embodichain/learning/sl/learner.py | 32 ++++++++++++++++++++++++++ embodichain/learning/sl/sources.py | 2 +- tests/learning/sl/test_trainer.py | 36 ++++++++++++++++++++++++++++++ 3 files changed, 69 insertions(+), 1 deletion(-) diff --git a/embodichain/learning/sl/learner.py b/embodichain/learning/sl/learner.py index 1ade37c9b5..33af0e32a9 100644 --- a/embodichain/learning/sl/learner.py +++ b/embodichain/learning/sl/learner.py @@ -133,6 +133,20 @@ def update(self, batch: SupervisedBatch) -> dict[str, float | int]: ) 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. @@ -142,6 +156,8 @@ def state_dict(self) -> dict[str, Any]: 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, } @@ -167,6 +183,22 @@ def load_state_dict(self, state: Mapping[str, Any]) -> None: 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 index bf4e7f9443..52fa06908e 100644 --- a/embodichain/learning/sl/sources.py +++ b/embodichain/learning/sl/sources.py @@ -112,7 +112,7 @@ def _gather(self, indices: torch.Tensor) -> SupervisedBatch: _map_tensors(self.batch.targets, gather), _map_tensors(self.batch.masks, gather), { - **self.batch.metadata, + **copy.deepcopy(self.batch.metadata), "source_id": self.source_id, "indices": indices.clone(), }, diff --git a/tests/learning/sl/test_trainer.py b/tests/learning/sl/test_trainer.py index 421d777561..d3b54da43f 100644 --- a/tests/learning/sl/test_trainer.py +++ b/tests/learning/sl/test_trainer.py @@ -124,3 +124,39 @@ def test_unified_cli_routes_sl_help(): ) 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 + ) From 8ed70f107a66a5de51f4848a906404a85a59c122 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 24 Oct 2026 18:50:38 +0000 Subject: [PATCH 11/21] docs(sl): link the standalone NeRD training integration --- .../api_reference/embodichain/embodichain.learning.sl.rst | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/docs/source/api_reference/embodichain/embodichain.learning.sl.rst b/docs/source/api_reference/embodichain/embodichain.learning.sl.rst index f016e712bb..490c9afd0b 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.sl.rst +++ b/docs/source/api_reference/embodichain/embodichain.learning.sl.rst @@ -123,3 +123,11 @@ embodichain.learning.sl.workflow .. autofunction:: train_from_config .. autofunction:: cli + +NeRD integration +---------------- + +.. toctree:: + :maxdepth: 1 + + embodichain.learning.dynamics.sl From 6e5a861c0864606155b4596b468cf62a952e126c Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 24 Oct 2026 19:04:03 +0000 Subject: [PATCH 12/21] fix(sl): bind scheduler and tensor-tree semantics during recovery --- embodichain/learning/sl/sources.py | 22 +++++++++- embodichain/learning/sl/trainer.py | 12 +++++ tests/learning/sl/test_trainer.py | 70 ++++++++++++++++++++++++++++++ 3 files changed, 102 insertions(+), 2 deletions(-) diff --git a/embodichain/learning/sl/sources.py b/embodichain/learning/sl/sources.py index 52fa06908e..5bf8c2b417 100644 --- a/embodichain/learning/sl/sources.py +++ b/embodichain/learning/sl/sources.py @@ -81,13 +81,31 @@ def snapshot(tensor: torch.Tensor) -> torch.Tensor: 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 {str(key): structure(item) for key, item in value.items()} + return { + "kind": "mapping", + "items": [ + [key_signature(key), structure(item)] + for key, item in value.items() + ], + } if isinstance(value, (tuple, list)): - return [structure(item) for item in value] + return { + "kind": "tuple" if isinstance(value, tuple) else "list", + "items": [structure(item) for item in value], + } return value digest.update( diff --git a/embodichain/learning/sl/trainer.py b/embodichain/learning/sl/trainer.py index d12eabfc43..b834ee5748 100644 --- a/embodichain/learning/sl/trainer.py +++ b/embodichain/learning/sl/trainer.py @@ -197,6 +197,11 @@ def save(self, path: str | Path) -> None: "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() ), @@ -224,6 +229,13 @@ def resume(self, path: str | Path) -> 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") self.learner.load_state_dict(state["learner"]) if self.scheduler is not None: self.scheduler.load_state_dict(state["scheduler"]) diff --git a/tests/learning/sl/test_trainer.py b/tests/learning/sl/test_trainer.py index d3b54da43f..77d1eb635a 100644 --- a/tests/learning/sl/test_trainer.py +++ b/tests/learning/sl/test_trainer.py @@ -19,6 +19,8 @@ from __future__ import annotations import copy +import random +import numpy as np import json import sys import subprocess @@ -160,3 +162,71 @@ def test_optimizer_incompatibility_rejected_before_model_mutation(): 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) + + reset_rng() + baseline = build(tmp_path / "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") + candidate.cfg.validation_interval = 1 + + def evaluate(learner): + random.random() + np.random.rand(5) + torch.rand(9) + 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) From d591e78ba0a966b85eef0a0c5219628372e72762 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 25 Oct 2026 00:40:14 +0000 Subject: [PATCH 13/21] fix(sl): reject duplicate NeRD acceptance seeds before startup --- scripts/experiments/supervised/nerd.md | 4 +- .../experiments/supervised/nerd_acceptance.py | 16 +++- .../supervised/test_nerd_acceptance.py | 75 +++++++++++++++++++ 3 files changed, 92 insertions(+), 3 deletions(-) create mode 100644 tests/scripts/experiments/supervised/test_nerd_acceptance.py diff --git a/scripts/experiments/supervised/nerd.md b/scripts/experiments/supervised/nerd.md index 16e162965f..e34afeeb83 100644 --- a/scripts/experiments/supervised/nerd.md +++ b/scripts/experiments/supervised/nerd.md @@ -55,7 +55,9 @@ The default matrix performs: compared with uninterrupted weights, optimizer, sampler, counters and RNG. `summary.json` is the machine-readable gate; the command exits nonzero if a gate -fails or fewer than five paired seeds are supplied. `report.md` contains time, +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. diff --git a/scripts/experiments/supervised/nerd_acceptance.py b/scripts/experiments/supervised/nerd_acceptance.py index 0decba9c17..130a0ea79b 100644 --- a/scripts/experiments/supervised/nerd_acceptance.py +++ b/scripts/experiments/supervised/nerd_acceptance.py @@ -275,6 +275,18 @@ def _state_equal(first: Any, second: Any) -> bool: 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__) @@ -289,7 +301,7 @@ def main() -> None: 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", default="0,1,2,3,4") + 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) @@ -317,7 +329,7 @@ def main() -> None: paired_noninferiority, ) - seeds = [int(value) for value in args.seeds.split(",")] + seeds = args.seeds sources = { "reference": str(args.reference_source.resolve()), "candidate": str(args.candidate_source.resolve()), diff --git a/tests/scripts/experiments/supervised/test_nerd_acceptance.py b/tests/scripts/experiments/supervised/test_nerd_acceptance.py new file mode 100644 index 0000000000..ab17351b44 --- /dev/null +++ b/tests/scripts/experiments/supervised/test_nerd_acceptance.py @@ -0,0 +1,75 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""CPU-only CLI preflight for independent-seed acceptance evidence.""" + +from __future__ import annotations + +from pathlib import Path +import subprocess +import sys + +import pytest + +from scripts.experiments.supervised import nerd_acceptance + + +@pytest.mark.parametrize( + ("seeds", "message"), + [ + ("0,0,0,0,0", "seeds must be distinct"), + ("0,1,2,3,00", "seeds must be distinct"), + ("0,1,2,3,invalid", "seeds must be comma-separated integers"), + ("", "seeds must be comma-separated integers"), + ], +) +def test_invalid_seed_list_fails_before_output_or_workers( + tmp_path: Path, seeds: str, message: str +) -> None: + runner = ( + Path(__file__).resolve().parents[4] + / "scripts/experiments/supervised/nerd_acceptance.py" + ) + output = tmp_path / "results" + # Missing artifacts must never be inspected or passed to a worker. + result = subprocess.run( + [ + sys.executable, + str(runner), + "--reference-source", + str(tmp_path / "missing-reference"), + "--datasets", + str(tmp_path / "missing-datasets"), + "--models", + str(tmp_path / "missing-models"), + "--output", + str(output), + "--seeds", + seeds, + ], + capture_output=True, + text=True, + timeout=10, + ) + assert result.returncode == 2 + assert message in result.stderr + assert not output.exists() + + +@pytest.mark.parametrize("seeds", [[0], [0, 1, 2, 3], [0, 1, 2, 3, 4]]) +def test_unique_seed_lists_preserve_short_and_full_runs(seeds: list[int]) -> None: + # Short runs remain available; the existing qualification gate returns pending. + assert nerd_acceptance._parse_seeds(",".join(map(str, seeds))) == seeds From a7668adf28512ca38a66c04df5bb3a7c9ab97273 Mon Sep 17 00:00:00 2001 From: Yueci Deng Date: Sun, 25 Oct 2026 15:17:07 +0800 Subject: [PATCH 14/21] Add shared SL updates, offline training and recovery (#857) --- agent_context/MAP.yaml | 12 + .../topics/learning-layers/learning-layers.md | 3 + .../supervised-learning.md | 36 +++ docs/architecture/curated.json | 43 ++- .../embodichain/embodichain.learning.rst | 15 +- .../embodichain/embodichain.learning.sl.rst | 125 ++++++++ embodichain/cli/main.py | 5 + embodichain/learning/__init__.py | 2 +- embodichain/learning/sl/__init__.py | 35 +++ embodichain/learning/sl/learner.py | 204 ++++++++++++ embodichain/learning/sl/sources.py | 169 ++++++++++ embodichain/learning/sl/trainer.py | 282 +++++++++++++++++ embodichain/learning/sl/types.py | 109 +++++++ embodichain/learning/sl/workflow.py | 78 +++++ tests/learning/sl/__init__.py | 19 ++ tests/learning/sl/test_learner.py | 107 +++++++ tests/learning/sl/test_trainer.py | 292 ++++++++++++++++++ tests/test_agent_context_map.py | 1 + 18 files changed, 1529 insertions(+), 8 deletions(-) create mode 100644 agent_context/topics/supervised-learning/supervised-learning.md create mode 100644 docs/source/api_reference/embodichain/embodichain.learning.sl.rst create mode 100644 embodichain/learning/sl/__init__.py create mode 100644 embodichain/learning/sl/learner.py create mode 100644 embodichain/learning/sl/sources.py create mode 100644 embodichain/learning/sl/trainer.py create mode 100644 embodichain/learning/sl/types.py create mode 100644 embodichain/learning/sl/workflow.py create mode 100644 tests/learning/sl/__init__.py create mode 100644 tests/learning/sl/test_learner.py create mode 100644 tests/learning/sl/test_trainer.py diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 2f3a5ff947..64966ac3d6 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -542,3 +542,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] + status: active diff --git a/agent_context/topics/learning-layers/learning-layers.md b/agent_context/topics/learning-layers/learning-layers.md index 7e609fb2f0..c0713cf071 100644 --- a/agent_context/topics/learning-layers/learning-layers.md +++ b/agent_context/topics/learning-layers/learning-layers.md @@ -24,6 +24,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/supervised-learning/supervised-learning.md b/agent_context/topics/supervised-learning/supervised-learning.md new file mode 100644 index 0000000000..3efd4cee39 --- /dev/null +++ b/agent_context/topics/supervised-learning/supervised-learning.md @@ -0,0 +1,36 @@ +# Supervised learning (SL) + +`embodichain.learning.sl` owns common supervised optimizer updates and offline +training orchestration. Models remain in `learning.models`; 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`. + +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. +Fresh runs require artifact-free output directories. Restoring an older +checkpoint into a directory with later metric history requires a fresh output +directory so prior results are retained. +Validation callbacks use frozen models and restored RNG; 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 0d05ee8810..f5652ba5d1 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…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\"\"\"" + "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\"\"\"" } ], "documentation": [ @@ -3368,6 +3368,35 @@ "label": "NeRD data and training workflow" } ] + }, + { + "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. Model-specific codecs, normalization and sequence unrolling remain with objectives and recipe factories. 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": [ @@ -5804,7 +5833,8 @@ "urdf-assembly", "convex-decomposition", "scene-services", - "neural-dynamics" + "neural-dynamics", + "supervised-learning" ], "edge_ids": [ "adapter-calls-compiler", @@ -6020,7 +6050,8 @@ "differentiable-collector", "observation-normalizer", "policy-evaluation", - "neural-dynamics" + "neural-dynamics", + "supervised-learning" ] }, { @@ -6360,7 +6391,8 @@ "rollout-buffer", "differentiable-trainer", "differentiable-collector", - "observation-normalizer" + "observation-normalizer", + "supervised-learning" ], "edge_ids": [ "cli-calls-training-assembly", @@ -6427,7 +6459,8 @@ "cli", "rl-routing", "learning-env", - "policy-evaluation" + "policy-evaluation", + "supervised-learning" ] }, { diff --git a/docs/source/api_reference/embodichain/embodichain.learning.rst b/docs/source/api_reference/embodichain/embodichain.learning.rst index f8c60d7ed6..1022a61fbe 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 dynamics Reinforcement Learning @@ -65,3 +65,14 @@ standalone IK layers do not import RL, simulation or TensorDict. :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..f016e712bb --- /dev/null +++ b/docs/source/api_reference/embodichain/embodichain.learning.sl.rst @@ -0,0 +1,125 @@ +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 diff --git a/embodichain/cli/main.py b/embodichain/cli/main.py index 315ec493a1..7c4522a3e6 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 b83c68b8a3..3dba1a64db 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", "dynamics"] +__all__ = ["layers", "models", "rl", "dynamics", "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/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..c7c1530cff --- /dev/null +++ b/embodichain/learning/sl/trainer.py @@ -0,0 +1,282 @@ +# ---------------------------------------------------------------------------- +# 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 .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.learner.model.to(self.cfg.device) + self._owned_artifacts: dict[Path, tuple[int, int, int, int]] = {} + self._fit_started = False + self._restored = False + + 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") + if not self._fit_started and not self._restored: + artifacts = [ + self.output / "metrics.jsonl", + self.output / "last.pt", + *self.output.glob("step_*.pt"), + ] + for path in artifacts: + if path.exists() and self._owned_artifacts.get( + path.resolve() + ) != self._artifact_stamp(path): + raise FileExistsError( + "SL output contains another run; restore a checkpoint or use a fresh output directory" + ) + self._fit_started = True + 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 + ): + rng = self._rng_state() + was_training = self.learner.model.training + try: + self.learner.model.eval() + with torch.no_grad(): + metrics.update( + { + f"validation/{key}": float(value) + for key, value in self.evaluate(self.learner).items() + } + ) + finally: + self._restore_rng(rng) + self.learner.model.train(was_training) + 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 _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) + + @staticmethod + def _artifact_stamp(path: Path) -> tuple[int, int, int, int]: + stat = path.stat() + return stat.st_dev, stat.st_ino, stat.st_size, stat.st_mtime_ns + + 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) + path.parent.mkdir(parents=True, exist_ok=True) + state = { + "version": 1, + "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() + ), + } + temporary = path.with_suffix(path.suffix + ".tmp") + torch.save(state, temporary) + temporary.replace(path) + self._owned_artifacts[path.resolve()] = self._artifact_stamp(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") + metrics_path = self.output / "metrics.jsonl" + if metrics_path.exists(): + last_line = "" + with metrics_path.open() as handle: + for line in handle: + if line.strip(): + last_line = line + if ( + last_line + and json.loads(last_line)["num_updates"] + > state["learner"]["num_updates"] + ): + raise ValueError( + "SL output history is ahead of this checkpoint; use a fresh output directory" + ) + 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._restored = True 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/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..b5c881e81e --- /dev/null +++ b/tests/learning/sl/test_trainer.py @@ -0,0 +1,292 @@ +# ---------------------------------------------------------------------------- +# 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) + + reset_rng() + baseline = build(tmp_path / "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") + candidate.cfg.validation_interval = 1 + + def evaluate(learner): + random.random() + np.random.rand(5) + torch.rand(9) + 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) + + +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 + + +def test_older_checkpoint_cannot_append_to_a_later_run_history(tmp_path): + first = build(tmp_path / "run") + first.fit(stop_after=2) + first.save(tmp_path / "step2.pt") + first.fit(stop_after=3) + candidate = build(tmp_path / "run") + before = copy.deepcopy(candidate.learner.model.state_dict()) + with pytest.raises(ValueError, match="history is ahead"): + candidate.resume(tmp_path / "step2.pt") + 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_replacement_of_owned_initial_snapshot_is_rejected(tmp_path): + first = build(tmp_path / "shared") + second = build(tmp_path / "shared") + path = first.output / "step_000000.pt" + first.save(path) + second.save(path) + with pytest.raises(FileExistsError, match="another run"): + first.fit(stop_after=1) diff --git a/tests/test_agent_context_map.py b/tests/test_agent_context_map.py index 1b9ac04d67..4e1472bae5 100644 --- a/tests/test_agent_context_map.py +++ b/tests/test_agent_context_map.py @@ -62,6 +62,7 @@ def test_map_registers_the_supported_context_domains() -> None: "motion-planning", "rl-learning", "learning-layers", + "supervised-learning", "neural-dynamics", "configclass-pattern", "randomization", From 6dc47c3b590741c55d406da4eb7c759a363cfe88 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 25 Oct 2026 09:41:24 +0000 Subject: [PATCH 15/21] fix(sl): preserve run identity, best models and NeRD recovery --- .../topics/neural-dynamics/neural-dynamics.md | 17 +- .../supervised-learning.md | 13 +- docs/source/features/neural_dynamics.md | 28 ++ .../data_pipeline/dynamics/generation.py | 30 +- embodichain/data_pipeline/dynamics/shards.py | 20 +- embodichain/learning/dynamics/sl_recipe.py | 1 + embodichain/learning/dynamics/trainer.py | 216 ++++++++++-- embodichain/learning/dynamics/workflow.py | 45 ++- .../learning/models/nerd/checkpoint.py | 20 +- embodichain/learning/sl/_artifacts.py | 147 ++++++++ embodichain/learning/sl/trainer.py | 101 +++--- .../ant/dynamics.train.full.yaml | 3 +- .../tasks/classic_control/ant/sl.train.yaml | 3 +- scripts/experiments/supervised/nerd.md | 35 ++ .../experiments/supervised/nerd_acceptance.py | 17 + tests/cli/test_dynamics.py | 169 +++++++++ .../data_pipeline/dynamics/test_generation.py | 128 +++++++ tests/data_pipeline/dynamics/test_shards.py | 26 ++ tests/learning/dynamics/test_trainer.py | 325 +++++++++++++++++- tests/learning/models/test_nerd.py | 9 +- tests/learning/sl/test_trainer.py | 93 ++++- 21 files changed, 1321 insertions(+), 125 deletions(-) create mode 100644 embodichain/learning/sl/_artifacts.py diff --git a/agent_context/topics/neural-dynamics/neural-dynamics.md b/agent_context/topics/neural-dynamics/neural-dynamics.md index 8db8d85b98..e220799479 100644 --- a/agent_context/topics/neural-dynamics/neural-dynamics.md +++ b/agent_context/topics/neural-dynamics/neural-dynamics.md @@ -24,8 +24,11 @@ functional Gym integration are separate consumers, not dependencies of training. A trajectory contains T+1 states and T controls/contact inputs. Splits belong to whole trajectories and windows cannot cross episodes or splits. Writers publish -only committed shards; resume advances candidate seeds past accepted and rejected -candidates. File handles and bounded caches belong to their reader/worker. +only committed shards. The manifest commits generation cursor and rejection +evidence with the shard; rejection-only commits preserve data fingerprints. +Legacy progress sidecars remain readable, while collector source/physical +identity checks still apply. File handles and bounded caches belong to their +reader/worker. The direct teacher stores its importer profile, asset/physical-array hashes, solver settings and actual Newton/Warp versions. Its @@ -49,9 +52,13 @@ Mass-matrix update intervals divide the macro substep count so resets can replay ## Training and validation -Training checkpoints preserve optimizer, scheduler, RNG and committed sample -progress. Portable exports contain tensors and JSON, including normalizers and -attention policy. Resume restores the same run; fine-tuning starts a new one. +`last.pt` preserves optimizer, scheduler, RNG, committed sample progress and one +owned CPU snapshot of the selected best model. New `best.pt` files are model-only +export snapshots; exact training recovery uses `last.pt`. Ordered input features +are part of recovery identity. Both trainers share run-directory/history admission +with [SL](../supervised-learning/supervised-learning.md). +Portable exports preserve saved floating precision, normalizers and attention +policy. Resume restores the same run; fine-tuning starts a new one. Only explicitly trusted legacy conversion may deserialize Python model objects. `DynamicsTrainer` delegates optimizer updates to `learning.sl.SupervisedLearner` diff --git a/agent_context/topics/supervised-learning/supervised-learning.md b/agent_context/topics/supervised-learning/supervised-learning.md index 3efd4cee39..098ce93ae8 100644 --- a/agent_context/topics/supervised-learning/supervised-learning.md +++ b/agent_context/topics/supervised-learning/supervised-learning.md @@ -22,11 +22,14 @@ 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. -Fresh runs require artifact-free output directories. Restoring an older -checkpoint into a directory with later metric history requires a fresh output -directory so prior results are retained. -Validation callbacks use frozen models and restored RNG; only development data -may select checkpoints. Acquisition and producer lifecycles remain application-owned. +`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 diff --git a/docs/source/features/neural_dynamics.md b/docs/source/features/neural_dynamics.md index c5b1c79d1f..bdd45421b4 100644 --- a/docs/source/features/neural_dynamics.md +++ b/docs/source/features/neural_dynamics.md @@ -130,6 +130,11 @@ Smoke presets use a small network and short budget to exercise the workflow. Full presets use the reference network scale and a substantially longer budget. Neither preset name is an accuracy qualification. +Collection commits its candidate cursor and rejection records in the same +manifest update as the shard. Rejection-only progress does not change training +data identity. Historical progress-file layouts remain readable; continuing a +collection still requires the recorded collector source and physical identity. + To train on a downloaded dataset, copy `dynamics.train.full.yaml`, then set `dataset` to the downloaded robot's `v1` directory and `output_dir` to a new training directory. Preserve its model feature dimensions and history length. @@ -139,10 +144,33 @@ then export its validation-selected `best.pt` with the `export` command. Training checkpoints retain model, optimizer, normalization, RNG and committed sampling progress. Continue the same run with `train --resume PATH/last.pt`. Initialize a new run from an inference export with `train --fine-tune PATH`. +Each run records its identity in `run.json`; fresh training and fine-tuning require +an artifact-free destination. Resuming into an occupied directory requires the +same run identity. If a destination has a different selected best model or newer +or incomplete metric history, use a fresh directory to preserve that evidence. +Historical checkpoints without a +run identity can also recover into a fresh directory. + +New `best.pt` files contain the validation-selected model and its own step, and +support export rather than exact optimizer recovery. `last.pt` retains the best +weights so recovery into another directory can recreate `best.pt` even without +later validation improvement. Export and reload preserve the saved floating +dtype and normalization. Input feature order must remain unchanged when resuming. Export destinations must be new directories. Paths written in configuration files resolve relative to the file; explicit CLI path overrides resolve relative to the working directory. +The same NeRD model/objective also uses the shared supervised trainer: + +```bash +embodichain train-sl --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/sl.train.yaml +embodichain train-sl --config extensions/embodichain_tasks/configs/tasks/classic_control/ant/sl.train.yaml +``` + +These bounded recipes use the common SL checkpoint format and resume through +`train-sl --resume PATH`. The historical `dynamics` CLI keeps its own checkpoint +format; both use the same supervised optimizer update and physical-target loss. + ## Use the exported model in PyTorch Load tensor weights and their recorded normalization buffers: diff --git a/embodichain/data_pipeline/dynamics/generation.py b/embodichain/data_pipeline/dynamics/generation.py index 49e2f8b1b3..954b07baec 100644 --- a/embodichain/data_pipeline/dynamics/generation.py +++ b/embodichain/data_pipeline/dynamics/generation.py @@ -193,12 +193,6 @@ def generate_dataset(config: Mapping[str, Any]) -> dict[str, Any]: cfg.validate() output = Path(cfg.output_dir) progress_path = output / "generation-progress.json" - progress = ( - json.loads(progress_path.read_text()) - if cfg.resume and progress_path.exists() - else {"next_candidate": 0, "rejected": 0, "rejections": []} - ) - progress.setdefault("rejections", []) started = time.perf_counter() with _NewtonTeacher( cfg.robot, @@ -225,6 +219,17 @@ def generate_dataset(config: Mapping[str, Any]) -> dict[str, Any]: 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( @@ -235,6 +240,12 @@ def generate_dataset(config: Mapping[str, Any]) -> dict[str, Any]: 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 @@ -292,6 +303,7 @@ def generate_dataset(config: Mapping[str, Any]) -> dict[str, Any]: 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()) @@ -315,6 +327,7 @@ def generate_dataset(config: Mapping[str, Any]) -> dict[str, Any]: } ) continue + progress["accepted"] = accepted + 1 writer.add_trajectory( f"{cfg.seed}:{candidate}", {k: v[row] for k, v in states.items()}, @@ -323,8 +336,9 @@ def generate_dataset(config: Mapping[str, Any]) -> dict[str, Any]: ) accepted += 1 writer.commit() - progress.update(next_candidate=cursor, accepted=accepted) - _write_json(progress_path, progress) + # 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()), diff --git a/embodichain/data_pipeline/dynamics/shards.py b/embodichain/data_pipeline/dynamics/shards.py index 0617787841..713953ed26 100644 --- a/embodichain/data_pipeline/dynamics/shards.py +++ b/embodichain/data_pipeline/dynamics/shards.py @@ -97,8 +97,15 @@ def load_manifest(root: str | Path, *, verify: bool = False) -> dict[str, Any]: def dataset_fingerprint(manifest: dict[str, Any]) -> str: - """Hash metadata, schema, immutable shard hashes and split assignments.""" - return hashlib.sha256(_json_bytes(manifest)).hexdigest() + """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: @@ -239,8 +246,15 @@ def add_trajectory( self.commit() def commit(self) -> None: - """Publish the current shard atomically, including a content checksum.""" + """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() diff --git a/embodichain/learning/dynamics/sl_recipe.py b/embodichain/learning/dynamics/sl_recipe.py index cec2ec19ed..7ca60f6ce8 100644 --- a/embodichain/learning/dynamics/sl_recipe.py +++ b/embodichain/learning/dynamics/sl_recipe.py @@ -204,6 +204,7 @@ def evaluate(current: SupervisedLearner) -> dict[str, float]: cfg, signature={ "model": model.cfg.to_dict(), + "model_input_layout": list(model.cfg.input_dims.items()), "objective": { "include_prefix": include_prefix, "numerical_profile": profile, diff --git a/embodichain/learning/dynamics/trainer.py b/embodichain/learning/dynamics/trainer.py index 87b3bec97e..8848ea0c69 100644 --- a/embodichain/learning/dynamics/trainer.py +++ b/embodichain/learning/dynamics/trainer.py @@ -33,6 +33,7 @@ from embodichain.utils import configclass from embodichain.data_pipeline.dynamics.shards import DynamicsWindowDataset from embodichain.learning.sl import SupervisedLearner +from embodichain.learning.sl._artifacts import _RunArtifacts from .supervised import NeRDObjective, make_nerd_batch from ._sample_order import _CheckpointedOrder, _OfficialShuffleOrder, _open_order from ._numerics import ( @@ -373,6 +374,14 @@ def evaluate_model( } +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. @@ -434,9 +443,12 @@ def __init__( 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) @@ -499,8 +511,17 @@ def __init__( 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 + return None, None + identity = None if isinstance(source, (str, Path)): path = Path(source).resolve() digest = hashlib.sha256() @@ -512,7 +533,7 @@ def _prepare_order( order, (_CheckpointedOrder, _OfficialShuffleOrder) ) and order.dataset_size != len(self.train_dataset): raise ValueError("Sample-order dataset size mismatch") - self._sample_order_identity = { + identity = { "path": str(path), "sha256": digest.hexdigest(), } @@ -520,7 +541,7 @@ def _prepare_order( 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 + return order, identity def _order_snapshot(self) -> Any: return ( @@ -542,6 +563,7 @@ 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 = ( @@ -572,8 +594,11 @@ def step(self) -> dict[str, float | int]: 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(), @@ -585,6 +610,8 @@ def save(self, path: str | Path) -> None: "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, @@ -599,9 +626,14 @@ def save(self, path: str | Path) -> None: 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" @@ -614,7 +646,7 @@ def resume(self, path: str | Path) -> None: for key, expected in ( ("dataset_fingerprint", self.train_dataset.fingerprint), ("train_cfg", self.cfg.to_dict()), - ("model_cfg", self.model.cfg.to_dict()), + ("model_cfg", _model_cfg_identity(self.model.cfg.to_dict())), ("manifest", self.manifest), ( "window_contract", @@ -627,12 +659,20 @@ def resume(self, path: str | Path) -> None: actual = ( saved_train_cfg if key == "train_cfg" - else saved_model_cfg if key == "model_cfg" else state[key] + 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 = ( @@ -645,11 +685,33 @@ def resume(self, path: str | Path) -> None: if not same: raise ValueError("Resume sample order mismatch") if isinstance(saved_order, dict): - self.sample_order = self._prepare_order(saved_order["path"]) - if self._sample_order_identity != saved_order: + order, order_identity = self._load_order(saved_order["path"]) + if order_identity != saved_order: raise ValueError("Resume sample order checksum mismatch") else: - self.sample_order = saved_order + 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 = ( @@ -657,9 +719,125 @@ def resume(self, path: str | Path) -> None: 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.""" @@ -670,20 +848,11 @@ def train(self, *, stop_after: int | None = None) -> dict[str, Any]: ) 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" - # Drop records newer than a restored checkpoint before appending. - if log_path.exists(): - temporary_log = log_path.with_suffix(".jsonl.pending") - with log_path.open() as source, temporary_log.open("w") as target: - for line in source: - try: - record = json.loads(line) - except json.JSONDecodeError: - break # Interrupted trailing write is not a committed record. - if record["step"] <= self.global_step: - target.write(line) - os.replace(temporary_log, log_path) with log_path.open("a") as log: while self.global_step < end: last = self.step() @@ -703,7 +872,12 @@ def train(self, *, stop_after: int | None = None) -> dict[str, Any]: score = float(validation["normalized_mse"]) if score < self.best_validation: self.best_validation = score - self.save(self.output_dir / "best.pt") + 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 diff --git a/embodichain/learning/dynamics/workflow.py b/embodichain/learning/dynamics/workflow.py index 4e152fe4be..817e500c79 100644 --- a/embodichain/learning/dynamics/workflow.py +++ b/embodichain/learning/dynamics/workflow.py @@ -56,6 +56,19 @@ def _arguments(value: dict[str, Any], allowed: set[str]) -> None: 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. @@ -117,6 +130,15 @@ def run_train(config: dict[str, Any]) -> dict[str, Any]: "dt" ) != manifest.get("dt"): raise ValueError("Fine-tuning requires compatible codec and timestep") + elif config.get("resume"): + from embodichain.learning.models.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( @@ -225,16 +247,23 @@ def run_eval(config: dict[str, Any]) -> dict[str, Any]: def run_export(checkpoint: str | Path, output: str | Path) -> dict[str, Any]: """Export a training checkpoint using the production model artifact writer.""" - import torch - from embodichain.learning.models.nerd import NeRDModel, NeRDModelCfg, export_model + from embodichain.learning.models.nerd import NeRDModelCfg, export_model + from embodichain.learning.models.nerd.checkpoint import _model_from_state from ._numerics import _checkpoint_numerical_profile - state = torch.load(checkpoint, map_location="cpu", weights_only=True) - if state.get("format") != "embodichain.nerd.training.v1": - raise ValueError("Expected a dynamics training checkpoint") - profile = _checkpoint_numerical_profile(state) - model = NeRDModel(NeRDModelCfg(**state["model_cfg"])) - model.load_state_dict(state["model"]) + state = _load_training_checkpoint(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, diff --git a/embodichain/learning/models/nerd/checkpoint.py b/embodichain/learning/models/nerd/checkpoint.py index f161a55d18..893a582e0f 100644 --- a/embodichain/learning/models/nerd/checkpoint.py +++ b/embodichain/learning/models/nerd/checkpoint.py @@ -17,6 +17,7 @@ """Versioned portable NeRD exports and explicitly trusted legacy conversion.""" from __future__ import annotations +from collections.abc import Mapping import hashlib import json import pickle @@ -61,6 +62,18 @@ def _digest(path: Path) -> str: 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. @@ -102,13 +115,8 @@ def load_model( raise ValueError("NeRD weights checksum mismatch") cfg = NeRDModelCfg(**manifest["model_config"]) _execution_metadata(manifest, cfg) - model = NeRDModel(cfg) state = torch.load(path / "weights.pt", map_location="cpu", weights_only=True) - # Preserve exported floating dtype instead of silently downcasting to FP32. - dtype = next(v.dtype for v in state.values() if v.is_floating_point()) - model.to(dtype=dtype) - model.load_state_dict(state, strict=True) - return model.to(device).eval(), manifest + return _model_from_state(cfg, state, device).eval(), manifest class _LegacyModule(nn.Module): 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/trainer.py b/embodichain/learning/sl/trainer.py index c7c1530cff..d3d46cc816 100644 --- a/embodichain/learning/sl/trainer.py +++ b/embodichain/learning/sl/trainer.py @@ -27,6 +27,7 @@ import numpy as np import torch from embodichain.utils import configclass +from ._artifacts import _RunArtifacts from .learner import SupervisedLearner from .sources import SupervisedSource @@ -101,10 +102,9 @@ def __init__( 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._owned_artifacts: dict[Path, tuple[int, int, int, int]] = {} - self._fit_started = False - self._restored = False + 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. @@ -122,20 +122,8 @@ def fit(self, *, stop_after: int | None = None) -> dict[str, Any]: ) if end < self.learner.num_updates: raise ValueError("stop_after precedes the restored update cursor") - if not self._fit_started and not self._restored: - artifacts = [ - self.output / "metrics.jsonl", - self.output / "last.pt", - *self.output.glob("step_*.pt"), - ] - for path in artifacts: - if path.exists() and self._owned_artifacts.get( - path.resolve() - ) != self._artifact_stamp(path): - raise FileExistsError( - "SL output contains another run; restore a checkpoint or use a fresh output directory" - ) - self._fit_started = True + 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") @@ -153,20 +141,7 @@ def fit(self, *, stop_after: int | None = None) -> dict[str, Any]: and self.cfg.validation_interval and step % self.cfg.validation_interval == 0 ): - rng = self._rng_state() - was_training = self.learner.model.training - try: - self.learner.model.eval() - with torch.no_grad(): - metrics.update( - { - f"validation/{key}": float(value) - for key, value in self.evaluate(self.learner).items() - } - ) - finally: - self._restore_rng(rng) - self.learner.model.train(was_training) + metrics.update(self._evaluate_step()) with (self.output / "metrics.jsonl").open("a") as handle: handle.write(json.dumps(metrics, allow_nan=False) + "\n") if ( @@ -182,6 +157,34 @@ def fit(self, *, stop_after: int | None = None) -> dict[str, Any]: "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 { @@ -198,11 +201,6 @@ def _restore_rng(self, state: Mapping[str, Any]) -> None: if state["cuda"] is not None: torch.cuda.set_rng_state(state["cuda"], self.cfg.device) - @staticmethod - def _artifact_stamp(path: Path) -> tuple[int, int, int, int]: - stat = path.stat() - return stat.st_dev, stat.st_ino, stat.st_size, stat.st_mtime_ns - def save(self, path: str | Path) -> None: """Atomically save a complete committed learning and sampling boundary. @@ -210,9 +208,11 @@ def save(self, path: str | Path) -> None: 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(), @@ -228,10 +228,16 @@ def save(self, path: str | Path) -> None: 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) - self._owned_artifacts[path.resolve()] = self._artifact_stamp(path) def resume(self, path: str | Path) -> None: """Restore the same run after validating data, recipe and configuration. @@ -259,24 +265,13 @@ def resume(self, path: str | Path) -> None: ) if state.get("scheduler_type") != scheduler_type: raise ValueError("Supervised checkpoint scheduler type differs") - metrics_path = self.output / "metrics.jsonl" - if metrics_path.exists(): - last_line = "" - with metrics_path.open() as handle: - for line in handle: - if line.strip(): - last_line = line - if ( - last_line - and json.loads(last_line)["num_updates"] - > state["learner"]["num_updates"] - ): - raise ValueError( - "SL output history is ahead of this checkpoint; use a fresh output directory" - ) + 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._restored = True + self._artifacts.adopt_resume(run_id) diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.train.full.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.train.full.yaml index 212b3e9517..bfa9485894 100644 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.train.full.yaml +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.train.full.yaml @@ -35,4 +35,5 @@ train: lr_update_interval: 5000 lr_offset: 1 log_interval: 500 - resident_cache_bytes: 4294967296 + # Full Ant data exceed the former 4 GiB hard cache budget; use bounded reads. + resident_cache_bytes: 0 diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/sl.train.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/sl.train.yaml index 57355107ac..d4b2ed6256 100644 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/sl.train.yaml +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/sl.train.yaml @@ -23,7 +23,8 @@ model: attention_backend: math data: path: outputs/datasets/neural_dynamics/nerd/ant/v1 - resident_cache_bytes: 4294967296 + # Keep full-corpus training on bounded reads unless a fitting budget is set. + resident_cache_bytes: 0 objective: include_prefix: false numerical_profile: reference diff --git a/scripts/experiments/supervised/nerd.md b/scripts/experiments/supervised/nerd.md index e34afeeb83..871cafed58 100644 --- a/scripts/experiments/supervised/nerd.md +++ b/scripts/experiments/supervised/nerd.md @@ -20,6 +20,34 @@ The existing `embodichain dynamics` CLI retains its historical training/export formats. New common-format checkpoints resume through `train-sl` with the same recipe and data identity. +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 the resume contract for both CLIs. + +`last.pt` is the complete resumable checkpoint. Automatically selected `best.pt` +files are model-only snapshots marked `resumable=False`, suitable for +`embodichain dynamics 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 Use PR #855 revision `a58de92a82b4e1c1c260083b5b1930173322aceb` as the @@ -54,6 +82,13 @@ The default matrix performs: - 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 diff --git a/scripts/experiments/supervised/nerd_acceptance.py b/scripts/experiments/supervised/nerd_acceptance.py index 130a0ea79b..583b15d09e 100644 --- a/scripts/experiments/supervised/nerd_acceptance.py +++ b/scripts/experiments/supervised/nerd_acceptance.py @@ -173,6 +173,7 @@ def _worker(request_path: Path) -> None: ) 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 = [ @@ -191,6 +192,7 @@ def _worker(request_path: Path) -> None: ) ) first = trainer.step() + last_update = first for handle in handles: handle.remove() captured.update( @@ -208,6 +210,21 @@ def _worker(request_path: Path) -> None: ) 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 ) diff --git a/tests/cli/test_dynamics.py b/tests/cli/test_dynamics.py index cf534b2da9..93570bfbd5 100644 --- a/tests/cli/test_dynamics.py +++ b/tests/cli/test_dynamics.py @@ -19,6 +19,7 @@ from __future__ import annotations import json +from copy import deepcopy import os from pathlib import Path import subprocess @@ -28,6 +29,7 @@ import numpy as np import pytest +import torch import yaml from embodichain.cli.dynamics import main @@ -172,6 +174,173 @@ def test_real_offline_cli_train_export_eval(tmp_path): assert report["split"] == "test" and report["windows"] == 3 +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +@pytest.mark.parametrize("model_snapshot", [False, True]) +def test_training_checkpoint_cli_export_preserves_dtype_statistics_and_predictions( + tmp_path, dtype, model_snapshot +): + from embodichain.learning.models.nerd import NeRDModel, NeRDModelCfg, load_model + + torch.set_num_threads(2) + model = ( + NeRDModel( + NeRDModelCfg( + input_dims={"states_embedding": 4, "joint_f": 2}, + output_dim=4, + n_embd=12, + n_head=3, + n_layer=1, + block_size=3, + history_length=3, + head_hidden_dims=(12,), + ) + ) + .to(dtype=dtype) + .eval() + ) + # Values beyond FP32 precision expose a downcast even if tensors were later + # converted back to FP64. Frozen normalizers must retain their precision too. + with torch.no_grad(): + for parameter in model.parameters(): + parameter.add_(1e-10) + model.set_normalization( + { + key: ( + torch.full((width,), 1 / 7, dtype=dtype), + torch.full((width,), 1 / 3, dtype=dtype), + ) + for key, width in model.cfg.input_dims.items() + }, + torch.full((4,), 1 / 11, dtype=dtype), + torch.full((4,), 1 / 13, dtype=dtype), + ) + checkpoint = tmp_path / "training.pt" + state = { + "format": "embodichain.nerd.training.v1", + "model_cfg": model.cfg.to_dict(), + "model": model.state_dict(), + "train_cfg": {"numerical_profile": "default"}, + "manifest": {"codec": {"robot": "cartpole"}, "dt": 1 / 60}, + "global_step": 7, + "dataset_fingerprint": "fixed-training-corpus", + } + if model_snapshot: + state.pop("train_cfg") + state.update( + resumable=False, numerical_profile="default", attention_backend="default" + ) + torch.save(state, checkpoint) + destination = tmp_path / "export" + main(["export", "--checkpoint", str(checkpoint), "--output", str(destination)]) + loaded, metadata = load_model(destination) + assert metadata["training"]["global_step"] == 7 + assert all(value.dtype == dtype for value in loaded.state_dict().values()) + for name, expected in model.state_dict().items(): + torch.testing.assert_close(loaded.state_dict()[name], expected, atol=0, rtol=0) + inputs = { + name: torch.randn(2, 3, width, dtype=dtype) + for name, width in model.cfg.input_dims.items() + } + torch.testing.assert_close(model(inputs), loaded(inputs), atol=0, rtol=0) + + +def test_fp64_fine_tune_resume_export_preserves_updates_and_input_order(tmp_path): + from embodichain.data_pipeline.dynamics.shards import DynamicsShardWriter + from embodichain.learning.dynamics.workflow import run_export, run_train + from embodichain.learning.models.nerd import ( + NeRDModel, + NeRDModelCfg, + export_model, + load_model, + ) + + torch.set_num_threads(2) + metadata = {"codec": {"robot": "cartpole"}, "dt": 1 / 60} + with DynamicsShardWriter(tmp_path / "data", metadata) as writer: + for split in ("train", "validation", "test"): + q = np.zeros((6, 2), np.float32) + q[:, 0] = np.arange(6) * 0.01 + qd = np.zeros_like(q) + qd[:, 0] = 0.6 + writer.add_trajectory( + split, + {"q": q, "qd": qd}, + {"joint_f": np.zeros((5, 2), np.float32)}, + split=split, + ) + model = NeRDModel( + NeRDModelCfg( + input_dims={"states_embedding": 4, "joint_f": 2}, + output_dim=4, + n_embd=12, + n_head=3, + n_layer=1, + block_size=3, + history_length=3, + head_hidden_dims=(12,), + ) + ).double() + model.input_rms["states_embedding"].mean.fill_(1 / 7) + export_model(tmp_path / "pretrained", model, metadata) + # Use the same JSON-shaped config as the CLI, including list-valued fields. + config = json.loads( + json.dumps( + { + "dataset": str(tmp_path / "data"), + "output_dir": str(tmp_path / "interrupted"), + "fine_tune": str(tmp_path / "pretrained"), + "model": model.cfg.to_dict(), + "train": { + "total_steps": 2, + "batch_size": 2, + "device": "cpu", + "validation_interval": 1, + "checkpoint_interval": 1, + }, + "stop_after": 1, + } + ) + ) + run_train(config) + checkpoint = tmp_path / "interrupted/last.pt" + saved = torch.load(checkpoint, weights_only=True) + assert all(value.dtype == torch.float64 for value in saved["model"].values()) + resumed_config = deepcopy(config) + resumed_config.pop("fine_tune") + resumed_config.update( + resume=str(checkpoint), output_dir=str(tmp_path / "resumed"), stop_after=2 + ) + reordered = deepcopy(resumed_config) + reordered["output_dir"] = str(tmp_path / "reordered") + reordered["model"]["input_dims"] = dict( + reversed(list(reordered["model"]["input_dims"].items())) + ) + with pytest.raises(ValueError, match="input_dims order"): + run_train(reordered) + assert list((tmp_path / "reordered").iterdir()) == [] + assert run_train(resumed_config)["global_step"] == 2 + uninterrupted = deepcopy(config) + uninterrupted.update(output_dir=str(tmp_path / "full"), stop_after=2) + run_train(uninterrupted) + expected = torch.load(tmp_path / "full/last.pt", weights_only=True) + actual = torch.load(tmp_path / "resumed/last.pt", weights_only=True) + for name, value in expected["model"].items(): + torch.testing.assert_close(actual["model"][name], value, atol=0, rtol=0) + assert actual["best_step"] == expected["best_step"] + run_export(tmp_path / "resumed/best.pt", tmp_path / "export") + loaded, _ = load_model(tmp_path / "export") + assert all(value.dtype == torch.float64 for value in loaded.state_dict().values()) + for name, value in actual["best_model"].items(): + torch.testing.assert_close(loaded.state_dict()[name], value, atol=0, rtol=0) + inputs = { + name: torch.randn(2, 3, width, dtype=torch.float64) + for name, width in model.cfg.input_dims.items() + } + prediction = loaded(inputs) + assert prediction.dtype == torch.float64 + assert torch.isfinite(prediction).all() + + def test_packaged_workflow_configs_decode_through_cli(tmp_path, monkeypatch): """Exercise the actual YAML loader/path resolver and model/train constructors.""" import embodichain.learning.dynamics.workflow as workflow diff --git a/tests/data_pipeline/dynamics/test_generation.py b/tests/data_pipeline/dynamics/test_generation.py index c04d60fd1d..8912f5f922 100644 --- a/tests/data_pipeline/dynamics/test_generation.py +++ b/tests/data_pipeline/dynamics/test_generation.py @@ -19,6 +19,7 @@ from __future__ import annotations import os +import json from pathlib import Path import subprocess import sys @@ -36,6 +37,133 @@ from embodichain.data_pipeline.dynamics.shards import load_manifest, trajectory_split +@pytest.fixture +def deterministic_generation(monkeypatch): + """Exercise persistence with reproducible accepted and rejected candidates.""" + from embodichain.data_pipeline.dynamics import generation + + class Teacher: + def __init__(self, robot, *, device, num_envs, **kwargs): + self.device = torch.device("cpu") + self.gravity_dir = torch.tensor([0.0, 0.0, -1.0]).expand(num_envs, -1) + self.manifest = { + "physics_signature": "fake-physical-v1", + "teacher_origin": {"collector": "fake"}, + } + + def reset(self, q, qd): + self.q, self.qd = q.clone(), qd.clone() + self.rejected = torch.tensor([int(seed) in {11, 13} for seed in q[:, 0]]) + self.index = 0 + + def snapshot(self): + q = self.q + self.index * 0.01 + q[self.rejected] = float("nan") + return {"q": q, "qd": self.qd.clone()}, {} + + def step(self, force): + self.index += 1 + + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + def sample(robot, seed, horizon, *args): + return torch.tensor([float(seed), 0.0]), torch.zeros(2), torch.zeros(horizon, 2) + + monkeypatch.setattr(generation, "_NewtonTeacher", Teacher) + monkeypatch.setattr(generation, "_sample", sample) + + +def test_resume_retains_rejections_when_interrupted_after_shard_manifest( + tmp_path, monkeypatch, deterministic_generation +): + from embodichain.data_pipeline.dynamics import shards + + config = { + "output_dir": str(tmp_path), + "seed": 11, + "trajectories": 2, + "num_envs": 2, + "horizon": 2, + "trajectories_per_shard": 1, + } + atomic_json = shards._atomic_json + + def interrupt_after_publish(path, value): + atomic_json(path, value) + if path.name == "manifest.json" and value["shards"]: + raise RuntimeError("interrupted-after-manifest-publication") + + monkeypatch.setattr(shards, "_atomic_json", interrupt_after_publish) + with pytest.raises(RuntimeError, match="interrupted-after-manifest-publication"): + generate_dataset(config) + committed = load_manifest(tmp_path, verify=True) + assert [ + row["id"] for shard in committed["shards"] for row in shard["trajectories"] + ] == ["11:1"] + # This legacy report can lag behind the committed shard; it cannot be the + # recovery authority once progress is included in the shard manifest. + (tmp_path / "generation-progress.json").write_text( + json.dumps({"next_candidate": 0, "rejected": 0, "rejections": []}) + ) + monkeypatch.setattr(shards, "_atomic_json", atomic_json) + result = generate_dataset({**config, "resume": True}) + assert result["trajectories"] == 2 + assert result["rejected"] == 2 + evidence = json.loads(Path(result["rejection_evidence"]).read_text()) + assert evidence["next_candidate"] == 4 + assert [record["id"] for record in evidence["rejections"]] == ["11:0", "11:2"] + final = load_manifest(tmp_path, verify=True) + assert final["metadata"] == committed["metadata"] + assert final["shards"][0] == committed["shards"][0] + baseline = tmp_path / "uninterrupted" + generate_dataset({**config, "output_dir": str(baseline)}) + expected = load_manifest(baseline, verify=True) + assert final["metadata"] == expected["metadata"] + for actual_shard, expected_shard in zip(final["shards"], expected["shards"]): + assert actual_shard["trajectories"] == expected_shard["trajectories"] + with ( + h5py.File(tmp_path / actual_shard["file"]) as actual_file, + h5py.File(baseline / expected_shard["file"]) as expected_file, + ): + for group in ("states", "transitions"): + for field in actual_file[f"0/{group}"]: + assert torch.equal( + torch.from_numpy(actual_file[f"0/{group}/{field}"][()]), + torch.from_numpy(expected_file[f"0/{group}/{field}"][()]), + ) + + +def test_resume_reads_legacy_generation_progress(tmp_path, deterministic_generation): + config = { + "output_dir": str(tmp_path), + "seed": 11, + "trajectories": 1, + "num_envs": 1, + "horizon": 2, + "trajectories_per_shard": 1, + "max_attempts": 1, + } + with pytest.raises(RuntimeError, match="Generation attempt limit"): + generate_dataset(config) + # Even a batch with no accepted trajectories advances a recoverable cursor. + manifest = load_manifest(tmp_path) + assert not manifest["shards"] + assert manifest["generation_progress"]["next_candidate"] == 1 + assert manifest["generation_progress"]["rejected"] == 1 + # Older collectors store progress only in the separate JSON report. + progress = manifest.pop("generation_progress") + (tmp_path / "manifest.json").write_text(json.dumps(manifest)) + (tmp_path / "generation-progress.json").write_text(json.dumps(progress)) + result = generate_dataset({**config, "resume": True, "max_attempts": 2}) + assert result["trajectories"] == result["rejected"] == 1 + evidence = json.loads(Path(result["rejection_evidence"]).read_text()) + assert [record["id"] for record in evidence["rejections"]] == ["11:0"] + + @pytest.mark.parametrize( "field,value", [ diff --git a/tests/data_pipeline/dynamics/test_shards.py b/tests/data_pipeline/dynamics/test_shards.py index 380a3401b0..ee82ed0ff1 100644 --- a/tests/data_pipeline/dynamics/test_shards.py +++ b/tests/data_pipeline/dynamics/test_shards.py @@ -108,6 +108,32 @@ def test_validation_and_integrity(tmp_path): load_manifest(tmp_path, verify=True) +def test_progress_only_publication_preserves_data_identity_and_pending_rows(tmp_path): + from embodichain.data_pipeline.dynamics.shards import dataset_fingerprint + + with DynamicsShardWriter(tmp_path, {}, trajectories_per_shard=4) as writer: + writer.add_trajectory("accepted", *trajectory(), split="train") + progress = {"next_candidate": 1, "accepted": 1, "rejected": 0, "rejections": []} + writer.manifest["generation_progress"] = progress + writer.commit() + first = load_manifest(tmp_path, verify=True) + assert first["shards"][0]["trajectories"][0]["id"] == "accepted" + legacy = { + key: value for key, value in first.items() if key != "generation_progress" + } + assert dataset_fingerprint(first) == dataset_fingerprint(legacy) + progress.update( + next_candidate=3, + rejected=2, + rejections=[{"id": "rejected-1"}, {"id": "rejected-2"}], + ) + writer.commit() + final = load_manifest(tmp_path, verify=True) + assert final["generation_progress"] == progress + assert final["shards"] == first["shards"] + assert dataset_fingerprint(final) == dataset_fingerprint(first) + + def test_stable_split_not_ingestion_order(): forward = {str(i): trajectory_split(str(i), 7) for i in range(100)} backward = {str(i): trajectory_split(str(i), 7) for i in reversed(range(100))} diff --git a/tests/learning/dynamics/test_trainer.py b/tests/learning/dynamics/test_trainer.py index e0def16c6c..6055215f27 100644 --- a/tests/learning/dynamics/test_trainer.py +++ b/tests/learning/dynamics/test_trainer.py @@ -20,6 +20,7 @@ import json import math +import random import subprocess import sys @@ -162,6 +163,8 @@ def test_shared_recipe_matches_legacy_updates_and_restores_sampler( steps=4, numerical_profile=numerical_profile, ) + legacy.cfg.lr_update_interval = 3 + legacy.cfg.lr_offset = 1 config = { "model": architecture().to_dict(), "data": {"path": str(tmp_path / "data")}, @@ -170,6 +173,8 @@ def test_shared_recipe_matches_legacy_updates_and_restores_sampler( "learning_rate": 0.005, "final_learning_rate": 0.002, "max_grad_norm": 1.0, + "lr_update_interval": 3, + "lr_offset": 1, }, "trainer": { "total_updates": 4, @@ -180,12 +185,48 @@ def test_shared_recipe_matches_legacy_updates_and_restores_sampler( } shared = build_supervised_trainer(config) shared.learner.model.load_state_dict(legacy.model.state_dict()) - for _ in range(4): - legacy.step() + legacy_rates = [legacy.step()["learning_rate"] for _ in range(4)] + shared_rates = [] + shared.learner.optimizer.register_step_pre_hook( + lambda optimizer, args, kwargs: shared_rates.append( + optimizer.param_groups[0]["lr"] + ) + ) shared.fit(stop_after=2) + reordered = build_supervised_trainer( + { + **config, + "model": { + **config["model"], + "input_dims": dict( + reversed(list(config["model"]["input_dims"].items())) + ), + }, + "trainer": {**config["trainer"], "output_dir": str(tmp_path / "reordered")}, + } + ) + before = { + key: value.clone() + for key, value in reordered.learner.model.state_dict().items() + } + with pytest.raises(ValueError, match="identity"): + reordered.resume(tmp_path / "shared/last.pt") + assert reordered.learner.num_updates == 0 + for key, value in before.items(): + torch.testing.assert_close( + reordered.learner.model.state_dict()[key], value, rtol=0, atol=0 + ) restored = build_supervised_trainer(config) restored.resume(tmp_path / "shared/last.pt") + restored.learner.optimizer.register_step_pre_hook( + lambda optimizer, args, kwargs: shared_rates.append( + optimizer.param_groups[0]["lr"] + ) + ) restored.fit() + assert shared_rates == pytest.approx(legacy_rates, rel=1e-15, abs=0) + assert shared_rates[0] == shared_rates[1] == shared_rates[2] + assert shared_rates[3] < shared_rates[2] for name, value in legacy.model.state_dict().items(): torch.testing.assert_close( restored.learner.model.state_dict()[name], value, atol=0, rtol=0 @@ -269,13 +310,16 @@ def test_training_improves_validation_and_export_matches(tmp_path): assert report["qualification"] == "offline_only" -def test_workflow_fine_tune_and_resume_contract(tmp_path): +@pytest.mark.parametrize("head_dims", [(12,), [12]]) +def test_workflow_fine_tune_and_resume_contract(tmp_path, head_dims): metadata = make_data(tmp_path / "data") + model_config = architecture().to_dict() + model_config["head_hidden_dims"] = head_dims cfg = { "dataset": str(tmp_path / "data"), "output_dir": str(tmp_path / "first"), "manifest": metadata, - "model": architecture().to_dict(), + "model": model_config, "train": { "total_steps": 2, "batch_size": 4, @@ -285,12 +329,35 @@ def test_workflow_fine_tune_and_resume_contract(tmp_path): } run_train(cfg) run_export(tmp_path / "first" / "last.pt", tmp_path / "export") + existing = { + path.name: path.read_bytes() + for path in (tmp_path / "first").iterdir() + if path.is_file() + } + for options in ({}, {"fine_tune": str(tmp_path / "export")}): + with pytest.raises(FileExistsError, match="another run"): + run_train({**cfg, **options}) + assert existing == { + path.name: path.read_bytes() + for path in (tmp_path / "first").iterdir() + if path.is_file() + } fine = { **cfg, "output_dir": str(tmp_path / "fine"), "fine_tune": str(tmp_path / "export"), } assert run_train(fine)["global_step"] == 2 + assert ( + run_train( + { + **cfg, + "output_dir": str(tmp_path / "fine-resume"), + "resume": str(tmp_path / "fine/last.pt"), + } + )["global_step"] + == 2 + ) mismatch = { **cfg, "output_dir": str(tmp_path / "bad"), @@ -301,6 +368,252 @@ def test_workflow_fine_tune_and_resume_contract(tmp_path): run_train(mismatch) +@pytest.mark.parametrize("operation", ["step", "train", "save"]) +def test_fresh_trainer_rejects_owned_output_before_update(tmp_path, operation): + make_data(tmp_path / "data") + original = build(tmp_path / "data", tmp_path / "run", steps=2) + original.train(stop_after=1) + before = { + path.name: path.read_bytes() + for path in (tmp_path / "run").iterdir() + if path.is_file() + } + candidate = build(tmp_path / "data", tmp_path / "run", steps=2) + parameters = { + key: value.clone() for key, value in candidate.model.state_dict().items() + } + sampler = candidate.sampler.get_state().clone() + with pytest.raises(FileExistsError, match="another run"): + if operation == "save": + candidate.save(tmp_path / "run/new.pt") + else: + getattr(candidate, operation)() + assert candidate.global_step == candidate.learner.num_updates == 0 + assert not candidate.optimizer.state + torch.testing.assert_close(candidate.sampler.get_state(), sampler, rtol=0, atol=0) + for key, value in parameters.items(): + torch.testing.assert_close( + candidate.model.state_dict()[key], value, rtol=0, atol=0 + ) + assert before == { + path.name: path.read_bytes() + for path in (tmp_path / "run").iterdir() + if path.is_file() + } + + +def test_resume_new_directory_preserves_best_export_without_new_improvement( + tmp_path, monkeypatch +): + make_data(tmp_path / "data") + scores = iter([1.0, 2.0, 3.0]) + monkeypatch.setattr( + "embodichain.learning.dynamics.trainer.evaluate_model", + lambda *args, **kwargs: {"normalized_mse": next(scores)}, + ) + original = build(tmp_path / "data", tmp_path / "original", steps=4) + original.cfg.validation_interval = 1 + original.train(stop_after=2) + saved_best = torch.load(tmp_path / "original/best.pt", weights_only=True) + assert saved_best["global_step"] == 1 + assert saved_best["resumable"] is False + assert any( + not torch.equal(value, original.model.state_dict()[key]) + for key, value in saved_best["model"].items() + ) + before = { + path.name: path.read_bytes() + for path in (tmp_path / "original").iterdir() + if path.is_file() + } + run_export(tmp_path / "original/best.pt", tmp_path / "original-export") + restored = build(tmp_path / "data", tmp_path / "resumed", steps=4) + restored.cfg.validation_interval = 1 + restored.resume(tmp_path / "original/last.pt") + assert list((tmp_path / "resumed").iterdir()) == [] + restored.train(stop_after=3) + restored_best = torch.load(tmp_path / "resumed/best.pt", weights_only=True) + assert restored_best["global_step"] == 1 + for key, value in saved_best["model"].items(): + torch.testing.assert_close(restored_best["model"][key], value, rtol=0, atol=0) + assert before == { + path.name: path.read_bytes() + for path in (tmp_path / "original").iterdir() + if path.is_file() + } + run_export(tmp_path / "resumed/best.pt", tmp_path / "resumed-export") + expected, _ = load_model(tmp_path / "original-export") + actual, _ = load_model(tmp_path / "resumed-export") + inputs, _, _ = encode_batch( + restored.validation_dataset.get_batch([0]), restored.codec, "cpu" + ) + torch.testing.assert_close(actual(inputs), expected(inputs), rtol=0, atol=0) + with pytest.raises(ValueError, match="export only"): + restored.resume(tmp_path / "resumed/best.pt") + + +def test_resume_preserves_future_best_after_interrupted_publication( + tmp_path, monkeypatch +): + make_data(tmp_path / "data") + scores = iter([1.0, 0.5]) + monkeypatch.setattr( + "embodichain.learning.dynamics.trainer.evaluate_model", + lambda *args, **kwargs: {"normalized_mse": next(scores)}, + ) + original = build(tmp_path / "data", tmp_path / "original", steps=4) + original.cfg.validation_interval = 1 + original.train(stop_after=1) + first_best = torch.load(tmp_path / "original/best.pt", weights_only=True) + publish = original._write_best_snapshot + + def interrupt_after_best() -> None: + publish() + raise RuntimeError("interrupted after publishing best") + + monkeypatch.setattr(original, "_write_best_snapshot", interrupt_after_best) + with pytest.raises(RuntimeError, match="publishing best"): + original.train(stop_after=2) + assert ( + torch.load(tmp_path / "original/best.pt", weights_only=True)["global_step"] == 2 + ) + assert ( + torch.load(tmp_path / "original/last.pt", weights_only=True)["global_step"] == 1 + ) + before = { + path.name: path.read_bytes() + for path in (tmp_path / "original").iterdir() + if path.is_file() + } + candidate = build(tmp_path / "data", tmp_path / "original", steps=4) + candidate.cfg.validation_interval = 1 + parameters = { + key: value.clone() for key, value in candidate.model.state_dict().items() + } + sampler, rng = candidate.sampler.get_state().clone(), torch.get_rng_state().clone() + python_rng, numpy_rng = random.getstate(), np.random.get_state() + with pytest.raises(ValueError, match="best.pt.*fresh output"): + candidate.resume(tmp_path / "original/last.pt") + assert candidate.global_step == 0 and not candidate.optimizer.state + 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) + assert random.getstate() == python_rng + actual_numpy = np.random.get_state() + assert actual_numpy[0] == numpy_rng[0] and actual_numpy[2:] == numpy_rng[2:] + np.testing.assert_array_equal(actual_numpy[1], numpy_rng[1]) + for key, value in parameters.items(): + torch.testing.assert_close( + candidate.model.state_dict()[key], value, rtol=0, atol=0 + ) + restored = build(tmp_path / "data", tmp_path / "fresh", steps=4) + restored.cfg.validation_interval = 1 + restored.resume(tmp_path / "original/last.pt") + restored.train(stop_after=1) + restored_best = torch.load(tmp_path / "fresh/best.pt", weights_only=True) + assert restored_best["global_step"] == 1 + for key, value in first_best["model"].items(): + torch.testing.assert_close(restored_best["model"][key], value, rtol=0, atol=0) + assert before == { + path.name: path.read_bytes() + for path in (tmp_path / "original").iterdir() + if path.is_file() + } + + +def test_resume_input_order_rejection_preserves_runtime_state(tmp_path): + make_data(tmp_path / "data") + original = build(tmp_path / "data", tmp_path / "original", steps=4) + original.train(stop_after=1) + candidate = build(tmp_path / "data", tmp_path / "candidate", steps=4) + candidate.model.cfg.input_dims = dict( + reversed(list(candidate.model.cfg.input_dims.items())) + ) + parameters = { + key: value.clone() for key, value in candidate.model.state_dict().items() + } + sampler, rng = candidate.sampler.get_state().clone(), torch.get_rng_state().clone() + with pytest.raises(ValueError, match="input_dims order"): + candidate.resume(tmp_path / "original/last.pt") + assert candidate.global_step == 0 + assert not candidate.optimizer.state + assert list((tmp_path / "candidate").iterdir()) == [] + 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 key, value in parameters.items(): + torch.testing.assert_close( + candidate.model.state_dict()[key], value, rtol=0, atol=0 + ) + + +def test_legacy_best_fallback_requires_matching_source_weights(tmp_path): + make_data(tmp_path / "data") + original = build(tmp_path / "data", tmp_path / "legacy", steps=4) + original.step() + original.best_validation = 1.0 + original.save(tmp_path / "legacy/best.pt") + original.step() + original.save(tmp_path / "legacy/last.pt") + # Construct the historical v1 layout from actual step-one and step-two + # checkpoints, retaining each original optimizer/RNG boundary. + for filename in ("best.pt", "last.pt"): + path = tmp_path / "legacy" / filename + state = torch.load(path, weights_only=True) + for key in ("run_id", "best_model", "best_step"): + state.pop(key) + torch.save(state, path) + expected = torch.load(tmp_path / "legacy/best.pt", weights_only=True) + restored = build(tmp_path / "data", tmp_path / "restored", steps=4) + restored.resume(tmp_path / "legacy/last.pt") + restored.train(stop_after=2) + best = torch.load(tmp_path / "restored/best.pt", weights_only=True) + assert best["global_step"] == 1 + assert best["resumable"] is False + for key, value in expected["model"].items(): + torch.testing.assert_close(best["model"][key], value, rtol=0, atol=0) + (tmp_path / "legacy/best.pt").unlink() + missing = build(tmp_path / "data", tmp_path / "missing", steps=4) + with pytest.raises(ValueError, match="matching best.pt"): + missing.resume(tmp_path / "legacy/last.pt") + assert missing.global_step == 0 and not missing.optimizer.state + assert list((tmp_path / "missing").iterdir()) == [] + + +def test_resume_foreign_directory_or_newer_history_preserves_artifacts(tmp_path): + make_data(tmp_path / "data") + source = build(tmp_path / "data", tmp_path / "source", steps=4) + source.step() + source.save(tmp_path / "source/step_1.pt") + source.train(stop_after=3) + foreign = build(tmp_path / "data", tmp_path / "foreign", steps=4) + foreign.train(stop_after=1) + for output, checkpoint, error, message in ( + ("foreign", "source/last.pt", FileExistsError, "another run"), + ("source", "source/step_1.pt", ValueError, "history is ahead"), + ): + before = { + path.name: path.read_bytes() + for path in (tmp_path / output).iterdir() + if path.is_file() + } + candidate = build(tmp_path / "data", tmp_path / output, steps=4) + parameters = { + key: value.clone() for key, value in candidate.model.state_dict().items() + } + with pytest.raises(error, match=message): + candidate.resume(tmp_path / checkpoint) + assert candidate.global_step == 0 + assert not candidate.optimizer.state + for key, value in parameters.items(): + torch.testing.assert_close( + candidate.model.state_dict()[key], value, rtol=0, atol=0 + ) + assert before == { + path.name: path.read_bytes() + for path in (tmp_path / output).iterdir() + if path.is_file() + } + + def test_resume_rejects_unknown_preprocessing_execution(tmp_path): make_data(tmp_path / "data") trainer = build(tmp_path / "data", tmp_path / "run", steps=2) @@ -457,8 +770,12 @@ def test_memory_mapped_order_resume_and_integrity(tmp_path): order[0, 0] = 1 np.save(path, order) invalid = build(tmp_path / "data", tmp_path / "invalid", steps=10) + sampler, rng = invalid.sampler.get_state().clone(), torch.get_rng_state().clone() with pytest.raises(ValueError, match="checksum"): invalid.resume(tmp_path / "run" / "last.pt") + assert invalid.sample_order is None and invalid._sample_order_identity is None + torch.testing.assert_close(invalid.sampler.get_state(), sampler, rtol=0, atol=0) + torch.testing.assert_close(torch.get_rng_state(), rng, rtol=0, atol=0) def test_checkpointed_order_resume_is_identical_to_uninterrupted_training(tmp_path): diff --git a/tests/learning/models/test_nerd.py b/tests/learning/models/test_nerd.py index 77ce6d093c..a3f6c29d41 100644 --- a/tests/learning/models/test_nerd.py +++ b/tests/learning/models/test_nerd.py @@ -105,15 +105,16 @@ def test_multistep_history_control_and_parameter_gradients() -> None: ) -def test_export_roundtrip_and_checksum(tmp_path: Path) -> None: - model = tiny().double().eval() +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_export_roundtrip_and_checksum(tmp_path: Path, dtype: torch.dtype) -> None: + model = tiny().to(dtype=dtype).eval() inputs = { - k: torch.randn(2, 3, d, dtype=torch.double) - for k, d in model.cfg.input_dims.items() + k: torch.randn(2, 3, d, dtype=dtype) for k, d in model.cfg.input_dims.items() } export_model(tmp_path, model, {"codec": {"robot": "cartpole"}, "dt": 1 / 60}) loaded, manifest = load_model(tmp_path) assert manifest["codec"]["robot"] == "cartpole" + assert all(value.dtype == dtype for value in loaded.state_dict().values()) torch.testing.assert_close(model(inputs), loaded(inputs), atol=0, rtol=0) with pytest.raises(FileExistsError): export_model(tmp_path, model, {}) diff --git a/tests/learning/sl/test_trainer.py b/tests/learning/sl/test_trainer.py index b5c881e81e..c1681efbad 100644 --- a/tests/learning/sl/test_trainer.py +++ b/tests/learning/sl/test_trainer.py @@ -198,20 +198,34 @@ 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} @@ -230,6 +244,17 @@ def evaluate(learner): 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): @@ -247,15 +272,22 @@ def test_fresh_run_cannot_overwrite_existing_artifacts(tmp_path): assert resumed.learner.num_updates == 3 -def test_older_checkpoint_cannot_append_to_a_later_run_history(tmp_path): +@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") - first.fit(stop_after=3) + 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 is ahead"): + 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 @@ -282,11 +314,60 @@ def test_preconstructed_trainer_rechecks_output_ownership_at_fit(tmp_path): assert second.learner.num_updates == 0 -def test_foreign_replacement_of_owned_initial_snapshot_is_rejected(tmp_path): +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) - second.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"): - first.fit(stop_after=1) + 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"] From d9cd707e09cf986e16ef5a56ed8792cd5caa39d8 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 25 Oct 2026 09:52:23 +0000 Subject: [PATCH 16/21] docs(sl): render NeRD adapters with explicit API directives --- .../embodichain.learning.dynamics.sl.rst | 34 ++++++++++++++++--- .../embodichain/embodichain.learning.rst | 2 +- 2 files changed, 30 insertions(+), 6 deletions(-) diff --git a/docs/source/api_reference/embodichain/embodichain.learning.dynamics.sl.rst b/docs/source/api_reference/embodichain/embodichain.learning.dynamics.sl.rst index e619e89692..eecc44a544 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.dynamics.sl.rst +++ b/docs/source/api_reference/embodichain/embodichain.learning.dynamics.sl.rst @@ -1,5 +1,5 @@ NeRD supervised-learning adapters -================================ +================================= NeRD uses the shared supervised optimizer update while retaining its physical codec, normalized loss and numerical execution policy. The standalone recipe @@ -7,8 +7,32 @@ connects immutable NeRD windows to the common trainer; its checkpoints use the SL format. Historical NeRD training checkpoints remain owned by the existing ``dynamics`` workflow. -.. automodule:: embodichain.learning.dynamics.supervised - :members: NeRDObjective, make_nerd_batch +Physical objective and batches +------------------------------ -.. automodule:: embodichain.learning.dynamics.sl_recipe - :members: NeRDWindowSource, build_supervised_trainer +.. currentmodule:: embodichain.learning.dynamics.supervised + +.. autosummary:: + + NeRDObjective + make_nerd_batch + +.. autoclass:: NeRDObjective + :members: + +.. autofunction:: make_nerd_batch + +Offline source and training recipe +---------------------------------- + +.. currentmodule:: embodichain.learning.dynamics.sl_recipe + +.. autosummary:: + + NeRDWindowSource + build_supervised_trainer + +.. autoclass:: NeRDWindowSource + :members: + +.. autofunction:: build_supervised_trainer diff --git a/docs/source/api_reference/embodichain/embodichain.learning.rst b/docs/source/api_reference/embodichain/embodichain.learning.rst index 1022a61fbe..cf20e4bfa9 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.rst +++ b/docs/source/api_reference/embodichain/embodichain.learning.rst @@ -1,5 +1,5 @@ embodichain.learning -============= +==================== .. automodule:: embodichain.learning .. rubric:: Submodules From 2e993020a0e127605df91e6cb8772d0ce11a0034 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 25 Oct 2026 10:28:36 +0000 Subject: [PATCH 17/21] fix(nerd): validate statistics and preserve unnormalized model precision --- embodichain/learning/dynamics/trainer.py | 11 ++ embodichain/learning/models/nerd/model.py | 4 + tests/cli/test_dynamics.py | 66 +++++++-- tests/learning/dynamics/test_trainer.py | 169 ++++++++++++++++++++++ tests/learning/models/test_nerd.py | 59 ++++++++ 5 files changed, 297 insertions(+), 12 deletions(-) diff --git a/embodichain/learning/dynamics/trainer.py b/embodichain/learning/dynamics/trainer.py index 8848ea0c69..5de8b08180 100644 --- a/embodichain/learning/dynamics/trainer.py +++ b/embodichain/learning/dynamics/trainer.py @@ -138,6 +138,9 @@ def normalized_mse( 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: @@ -389,6 +392,10 @@ class DynamicsTrainer: 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__( @@ -436,6 +443,10 @@ def __init__( 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": diff --git a/embodichain/learning/models/nerd/model.py b/embodichain/learning/models/nerd/model.py index 3abd86c07b..4a36843da5 100644 --- a/embodichain/learning/models/nerd/model.py +++ b/embodichain/learning/models/nerd/model.py @@ -206,6 +206,8 @@ class NeRDModel(nn.Module): 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. """ @@ -275,6 +277,7 @@ 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: @@ -282,6 +285,7 @@ def forward(self, inputs: Mapping[str, Tensor]) -> Tensor: 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 ) diff --git a/tests/cli/test_dynamics.py b/tests/cli/test_dynamics.py index 93570bfbd5..61b61e3e00 100644 --- a/tests/cli/test_dynamics.py +++ b/tests/cli/test_dynamics.py @@ -244,9 +244,14 @@ def test_training_checkpoint_cli_export_preserves_dtype_statistics_and_predictio torch.testing.assert_close(model(inputs), loaded(inputs), atol=0, rtol=0) -def test_fp64_fine_tune_resume_export_preserves_updates_and_input_order(tmp_path): +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +@pytest.mark.parametrize("normalize_input", [False, True]) +@pytest.mark.parametrize("normalize_output", [False, True]) +def test_fine_tune_resume_export_preserves_dtype_updates_and_input_order( + tmp_path, dtype, normalize_input, normalize_output +): from embodichain.data_pipeline.dynamics.shards import DynamicsShardWriter - from embodichain.learning.dynamics.workflow import run_export, run_train + from embodichain.learning.dynamics.workflow import run_train from embodichain.learning.models.nerd import ( NeRDModel, NeRDModelCfg, @@ -278,8 +283,10 @@ def test_fp64_fine_tune_resume_export_preserves_updates_and_input_order(tmp_path block_size=3, history_length=3, head_hidden_dims=(12,), + normalize_input=normalize_input, + normalize_output=normalize_output, ) - ).double() + ).to(dtype=dtype) model.input_rms["states_embedding"].mean.fill_(1 / 7) export_model(tmp_path / "pretrained", model, metadata) # Use the same JSON-shaped config as the CLI, including list-valued fields. @@ -301,10 +308,34 @@ def test_fp64_fine_tune_resume_export_preserves_updates_and_input_order(tmp_path } ) ) - run_train(config) + + def train_cli(value, label): + path = tmp_path / f"{label}.json" + path.write_text(json.dumps(value)) + main(["train", "--config", str(path)]) + + def evaluate_cli(model_path, label): + path = tmp_path / f"{label}.json" + output = tmp_path / f"{label}-result.json" + path.write_text( + json.dumps( + { + "mode": "offline", + "dataset": str(tmp_path / "data"), + "model": str(model_path), + "device": "cpu", + "output": str(output), + } + ) + ) + main(["eval", "--config", str(path)]) + assert json.loads(output.read_text())["windows"] == 3 + + evaluate_cli(tmp_path / "pretrained", "pretrained-eval") + train_cli(config, "fine-tune") checkpoint = tmp_path / "interrupted/last.pt" saved = torch.load(checkpoint, weights_only=True) - assert all(value.dtype == torch.float64 for value in saved["model"].values()) + assert all(value.dtype == dtype for value in saved["model"].values()) resumed_config = deepcopy(config) resumed_config.pop("fine_tune") resumed_config.update( @@ -318,27 +349,38 @@ def test_fp64_fine_tune_resume_export_preserves_updates_and_input_order(tmp_path with pytest.raises(ValueError, match="input_dims order"): run_train(reordered) assert list((tmp_path / "reordered").iterdir()) == [] - assert run_train(resumed_config)["global_step"] == 2 + train_cli(resumed_config, "resume") uninterrupted = deepcopy(config) uninterrupted.update(output_dir=str(tmp_path / "full"), stop_after=2) - run_train(uninterrupted) + train_cli(uninterrupted, "uninterrupted") expected = torch.load(tmp_path / "full/last.pt", weights_only=True) actual = torch.load(tmp_path / "resumed/last.pt", weights_only=True) + assert actual["global_step"] == 2 for name, value in expected["model"].items(): torch.testing.assert_close(actual["model"][name], value, atol=0, rtol=0) assert actual["best_step"] == expected["best_step"] - run_export(tmp_path / "resumed/best.pt", tmp_path / "export") + main( + [ + "export", + "--checkpoint", + str(tmp_path / "resumed/best.pt"), + "--output", + str(tmp_path / "export"), + ] + ) loaded, _ = load_model(tmp_path / "export") - assert all(value.dtype == torch.float64 for value in loaded.state_dict().values()) + assert all(value.dtype == dtype for value in loaded.state_dict().values()) + assert loaded.cfg.normalize_input is normalize_input + assert loaded.cfg.normalize_output is normalize_output for name, value in actual["best_model"].items(): torch.testing.assert_close(loaded.state_dict()[name], value, atol=0, rtol=0) inputs = { - name: torch.randn(2, 3, width, dtype=torch.float64) - for name, width in model.cfg.input_dims.items() + name: torch.randn(2, 3, width) for name, width in model.cfg.input_dims.items() } prediction = loaded(inputs) - assert prediction.dtype == torch.float64 + assert prediction.dtype == dtype assert torch.isfinite(prediction).all() + evaluate_cli(tmp_path / "export", "resumed-eval") def test_packaged_workflow_configs_decode_through_cli(tmp_path, monkeypatch): diff --git a/tests/learning/dynamics/test_trainer.py b/tests/learning/dynamics/test_trainer.py index 6055215f27..ad58335428 100644 --- a/tests/learning/dynamics/test_trainer.py +++ b/tests/learning/dynamics/test_trainer.py @@ -21,6 +21,7 @@ import json import math import random +from copy import deepcopy import subprocess import sys @@ -41,7 +42,9 @@ compute_statistics, encode_batch, evaluate_model, + load_statistics, normalized_mse, + save_statistics, ) from embodichain.learning.dynamics.workflow import run_train, run_export, run_eval from embodichain.learning.dynamics.qualification import ( @@ -149,6 +152,130 @@ def test_statistics_only_training_windows(tmp_path): ) +@pytest.fixture +def statistics_sources(tmp_path): + """Provide distinct immutable datasets with the same numeric field schema.""" + metadata = make_data(tmp_path / "own") + with DynamicsShardWriter(tmp_path / "foreign", metadata) as writer: + writer.add_trajectory( + "foreign", + { + "q": np.full((6, 2), 17.0, np.float32), + "qd": np.full((6, 2), 3.0, np.float32), + }, + {"joint_f": np.zeros((5, 2), np.float32)}, + split="train", + ) + train = DynamicsWindowDataset(tmp_path / "own", "train", 3) + validation = DynamicsWindowDataset(tmp_path / "own", "validation", 3) + foreign = DynamicsWindowDataset(tmp_path / "foreign", "train", 3) + assert train.manifest["fields"] == foreign.manifest["fields"] + assert train.fingerprint != foreign.fingerprint + try: + yield train, validation, foreign, NeRDStateCodec("cartpole") + finally: + for reader in (train, validation, foreign): + reader.close() + + +@pytest.mark.parametrize("origin", ["foreign", "missing"]) +@pytest.mark.parametrize("initialize_normalization", [True, False]) +def test_supplied_statistics_identity_rejected_before_mutation( + tmp_path, statistics_sources, origin, initialize_normalization +): + train, validation, foreign, codec = statistics_sources + statistics = compute_statistics(foreign if origin == "foreign" else train, codec) + if origin == "missing": + statistics.pop("dataset_fingerprint") + path = tmp_path / "statistics.pt" + save_statistics(path, statistics) + original_file = path.read_bytes() + loaded = load_statistics(path) + original_statistics = deepcopy(loaded) + model = NeRDModel(architecture()) + original_model = {key: value.clone() for key, value in model.state_dict().items()} + original_config = deepcopy(model.cfg.to_dict()) + # A rejected constructor must not switch the attention policy or reset any + # global RNG while admitting an unrelated normalization artifact. + cfg = DynamicsTrainCfg( + total_steps=2, batch_size=4, numerical_profile="reference", seed=97 + ) + random.seed(101) + np.random.seed(102) + torch.manual_seed(103) + rng = (random.getstate(), np.random.get_state(), torch.get_rng_state().clone()) + output = tmp_path / "run" + with pytest.raises(ValueError, match="statistics/dataset fingerprint mismatch"): + DynamicsTrainer( + model, + codec, + train, + validation, + cfg, + output, + manifest=train.manifest["metadata"], + statistics=loaded, + initialize_normalization=initialize_normalization, + ) + assert not output.exists() + assert model.cfg.to_dict() == original_config + for key, value in original_model.items(): + torch.testing.assert_close(model.state_dict()[key], value, atol=0, rtol=0) + assert random.getstate() == rng[0] + numpy_rng = np.random.get_state() + assert numpy_rng[0] == rng[1][0] and numpy_rng[2:] == rng[1][2:] + np.testing.assert_array_equal(numpy_rng[1], rng[1][1]) + torch.testing.assert_close(torch.get_rng_state(), rng[2], atol=0, rtol=0) + assert loaded.keys() == original_statistics.keys() + for group in ("inputs", "output"): + values = loaded[group] if group == "inputs" else {"output": loaded[group]} + expected = ( + original_statistics[group] + if group == "inputs" + else {"output": original_statistics[group]} + ) + for field in values: + for moment in ("mean", "var"): + torch.testing.assert_close( + values[field][moment], expected[field][moment], atol=0, rtol=0 + ) + assert path.read_bytes() == original_file + + +@pytest.mark.parametrize("initialize_normalization", [True, False]) +def test_matching_supplied_statistics_preserve_normalization_and_update( + tmp_path, statistics_sources, initialize_normalization +): + train, validation, _, codec = statistics_sources + baseline = build(tmp_path / "own", tmp_path / "baseline", steps=2) + path = tmp_path / "matching-statistics.pt" + save_statistics(path, baseline.statistics) + loaded = load_statistics(path) + if not initialize_normalization: + # Explicitly skipping initialization must retain existing model buffers, + # even if matching-corpus statistics were supplied for its record. + baseline.model.input_rms["states_embedding"].mean[0] = 0.5 + model = deepcopy(baseline.model) + candidate = DynamicsTrainer( + model, + codec, + train, + validation, + baseline.cfg.copy(), + tmp_path / "matching", + manifest=train.manifest["metadata"], + statistics=loaded, + initialize_normalization=initialize_normalization, + ) + for key, value in baseline.model.state_dict().items(): + torch.testing.assert_close(model.state_dict()[key], value, atol=0, rtol=0) + assert candidate.step() == baseline.step() + for key, value in baseline.model.state_dict().items(): + torch.testing.assert_close(model.state_dict()[key], value, atol=0, rtol=0) + baseline.train_dataset.close() + baseline.validation_dataset.close() + + @pytest.mark.parametrize("numerical_profile", ["default", "reference"]) def test_shared_recipe_matches_legacy_updates_and_restores_sampler( tmp_path, numerical_profile @@ -642,6 +769,48 @@ def test_official_loss_order_and_mask(): assert torch.count_nonzero(prediction.grad[:, 1]) == 0 +@pytest.mark.parametrize("masked", [False, True]) +def test_unnormalized_fp64_loss_backward_matches_analytic_oracle(masked): + prediction = torch.tensor( + [[[0.25, -0.5], [1.25, 0.125]]], dtype=torch.float64, requires_grad=True + ) + target = torch.tensor([[[0.125, -0.25], [0.5, 0.75]]], dtype=torch.float32) + mask = torch.tensor([[True, False]]) if masked else None + actual = normalized_mse(prediction, target, None, mask) + # Binary fractions yield exact independent values and derivatives: + # residuals=(1/8,-1/4,3/4,-5/8), gradient=2*residual/valid scalar count. + expected_loss = 5 / 128 if masked else 33 / 128 + expected_gradient = torch.tensor( + ( + [[[0.125, -0.25], [0.0, 0.0]]] + if masked + else [[[0.0625, -0.125], [0.375, -0.3125]]] + ), + dtype=torch.float64, + ) + assert actual.dtype == torch.float64 + assert actual.item() == expected_loss + actual.backward() + torch.testing.assert_close(prediction.grad, expected_gradient, atol=0, rtol=0) + assert target.dtype == torch.float32 + + +def test_bfloat16_predictions_do_not_quantize_float32_targets(): + prediction = torch.tensor([[1.0, 2.0]], dtype=torch.bfloat16, requires_grad=True) + # Both differences disappear if labels are incorrectly converted to BF16. + target = torch.tensor([[1.0 + 2**-12, 2.0 + 2**-10]], dtype=torch.float32) + oracle = prediction.detach().float().requires_grad_() + expected = (oracle - target).square().mean() + expected.backward() + with torch.autocast("cpu", dtype=torch.bfloat16): + actual = normalized_mse(prediction, target, None) + assert actual.dtype == torch.float32 + assert actual.item() > 0 + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + actual.backward() + torch.testing.assert_close(prediction.grad, oracle.grad.bfloat16(), atol=0, rtol=0) + + @pytest.mark.parametrize( "device", ["cpu", pytest.param("cuda:0", marks=pytest.mark.gpu)] ) diff --git a/tests/learning/models/test_nerd.py b/tests/learning/models/test_nerd.py index a3f6c29d41..577b3d6e3e 100644 --- a/tests/learning/models/test_nerd.py +++ b/tests/learning/models/test_nerd.py @@ -63,6 +63,65 @@ def test_normalizers_and_parameters_follow_dtype() -> None: ) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +@pytest.mark.parametrize("normalize_input", [False, True]) +@pytest.mark.parametrize("normalize_output", [False, True]) +def test_float32_features_follow_model_dtype_without_changing_math_or_gradients( + dtype: torch.dtype, normalize_input: bool, normalize_output: bool +) -> None: + torch.set_num_threads(2) + model = tiny().to(dtype=dtype).eval() + model.cfg.normalize_input = normalize_input + model.cfg.normalize_output = normalize_output + model.set_normalization( + { + name: ( + torch.full((width,), 1 / 7, dtype=dtype), + torch.full((width,), 1 / 3, dtype=dtype), + ) + for name, width in model.cfg.input_dims.items() + }, + torch.full((4,), 1 / 11, dtype=dtype), + torch.full((4,), 1 / 13, dtype=dtype), + ) + inputs = { + name: torch.randn(2, 3, width, requires_grad=True) + for name, width in model.cfg.input_dims.items() + } + references = { + name: value.detach().clone().requires_grad_() for name, value in inputs.items() + } + original = {name: value.detach().clone() for name, value in inputs.items()} + buffers = {name: value.clone() for name, value in model.named_buffers()} + # Independent old network/normalizer math with an explicit boundary cast. + values = [] + for name in model.cfg.input_dims: + value = references[name].to(dtype=dtype) + if normalize_input: + normalizer = model.input_rms[name] + value = (value - normalizer.mean) / torch.sqrt(normalizer.var + 1e-5) + values.append(value) + expected = model.model(model.transformer_model(torch.cat(values, -1))) + if model.cfg.normalize_output: + expected = ( + expected * torch.sqrt(model.output_rms.var + 1e-5) + model.output_rms.mean + ) + actual = model(inputs) + assert actual.dtype == dtype + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + actual.square().mean().backward() + gradients = {name: value.grad.clone() for name, value in model.named_parameters()} + model.zero_grad(set_to_none=True) + expected.square().mean().backward() + for name, value in model.named_parameters(): + torch.testing.assert_close(value.grad, gradients[name], atol=0, rtol=0) + for name, value in inputs.items(): + torch.testing.assert_close(value.grad, references[name].grad, atol=0, rtol=0) + torch.testing.assert_close(value, original[name], atol=0, rtol=0) + for name, value in model.named_buffers(): + torch.testing.assert_close(value, buffers[name], atol=0, rtol=0) + + def test_variable_prefix_matches_individual_unpadded_calls() -> None: model = tiny().eval() inputs = {k: torch.randn(3, 10, d) for k, d in model.cfg.input_dims.items()} From 048d1474ca6c3a72f1303850864ff7b24c568beb Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 25 Oct 2026 14:10:17 +0000 Subject: [PATCH 18/21] refactor(models): move NeRD into extension and downloadable showcase --- .github/workflows/main.yml | 2 + agent_context/MAP.yaml | 29 +- .../topics/data-pipeline/data-pipeline.md | 2 +- .../topics/learning-layers/learning-layers.md | 12 +- .../topics/neural-dynamics/neural-dynamics.md | 139 ++++---- .../supervised-learning.md | 10 +- docs/architecture/curated.json | 109 +++++- docs/scripts/build_architecture.py | 2 + docs/scripts/check_api_docs.py | 4 + .../embodichain/embodichain.data_pipeline.rst | 1 - .../embodichain.learning.dynamics.sl.rst | 38 -- .../embodichain/embodichain.learning.rst | 3 +- .../embodichain/embodichain.learning.sl.rst | 12 +- docs/source/api_reference/index.rst | 8 +- docs/source/api_reference/model_catalogs.rst | 254 +++++++++++++ docs/source/api_reference/public_api.rst | 237 ------------ docs/source/conf.py | 3 + docs/source/features/neural_dynamics.md | 236 ++++++------ docs/source/guides/neural_dynamics.md | 7 +- embodichain/cli/main.py | 5 - embodichain/data/nerd/ATTRIBUTION.md | 21 -- embodichain/data/nerd/LICENSE.md | 176 --------- embodichain/data/nerd/ant.xml | 92 ----- embodichain/data/nerd/cartpole_single.urdf | 82 ----- embodichain/data_pipeline/__init__.py | 22 +- embodichain/learning/__init__.py | 2 +- examples/README.md | 8 + examples/learning/nerd/README.md | 91 +++++ .../learning/nerd}/__init__.py | 2 +- .../learning/nerd/__main__.py | 113 +++++- examples/learning/nerd/_resources.py | 337 ++++++++++++++++++ .../learning/nerd/_support}/__init__.py | 0 .../learning/nerd/_support/_acceptance.py | 111 ++++-- .../_support/_benchmark_dynamics_input.py | 10 +- .../_support/_benchmark_dynamics_sampling.py | 4 +- .../_support/_benchmark_dynamics_training.py | 8 +- .../nerd/_support/_compare_official_model.py | 12 +- .../learning/nerd/_support/_comparison.py | 33 +- .../learning/nerd/_support}/_evaluation.py | 0 .../learning/nerd/_support/_generation.py | 4 +- .../nerd/_support}/_historical_physics.py | 2 +- .../nerd/_support/_matched_small_budget.py | 10 +- .../learning/nerd/_support}/_newton.py | 12 +- .../_support/_official_evaluation_data.py | 60 +++- .../nerd/_support/_official_math_rollout.py | 11 +- .../_support/_official_model_reference.py | 0 .../nerd/_support/_official_small_budget.py | 0 .../_official_small_budget_evaluate.py | 0 .../_official_small_budget_rollout.py | 5 +- .../nerd/_support}/_portable_reference.py | 4 +- .../learning/nerd/_support/_qualification.py | 2 +- .../learning/nerd/_support/_reference.py | 22 +- .../nerd/_support}/_reference_assets.py | 6 +- .../nerd/_support/_reference_rollout.py | 15 +- .../nerd/_support}/_reference_traces.py | 6 +- .../learning/nerd/_support/_report.py | 8 +- .../learning/nerd/_support}/_runtime.py | 0 .../learning/nerd/_support}/_sample_order.py | 0 .../learning/nerd/_support/_task_values.py | 2 +- .../learning/nerd/_support/_trainer.py | 279 +-------------- examples/learning/nerd/_support/_worker.py | 61 ++++ .../learning/nerd/_support/_workflow.py | 90 ++++- examples/learning/nerd/configs/ant/README.md | 40 +++ .../nerd/configs/ant/dynamics.eval.full.yaml | 8 + .../configs/ant/dynamics.eval.pretrained.yaml | 8 + .../nerd/configs/ant/dynamics.eval.smoke.yaml | 8 + .../configs}/ant/dynamics.generate.full.yaml | 2 +- .../configs}/ant/dynamics.generate.smoke.yaml | 2 +- .../ant/dynamics.generate.warp-reference.yaml | 2 +- .../configs}/ant/dynamics.train.full.yaml | 4 +- .../configs}/ant/dynamics.train.smoke.yaml | 4 +- .../learning/nerd/configs}/ant/sl.train.yaml | 2 +- .../learning/nerd/configs/cartpole/README.md | 41 +++ .../configs/cartpole/dynamics.eval.full.yaml | 8 + .../cartpole/dynamics.eval.pretrained.yaml | 8 + .../configs/cartpole/dynamics.eval.smoke.yaml | 8 + .../cartpole/dynamics.generate.full.yaml | 2 +- .../cartpole/dynamics.generate.smoke.yaml | 2 +- .../dynamics.generate.warp-reference.yaml | 2 +- .../cartpole/dynamics.train.full.yaml | 4 +- .../cartpole/dynamics.train.smoke.yaml | 4 +- .../nerd/configs}/cartpole/sl.train.yaml | 2 +- examples/learning/nerd/research/README.md | 65 ++++ .../learning/nerd/research/acceptance.md | 58 +-- extensions/embodichain_models/README.md | 25 ++ .../embodichain_models}/__init__.py | 13 +- .../embodichain_models}/dynamics/__init__.py | 24 +- .../dynamics}/nerd/THIRD_PARTY_LICENSES | 0 .../dynamics/nerd/__init__.py | 58 +++ .../dynamics/nerd}/_numerics.py | 0 .../dynamics/nerd}/_resident.py | 0 .../dynamics}/nerd/checkpoint.py | 0 .../dynamics/nerd}/codec.py | 0 .../embodichain_models/dynamics/nerd/data.py | 0 .../dynamics/nerd/integration.py | 299 ++++++++++++++++ .../dynamics}/nerd/model.py | 0 .../dynamics/nerd}/sl_recipe.py | 12 +- .../dynamics/nerd}/supervised.py | 7 +- .../tasks/classic_control/ant/README.md | 25 -- .../ant/dynamics.eval.full.yaml | 8 - .../ant/dynamics.eval.pretrained.yaml | 8 - .../ant/dynamics.eval.smoke.yaml | 8 - .../tasks/classic_control/cartpole/README.md | 25 -- .../cartpole/dynamics.eval.full.yaml | 8 - .../cartpole/dynamics.eval.pretrained.yaml | 8 - .../cartpole/dynamics.eval.smoke.yaml | 8 - scripts/experiments/neural_dynamics/README.md | 51 --- setup.py | 10 +- tests/ci/test_workflows.py | 10 +- tests/docs/test_build_architecture.py | 27 +- tests/docs/test_check_api_docs.py | 29 +- tests/examples/learning/__init__.py | 21 ++ tests/examples/learning/nerd/__init__.py | 21 ++ .../learning/nerd/test_acceptance.py} | 4 +- .../learning/nerd}/test_comparison.py | 35 +- .../nerd}/test_dynamics_runtime_identity.py | 4 +- .../learning/nerd/test_entrypoint.py} | 152 ++++---- .../learning/nerd}/test_generation.py | 20 +- .../learning/nerd/test_metrics.py} | 4 +- .../learning/nerd}/test_reference.py | 122 ++++--- .../learning/nerd}/test_reference_rollout.py | 22 +- .../learning/nerd}/test_reference_traces.py | 4 +- .../learning/nerd}/test_report.py | 16 +- .../examples/learning/nerd/test_resources.py | 212 +++++++++++ .../learning/nerd}/test_sample_order.py | 4 +- .../nerd}/test_small_budget_observation.py | 6 +- .../learning/nerd/test_task_values.py} | 2 +- tests/models/__init__.py | 21 ++ tests/models/dynamics/__init__.py | 21 ++ tests/models/dynamics/nerd/__init__.py | 21 ++ .../dynamics/nerd}/_nerd_reference_check.py | 4 +- .../nerd}/data/ant-official-frame.json | 0 .../nerd}/data/ant-official-frame.npz | Bin .../dynamics/nerd}/test_codec.py | 2 +- .../dynamics/nerd/test_data.py} | 4 +- .../dynamics/nerd/test_model.py} | 10 +- .../dynamics/nerd/test_training.py} | 63 ++-- tests/test_benchmark_package.py | 1 + tests/test_example_package.py | 1 + tests/test_main.py | 1 + tests/test_models_package.py | 150 ++++++++ tests/test_release_metadata.py | 2 + 142 files changed, 2950 insertions(+), 1783 deletions(-) delete mode 100644 docs/source/api_reference/embodichain/embodichain.learning.dynamics.sl.rst create mode 100644 docs/source/api_reference/model_catalogs.rst delete mode 100644 embodichain/data/nerd/ATTRIBUTION.md delete mode 100644 embodichain/data/nerd/LICENSE.md delete mode 100644 embodichain/data/nerd/ant.xml delete mode 100644 embodichain/data/nerd/cartpole_single.urdf create mode 100644 examples/learning/nerd/README.md rename {embodichain/learning/dynamics => examples/learning/nerd}/__init__.py (93%) rename embodichain/cli/dynamics.py => examples/learning/nerd/__main__.py (64%) create mode 100644 examples/learning/nerd/_resources.py rename {embodichain/data_pipeline/dynamics => examples/learning/nerd/_support}/__init__.py (100%) rename scripts/experiments/supervised/nerd_acceptance.py => examples/learning/nerd/_support/_acceptance.py (84%) rename scripts/benchmark/learning/dynamics_input.py => examples/learning/nerd/_support/_benchmark_dynamics_input.py (96%) rename scripts/benchmark/learning/dynamics_sampling.py => examples/learning/nerd/_support/_benchmark_dynamics_sampling.py (97%) rename scripts/benchmark/learning/dynamics_training.py => examples/learning/nerd/_support/_benchmark_dynamics_training.py (97%) rename scripts/experiments/neural_dynamics/compare_official_model.py => examples/learning/nerd/_support/_compare_official_model.py (96%) rename embodichain/learning/dynamics/comparison.py => examples/learning/nerd/_support/_comparison.py (95%) rename {embodichain/compute/dynamics => examples/learning/nerd/_support}/_evaluation.py (100%) rename embodichain/data_pipeline/dynamics/generation.py => examples/learning/nerd/_support/_generation.py (99%) rename {scripts/experiments/neural_dynamics => examples/learning/nerd/_support}/_historical_physics.py (99%) rename scripts/experiments/neural_dynamics/matched_small_budget.py => examples/learning/nerd/_support/_matched_small_budget.py (97%) rename {embodichain/data_pipeline/dynamics => examples/learning/nerd/_support}/_newton.py (97%) rename scripts/experiments/neural_dynamics/official_evaluation_data.py => examples/learning/nerd/_support/_official_evaluation_data.py (90%) rename scripts/experiments/neural_dynamics/official_math_rollout.py => examples/learning/nerd/_support/_official_math_rollout.py (97%) rename scripts/experiments/neural_dynamics/official_model_reference.py => examples/learning/nerd/_support/_official_model_reference.py (100%) rename scripts/experiments/neural_dynamics/official_small_budget.py => examples/learning/nerd/_support/_official_small_budget.py (100%) rename scripts/experiments/neural_dynamics/official_small_budget_evaluate.py => examples/learning/nerd/_support/_official_small_budget_evaluate.py (100%) rename scripts/experiments/neural_dynamics/official_small_budget_rollout.py => examples/learning/nerd/_support/_official_small_budget_rollout.py (98%) rename {scripts/experiments/neural_dynamics => examples/learning/nerd/_support}/_portable_reference.py (98%) rename embodichain/learning/dynamics/qualification.py => examples/learning/nerd/_support/_qualification.py (99%) rename embodichain/learning/dynamics/reference.py => examples/learning/nerd/_support/_reference.py (98%) rename {embodichain/data_pipeline/dynamics => examples/learning/nerd/_support}/_reference_assets.py (97%) rename embodichain/learning/dynamics/reference_rollout.py => examples/learning/nerd/_support/_reference_rollout.py (98%) rename {embodichain/learning/dynamics => examples/learning/nerd/_support}/_reference_traces.py (97%) rename embodichain/learning/dynamics/report.py => examples/learning/nerd/_support/_report.py (99%) rename {embodichain/learning/dynamics => examples/learning/nerd/_support}/_runtime.py (100%) rename {embodichain/learning/dynamics => examples/learning/nerd/_support}/_sample_order.py (100%) rename embodichain/compute/dynamics/reference.py => examples/learning/nerd/_support/_task_values.py (98%) rename embodichain/learning/dynamics/trainer.py => examples/learning/nerd/_support/_trainer.py (73%) create mode 100644 examples/learning/nerd/_support/_worker.py rename embodichain/learning/dynamics/workflow.py => examples/learning/nerd/_support/_workflow.py (73%) create mode 100644 examples/learning/nerd/configs/ant/README.md create mode 100644 examples/learning/nerd/configs/ant/dynamics.eval.full.yaml create mode 100644 examples/learning/nerd/configs/ant/dynamics.eval.pretrained.yaml create mode 100644 examples/learning/nerd/configs/ant/dynamics.eval.smoke.yaml rename {extensions/embodichain_tasks/configs/tasks/classic_control => examples/learning/nerd/configs}/ant/dynamics.generate.full.yaml (85%) rename {extensions/embodichain_tasks/configs/tasks/classic_control => examples/learning/nerd/configs}/ant/dynamics.generate.smoke.yaml (85%) rename {extensions/embodichain_tasks/configs/tasks/classic_control => examples/learning/nerd/configs}/ant/dynamics.generate.warp-reference.yaml (85%) rename {extensions/embodichain_tasks/configs/tasks/classic_control => examples/learning/nerd/configs}/ant/dynamics.train.full.yaml (87%) rename {extensions/embodichain_tasks/configs/tasks/classic_control => examples/learning/nerd/configs}/ant/dynamics.train.smoke.yaml (84%) rename {extensions/embodichain_tasks/configs/tasks/classic_control => examples/learning/nerd/configs}/ant/sl.train.yaml (92%) create mode 100644 examples/learning/nerd/configs/cartpole/README.md create mode 100644 examples/learning/nerd/configs/cartpole/dynamics.eval.full.yaml create mode 100644 examples/learning/nerd/configs/cartpole/dynamics.eval.pretrained.yaml create mode 100644 examples/learning/nerd/configs/cartpole/dynamics.eval.smoke.yaml rename {extensions/embodichain_tasks/configs/tasks/classic_control => examples/learning/nerd/configs}/cartpole/dynamics.generate.full.yaml (84%) rename {extensions/embodichain_tasks/configs/tasks/classic_control => examples/learning/nerd/configs}/cartpole/dynamics.generate.smoke.yaml (84%) rename {extensions/embodichain_tasks/configs/tasks/classic_control => examples/learning/nerd/configs}/cartpole/dynamics.generate.warp-reference.yaml (83%) rename {extensions/embodichain_tasks/configs/tasks/classic_control => examples/learning/nerd/configs}/cartpole/dynamics.train.full.yaml (83%) rename {extensions/embodichain_tasks/configs/tasks/classic_control => examples/learning/nerd/configs}/cartpole/dynamics.train.smoke.yaml (81%) rename {extensions/embodichain_tasks/configs/tasks/classic_control => examples/learning/nerd/configs}/cartpole/sl.train.yaml (91%) create mode 100644 examples/learning/nerd/research/README.md rename scripts/experiments/supervised/nerd.md => examples/learning/nerd/research/acceptance.md (65%) create mode 100644 extensions/embodichain_models/README.md rename {embodichain/learning/models/nerd => extensions/embodichain_models/embodichain_models}/__init__.py (69%) rename {embodichain/compute => extensions/embodichain_models/embodichain_models}/dynamics/__init__.py (63%) rename {embodichain/learning/models => extensions/embodichain_models/embodichain_models/dynamics}/nerd/THIRD_PARTY_LICENSES (100%) create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/nerd/__init__.py rename {embodichain/learning/dynamics => extensions/embodichain_models/embodichain_models/dynamics/nerd}/_numerics.py (100%) rename {embodichain/data_pipeline/dynamics => extensions/embodichain_models/embodichain_models/dynamics/nerd}/_resident.py (100%) rename {embodichain/learning/models => extensions/embodichain_models/embodichain_models/dynamics}/nerd/checkpoint.py (100%) rename {embodichain/compute/dynamics => extensions/embodichain_models/embodichain_models/dynamics/nerd}/codec.py (100%) rename embodichain/data_pipeline/dynamics/shards.py => extensions/embodichain_models/embodichain_models/dynamics/nerd/data.py (100%) create mode 100644 extensions/embodichain_models/embodichain_models/dynamics/nerd/integration.py rename {embodichain/learning/models => extensions/embodichain_models/embodichain_models/dynamics}/nerd/model.py (100%) rename {embodichain/learning/dynamics => extensions/embodichain_models/embodichain_models/dynamics/nerd}/sl_recipe.py (95%) rename {embodichain/learning/dynamics => extensions/embodichain_models/embodichain_models/dynamics/nerd}/supervised.py (93%) delete mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/ant/README.md delete mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.full.yaml delete mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.pretrained.yaml delete mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.smoke.yaml delete mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/README.md delete mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.full.yaml delete mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.pretrained.yaml delete mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.smoke.yaml delete mode 100644 scripts/experiments/neural_dynamics/README.md create mode 100644 tests/examples/learning/__init__.py create mode 100644 tests/examples/learning/nerd/__init__.py rename tests/{scripts/experiments/supervised/test_nerd_acceptance.py => examples/learning/nerd/test_acceptance.py} (94%) rename tests/{learning/dynamics => examples/learning/nerd}/test_comparison.py (84%) rename tests/{learning/dynamics => examples/learning/nerd}/test_dynamics_runtime_identity.py (97%) rename tests/{cli/test_dynamics.py => examples/learning/nerd/test_entrypoint.py} (80%) rename tests/{data_pipeline/dynamics => examples/learning/nerd}/test_generation.py (94%) rename tests/{compute/dynamics/test_dynamics_metrics.py => examples/learning/nerd/test_metrics.py} (92%) rename tests/{learning/dynamics => examples/learning/nerd}/test_reference.py (89%) rename tests/{learning/dynamics => examples/learning/nerd}/test_reference_rollout.py (96%) rename tests/{learning/dynamics => examples/learning/nerd}/test_reference_traces.py (98%) rename tests/{learning/dynamics => examples/learning/nerd}/test_report.py (96%) create mode 100644 tests/examples/learning/nerd/test_resources.py rename tests/{learning/dynamics => examples/learning/nerd}/test_sample_order.py (98%) rename tests/{learning/dynamics => examples/learning/nerd}/test_small_budget_observation.py (97%) rename tests/{compute/dynamics/test_dynamics_task_reference.py => examples/learning/nerd/test_task_values.py} (97%) create mode 100644 tests/models/__init__.py create mode 100644 tests/models/dynamics/__init__.py create mode 100644 tests/models/dynamics/nerd/__init__.py rename tests/{learning/models => models/dynamics/nerd}/_nerd_reference_check.py (99%) rename tests/{compute/dynamics => models/dynamics/nerd}/data/ant-official-frame.json (100%) rename tests/{compute/dynamics => models/dynamics/nerd}/data/ant-official-frame.npz (100%) rename tests/{compute/dynamics => models/dynamics/nerd}/test_codec.py (99%) rename tests/{data_pipeline/dynamics/test_shards.py => models/dynamics/nerd/test_data.py} (99%) rename tests/{learning/models/test_nerd.py => models/dynamics/nerd/test_model.py} (97%) rename tests/{learning/dynamics/test_trainer.py => models/dynamics/nerd/test_training.py} (95%) create mode 100644 tests/test_models_package.py diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index e20041da31..a5859e1c4d 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", } with ZipFile(wheels[0]) as wheel: members = set(wheel.namelist()) @@ -603,6 +604,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 64966ac3d6..18b44d5967 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -270,22 +270,23 @@ topics: status: active - id: neural-dynamics title: Neural Dynamics Models and Training - aliases: [neural dynamics, neural robot dynamics, NeRD, 神经动力学, 动力学模型训练] + aliases: [neural dynamics, neural robot dynamics, NeRD, 神经动力学, 动力学模型训练, NeRD 示例] keywords: [NeRDStateCodec, NeRDModel, DynamicsGenerationCfg, - DynamicsShardWriter, DynamicsWindowDataset, DynamicsTrainer, + DynamicsShardWriter, DynamicsWindowDataset, DynamicsTrainer, embodichain_models, + learning/nerd, nerd reference assets, 官方训练精度对比] paths: [topics/neural-dynamics/neural-dynamics.md] source_of_truth: - - embodichain/compute/dynamics/ - - embodichain/learning/models/nerd/ - - embodichain/learning/dynamics/ - - embodichain/data_pipeline/dynamics/ - - embodichain/cli/dynamics.py - watch_paths: [embodichain/compute/dynamics/, - embodichain/learning/models/nerd/, embodichain/learning/dynamics/, - embodichain/data_pipeline/dynamics/, - tests/compute/dynamics/, tests/learning/dynamics/, tests/data_pipeline/dynamics/] - related_topics: [learning-layers, data-pipeline] + - 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/, + tests/models/dynamics/nerd/, tests/examples/learning/nerd/, + tests/test_models_package.py, 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 @@ -502,7 +503,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 @@ -552,5 +553,5 @@ topics: - embodichain/learning/sl/ - embodichain/cli/main.py watch_paths: [embodichain/learning/sl/, tests/learning/sl/] - related_topics: [learning-layers, data-pipeline, rl-learning] + 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 5db12b3faa..612ba903f2 100644 --- a/agent_context/topics/data-pipeline/data-pipeline.md +++ b/agent_context/topics/data-pipeline/data-pipeline.md @@ -23,7 +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` | -| Dynamics teacher generation and HDF5 windows | `embodichain/data_pipeline/dynamics/`; contracts in [neural dynamics](../neural-dynamics/neural-dynamics.md) | +| 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 c0713cf071..e1d5ff3c23 100644 --- a/agent_context/topics/learning-layers/learning-layers.md +++ b/agent_context/topics/learning-layers/learning-layers.md @@ -9,11 +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. `learning/models/nerd/` owns the NeRD causal model and portable weights; -offline dynamics training is in `learning/dynamics/`. See -[neural dynamics](../neural-dynamics/neural-dynamics.md) for dataset and checkpoint -contracts. 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. diff --git a/agent_context/topics/neural-dynamics/neural-dynamics.md b/agent_context/topics/neural-dynamics/neural-dynamics.md index e220799479..f397bad326 100644 --- a/agent_context/topics/neural-dynamics/neural-dynamics.md +++ b/agent_context/topics/neural-dynamics/neural-dynamics.md @@ -2,81 +2,88 @@ ## Entry points and ownership -| Request | Owner (relative to `embodichain/`) | +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). + +| Request | Owner | |---|---| -| Tensor state, contact and history codecs | `compute/dynamics/codec.py` | -| Model and portable checkpoints | `learning/models/nerd/` | -| Analytical teacher generation | `data_pipeline/dynamics/generation.py`; private `_newton.py` owns its Newton model | -| Reference assets and physical profiles | `data_pipeline/dynamics/_reference_assets.py`; packaged files in `data/nerd/` | -| Committed shards and sequence windows | `data_pipeline/dynamics/shards.py` | -| Training, resume, evaluation and export | `learning/dynamics/trainer.py`, `workflow.py` | -| Shared SL objective and encoded batches | `learning/dynamics/supervised.py` | -| Common offline trainer recipe | `learning/dynamics/sl_recipe.py`; `embodichain train-sl` | -| Original-repository comparisons | `learning/dynamics/reference.py`, `comparison.py`, `report.py` | -| Command dispatch | `cli/dynamics.py`: `embodichain dynamics` | +| 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/` | -Model, codec, training and offline evaluation do not import simulation. Generation -uses Newton's built-in Featherstone/MuJoCo solver directly without DexSim, its -custom provider registry or Scene APIs. Native learned-solver deployment and -functional Gym integration are separate consumers, not dependencies of training. +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. -## Data and physical identity +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. -A trajectory contains T+1 states and T controls/contact inputs. Splits belong to -whole trajectories and windows cannot cross episodes or splits. Writers publish -only committed shards. The manifest commits generation cursor and rejection -evidence with the shard; rejection-only commits preserve data fingerprints. -Legacy progress sidecars remain readable, while collector source/physical -identity checks still apply. File handles and bounded caches belong to their -reader/worker. +## Resource and physical identity -The direct teacher stores its importer profile, asset/physical-array hashes, -solver settings and actual Newton/Warp versions. Its -`nerd-direct-newton-builder-v1` signature is distinct from Scene-built datasets. -Existing published Scene-generated data remains readable without simulation. -`warp_reference` reconstructs measured legacy asset parameters and the coupled -Ant Featherstone friction mapping; current Newton still owns integration. -Matching a profile name alone does not establish matching physics across -model-building paths, versions or solvers. +`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. -Newton COM velocities and NeRD twists differ in ordering and reference point. -The codec owns transport, relative targets use the frame at t, and frozen train -normalizers are reused for evaluation. Ant uses 25 ordered plane-contact slots; -Cartpole is contact-free. Controls are generalized efforts with zero free-base -force. Physical records cannot be relabeled through training manifest overrides. +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 collector keeps one bounded batch of history on the model device, then writes -owned CPU arrays. CUDA teachers use graph capture by default where supported; -CPU MuJoCo execution does not capture. Snapshots retain independent ownership. -Mass-matrix update intervals divide the macro substep count so resets can replay. +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. -## Training and validation +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. -`last.pt` preserves optimizer, scheduler, RNG, committed sample progress and one -owned CPU snapshot of the selected best model. New `best.pt` files are model-only -export snapshots; exact training recovery uses `last.pt`. Ordered input features -are part of recovery identity. Both trainers share run-directory/history admission -with [SL](../supervised-learning/supervised-learning.md). -Portable exports preserve saved floating precision, normalizers and attention -policy. Resume restores the same run; fine-tuning starts a new one. -Only explicitly trusted legacy conversion may deserialize Python model objects. +## Training and validation -`DynamicsTrainer` delegates optimizer updates to `learning.sl.SupervisedLearner` -while retaining its historical sampling, numerical profiles, LR schedule and -training-checkpoint format. The standalone `sl_recipe.build_supervised_trainer` -instead composes the common trainer and uses SL checkpoints. Both reuse the -same codec and normalized physical-target objective. Models and datasets do not -depend on either trainer. The paired acceptance runner under -`scripts/experiments/supervised/nerd_acceptance.py` compares independent reference -and candidate processes; fixed-budget parity is separate from full convergence. +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. -Optional shared sample-order/statistics artifacts support controlled original -trainer comparisons; equal seeds alone do not prove identical batches or physics. -Original Warp references run in isolated supported environments. Keep virtualenv -executable paths intact instead of resolving their interpreter symlinks. +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. -Tests cover codecs, model conversion, bounded shards, reset/replay and partial -collection batches, training/resume and the actual offline CLI. Public usage is -in [Neural Dynamics Models](../../../docs/source/features/neural_dynamics.md). -Fixed-set transition error, live rollout error and paper protocol accuracy are -separate claims; this workflow does not establish their equivalence. +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. Tests are in `tests/models/dynamics/nerd/`, +`tests/examples/learning/nerd/`, `tests/test_models_package.py` and the common SL +suite. 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 index 098ce93ae8..982c189bd7 100644 --- a/agent_context/topics/supervised-learning/supervised-learning.md +++ b/agent_context/topics/supervised-learning/supervised-learning.md @@ -1,8 +1,10 @@ # Supervised learning (SL) `embodichain.learning.sl` owns common supervised optimizer updates and offline -training orchestration. Models remain in `learning.models`; numerical state and -geometry operations remain in `compute`; source storage belongs to data adapters. +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 @@ -17,6 +19,10 @@ geometry operations remain in `compute`; source storage belongs to data adapters - `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. diff --git a/docs/architecture/curated.json b/docs/architecture/curated.json index f5652ba5d1..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": [ @@ -3342,20 +3342,20 @@ }, { "id": "neural-dynamics", - "label": "NeRD offline workflow", + "label": "NeRD model extension", "kind": "package", "topic_ids": [ "neural-dynamics" ], - "summary": "Generates analytical trajectories and trains, evaluates and exports causal neural dynamics models.", - "details": "State/contact codecs live in compute/dynamics, bounded trajectory collection and HDF5 storage in data_pipeline/dynamics, NeRD networks in learning/models/nerd, and offline learning in learning/dynamics. The independent workflow uses built-in Newton teachers without initializing DexSim; live Scene solver integration is a separate consumer.", + "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": [ - "Model training and offline evaluation do not own a simulation World.", - "A collector profile records its physical domain; exported weights do not establish equivalence across teachers." + "The core does not import concrete model implementations.", + "Data/model/reference-asset payloads are external resources; numerical and physical identities remain explicit." ], "evidence": [ { - "path": "embodichain/learning/models/nerd/model.py", + "path": "extensions/embodichain_models/embodichain_models/dynamics/nerd/model.py", "symbol": "NeRDModel", "start_line": 201, "end_line": 201, @@ -3363,12 +3363,45 @@ } ], "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)", @@ -3377,7 +3410,7 @@ "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. Model-specific codecs, normalization and sequence unrolling remain with objectives and recipe factories. Training does not create a simulation World or data producer.", + "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." @@ -5739,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": [ @@ -5834,7 +5903,8 @@ "convex-decomposition", "scene-services", "neural-dynamics", - "supervised-learning" + "supervised-learning", + "nerd-showcase" ], "edge_ids": [ "adapter-calls-compiler", @@ -5920,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": [ { @@ -6051,7 +6123,8 @@ "observation-normalizer", "policy-evaluation", "neural-dynamics", - "supervised-learning" + "supervised-learning", + "nerd-showcase" ] }, { @@ -6361,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", @@ -6392,7 +6465,9 @@ "differentiable-trainer", "differentiable-collector", "observation-normalizer", - "supervised-learning" + "supervised-learning", + "neural-dynamics", + "nerd-showcase" ], "edge_ids": [ "cli-calls-training-assembly", @@ -6419,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": [ { @@ -6460,7 +6537,9 @@ "rl-routing", "learning-env", "policy-evaluation", - "supervised-learning" + "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.data_pipeline.rst b/docs/source/api_reference/embodichain/embodichain.data_pipeline.rst index 52476fc8ea..bfe2e2bf72 100644 --- a/docs/source/api_reference/embodichain/embodichain.data_pipeline.rst +++ b/docs/source/api_reference/embodichain/embodichain.data_pipeline.rst @@ -22,7 +22,6 @@ storage for LeRobot datasets on Python 3.10--3.12). depth_video engine recording - dynamics runtime Datasets diff --git a/docs/source/api_reference/embodichain/embodichain.learning.dynamics.sl.rst b/docs/source/api_reference/embodichain/embodichain.learning.dynamics.sl.rst deleted file mode 100644 index eecc44a544..0000000000 --- a/docs/source/api_reference/embodichain/embodichain.learning.dynamics.sl.rst +++ /dev/null @@ -1,38 +0,0 @@ -NeRD supervised-learning adapters -================================= - -NeRD uses the shared supervised optimizer update while retaining its physical -codec, normalized loss and numerical execution policy. The standalone recipe -connects immutable NeRD windows to the common trainer; its checkpoints use the -SL format. Historical NeRD training checkpoints remain owned by the existing -``dynamics`` workflow. - -Physical objective and batches ------------------------------- - -.. currentmodule:: embodichain.learning.dynamics.supervised - -.. autosummary:: - - NeRDObjective - make_nerd_batch - -.. autoclass:: NeRDObjective - :members: - -.. autofunction:: make_nerd_batch - -Offline source and training recipe ----------------------------------- - -.. currentmodule:: embodichain.learning.dynamics.sl_recipe - -.. autosummary:: - - NeRDWindowSource - build_supervised_trainer - -.. autoclass:: NeRDWindowSource - :members: - -.. autofunction:: build_supervised_trainer diff --git a/docs/source/api_reference/embodichain/embodichain.learning.rst b/docs/source/api_reference/embodichain/embodichain.learning.rst index cf20e4bfa9..662af80727 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.rst +++ b/docs/source/api_reference/embodichain/embodichain.learning.rst @@ -10,7 +10,6 @@ embodichain.learning models rl sl - dynamics Reinforcement Learning ---------------------- @@ -60,6 +59,8 @@ 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 diff --git a/docs/source/api_reference/embodichain/embodichain.learning.sl.rst b/docs/source/api_reference/embodichain/embodichain.learning.sl.rst index 490c9afd0b..47aa97a49d 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.sl.rst +++ b/docs/source/api_reference/embodichain/embodichain.learning.sl.rst @@ -124,10 +124,10 @@ embodichain.learning.sl.workflow .. autofunction:: cli -NeRD integration ----------------- +Model-specific recipes +---------------------- -.. toctree:: - :maxdepth: 1 - - embodichain.learning.dynamics.sl +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/api_reference/public_api.rst b/docs/source/api_reference/public_api.rst index 98eb579353..095a621c3d 100644 --- a/docs/source/api_reference/public_api.rst +++ b/docs/source/api_reference/public_api.rst @@ -22,56 +22,6 @@ decoding and compiling the constrained Task Program schema surface. compile_mllm_task_program decode_mllm_task_program -embodichain.compute.dynamics ----------------------------- - -Simulation-independent NeRD coordinate, contact and history transformations. -State derivatives retain their graphs; discrete contact and orientation decisions -use branch-local derivatives. - -.. currentmodule:: embodichain.compute.dynamics - -.. autosummary:: - - reference_task_values - NeRDStateCodec - quat_multiply - quat_rotate - quat_inverse - plane_contact_features - append_history - reset_history - -embodichain.compute.dynamics.codec ----------------------------------- - -Convert public Newton COM velocities to NeRD spatial twists, encode relative -targets in the current body frame, and update per-world recurrent windows. - -.. currentmodule:: embodichain.compute.dynamics.codec - -.. autosummary:: - - NeRDStateCodec - quat_multiply - quat_rotate - quat_inverse - plane_contact_features - append_history - reset_history - -embodichain.compute.dynamics.reference --------------------------------------- - -Shared Cartpole and Ant observation, reward and failure calculations used by -registered tasks and rollout comparison, with explicit quaternion validation. - -.. currentmodule:: embodichain.compute.dynamics.reference - -.. autosummary:: - - reference_task_values - embodichain.data.assets.demo_assets ----------------------------------- @@ -287,37 +237,6 @@ embodichain.data_pipeline.depth_video DEPTH_MILLIMETER_UNIT DEPTH_QMAX -embodichain.data_pipeline.dynamics.generation ---------------------------------------------- - -Generate bounded batches of Cartpole/Ant trajectories with Newton's built-in -analytical solvers. This offline collector does not initialize DexSim or register -a learned simulation solver. Committed shards support resumable generation. - -.. currentmodule:: embodichain.data_pipeline.dynamics.generation - -.. autosummary:: - - DynamicsGenerationCfg - generate_dataset - -embodichain.data_pipeline.dynamics.shards ------------------------------------------ - -Committed HDF5 trajectory shards and bounded sequence-window reads. State -sequences have T+1 entries, controls have T; windows respect trajectory and split -boundaries. Resume uses only manifest-committed shards. - -.. currentmodule:: embodichain.data_pipeline.dynamics.shards - -.. autosummary:: - - DynamicsShardWriter - DynamicsWindowDataset - trajectory_split - load_manifest - dataset_fingerprint - embodichain.gen_sim.env ----------------------- @@ -1736,162 +1655,6 @@ embodichain.lab.sim.motion.planners.trapezoidal_planner TrapezoidalPlanner TrapezoidalPlannerCfg -embodichain.learning.dynamics.comparison ----------------------------------------- - -Prepare and execute matched-data trainer comparisons with immutable source, -initialization, statistics and sample-order identities. Completion status reflects -which experiments actually ran. - -.. currentmodule:: embodichain.learning.dynamics.comparison - -.. autosummary:: - - prepare_comparison - run_comparison - run_candidate - -embodichain.learning.dynamics.qualification -------------------------------------------- - -Compare independent forward/gradient/update snapshots and evaluate paired -seed-level non-inferiority. Fewer than five seeds remains pending qualification. - -.. currentmodule:: embodichain.learning.dynamics.qualification - -.. autosummary:: - - capture_training_step - compare_training_steps - paired_noninferiority - -embodichain.learning.dynamics.reference ---------------------------------------- - -Explicit adapters for a trusted, pinned official source checkout. Reference -execution is isolated from the production runtime; record compatibility transport -and source identity rather than relabeling migrated code as an official baseline. - -.. currentmodule:: embodichain.learning.dynamics.reference - -.. autosummary:: - - OFFICIAL_COMMIT - reference_identity - run_reference - export_legacy_dataset - import_legacy_datasets - generate_official_dataset - repair_gravity_transport - -embodichain.learning.dynamics.reference_rollout ------------------------------------------------ - -Evaluate portable weights in the pinned original Warp environment and neural -integrator. This explicitly trusted, isolated reference process measures -historical-domain training quality; it does not deploy historical physics into -current Newton. Physical provenance, shared teacher traces and portable-network -equivalence are retained with failure and policy-return evidence. - -.. currentmodule:: embodichain.learning.dynamics.reference_rollout - -.. autosummary:: - - evaluate_reference_rollouts - -embodichain.learning.dynamics.report ------------------------------------- - -Write machine-readable and human-readable comparison evidence with explicit -complete, failed and pending qualification status. - -.. currentmodule:: embodichain.learning.dynamics.report - -.. autosummary:: - - write_comparison_report - annotate_rollout_evidence - -embodichain.learning.dynamics.trainer -------------------------------------- - -Single-device causal dynamics training with frozen train-only normalization, -committed sampling progress, optimizer/RNG resume and independent evaluation. - -.. currentmodule:: embodichain.learning.dynamics.trainer - -.. autosummary:: - - DynamicsTrainCfg - DynamicsTrainer - encode_batch - normalized_mse - compute_statistics - load_statistics - save_statistics - evaluate_model - -embodichain.learning.dynamics.workflow --------------------------------------- - -Offline train, evaluate and export workflows validate dataset/model contracts -without importing or initializing the simulation runtime. - -.. currentmodule:: embodichain.learning.dynamics.workflow - -.. autosummary:: - - run_train - run_eval - run_export - -embodichain.learning.models.nerd --------------------------------- - -Portable causal NeRD networks, registered normalization buffers and versioned -weights. Evaluation preserves gradients; legacy Python-object conversion requires -explicit trust and is separate from normal runtime loading. - -.. currentmodule:: embodichain.learning.models.nerd - -.. autosummary:: - - NeRDModelCfg - NeRDModel - FeatureNormalizer - export_model - load_model - convert_legacy_checkpoint - -embodichain.learning.models.nerd.checkpoint -------------------------------------------- - -Export tensor state and JSON metadata with checksums; load portable models or -convert an explicitly trusted official Transformer checkpoint. - -.. currentmodule:: embodichain.learning.models.nerd.checkpoint - -.. autosummary:: - - export_model - load_model - convert_legacy_checkpoint - -embodichain.learning.models.nerd.model --------------------------------------- - -Causal Transformer architecture matching official NeRD parameter names. -Right-padded batches retain valid-prefix positional semantics and per-row output -selection; normalizer statistics move with the network. - -.. currentmodule:: embodichain.learning.models.nerd.model - -.. autosummary:: - - NeRDModelCfg - NeRDModel - FeatureNormalizer - embodichain.utils.warp.kinematics.trapezoidal_warp -------------------------------------------------- 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 index bdd45421b4..2300b89876 100644 --- a/docs/source/features/neural_dynamics.md +++ b/docs/source/features/neural_dynamics.md @@ -1,15 +1,16 @@ # Neural Dynamics Models -A neural dynamics model predicts how a robot's physical state changes under -applied joint efforts and contact. EmbodiChain provides a simulation-independent NeRD model, trajectory -storage, training, offline evaluation and portable export workflow. Data generation -uses Newton's built-in analytical solvers directly and does not require a custom -DexSim build. - -[NeRD (Neural Robot Dynamics)](https://neural-robot-dynamics.github.io/) is the -first implementation. Its causal Transformer consumes a short history of state, -efforts and contact features, then predicts the next state. The current reference -systems are contact-free Cartpole and free-base Ant on a static plane. +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 @@ -43,156 +44,131 @@ or [Ant video](https://huggingface.co/DexForceAI/embodichain_model/resolve/main/ The recordings visualize the released reference models; they are not a claim of paper-level accuracy or equivalent behavior in arbitrary robot scenes. -## Published datasets and models +## Pinned resources -Research datasets live in -[DexForceAI/embodichain_datasets](https://huggingface.co/datasets/DexForceAI/embodichain_datasets). -Deployable models live in -[DexForceAI/embodichain_model](https://huggingface.co/DexForceAI/embodichain_model). -Both repositories organize assets by model family, implementation and robot: +The showcase prepares these resources at fixed revisions: -| Reference system | Path in each repository | Deployment profile | +| Resource | Repository | Revision | | --- | --- | --- | -| Cartpole | `neural_dynamics/nerd/cartpole/v1` | Contact-free, `warp_reference`, 1/60 s | -| Ant | `neural_dynamics/nerd/ant/v1` | Static plane, `warp_reference`, 1/60 s | +| 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` | -Dataset directories contain a manifest and HDF5 shards with explicit -train/validation/test trajectory splits. Each trajectory stores T+1 states and T -efforts/contact inputs. Windows cannot cross trajectory or split boundaries. -Model directories contain `weights.pt` and `manifest.json`, including network -configuration, normalization buffers, codec, timestep and physical identity. -These are inference exports; resumable training checkpoints are local training -outputs, not part of the published model package. +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. -The dataset and model cards describe the generating solver, supported dependency -versions, data size and evaluation limits. In particular, long-horizon Ant -predictions can drift. Assess the model against its matching analytical solver -for the intended motion and prediction horizon before using it downstream. +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. -## Download and evaluate +## Prepare and evaluate -Run from the repository root. The Hugging Face client is an EmbodiChain -dependency. Download one matching dataset/model pair: +These commands also work from an installed wheel outside the checkout: -```python -from huggingface_hub import snapshot_download - -robot = "cartpole" # or "ant" -prefix = f"neural_dynamics/nerd/{robot}/v1" -snapshot_download( - repo_id="DexForceAI/embodichain_datasets", - repo_type="dataset", - allow_patterns=[f"{prefix}/*"], - local_dir="outputs/datasets", -) -snapshot_download( - repo_id="DexForceAI/embodichain_model", - allow_patterns=[f"{prefix}/*"], - local_dir="outputs/models", -) +```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 ``` -Evaluate the saved test transitions without starting a simulator: +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`. -```bash -embodichain dynamics eval --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.pretrained.yaml -``` +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. -The pretrained configuration reads the downloaded paths above. Replace -`cartpole` with `ant` for Ant. Use `--device cpu` or `--device cuda:0` to choose -the tensor device. Offline evaluation measures held-out transition prediction; -it does not measure accumulated rollout error. +## Train, recover and export -The released datasets were collected through the matching Newton/Scene -integration. The independent collector below builds a new Newton model directly; -its manifests use a separate `nerd-direct-newton-builder-v1` physical identity. -Both datasets share the same state/contact schema and can be trained and evaluated -here. Do not treat different model-building paths or solver versions as identical -physical systems, or use offline transition error as a live rollout measurement. +After preparation, train through the default common SL recipe: -## Generate data, train and export +```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 +``` -The workflow configurations are under -`extensions/embodichain_tasks/configs/tasks/classic_control//`. -They configure `embodichain dynamics` commands and are separate from runnable -Gym deployment configs. +`--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. -For a short local workflow, use the smoke presets: +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 dynamics generate --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.smoke.yaml -embodichain dynamics train --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.train.smoke.yaml -embodichain dynamics export --checkpoint outputs/dynamics/cartpole/smoke/train/best.pt --output outputs/dynamics/cartpole/smoke/export -embodichain dynamics eval --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.smoke.yaml +embodichain train-sl --config /path/to/recipe.yaml +embodichain train-sl --config /path/to/recipe.yaml --resume /path/to/train/last.pt ``` -Smoke presets use a small network and short budget to exercise the workflow. -Full presets use the reference network scale and a substantially longer budget. -Neither preset name is an accuracy qualification. - -Collection commits its candidate cursor and rejection records in the same -manifest update as the shard. Rejection-only progress does not change training -data identity. Historical progress-file layouts remain readable; continuing a -collection still requires the recorded collector source and physical identity. - -To train on a downloaded dataset, copy `dynamics.train.full.yaml`, then set -`dataset` to the downloaded robot's `v1` directory and `output_dir` to a new -training directory. Preserve its model feature dimensions and history length. -Run the copied configuration with `embodichain dynamics train --config PATH`, -then export its validation-selected `best.pt` with the `export` command. - -Training checkpoints retain model, optimizer, normalization, RNG and committed -sampling progress. Continue the same run with `train --resume PATH/last.pt`. -Initialize a new run from an inference export with `train --fine-tune PATH`. -Each run records its identity in `run.json`; fresh training and fine-tuning require -an artifact-free destination. Resuming into an occupied directory requires the -same run identity. If a destination has a different selected best model or newer -or incomplete metric history, use a fresh directory to preserve that evidence. -Historical checkpoints without a -run identity can also recover into a fresh directory. - -New `best.pt` files contain the validation-selected model and its own step, and -support export rather than exact optimizer recovery. `last.pt` retains the best -weights so recovery into another directory can recreate `best.pt` even without -later validation improvement. Export and reload preserve the saved floating -dtype and normalization. Input feature order must remain unchanged when resuming. -Export destinations must be new directories. Paths written in configuration -files resolve relative to the file; explicit CLI path overrides resolve relative -to the working directory. - -The same NeRD model/objective also uses the shared supervised trainer: +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 train-sl --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/sl.train.yaml -embodichain train-sl --config extensions/embodichain_tasks/configs/tasks/classic_control/ant/sl.train.yaml +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 ``` -These bounded recipes use the common SL checkpoint format and resume through -`train-sl --resume PATH`. The historical `dynamics` CLI keeps its own checkpoint -format; both use the same supervised optimizer update and physical-target loss. +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. -## Use the exported model in PyTorch +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. -Load tensor weights and their recorded normalization buffers: +## Use the model in PyTorch ```python -from embodichain.learning.models.nerd import load_model +from embodichain_models.dynamics.nerd import load_model model, manifest = load_model( - "outputs/models/neural_dynamics/nerd/cartpole/v1", device="cpu" + "outputs/nerd-resources/models/neural_dynamics/nerd/cartpole/v1", + device="cpu", ) ``` -`NeRDStateCodec` and `encode_batch` construct features from the dataset's state, -effort and contact fields. The model accepts feature windows and valid history -lengths; the training and evaluation workflows use this same preprocessing. - -Live Scene stepping, `solver_type: nerd`, and differentiable Gym integration are -provided separately by the [neural dynamics solver PR](https://github.com/DexForce/EmbodiChain/pull/822). -The comparison videos above were produced using that integration; they are -reference recordings, not a command available in this offline workflow. - -Keep state transforms in {mod}`embodichain.compute.dynamics`, model/checkpoint -code in {mod}`embodichain.learning.models.nerd`, training in -{mod}`embodichain.learning.dynamics.workflow`, and trajectory storage/generation in -{mod}`embodichain.data_pipeline.dynamics`. +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 index 74fcac64e4..c0d657ff21 100644 --- a/docs/source/guides/neural_dynamics.md +++ b/docs/source/guides/neural_dynamics.md @@ -4,5 +4,8 @@ orphan: true # Neural dynamics -The feature overview, published datasets and models, comparison videos, and -workflow commands are now in {doc}`/features/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/embodichain/cli/main.py b/embodichain/cli/main.py index 7c4522a3e6..0c90341580 100644 --- a/embodichain/cli/main.py +++ b/embodichain/cli/main.py @@ -117,11 +117,6 @@ class Command: target="embodichain.cli.dataset:main", help="Validate, split, and recover recorded expert datasets.", ), - Command( - name="dynamics", - target="embodichain.cli.dynamics:main", - help="Generate data, train and evaluate neural dynamics models.", - ), Command( name="preview_lerobot_data", target="embodichain.lab.scripts.preview_lerobot_data:cli", diff --git a/embodichain/data/nerd/ATTRIBUTION.md b/embodichain/data/nerd/ATTRIBUTION.md deleted file mode 100644 index b791343085..0000000000 --- a/embodichain/data/nerd/ATTRIBUTION.md +++ /dev/null @@ -1,21 +0,0 @@ -# NeRD reference robot assets - -`cartpole_single.urdf` and `ant.xml` are unmodified copies from -NVlabs/neural-robot-dynamics, commit -`476dc67fccc6a99be293bf611f897b1c06f13e30`, under -`envs/newton_envs/assets/`. The upstream repository license is reproduced in -`LICENSE.md` (Apache License 2.0; NVIDIA Corporation and affiliates, 2025). - -Source: https://github.com/NVlabs/neural-robot-dynamics/tree/476dc67fccc6a99be293bf611f897b1c06f13e30/envs/newton_envs/assets - -EmbodiChain interprets these assets using the installed Newton public importer. -The independent data collector finalizes an analytical Newton model directly; -Scene-based neural solver integration is a separate consumer of the same asset -interpretation. Neither path claims to reproduce legacy Warp 1.5.1 physics. - -The host supplies generalized forces and its own solver configuration rather -than installing the source motor actuators. Dataset manifests identify the -importer profile, asset hash, actual parameter hashes and solver versions. -`warp_reference` reconstructs measured legacy mass, inertia and contact-parameter -semantics; it does not select the historical simulator or guarantee identical -trajectories across solver versions. diff --git a/embodichain/data/nerd/LICENSE.md b/embodichain/data/nerd/LICENSE.md deleted file mode 100644 index 2bb9ad240f..0000000000 --- a/embodichain/data/nerd/LICENSE.md +++ /dev/null @@ -1,176 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS \ No newline at end of file diff --git a/embodichain/data/nerd/ant.xml b/embodichain/data/nerd/ant.xml deleted file mode 100644 index 04c237b0d7..0000000000 --- a/embodichain/data/nerd/ant.xml +++ /dev/null @@ -1,92 +0,0 @@ - - - - - - - - - - - - - \ No newline at end of file diff --git a/embodichain/data/nerd/cartpole_single.urdf b/embodichain/data/nerd/cartpole_single.urdf deleted file mode 100644 index 3f5981a495..0000000000 --- a/embodichain/data/nerd/cartpole_single.urdf +++ /dev/null @@ -1,82 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/embodichain/data_pipeline/__init__.py b/embodichain/data_pipeline/__init__.py index 9597ede6bd..8e3ea1615f 100644 --- a/embodichain/data_pipeline/__init__.py +++ b/embodichain/data_pipeline/__init__.py @@ -27,25 +27,9 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from . import ( - configuration, - datasets, - depth_video, - dynamics, - engine, - recording, - runtime, - ) - -__all__ = [ - "configuration", - "datasets", - "depth_video", - "dynamics", - "engine", - "recording", - "runtime", -] + from . import configuration, datasets, depth_video, engine, recording, runtime + +__all__ = ["configuration", "datasets", "depth_video", "engine", "recording", "runtime"] def __getattr__(name: str) -> ModuleType: diff --git a/embodichain/learning/__init__.py b/embodichain/learning/__init__.py index 3dba1a64db..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", "dynamics", "sl"] +__all__ = ["layers", "models", "rl", "sl"] def __getattr__(name: str) -> ModuleType: 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/embodichain/learning/dynamics/__init__.py b/examples/learning/nerd/__init__.py similarity index 93% rename from embodichain/learning/dynamics/__init__.py rename to examples/learning/nerd/__init__.py index b1f8e18a59..051f680fe4 100644 --- a/embodichain/learning/dynamics/__init__.py +++ b/examples/learning/nerd/__init__.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Simulation-independent neural dynamics workflows.""" +"""NeRD model and showcase support.""" from __future__ import annotations diff --git a/embodichain/cli/dynamics.py b/examples/learning/nerd/__main__.py similarity index 64% rename from embodichain/cli/dynamics.py rename to examples/learning/nerd/__main__.py index 6f98492a14..7bf671174f 100644 --- a/embodichain/cli/dynamics.py +++ b/examples/learning/nerd/__main__.py @@ -21,6 +21,7 @@ import argparse from collections.abc import Sequence import json +import os from pathlib import Path from typing import Any @@ -78,16 +79,27 @@ def _resolve_paths( def main(argv: Sequence[str] | None = None) -> None: - """Dispatch one dynamics operation; offline paths never import simulation.""" - parser = argparse.ArgumentParser(prog="embodichain dynamics", description=__doc__) + """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, required=True) + 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( @@ -161,13 +173,34 @@ def main(argv: Sequence[str] | None = None) -> None: ) 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 embodichain.learning.dynamics.workflow import run_export + from examples.learning.nerd._support._workflow import run_export result = run_export(checkpoint, output) else: @@ -180,15 +213,77 @@ def main(argv: Sequence[str] | None = None) -> None: raise ValueError( "Conversion manifest requires codec mapping and dt" ) - from embodichain.learning.models.nerd import convert_legacy_checkpoint + 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: - path = args.config.expanduser().resolve() + 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} @@ -213,19 +308,19 @@ def main(argv: Sequence[str] | None = None) -> None: ) config = _resolve_paths(config, path.parent, args.operation) if args.operation == "generate": - from embodichain.data_pipeline.dynamics.generation import ( + from examples.learning.nerd._support._generation import ( generate_dataset, ) result = generate_dataset(config) elif args.operation == "train": - from embodichain.learning.dynamics.workflow import run_train + from examples.learning.nerd._support._workflow import run_train result = run_train(config) else: mode = config.pop("mode", "offline") if mode == "offline": - from embodichain.learning.dynamics.workflow import run_eval + from examples.learning.nerd._support._workflow import run_eval result = run_eval(config) else: diff --git a/examples/learning/nerd/_resources.py b/examples/learning/nerd/_resources.py new file mode 100644 index 0000000000..7ca427a998 --- /dev/null +++ b/examples/learning/nerd/_resources.py @@ -0,0 +1,337 @@ +# ---------------------------------------------------------------------------- +# 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, + "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/embodichain/data_pipeline/dynamics/__init__.py b/examples/learning/nerd/_support/__init__.py similarity index 100% rename from embodichain/data_pipeline/dynamics/__init__.py rename to examples/learning/nerd/_support/__init__.py diff --git a/scripts/experiments/supervised/nerd_acceptance.py b/examples/learning/nerd/_support/_acceptance.py similarity index 84% rename from scripts/experiments/supervised/nerd_acceptance.py rename to examples/learning/nerd/_support/_acceptance.py index 583b15d09e..db2cae354c 100644 --- a/scripts/experiments/supervised/nerd_acceptance.py +++ b/examples/learning/nerd/_support/_acceptance.py @@ -17,7 +17,7 @@ """Compare independent NeRD trainers on immutable published data. Run: - python scripts/experiments/supervised/nerd_acceptance.py \ + 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 @@ -30,6 +30,7 @@ import argparse import hashlib +import importlib import json import os from pathlib import Path @@ -53,16 +54,67 @@ def _sha(path: Path) -> str: 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]) - modules = ( - "embodichain/learning/dynamics/trainer.py", - "embodichain/learning/dynamics/supervised.py", - "embodichain/learning/sl/learner.py", - "embodichain/learning/sl/types.py", - ) + 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(), @@ -81,21 +133,23 @@ def _worker(request_path: Path) -> None: and _source_identity(Path(request["source"])) != request["source_identity"] ): raise RuntimeError("Trainer source changed after protocol freeze") - sys.path.insert(0, request["source"]) + source = Path(request["source"]) + sys.path[:0] = _source_paths(source) import torch - from embodichain.compute.dynamics import NeRDStateCodec - from embodichain.data_pipeline.dynamics.shards import ( - DynamicsWindowDataset, - dataset_fingerprint, - load_manifest, - ) - from embodichain.learning.dynamics.trainer import ( - DynamicsTrainCfg, - DynamicsTrainer, - evaluate_model, - ) - from embodichain.learning.models.nerd import NeRDModel, load_model + 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() @@ -254,7 +308,7 @@ def _run(request: dict[str, Any]) -> dict[str, Any]: _write(request_path, request) environment = dict(os.environ) environment.update( - PYTHONPATH=request["source"], + 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"]), @@ -310,7 +364,7 @@ def main() -> None: 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[3] + "--candidate-source", type=Path, default=Path(__file__).resolve().parents[4] ) parser.add_argument("--datasets", type=Path) parser.add_argument("--models", type=Path) @@ -339,12 +393,13 @@ def main() -> None: ) import torch - sys.path.insert(0, str(args.candidate_source.resolve())) - from embodichain.data_pipeline.dynamics.shards import load_manifest - from embodichain.learning.dynamics.qualification import ( - compare_training_steps, - paired_noninferiority, - ) + 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 = { diff --git a/scripts/benchmark/learning/dynamics_input.py b/examples/learning/nerd/_support/_benchmark_dynamics_input.py similarity index 96% rename from scripts/benchmark/learning/dynamics_input.py rename to examples/learning/nerd/_support/_benchmark_dynamics_input.py index 89772bd6ff..3bfc12c962 100644 --- a/scripts/benchmark/learning/dynamics_input.py +++ b/examples/learning/nerd/_support/_benchmark_dynamics_input.py @@ -16,7 +16,7 @@ """Compare individual and grouped dynamics window I/O on real and stress data. -Run: python -m scripts.benchmark.learning.dynamics_input --dataset DATASET +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 @@ -70,7 +70,7 @@ def create_stress_dataset( ) -> Path: """Repeat one real trajectory into independent shards for bounded-I/O stress.""" import h5py - from embodichain.data_pipeline.dynamics.shards import ( + from embodichain_models.dynamics.nerd.data import ( load_manifest, DynamicsShardWriter, dataset_fingerprint, @@ -124,9 +124,9 @@ def benchmark_inputs( import numpy as np import torch from torch.utils.data import default_collate - from embodichain.compute.dynamics import NeRDStateCodec - from embodichain.data_pipeline.dynamics.shards import DynamicsWindowDataset - from embodichain.learning.dynamics.trainer import encode_batch + 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) diff --git a/scripts/benchmark/learning/dynamics_sampling.py b/examples/learning/nerd/_support/_benchmark_dynamics_sampling.py similarity index 97% rename from scripts/benchmark/learning/dynamics_sampling.py rename to examples/learning/nerd/_support/_benchmark_dynamics_sampling.py index 62360ee3a0..aa3c8837c8 100644 --- a/scripts/benchmark/learning/dynamics_sampling.py +++ b/examples/learning/nerd/_support/_benchmark_dynamics_sampling.py @@ -16,7 +16,7 @@ """Measure exact minibatch-order storage and replay with bounded RNG checkpoints. -Run: python -m scripts.benchmark.learning.dynamics_sampling --output-dir PATH +Run: python -m embodichain run-example learning/nerd/_support/_benchmark_dynamics_sampling --output-dir PATH """ from __future__ import annotations @@ -33,7 +33,7 @@ import psutil import torch -from embodichain.learning.dynamics._sample_order import ( +from examples.learning.nerd._support._sample_order import ( _open_order, _write_checkpointed_order, ) diff --git a/scripts/benchmark/learning/dynamics_training.py b/examples/learning/nerd/_support/_benchmark_dynamics_training.py similarity index 97% rename from scripts/benchmark/learning/dynamics_training.py rename to examples/learning/nerd/_support/_benchmark_dynamics_training.py index 69dd5e9c65..197ac5dd82 100644 --- a/scripts/benchmark/learning/dynamics_training.py +++ b/examples/learning/nerd/_support/_benchmark_dynamics_training.py @@ -16,7 +16,7 @@ """Measure official-width resident NeRD optimizer updates and budget lower bounds. -Run: python -m scripts.benchmark.learning.dynamics_training --source /path/to/main +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, @@ -55,9 +55,9 @@ def benchmark_training( if min(batch_size, history, warmup, repeats) < 1: raise ValueError("Workload sizes must be positive") import torch - from embodichain.learning.dynamics.reference import reference_identity - from embodichain.learning.dynamics.trainer import normalized_mse - from embodichain.learning.models.nerd import NeRDModel, NeRDModelCfg + 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(): diff --git a/scripts/experiments/neural_dynamics/compare_official_model.py b/examples/learning/nerd/_support/_compare_official_model.py similarity index 96% rename from scripts/experiments/neural_dynamics/compare_official_model.py rename to examples/learning/nerd/_support/_compare_official_model.py index 81404b1979..e8af2aa9c1 100644 --- a/scripts/experiments/neural_dynamics/compare_official_model.py +++ b/examples/learning/nerd/_support/_compare_official_model.py @@ -52,18 +52,18 @@ def main() -> None: import numpy as np import torch import warp as wp - from embodichain.compute.dynamics import NeRDStateCodec - from embodichain.learning.models.nerd import load_model - from embodichain.learning.dynamics.reference import ( + 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 embodichain.learning.dynamics.reference_rollout import ( + from examples.learning.nerd._support._reference_rollout import ( _cpu_stream_compatibility, ) - from embodichain.learning.dynamics.trainer import encode_batch - from embodichain.data_pipeline.dynamics.shards import DynamicsWindowDataset + from examples.learning.nerd._support._trainer import encode_batch + from embodichain_models.dynamics.nerd.data import DynamicsWindowDataset source = reference_identity(args.source) if ( diff --git a/embodichain/learning/dynamics/comparison.py b/examples/learning/nerd/_support/_comparison.py similarity index 95% rename from embodichain/learning/dynamics/comparison.py rename to examples/learning/nerd/_support/_comparison.py index 74b5f5414d..ac214e1812 100644 --- a/embodichain/learning/dynamics/comparison.py +++ b/examples/learning/nerd/_support/_comparison.py @@ -27,7 +27,7 @@ import sys from typing import Any -from .reference import OFFICIAL_COMMIT, _hash, reference_identity +from ._reference import OFFICIAL_COMMIT, _hash, reference_identity __all__ = ["prepare_comparison", "run_comparison", "run_candidate"] _TEACHERS = ("official_warp", "current_newton") @@ -55,10 +55,10 @@ def prepare_comparison(config: dict[str, Any]) -> Path: """ import numpy as np import torch - from embodichain.compute.dynamics import NeRDStateCodec - from embodichain.data_pipeline.dynamics.shards import DynamicsWindowDataset - from embodichain.learning.models.nerd import NeRDModel, NeRDModelCfg, export_model - from .trainer import ( + 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, @@ -210,12 +210,14 @@ def prepare_comparison(config: dict[str, Any]) -> Path: spec_path = seed_dir / f"{trainer}.json" _json(spec_path, spec, exclusive=True) artifacts[str(spec_path.relative_to(root))] = _hash(spec_path) - module = ( - "embodichain.learning.dynamics.reference" - if trainer == "official" - else "embodichain.learning.dynamics.comparison" - ) - command = [python, "-m", module, "--spec", str(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: @@ -256,6 +258,7 @@ def prepare_comparison(config: dict[str, Any]) -> Path: "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, @@ -303,11 +306,11 @@ def run_candidate( ) -> dict[str, Any]: """Execute the migrated trainer on the same frozen run specification.""" import torch - from embodichain.data_pipeline.dynamics.shards import ( + from embodichain_models.dynamics.nerd.data import ( load_manifest, dataset_fingerprint, ) - from .workflow import run_train, run_export, run_eval + 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") @@ -436,7 +439,9 @@ def run_comparison( if plan["reference_dependencies"]: pythonpaths.append(plan["reference_dependencies"]) pythonpaths.append(plan["source"]["root"]) - pythonpaths.append(str(Path(__file__).resolve().parents[3])) + root = Path(__file__).resolve().parents[4] + if (root / "pyproject.toml").is_file(): + pythonpaths.extend((str(root / "extensions/embodichain_models"), str(root))) if environment.get("PYTHONPATH"): pythonpaths.append(environment["PYTHONPATH"]) environment["PYTHONPATH"] = os.pathsep.join(pythonpaths) diff --git a/embodichain/compute/dynamics/_evaluation.py b/examples/learning/nerd/_support/_evaluation.py similarity index 100% rename from embodichain/compute/dynamics/_evaluation.py rename to examples/learning/nerd/_support/_evaluation.py diff --git a/embodichain/data_pipeline/dynamics/generation.py b/examples/learning/nerd/_support/_generation.py similarity index 99% rename from embodichain/data_pipeline/dynamics/generation.py rename to examples/learning/nerd/_support/_generation.py index 954b07baec..d9d8b6c74c 100644 --- a/embodichain/data_pipeline/dynamics/generation.py +++ b/examples/learning/nerd/_support/_generation.py @@ -32,8 +32,8 @@ import numpy as np import torch -from embodichain.compute.dynamics import quat_rotate -from embodichain.data_pipeline.dynamics.shards import DynamicsShardWriter +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 diff --git a/scripts/experiments/neural_dynamics/_historical_physics.py b/examples/learning/nerd/_support/_historical_physics.py similarity index 99% rename from scripts/experiments/neural_dynamics/_historical_physics.py rename to examples/learning/nerd/_support/_historical_physics.py index f75718a722..6090c5d8d6 100644 --- a/scripts/experiments/neural_dynamics/_historical_physics.py +++ b/examples/learning/nerd/_support/_historical_physics.py @@ -23,7 +23,7 @@ import math from pathlib import Path from typing import Any -from official_evaluation_data import _hash +from ._official_evaluation_data import _hash __all__: list[str] = [] diff --git a/scripts/experiments/neural_dynamics/matched_small_budget.py b/examples/learning/nerd/_support/_matched_small_budget.py similarity index 97% rename from scripts/experiments/neural_dynamics/matched_small_budget.py rename to examples/learning/nerd/_support/_matched_small_budget.py index ec4d2d8d35..e53584ae96 100644 --- a/scripts/experiments/neural_dynamics/matched_small_budget.py +++ b/examples/learning/nerd/_support/_matched_small_budget.py @@ -35,21 +35,21 @@ import numpy as np import torch -from embodichain.compute.dynamics import NeRDStateCodec -from embodichain.data_pipeline.dynamics.shards import ( +from embodichain_models.dynamics.nerd.codec import NeRDStateCodec +from embodichain_models.dynamics.nerd.data import ( DynamicsWindowDataset, dataset_fingerprint, load_manifest, ) -from embodichain.learning.dynamics.trainer import ( +from examples.learning.nerd._support._trainer import ( DynamicsTrainCfg, DynamicsTrainer, encode_batch, evaluate_model, normalized_mse, ) -from embodichain.learning.dynamics.workflow import run_export -from embodichain.learning.models.nerd import convert_legacy_checkpoint, load_model +from examples.learning.nerd._support._workflow import run_export +from embodichain_models.dynamics.nerd import convert_legacy_checkpoint, load_model __all__: list[str] = [] diff --git a/embodichain/data_pipeline/dynamics/_newton.py b/examples/learning/nerd/_support/_newton.py similarity index 97% rename from embodichain/data_pipeline/dynamics/_newton.py rename to examples/learning/nerd/_support/_newton.py index c09a147d27..2e684d2ea3 100644 --- a/embodichain/data_pipeline/dynamics/_newton.py +++ b/examples/learning/nerd/_support/_newton.py @@ -29,7 +29,7 @@ import torch import warp as wp -from embodichain.compute.dynamics import plane_contact_features, quat_rotate +from embodichain_models.dynamics.nerd.codec import plane_contact_features, quat_rotate from ._reference_assets import ( _ant_reference_friction, _import_reference, @@ -235,10 +235,14 @@ def __init__( "source_physical_array_sha256": arrays, "ground_material": ground, "collector_source_sha256": { - name: hashlib.sha256( - Path(__file__).with_name(name).read_bytes() + logical: hashlib.sha256( + Path(__file__).with_name(filename).read_bytes() ).hexdigest() - for name in ("_newton.py", "_reference_assets.py", "generation.py") + for logical, filename in ( + ("_newton.py", "_newton.py"), + ("_reference_assets.py", "_reference_assets.py"), + ("generation.py", "_generation.py"), + ) }, "gravity": gravity, "codec": { diff --git a/scripts/experiments/neural_dynamics/official_evaluation_data.py b/examples/learning/nerd/_support/_official_evaluation_data.py similarity index 90% rename from scripts/experiments/neural_dynamics/official_evaluation_data.py rename to examples/learning/nerd/_support/_official_evaluation_data.py index b409337763..c63c9b01a5 100644 --- a/scripts/experiments/neural_dynamics/official_evaluation_data.py +++ b/examples/learning/nerd/_support/_official_evaluation_data.py @@ -48,18 +48,61 @@ def _hash(path: Path) -> str: 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, directory / (name + ".py") + package.__name__ + "." + name, path ) module = importlib.util.module_from_spec(spec) sys.modules[spec.name] = module - spec.loader.exec_module(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) @@ -206,11 +249,10 @@ def main() -> None: "Use supported Python3.8.20/Torch2.2.2/NumPy1.24.4 and robot-pinned Warp" ) torch.set_num_threads(1) - transport = Path( - cfg.get( - "transport_root", - Path(__file__).resolve().parents[3] / "embodichain/compute/dynamics", - ) + transport = ( + Path(cfg["transport_root"]) + if "transport_root" in cfg + else _default_transport_root() ) math_module, observations = _transport(transport) set_random_seed(seed) @@ -367,8 +409,8 @@ def main() -> None: }, "canonical_transport_modules": { name: { - "path": str(transport / (name + ".py")), - "sha256": _hash(transport / (name + ".py")), + "path": str(_transport_paths(transport)[name]), + "sha256": _hash(_transport_paths(transport)[name]), } for name in ("codec", "reference") }, diff --git a/scripts/experiments/neural_dynamics/official_math_rollout.py b/examples/learning/nerd/_support/_official_math_rollout.py similarity index 97% rename from scripts/experiments/neural_dynamics/official_math_rollout.py rename to examples/learning/nerd/_support/_official_math_rollout.py index 0f28e8e237..e7a192e738 100644 --- a/scripts/experiments/neural_dynamics/official_math_rollout.py +++ b/examples/learning/nerd/_support/_official_math_rollout.py @@ -34,14 +34,14 @@ from contextlib import nullcontext from typing import Any -from official_evaluation_data import _canonical, _hash, _transport -from _portable_reference import ( +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 +from ._historical_physics import _verify_historical_physics __all__: list[str] = [] @@ -92,7 +92,8 @@ def main() -> None: 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", args.transport_root / "_evaluation.py" + "_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) @@ -337,7 +338,7 @@ def controls(q: Any, qd: Any) -> Any: for name in ( "_portable_reference.py", "_historical_physics.py", - "official_evaluation_data.py", + "_official_evaluation_data.py", ) }, "teacher_manifest_sha256": _hash(args.teacher / "manifest.json"), diff --git a/scripts/experiments/neural_dynamics/official_model_reference.py b/examples/learning/nerd/_support/_official_model_reference.py similarity index 100% rename from scripts/experiments/neural_dynamics/official_model_reference.py rename to examples/learning/nerd/_support/_official_model_reference.py diff --git a/scripts/experiments/neural_dynamics/official_small_budget.py b/examples/learning/nerd/_support/_official_small_budget.py similarity index 100% rename from scripts/experiments/neural_dynamics/official_small_budget.py rename to examples/learning/nerd/_support/_official_small_budget.py diff --git a/scripts/experiments/neural_dynamics/official_small_budget_evaluate.py b/examples/learning/nerd/_support/_official_small_budget_evaluate.py similarity index 100% rename from scripts/experiments/neural_dynamics/official_small_budget_evaluate.py rename to examples/learning/nerd/_support/_official_small_budget_evaluate.py diff --git a/scripts/experiments/neural_dynamics/official_small_budget_rollout.py b/examples/learning/nerd/_support/_official_small_budget_rollout.py similarity index 98% rename from scripts/experiments/neural_dynamics/official_small_budget_rollout.py rename to examples/learning/nerd/_support/_official_small_budget_rollout.py index f44fe077e9..1e85dfbe54 100644 --- a/scripts/experiments/neural_dynamics/official_small_budget_rollout.py +++ b/examples/learning/nerd/_support/_official_small_budget_rollout.py @@ -33,7 +33,7 @@ import sys from typing import Any -from official_evaluation_data import _canonical, _hash, _transport +from ._official_evaluation_data import _canonical, _hash, _transport, _transport_paths __all__: list[str] = [] @@ -79,7 +79,8 @@ def main() -> None: 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", args.transport_root / "_evaluation.py" + "_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) diff --git a/scripts/experiments/neural_dynamics/_portable_reference.py b/examples/learning/nerd/_support/_portable_reference.py similarity index 98% rename from scripts/experiments/neural_dynamics/_portable_reference.py rename to examples/learning/nerd/_support/_portable_reference.py index c6cffc09ce..9deba5b13f 100644 --- a/scripts/experiments/neural_dynamics/_portable_reference.py +++ b/examples/learning/nerd/_support/_portable_reference.py @@ -25,7 +25,7 @@ import subprocess from typing import Any -from official_evaluation_data import _hash +from ._official_evaluation_data import _hash, _transport_paths __all__: list[str] = [] COMMIT = "370e89d2b8ce68aa3578aa7812859c681b7502bc" @@ -78,7 +78,7 @@ def verify_source( verified[name] = _hash(source / name) verified_transport = {} for name, identity in metadata["canonical_transport_modules"].items(): - actual = _hash(transport / (name + ".py")) + actual = _hash(_transport_paths(transport)[name]) if actual != identity["sha256"]: raise ValueError("Frozen postprocessing module checksum mismatch: " + name) verified_transport[name] = actual diff --git a/embodichain/learning/dynamics/qualification.py b/examples/learning/nerd/_support/_qualification.py similarity index 99% rename from embodichain/learning/dynamics/qualification.py rename to examples/learning/nerd/_support/_qualification.py index 76ad6a6f1d..e8113e3ae5 100644 --- a/embodichain/learning/dynamics/qualification.py +++ b/examples/learning/nerd/_support/_qualification.py @@ -24,7 +24,7 @@ import torch from torch import Tensor -from .trainer import normalized_mse +from ._trainer import normalized_mse __all__ = ["capture_training_step", "compare_training_steps", "paired_noninferiority"] diff --git a/embodichain/learning/dynamics/reference.py b/examples/learning/nerd/_support/_reference.py similarity index 98% rename from embodichain/learning/dynamics/reference.py rename to examples/learning/nerd/_support/_reference.py index 6cfdff3dbb..1d889374ab 100644 --- a/embodichain/learning/dynamics/reference.py +++ b/examples/learning/nerd/_support/_reference.py @@ -299,7 +299,7 @@ def run_reference( """ if not trust_external_code: raise ValueError("Official source execution requires trust_external_code=True") - from ._numerics import _validate_numerical_profile + from embodichain_models.dynamics.nerd._numerics import _validate_numerical_profile if ( _validate_numerical_profile( @@ -316,10 +316,10 @@ def run_reference( import numpy as np from algorithms.vanilla_trainer import VanillaTrainer from utils.time_report import TimeReport - from embodichain.compute.dynamics import NeRDStateCodec - from embodichain.data_pipeline.dynamics.shards import DynamicsWindowDataset - from embodichain.learning.models.nerd import load_model, export_model - from .trainer import ( + 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, @@ -587,8 +587,8 @@ def export_legacy_dataset( import h5py import numpy as np import torch - from embodichain.compute.dynamics import NeRDStateCodec - from embodichain.data_pipeline.dynamics.shards import ( + from embodichain_models.dynamics.nerd.codec import NeRDStateCodec + from embodichain_models.dynamics.nerd.data import ( load_manifest, dataset_fingerprint, ) @@ -693,8 +693,8 @@ def import_legacy_datasets( import h5py import numpy as np import torch - from embodichain.compute.dynamics import NeRDStateCodec - from embodichain.data_pipeline.dynamics.shards import ( + from embodichain_models.dynamics.nerd.codec import NeRDStateCodec + from embodichain_models.dynamics.nerd.data import ( DynamicsShardWriter, load_manifest, dataset_fingerprint, @@ -709,7 +709,7 @@ def import_legacy_datasets( "Preserved legacy velocity requires the source world frame without rotation" ) codec = NeRDStateCodec(**metadata["codec"]) - from embodichain.compute.dynamics import quat_rotate, quat_multiply + from embodichain_models.dynamics.nerd.codec import quat_rotate, quat_multiply rotation = ( None @@ -974,7 +974,7 @@ def generate_official_dataset( import math import time import xml.etree.ElementTree as ET - from embodichain.compute.dynamics import quat_multiply + from embodichain_models.dynamics.nerd.codec import quat_multiply robot = str(spec["robot"]).lower() if robot not in {"cartpole", "ant"}: diff --git a/embodichain/data_pipeline/dynamics/_reference_assets.py b/examples/learning/nerd/_support/_reference_assets.py similarity index 97% rename from embodichain/data_pipeline/dynamics/_reference_assets.py rename to examples/learning/nerd/_support/_reference_assets.py index 6928441396..618ec18d92 100644 --- a/embodichain/data_pipeline/dynamics/_reference_assets.py +++ b/examples/learning/nerd/_support/_reference_assets.py @@ -73,13 +73,13 @@ def _ant_reference_friction() -> dict[str, Any]: def _asset(robot: str) -> Path: if robot not in _ASSETS: raise ValueError(f"Unsupported NeRD reference robot: {robot}") - import embodichain + from .._resources import prepare_reference_assets - return Path(embodichain.__file__).parent / "data" / "nerd" / _ASSETS[robot] + return prepare_reference_assets() / _ASSETS[robot] def reference_asset_hash(robot: str) -> str: - """Return the SHA256 of the exact packaged reference robot source.""" + """Return the SHA256 of the pinned downloaded reference robot source.""" return hashlib.sha256(_asset(robot).read_bytes()).hexdigest() diff --git a/embodichain/learning/dynamics/reference_rollout.py b/examples/learning/nerd/_support/_reference_rollout.py similarity index 98% rename from embodichain/learning/dynamics/reference_rollout.py rename to examples/learning/nerd/_support/_reference_rollout.py index 3226dd71f7..a1f9975a3d 100644 --- a/embodichain/learning/dynamics/reference_rollout.py +++ b/examples/learning/nerd/_support/_reference_rollout.py @@ -32,7 +32,7 @@ import time from typing import Any from collections.abc import Mapping, Sequence -from .reference import OFFICIAL_COMMIT, reference_identity, _hash +from ._reference import OFFICIAL_COMMIT, reference_identity, _hash __all__ = ["evaluate_reference_rollouts"] @@ -266,7 +266,7 @@ def _verify_historical_physics( def _canonical_states(states: Any, codec: Any, rotation: Any) -> tuple[Any, Any]: import torch - from embodichain.compute.dynamics import quat_rotate, quat_multiply + 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 :]) @@ -352,7 +352,7 @@ def __init__( def __call__(self, q: Any, qd: Any) -> Any: import torch - from embodichain.compute.dynamics import reference_task_values + from ._task_values import reference_task_values observation, _, _ = reference_task_values(self.robot, q, qd) finite = torch.isfinite(observation).all(-1) @@ -462,14 +462,15 @@ def evaluate_reference_rollouts( JOINT_ACT_SCALE, ) from utils.python_utils import set_random_seed - from embodichain.compute.dynamics import NeRDStateCodec, reference_task_values - from embodichain.compute.dynamics._evaluation import ( + 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.learning.models.nerd import load_model - from .reference import _restore_official_model + 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))) diff --git a/embodichain/learning/dynamics/_reference_traces.py b/examples/learning/nerd/_support/_reference_traces.py similarity index 97% rename from embodichain/learning/dynamics/_reference_traces.py rename to examples/learning/nerd/_support/_reference_traces.py index def714dcfc..0fd582e14f 100644 --- a/embodichain/learning/dynamics/_reference_traces.py +++ b/examples/learning/nerd/_support/_reference_traces.py @@ -27,9 +27,9 @@ import numpy as np import torch -from embodichain.compute.dynamics._evaluation import _teacher_row_sha256 -from embodichain.data_pipeline.dynamics.shards import load_manifest -from .reference import OFFICIAL_COMMIT +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] = [] diff --git a/embodichain/learning/dynamics/report.py b/examples/learning/nerd/_support/_report.py similarity index 99% rename from embodichain/learning/dynamics/report.py rename to examples/learning/nerd/_support/_report.py index 406fe33780..3eeaa3d67c 100644 --- a/embodichain/learning/dynamics/report.py +++ b/examples/learning/nerd/_support/_report.py @@ -24,9 +24,9 @@ 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 +from ._comparison import _read_plan, _json +from ._qualification import paired_noninferiority +from ._reference import OFFICIAL_COMMIT, _hash __all__ = ["write_comparison_report", "annotate_rollout_evidence"] @@ -369,7 +369,7 @@ def _dataset_evidence(teacher: str, declaration: dict[str, Any]) -> dict[str, An """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.data_pipeline.dynamics.shards import ( + from embodichain_models.dynamics.nerd.data import ( dataset_fingerprint, load_manifest, ) diff --git a/embodichain/learning/dynamics/_runtime.py b/examples/learning/nerd/_support/_runtime.py similarity index 100% rename from embodichain/learning/dynamics/_runtime.py rename to examples/learning/nerd/_support/_runtime.py diff --git a/embodichain/learning/dynamics/_sample_order.py b/examples/learning/nerd/_support/_sample_order.py similarity index 100% rename from embodichain/learning/dynamics/_sample_order.py rename to examples/learning/nerd/_support/_sample_order.py diff --git a/embodichain/compute/dynamics/reference.py b/examples/learning/nerd/_support/_task_values.py similarity index 98% rename from embodichain/compute/dynamics/reference.py rename to examples/learning/nerd/_support/_task_values.py index 0b332e2b6b..33bcbf5ac6 100644 --- a/embodichain/compute/dynamics/reference.py +++ b/examples/learning/nerd/_support/_task_values.py @@ -25,7 +25,7 @@ from __future__ import annotations import torch from torch import Tensor -from .codec import quat_rotate +from embodichain_models.dynamics.nerd.codec import quat_rotate __all__ = ["reference_task_values"] diff --git a/embodichain/learning/dynamics/trainer.py b/examples/learning/nerd/_support/_trainer.py similarity index 73% rename from embodichain/learning/dynamics/trainer.py rename to examples/learning/nerd/_support/_trainer.py index 5de8b08180..44c6b93d9f 100644 --- a/embodichain/learning/dynamics/trainer.py +++ b/examples/learning/nerd/_support/_trainer.py @@ -21,7 +21,6 @@ import json import hashlib import math -import os from pathlib import Path import random from typing import Any @@ -31,19 +30,29 @@ from torch import Tensor from embodichain.utils import configclass -from embodichain.data_pipeline.dynamics.shards import DynamicsWindowDataset +from embodichain_models.dynamics.nerd.data import DynamicsWindowDataset from embodichain.learning.sl import SupervisedLearner from embodichain.learning.sl._artifacts import _RunArtifacts -from .supervised import NeRDObjective, make_nerd_batch +from embodichain_models.dynamics.nerd.supervised import NeRDObjective, make_nerd_batch from ._sample_order import _CheckpointedOrder, _OfficialShuffleOrder, _open_order -from ._numerics import ( +from embodichain_models.dynamics.nerd._numerics import ( _attention_context, _checkpoint_numerical_profile, _clip_grad_norm_, _validate_numerical_profile, ) -_PREPROCESSING_EXECUTION = "nerd-shared-frame-scripted-cuda-warmup-v2" +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", @@ -88,166 +97,6 @@ class DynamicsTrainCfg: numerical_profile: str = "default" -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 _capture_rng_state(device: str) -> dict[str, Any]: numpy_state = np.random.get_state() return { @@ -277,106 +126,6 @@ def _restore_rng_state(state: dict[str, Any]) -> None: torch.cuda.set_rng_state_all(state["cuda"]) -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=torch.float32) - qd = batch["states"]["qd"].to(device=device, dtype=torch.float32) - 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, - } - - def _model_cfg_identity(cfg: dict[str, Any]) -> dict[str, Any]: """Compare legal head-width sequences without changing feature order.""" identity = dict(cfg) 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/embodichain/learning/dynamics/workflow.py b/examples/learning/nerd/_support/_workflow.py similarity index 73% rename from embodichain/learning/dynamics/workflow.py rename to examples/learning/nerd/_support/_workflow.py index 817e500c79..c741ded2f9 100644 --- a/embodichain/learning/dynamics/workflow.py +++ b/examples/learning/nerd/_support/_workflow.py @@ -94,13 +94,13 @@ def run_train(config: dict[str, Any]) -> dict[str, Any]: }, ) import torch - from embodichain.compute.dynamics import NeRDStateCodec - from embodichain.learning.models.nerd import NeRDModel, NeRDModelCfg, load_model - from embodichain.data_pipeline.dynamics.shards import ( + 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 ._trainer import DynamicsTrainCfg, DynamicsTrainer, load_statistics from ._runtime import _training_runtime_identity if config.get("resume") and config.get("statistics"): @@ -131,7 +131,7 @@ def run_train(config: dict[str, Any]) -> dict[str, Any]: ) != manifest.get("dt"): raise ValueError("Fine-tuning requires compatible codec and timestep") elif config.get("resume"): - from embodichain.learning.models.nerd.checkpoint import _model_from_state + 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, @@ -206,10 +206,10 @@ def run_eval(config: dict[str, Any]) -> dict[str, Any]: "include_prefix", }, ) - from embodichain.compute.dynamics import NeRDStateCodec - from embodichain.learning.models.nerd import load_model - from embodichain.data_pipeline.dynamics.shards import DynamicsWindowDataset - from .trainer import evaluate_model + 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) @@ -246,12 +246,74 @@ def run_eval(config: dict[str, Any]) -> dict[str, Any]: def run_export(checkpoint: str | Path, output: str | Path) -> dict[str, Any]: - """Export a training checkpoint using the production model artifact writer.""" - from embodichain.learning.models.nerd import NeRDModelCfg, export_model - from embodichain.learning.models.nerd.checkpoint import _model_from_state - from ._numerics import _checkpoint_numerical_profile + """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 - state = _load_training_checkpoint(checkpoint) + # 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 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/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.full.yaml b/examples/learning/nerd/configs/ant/dynamics.generate.full.yaml similarity index 85% rename from extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.full.yaml rename to examples/learning/nerd/configs/ant/dynamics.generate.full.yaml index da72a9f087..1da63e9fb5 100644 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.full.yaml +++ b/examples/learning/nerd/configs/ant/dynamics.generate.full.yaml @@ -1,6 +1,6 @@ # Dynamics workflow configuration; not a runnable Gym deployment. robot: ant -output_dir: ../../../../../../outputs/dynamics/ant/full/data +output_dir: ../../../../../outputs/dynamics/ant/full/data device: cuda:0 num_envs: 512 trajectories: 125000 diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.smoke.yaml b/examples/learning/nerd/configs/ant/dynamics.generate.smoke.yaml similarity index 85% rename from extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.smoke.yaml rename to examples/learning/nerd/configs/ant/dynamics.generate.smoke.yaml index c134f5c0db..06394bd6af 100644 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.smoke.yaml +++ b/examples/learning/nerd/configs/ant/dynamics.generate.smoke.yaml @@ -1,6 +1,6 @@ # Dynamics workflow configuration; not a runnable Gym deployment. robot: ant -output_dir: ../../../../../../outputs/dynamics/ant/smoke/data +output_dir: ../../../../../outputs/dynamics/ant/smoke/data device: cpu num_envs: 4 trajectories: 64 diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.warp-reference.yaml b/examples/learning/nerd/configs/ant/dynamics.generate.warp-reference.yaml similarity index 85% rename from extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.warp-reference.yaml rename to examples/learning/nerd/configs/ant/dynamics.generate.warp-reference.yaml index 15ff72bec6..6db1478a8c 100644 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.warp-reference.yaml +++ b/examples/learning/nerd/configs/ant/dynamics.generate.warp-reference.yaml @@ -1,6 +1,6 @@ # Bounded parameter-alignment collection; not full accuracy qualification. robot: ant -output_dir: ../../../../../../outputs/dynamics/ant/warp-reference/data +output_dir: ../../../../../outputs/dynamics/ant/warp-reference/data device: cpu num_envs: 64 trajectories: 512 diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.train.full.yaml b/examples/learning/nerd/configs/ant/dynamics.train.full.yaml similarity index 87% rename from extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.train.full.yaml rename to examples/learning/nerd/configs/ant/dynamics.train.full.yaml index bfa9485894..cf9d1df029 100644 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.train.full.yaml +++ b/examples/learning/nerd/configs/ant/dynamics.train.full.yaml @@ -1,6 +1,6 @@ # Dynamics workflow configuration; not a runnable Gym deployment. -dataset: ../../../../../../outputs/dynamics/ant/full/data -output_dir: ../../../../../../outputs/dynamics/ant/full/train +dataset: ../../../../../outputs/dynamics/ant/full/data +output_dir: ../../../../../outputs/dynamics/ant/full/train model: input_dims: states_embedding: 29 diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.train.smoke.yaml b/examples/learning/nerd/configs/ant/dynamics.train.smoke.yaml similarity index 84% rename from extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.train.smoke.yaml rename to examples/learning/nerd/configs/ant/dynamics.train.smoke.yaml index 8533943066..7e8b697d89 100644 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.train.smoke.yaml +++ b/examples/learning/nerd/configs/ant/dynamics.train.smoke.yaml @@ -1,6 +1,6 @@ # Dynamics workflow configuration; not a runnable Gym deployment. -dataset: ../../../../../../outputs/dynamics/ant/smoke/data -output_dir: ../../../../../../outputs/dynamics/ant/smoke/train +dataset: ../../../../../outputs/dynamics/ant/smoke/data +output_dir: ../../../../../outputs/dynamics/ant/smoke/train model: input_dims: states_embedding: 29 diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/sl.train.yaml b/examples/learning/nerd/configs/ant/sl.train.yaml similarity index 92% rename from extensions/embodichain_tasks/configs/tasks/classic_control/ant/sl.train.yaml rename to examples/learning/nerd/configs/ant/sl.train.yaml index d4b2ed6256..f794bc4c26 100644 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/sl.train.yaml +++ b/examples/learning/nerd/configs/ant/sl.train.yaml @@ -1,5 +1,5 @@ # Standalone SL recipe; run from the repository root after downloading data. -factory: embodichain.learning.dynamics.sl_recipe:build_supervised_trainer +factory: embodichain_models.dynamics.nerd.sl_recipe:build_supervised_trainer model: input_dims: states_embedding: 29 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/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.full.yaml b/examples/learning/nerd/configs/cartpole/dynamics.generate.full.yaml similarity index 84% rename from extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.full.yaml rename to examples/learning/nerd/configs/cartpole/dynamics.generate.full.yaml index 8ee89e2b30..02b2b20ea0 100644 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.full.yaml +++ b/examples/learning/nerd/configs/cartpole/dynamics.generate.full.yaml @@ -1,6 +1,6 @@ # Dynamics workflow configuration; not a runnable Gym deployment. robot: cartpole -output_dir: ../../../../../../outputs/dynamics/cartpole/full/data +output_dir: ../../../../../outputs/dynamics/cartpole/full/data device: cuda:0 num_envs: 512 trajectories: 62500 diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.smoke.yaml b/examples/learning/nerd/configs/cartpole/dynamics.generate.smoke.yaml similarity index 84% rename from extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.smoke.yaml rename to examples/learning/nerd/configs/cartpole/dynamics.generate.smoke.yaml index c5e4d1431b..9efc7d69b3 100644 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.smoke.yaml +++ b/examples/learning/nerd/configs/cartpole/dynamics.generate.smoke.yaml @@ -1,6 +1,6 @@ # Dynamics workflow configuration; not a runnable Gym deployment. robot: cartpole -output_dir: ../../../../../../outputs/dynamics/cartpole/smoke/data +output_dir: ../../../../../outputs/dynamics/cartpole/smoke/data device: cpu num_envs: 4 trajectories: 64 diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.warp-reference.yaml b/examples/learning/nerd/configs/cartpole/dynamics.generate.warp-reference.yaml similarity index 83% rename from extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.warp-reference.yaml rename to examples/learning/nerd/configs/cartpole/dynamics.generate.warp-reference.yaml index a7fd7e8b38..d15857ec8a 100644 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.warp-reference.yaml +++ b/examples/learning/nerd/configs/cartpole/dynamics.generate.warp-reference.yaml @@ -1,6 +1,6 @@ # Bounded parameter-alignment collection; not full accuracy qualification. robot: cartpole -output_dir: ../../../../../../outputs/dynamics/cartpole/warp-reference/data +output_dir: ../../../../../outputs/dynamics/cartpole/warp-reference/data device: cpu num_envs: 64 trajectories: 512 diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.train.full.yaml b/examples/learning/nerd/configs/cartpole/dynamics.train.full.yaml similarity index 83% rename from extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.train.full.yaml rename to examples/learning/nerd/configs/cartpole/dynamics.train.full.yaml index cb287d7f1a..1a1d8f6f93 100644 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.train.full.yaml +++ b/examples/learning/nerd/configs/cartpole/dynamics.train.full.yaml @@ -1,6 +1,6 @@ # Dynamics workflow configuration; not a runnable Gym deployment. -dataset: ../../../../../../outputs/dynamics/cartpole/full/data -output_dir: ../../../../../../outputs/dynamics/cartpole/full/train +dataset: ../../../../../outputs/dynamics/cartpole/full/data +output_dir: ../../../../../outputs/dynamics/cartpole/full/train model: input_dims: states_embedding: 4 diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.train.smoke.yaml b/examples/learning/nerd/configs/cartpole/dynamics.train.smoke.yaml similarity index 81% rename from extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.train.smoke.yaml rename to examples/learning/nerd/configs/cartpole/dynamics.train.smoke.yaml index cf1815fde7..a9d1543bf9 100644 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.train.smoke.yaml +++ b/examples/learning/nerd/configs/cartpole/dynamics.train.smoke.yaml @@ -1,6 +1,6 @@ # Dynamics workflow configuration; not a runnable Gym deployment. -dataset: ../../../../../../outputs/dynamics/cartpole/smoke/data -output_dir: ../../../../../../outputs/dynamics/cartpole/smoke/train +dataset: ../../../../../outputs/dynamics/cartpole/smoke/data +output_dir: ../../../../../outputs/dynamics/cartpole/smoke/train model: input_dims: states_embedding: 4 diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/sl.train.yaml b/examples/learning/nerd/configs/cartpole/sl.train.yaml similarity index 91% rename from extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/sl.train.yaml rename to examples/learning/nerd/configs/cartpole/sl.train.yaml index 8b6b323d88..8fa869c445 100644 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/sl.train.yaml +++ b/examples/learning/nerd/configs/cartpole/sl.train.yaml @@ -1,5 +1,5 @@ # Standalone SL recipe; run from the repository root after downloading data. -factory: embodichain.learning.dynamics.sl_recipe:build_supervised_trainer +factory: embodichain_models.dynamics.nerd.sl_recipe:build_supervised_trainer model: input_dims: states_embedding: 4 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/scripts/experiments/supervised/nerd.md b/examples/learning/nerd/research/acceptance.md similarity index 65% rename from scripts/experiments/supervised/nerd.md rename to examples/learning/nerd/research/acceptance.md index 871cafed58..1a02b8a37a 100644 --- a/scripts/experiments/supervised/nerd.md +++ b/examples/learning/nerd/research/acceptance.md @@ -1,24 +1,28 @@ # NeRD supervised-learning acceptance -The production NeRD trainer delegates optimizer updates to -`SupervisedLearner`. Its existing data, codec, numerical execution, learning-rate -schedule, export and checkpoint contracts remain authoritative. The separate -`train-sl` recipe builds the common offline trainer and uses the common SL -checkpoint format. - -Run a bounded standalone SL recipe from the repository root after downloading -the published dataset into `outputs/datasets`: +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 train-sl --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/sl.train.yaml -embodichain train-sl --config extensions/embodichain_tasks/configs/tasks/classic_control/ant/sl.train.yaml +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 ``` -These 1,000-update recipes demonstrate the public entry point; they do not claim -full-budget convergence. Configure the output directory explicitly for each run. -The existing `embodichain dynamics` CLI retains its historical training/export -formats. New common-format checkpoints resume through `train-sl` with the same -recipe and data identity. +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 @@ -30,12 +34,14 @@ 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 the resume contract for both CLIs. +destination. Input feature order is part of recovery identity for both trainer +formats. -`last.pt` is the complete resumable checkpoint. Automatically selected `best.pt` -files are model-only snapshots marked `resumable=False`, suitable for -`embodichain dynamics export`. They retain the best model's own step and frozen -normalizers; resume from `last.pt`. The full checkpoint retains one owned CPU +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 @@ -50,6 +56,13 @@ has a finite best score but no embedded best weights, its matching source-side ## 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 @@ -58,8 +71,9 @@ independent reference source. Use published dataset revision records model/dataset manifest hashes, weights hashes and source identities. ```bash -CUDA_VISIBLE_DEVICES=1 python scripts/experiments/supervised/nerd_acceptance.py \ - --reference-source /path/to/pr855-worktree \ +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 diff --git a/extensions/embodichain_models/README.md b/extensions/embodichain_models/README.md new file mode 100644 index 0000000000..7016131911 --- /dev/null +++ b/extensions/embodichain_models/README.md @@ -0,0 +1,25 @@ +# 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. + +Organize implementations by domain and model type. The current package is +`dynamics/nerd`; future contributions can use policy families for BC, +`world_models` for RWM, or `dynamics` for RigidFormer-style models as their +contracts settle. All use the shared SL core. Add a family when its +implementation lands, without empty placeholder directories or a new registry. diff --git a/embodichain/learning/models/nerd/__init__.py b/extensions/embodichain_models/embodichain_models/__init__.py similarity index 69% rename from embodichain/learning/models/nerd/__init__.py rename to extensions/embodichain_models/embodichain_models/__init__.py index 1d3ba033ea..addbcd4951 100644 --- a/embodichain/learning/models/nerd/__init__.py +++ b/extensions/embodichain_models/embodichain_models/__init__.py @@ -14,17 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""NeRD networks, portable exports and trusted checkpoint conversion.""" +"""Official model implementations bundled in the EmbodiChain distribution.""" from __future__ import annotations -from .model import NeRDModelCfg, NeRDModel, FeatureNormalizer -from .checkpoint import export_model, load_model, convert_legacy_checkpoint -__all__ = [ - "NeRDModelCfg", - "NeRDModel", - "FeatureNormalizer", - "export_model", - "load_model", - "convert_legacy_checkpoint", -] +__all__: list[str] = [] diff --git a/embodichain/compute/dynamics/__init__.py b/extensions/embodichain_models/embodichain_models/dynamics/__init__.py similarity index 63% rename from embodichain/compute/dynamics/__init__.py rename to extensions/embodichain_models/embodichain_models/dynamics/__init__.py index ffc3c550f6..6fd00181de 100644 --- a/embodichain/compute/dynamics/__init__.py +++ b/extensions/embodichain_models/embodichain_models/dynamics/__init__.py @@ -14,28 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Simulation-independent learned dynamics mathematics.""" +"""Official learned dynamics models, imported through their named subpackages.""" from __future__ import annotations -from .codec import ( - NeRDStateCodec, - quat_multiply, - quat_rotate, - quat_inverse, - plane_contact_features, - append_history, - reset_history, -) -from .reference import reference_task_values - -__all__ = [ - "reference_task_values", - "NeRDStateCodec", - "quat_multiply", - "quat_rotate", - "quat_inverse", - "plane_contact_features", - "append_history", - "reset_history", -] +__all__: list[str] = [] diff --git a/embodichain/learning/models/nerd/THIRD_PARTY_LICENSES b/extensions/embodichain_models/embodichain_models/dynamics/nerd/THIRD_PARTY_LICENSES similarity index 100% rename from embodichain/learning/models/nerd/THIRD_PARTY_LICENSES rename to extensions/embodichain_models/embodichain_models/dynamics/nerd/THIRD_PARTY_LICENSES 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/embodichain/learning/dynamics/_numerics.py b/extensions/embodichain_models/embodichain_models/dynamics/nerd/_numerics.py similarity index 100% rename from embodichain/learning/dynamics/_numerics.py rename to extensions/embodichain_models/embodichain_models/dynamics/nerd/_numerics.py diff --git a/embodichain/data_pipeline/dynamics/_resident.py b/extensions/embodichain_models/embodichain_models/dynamics/nerd/_resident.py similarity index 100% rename from embodichain/data_pipeline/dynamics/_resident.py rename to extensions/embodichain_models/embodichain_models/dynamics/nerd/_resident.py diff --git a/embodichain/learning/models/nerd/checkpoint.py b/extensions/embodichain_models/embodichain_models/dynamics/nerd/checkpoint.py similarity index 100% rename from embodichain/learning/models/nerd/checkpoint.py rename to extensions/embodichain_models/embodichain_models/dynamics/nerd/checkpoint.py diff --git a/embodichain/compute/dynamics/codec.py b/extensions/embodichain_models/embodichain_models/dynamics/nerd/codec.py similarity index 100% rename from embodichain/compute/dynamics/codec.py rename to extensions/embodichain_models/embodichain_models/dynamics/nerd/codec.py diff --git a/embodichain/data_pipeline/dynamics/shards.py b/extensions/embodichain_models/embodichain_models/dynamics/nerd/data.py similarity index 100% rename from embodichain/data_pipeline/dynamics/shards.py rename to extensions/embodichain_models/embodichain_models/dynamics/nerd/data.py 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/embodichain/learning/models/nerd/model.py b/extensions/embodichain_models/embodichain_models/dynamics/nerd/model.py similarity index 100% rename from embodichain/learning/models/nerd/model.py rename to extensions/embodichain_models/embodichain_models/dynamics/nerd/model.py diff --git a/embodichain/learning/dynamics/sl_recipe.py b/extensions/embodichain_models/embodichain_models/dynamics/nerd/sl_recipe.py similarity index 95% rename from embodichain/learning/dynamics/sl_recipe.py rename to extensions/embodichain_models/embodichain_models/dynamics/nerd/sl_recipe.py index 7ca60f6ce8..8616ff480f 100644 --- a/embodichain/learning/dynamics/sl_recipe.py +++ b/extensions/embodichain_models/embodichain_models/dynamics/nerd/sl_recipe.py @@ -26,15 +26,15 @@ import numpy as np import torch -from embodichain.compute.dynamics import NeRDStateCodec -from embodichain.data_pipeline.dynamics.shards import DynamicsWindowDataset -from embodichain.learning.models.nerd import NeRDModel, NeRDModelCfg +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 .trainer import ( +from .integration import ( _PREPROCESSING_EXECUTION, compute_statistics, encode_batch, @@ -93,8 +93,8 @@ def build_supervised_trainer(config: dict[str, Any]) -> SupervisedTrainer: Returns: Common trainer with a NeRD objective and owned window reader references. - Its checkpoints use the common SL format; the existing ``dynamics`` CLI - continues reading and writing historical NeRD training checkpoints. + 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() diff --git a/embodichain/learning/dynamics/supervised.py b/extensions/embodichain_models/embodichain_models/dynamics/nerd/supervised.py similarity index 93% rename from embodichain/learning/dynamics/supervised.py rename to extensions/embodichain_models/embodichain_models/dynamics/nerd/supervised.py index 8a4583ad56..80acc8879c 100644 --- a/embodichain/learning/dynamics/supervised.py +++ b/extensions/embodichain_models/embodichain_models/dynamics/nerd/supervised.py @@ -53,9 +53,8 @@ def compute_loss( Returns: Graph-connected loss for one shared supervised update. """ - # Keep the historical public helper and its arithmetic as the owner; - # importing lazily also permits DynamicsTrainer to use this adapter. - from .trainer import _model_variance, normalized_mse + # Share the physical-target reduction with offline evaluation. + from .integration import _model_variance, normalized_mse prediction = model(batch.inputs) loss = normalized_mse( @@ -87,7 +86,7 @@ def make_nerd_batch( Returns: Encoded history inputs, physical increment targets and validity mask. """ - from .trainer import encode_batch + from .integration import encode_batch inputs, targets, valid = encode_batch(batch, codec, device) return SupervisedBatch( diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/README.md b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/README.md deleted file mode 100644 index d7d18c5508..0000000000 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/README.md +++ /dev/null @@ -1,25 +0,0 @@ -# Ant neural dynamics workflows - -These configuration files run the offline `embodichain dynamics` workflow; -they are not runnable Gym task deployments. - -```bash -embodichain dynamics generate --config extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.smoke.yaml -embodichain dynamics train --config extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.train.smoke.yaml -embodichain dynamics export --checkpoint outputs/dynamics/ant/smoke/train/best.pt --output outputs/dynamics/ant/smoke/export -embodichain dynamics eval --config extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.smoke.yaml -``` - -Generation calls Newton's built-in analytical solver directly. It does not -initialize DexSim or require its custom solver/Scene extension. The `full` and -`warp-reference` presets retain the recorded model/data scales and parameter -profiles; a profile does not promise identical dynamics across Newton versions. - -The pretrained evaluation preset reads the published data/model directories -shown in Features → Neural Dynamics Models. Downloaded datasets retain their -original teacher provenance. New direct-Newton datasets use a distinct physical -signature; replacing provenance fields does not align different teachers. - -See [Neural Dynamics Models](../../../../../../../docs/source/features/neural_dynamics.md) -for downloads, resume, fine-tuning and portable exports. Live neural simulation -and differentiable task integration are tracked separately in PR #822. diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.full.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.full.yaml deleted file mode 100644 index 2e4a7e59ba..0000000000 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.full.yaml +++ /dev/null @@ -1,8 +0,0 @@ -# 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/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.pretrained.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.pretrained.yaml deleted file mode 100644 index 9a1c928796..0000000000 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.pretrained.yaml +++ /dev/null @@ -1,8 +0,0 @@ -# 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/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.smoke.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.smoke.yaml deleted file mode 100644 index 0caca1bfb6..0000000000 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.smoke.yaml +++ /dev/null @@ -1,8 +0,0 @@ -# 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/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/README.md b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/README.md deleted file mode 100644 index 5e0bac7331..0000000000 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/README.md +++ /dev/null @@ -1,25 +0,0 @@ -# Cartpole neural dynamics workflows - -These configuration files run the offline `embodichain dynamics` workflow; -they are not runnable Gym task deployments. - -```bash -embodichain dynamics generate --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.smoke.yaml -embodichain dynamics train --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.train.smoke.yaml -embodichain dynamics export --checkpoint outputs/dynamics/cartpole/smoke/train/best.pt --output outputs/dynamics/cartpole/smoke/export -embodichain dynamics eval --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.smoke.yaml -``` - -Generation calls Newton's built-in analytical solver directly. It does not -initialize DexSim or require its custom solver/Scene extension. The `full` and -`warp-reference` presets retain the recorded model/data scales and parameter -profiles; a profile does not promise identical dynamics across Newton versions. - -The pretrained evaluation preset reads the published data/model directories -shown in Features → Neural Dynamics Models. Downloaded datasets retain their -original teacher provenance. New direct-Newton datasets use a distinct physical -signature; replacing provenance fields does not align different teachers. - -See [Neural Dynamics Models](../../../../../../../docs/source/features/neural_dynamics.md) -for downloads, resume, fine-tuning and portable exports. Live neural simulation -and differentiable task integration are tracked separately in PR #822. diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.full.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.full.yaml deleted file mode 100644 index 7cd4eef565..0000000000 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.full.yaml +++ /dev/null @@ -1,8 +0,0 @@ -# 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/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.pretrained.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.pretrained.yaml deleted file mode 100644 index 73b9c4926b..0000000000 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.pretrained.yaml +++ /dev/null @@ -1,8 +0,0 @@ -# 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/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.smoke.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.smoke.yaml deleted file mode 100644 index 6df395f4cb..0000000000 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.smoke.yaml +++ /dev/null @@ -1,8 +0,0 @@ -# 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/scripts/experiments/neural_dynamics/README.md b/scripts/experiments/neural_dynamics/README.md deleted file mode 100644 index 0cf9daa0ed..0000000000 --- a/scripts/experiments/neural_dynamics/README.md +++ /dev/null @@ -1,51 +0,0 @@ -# NeRD reference experiments - -These command-line tools compare an upstream NeRD implementation with the -EmbodiChain implementation. The normal dataset, training, export -workflow is documented in -[Neural Dynamics Models](../../../docs/source/features/neural_dynamics.md). -Downloadable datasets and inference models are linked there. Training logs, -checkpoints used for investigation, per-run reports and machine-specific plans -belong in an external output directory rather than this source tree. - -## Environments - -Run the unchanged upstream source in its supported Python/Torch environment. -Use the Warp release required by the selected upstream model. Run current -Newton collection and EmbodiChain training in a separate EmbodiChain environment. -The original-runtime helpers avoid importing the installed simulation SDK. -Use explicit interpreter and source paths; do not resolve a virtualenv executable -symlink to its base interpreter. - -External source/checkpoint execution requires the tools' explicit trust option. -Pin the source revision and keep any portability adapter separate from the -upstream numerical implementation. Use each tool's `--help` for its arguments. - -## Tool ownership - -`_portable_reference.py` and `_historical_physics.py` are private support modules -for the isolated original runtime. - -## Comparison contracts - -- Separate model/trainer changes, runtime kernels, coordinate transport and - analytical-solver versions instead of assigning every difference to migration. -- Preserve original sampling, input/output normalization, loss and model - selection when evaluating training-port parity. Equal random seeds alone do - not prove the same consumed sample sequence. -- Replay exactly the same realized open-loop efforts for predictive accuracy. - With feedback control, label comparisons that use each model's own state as - closed-loop behavior rather than same-action transition accuracy. -- Keep current Newton deployments and historical Warp models in their recorded - physical domains. Do not replace physics signatures or normalizers to make an - incompatible export load. -- Retain all requested trajectories and report nonfinite outputs and task - failures. Separate endpoint error from trajectory-time average error. -- Use the matching analytical solver as each model's accuracy reference. A - distance to another neural model is model similarity, not ground-truth error. -- State dataset size, update budget, evaluation population and horizon - independently. A README example or a complete optimizer budget does not prove - reproduction of a paper's training-data scale or final accuracy. - -These tools produce their outputs under caller-supplied directories. They do not -change the production numerical policy or provide a blanket accuracy acceptance. diff --git a/setup.py b/setup.py index c0922dd293..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,8 +184,7 @@ def main(): package_dir=get_package_dir(), package_data={ "embodichain": ["VERSION"], - "embodichain.data.nerd": ["*.urdf", "*.xml", "*.md"], - "embodichain.learning.models.nerd": ["THIRD_PARTY_LICENSES"], + "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..a726cfba6f 100644 --- a/tests/ci/test_workflows.py +++ b/tests/ci/test_workflows.py @@ -343,13 +343,14 @@ def test_test_job_restores_git_access_before_running_selected_tests( [ None, "missing_robots", + "missing_models", "missing_entry_point", "wrong_entry_point_group", "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.""" @@ -382,6 +383,7 @@ def test_release_guard_validates_bundled_robot_catalog( "embodichain_robots/__init__.py": "", "embodichain_robots/_catalog.py": "ROBOT_CONFIGS = {}\n", "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 +394,12 @@ 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 == "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/examples/learning/__init__.py b/tests/examples/learning/__init__.py new file mode 100644 index 0000000000..051f680fe4 --- /dev/null +++ b/tests/examples/learning/__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/tests/examples/learning/nerd/__init__.py b/tests/examples/learning/nerd/__init__.py new file mode 100644 index 0000000000..051f680fe4 --- /dev/null +++ b/tests/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/tests/scripts/experiments/supervised/test_nerd_acceptance.py b/tests/examples/learning/nerd/test_acceptance.py similarity index 94% rename from tests/scripts/experiments/supervised/test_nerd_acceptance.py rename to tests/examples/learning/nerd/test_acceptance.py index ab17351b44..37ecc24df0 100644 --- a/tests/scripts/experiments/supervised/test_nerd_acceptance.py +++ b/tests/examples/learning/nerd/test_acceptance.py @@ -24,7 +24,7 @@ import pytest -from scripts.experiments.supervised import nerd_acceptance +from examples.learning.nerd._support import _acceptance as nerd_acceptance @pytest.mark.parametrize( @@ -41,7 +41,7 @@ def test_invalid_seed_list_fails_before_output_or_workers( ) -> None: runner = ( Path(__file__).resolve().parents[4] - / "scripts/experiments/supervised/nerd_acceptance.py" + / "examples/learning/nerd/_support/_acceptance.py" ) output = tmp_path / "results" # Missing artifacts must never be inspected or passed to a worker. diff --git a/tests/learning/dynamics/test_comparison.py b/tests/examples/learning/nerd/test_comparison.py similarity index 84% rename from tests/learning/dynamics/test_comparison.py rename to tests/examples/learning/nerd/test_comparison.py index c7da565317..6eece29673 100644 --- a/tests/learning/dynamics/test_comparison.py +++ b/tests/examples/learning/nerd/test_comparison.py @@ -19,25 +19,26 @@ from __future__ import annotations import json +import os from pathlib import Path +import subprocess import numpy as np import pytest import torch -from embodichain.data_pipeline.dynamics.shards import DynamicsShardWriter -from embodichain.learning.dynamics.comparison import ( +from embodichain_models.dynamics.nerd.data import DynamicsShardWriter +from examples.learning.nerd._support._comparison import ( prepare_comparison, run_comparison, - run_candidate, _read_plan, ) -from embodichain.learning.dynamics.report import write_comparison_report +from examples.learning.nerd._support._report import write_comparison_report def configuration(tmp_path, monkeypatch): - import embodichain.learning.dynamics.comparison as comparison - from embodichain.learning.dynamics.reference import OFFICIAL_COMMIT + import examples.learning.nerd._support._comparison as comparison + from examples.learning.nerd._support._reference import OFFICIAL_COMMIT monkeypatch.setattr( comparison, @@ -128,14 +129,32 @@ def test_candidate_worker_uses_frozen_artifacts(tmp_path, monkeypatch): for cell in plan["cells"] if cell["status"] == "prepared" and cell["trainer"] == "embodichain" ) - result = run_candidate(json.loads(Path(cell["spec"]).read_text())) + repository = Path(__file__).resolve().parents[4] + environment = { + **os.environ, + "CUDA_VISIBLE_DEVICES": "", + "OMP_NUM_THREADS": "1", + "PYTHONPATH": os.pathsep.join( + (str(repository / "extensions/embodichain_models"), str(repository)) + ), + } + process = subprocess.run( + cell["command"], + cwd=tmp_path, + env=environment, + capture_output=True, + text=True, + timeout=30, + ) + assert process.returncode == 0, process.stdout + process.stderr + result = json.loads((Path(cell["output_dir"]) / "result.json").read_text()) assert result["completed_updates"] == 2 assert result["test"]["split"] == "test" assert result["qualification"] == "matched_training_run_only" def test_rollout_attachment_preserves_measurements_and_requires_identity(tmp_path): - from embodichain.learning.dynamics.report import annotate_rollout_evidence + from examples.learning.nerd._support._report import annotate_rollout_evidence source = tmp_path / "source.json" source.write_text( diff --git a/tests/learning/dynamics/test_dynamics_runtime_identity.py b/tests/examples/learning/nerd/test_dynamics_runtime_identity.py similarity index 97% rename from tests/learning/dynamics/test_dynamics_runtime_identity.py rename to tests/examples/learning/nerd/test_dynamics_runtime_identity.py index 15042ae47e..a58d1e10a3 100644 --- a/tests/learning/dynamics/test_dynamics_runtime_identity.py +++ b/tests/examples/learning/nerd/test_dynamics_runtime_identity.py @@ -26,7 +26,7 @@ import pytest import torch -from embodichain.learning.dynamics import _runtime +from examples.learning.nerd._support import _runtime def test_cpu_identity_reflects_actual_parameters_without_side_effects(monkeypatch): @@ -175,7 +175,7 @@ def test_fresh_process_identity_does_not_import_simulation(): code = """ import sys import torch -from embodichain.learning.dynamics._runtime import _training_runtime_identity +from examples.learning.nerd._support._runtime import _training_runtime_identity _training_runtime_identity(torch.nn.Linear(1, 1)) assert not any(name.split('.')[0] in {'warp', 'dexsim', 'newton'} for name in sys.modules) assert not any(name.startswith('embodichain.lab.sim') for name in sys.modules) diff --git a/tests/cli/test_dynamics.py b/tests/examples/learning/nerd/test_entrypoint.py similarity index 80% rename from tests/cli/test_dynamics.py rename to tests/examples/learning/nerd/test_entrypoint.py index 61b61e3e00..1b2e0368da 100644 --- a/tests/cli/test_dynamics.py +++ b/tests/examples/learning/nerd/test_entrypoint.py @@ -32,18 +32,18 @@ import torch import yaml -from embodichain.cli.dynamics import main +from examples.learning.nerd.__main__ import main def test_help_and_dispatcher_are_simulation_free(): - code = "from embodichain.cli.main import main; import sys\ntry: main(['dynamics','--help'])\nexcept SystemExit as exc: assert exc.code == 0\nassert not any(k == 'dexsim' or k.startswith('embodichain.lab.sim') for k in sys.modules)" + code = "from examples.learning.nerd.__main__ import main; import sys\ntry: main(['--help'])\nexcept SystemExit as exc: assert exc.code == 0\nassert not any(k == 'dexsim' or k.startswith('embodichain.lab.sim') for k in sys.modules)" subprocess.run( [sys.executable, "-c", code], check=True, capture_output=True, text=True ) def test_generate_routing_and_path_resolution(tmp_path, monkeypatch): - host = ModuleType("embodichain.data_pipeline.dynamics.generation") + host = ModuleType("examples.learning.nerd._support._generation") host.generate_dataset = Mock(return_value={"status": "ok"}) monkeypatch.setitem(sys.modules, host.__name__, host) config = tmp_path / "generate.yaml" @@ -54,6 +54,21 @@ def test_generate_routing_and_path_resolution(tmp_path, monkeypatch): assert received["seed"] == 7 and received["resume"] is True +def test_default_evaluation_writes_under_current_directory(tmp_path, monkeypatch): + host = ModuleType("examples.learning.nerd._support._workflow") + host.run_eval = Mock(return_value={"status": "ok"}) + monkeypatch.setitem(sys.modules, host.__name__, host) + monkeypatch.chdir(tmp_path) + resources = tmp_path / "resources" + prefix = Path("neural_dynamics/nerd/cartpole/v1") + for kind in ("datasets", "models"): + (resources / kind / prefix).mkdir(parents=True) + main(["eval", "--resource-root", str(resources), "--device", "cpu"]) + received = host.run_eval.call_args.args[0] + assert received["output"] == str(tmp_path / "outputs/nerd/cartpole/offline.json") + assert received["dataset"] == str(resources / "datasets" / prefix) + + def test_invalid_paths_and_conversion_requires_explicit_trust(tmp_path): config = tmp_path / "bad.yaml" config.write_text("dataset: missing\n") @@ -76,7 +91,7 @@ def test_invalid_paths_and_conversion_requires_explicit_trust(tmp_path): def test_conversion_manifest_and_trust_reach_converter(tmp_path, monkeypatch): - import embodichain.learning.models.nerd as nerd + import embodichain_models.dynamics.nerd as nerd checkpoint = tmp_path / "old.pt" checkpoint.write_bytes(b"unused mock fixture") @@ -100,8 +115,9 @@ def test_conversion_manifest_and_trust_reach_converter(tmp_path, monkeypatch): assert convert.call_args.args[2]["codec"]["robot"] == "cartpole" -def test_real_offline_cli_train_export_eval(tmp_path): - from embodichain.data_pipeline.dynamics.shards import DynamicsShardWriter +@pytest.mark.parametrize("trainer_kind", ["legacy", "sl"]) +def test_real_offline_cli_train_export_eval(tmp_path, trainer_kind): + from embodichain_models.dynamics.nerd.data import DynamicsShardWriter metadata = {"codec": {"robot": "cartpole"}, "dt": 1 / 60} with DynamicsShardWriter(tmp_path / "data", metadata) as writer: @@ -136,6 +152,20 @@ def test_real_offline_cli_train_export_eval(tmp_path): "validation_interval": 1, }, } + if trainer_kind == "sl": + training = { + "factory": "embodichain_models.dynamics.nerd.sl_recipe:build_supervised_trainer", + "model": training["model"], + "data": {"path": str(tmp_path / "data")}, + "trainer": { + "total_updates": 2, + "batch_size": 2, + "device": "cpu", + "validation_interval": 1, + "checkpoint_interval": 0, + "output_dir": str(tmp_path / "train"), + }, + } train_cfg = tmp_path / "train.json" train_cfg.write_text(json.dumps(training)) evaluation = { @@ -148,10 +178,10 @@ def test_real_offline_cli_train_export_eval(tmp_path): eval_cfg = tmp_path / "eval.yaml" eval_cfg.write_text(yaml.safe_dump(evaluation)) code = """import sys -from embodichain.cli.main import main -main(['dynamics','train','--config',sys.argv[1]]) -main(['dynamics','export','--checkpoint',sys.argv[2],'--output',sys.argv[3]]) -main(['dynamics','eval','--config',sys.argv[4]]) +from examples.learning.nerd.__main__ import main +main(['train','--config',sys.argv[1]]) +main(['export','--checkpoint',sys.argv[2],'--output',sys.argv[3]]) +main(['eval','--config',sys.argv[4]]) assert not any(k == 'dexsim' or k.startswith('embodichain.lab.sim') for k in sys.modules) """ env = {**os.environ, "OMP_NUM_THREADS": "1", "MKL_NUM_THREADS": "1"} @@ -161,7 +191,9 @@ def test_real_offline_cli_train_export_eval(tmp_path): "-c", code, str(train_cfg), - str(tmp_path / "train" / "best.pt"), + str( + tmp_path / "train" / ("last.pt" if trainer_kind == "sl" else "best.pt") + ), str(tmp_path / "export"), str(eval_cfg), ], @@ -172,6 +204,16 @@ def test_real_offline_cli_train_export_eval(tmp_path): assert result.returncode == 0, result.stdout + result.stderr report = json.loads((tmp_path / "reports/evaluation/result.json").read_text()) assert report["split"] == "test" and report["windows"] == 3 + if trainer_kind == "sl": + from embodichain_models.dynamics.nerd import load_model + + model, _ = load_model(tmp_path / "export") + assert list(model.cfg.input_dims) == ["states_embedding", "joint_f"] + state = torch.load(tmp_path / "train/last.pt", weights_only=False) + for name, value in model.state_dict().items(): + torch.testing.assert_close( + value, state["learner"]["model"][name], rtol=0, atol=0 + ) @pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) @@ -179,7 +221,7 @@ def test_real_offline_cli_train_export_eval(tmp_path): def test_training_checkpoint_cli_export_preserves_dtype_statistics_and_predictions( tmp_path, dtype, model_snapshot ): - from embodichain.learning.models.nerd import NeRDModel, NeRDModelCfg, load_model + from embodichain_models.dynamics.nerd import NeRDModel, NeRDModelCfg, load_model torch.set_num_threads(2) model = ( @@ -250,9 +292,9 @@ def test_training_checkpoint_cli_export_preserves_dtype_statistics_and_predictio def test_fine_tune_resume_export_preserves_dtype_updates_and_input_order( tmp_path, dtype, normalize_input, normalize_output ): - from embodichain.data_pipeline.dynamics.shards import DynamicsShardWriter - from embodichain.learning.dynamics.workflow import run_train - from embodichain.learning.models.nerd import ( + from embodichain_models.dynamics.nerd.data import DynamicsShardWriter + from examples.learning.nerd._support._workflow import run_train + from embodichain_models.dynamics.nerd import ( NeRDModel, NeRDModelCfg, export_model, @@ -385,11 +427,11 @@ def evaluate_cli(model_path, label): def test_packaged_workflow_configs_decode_through_cli(tmp_path, monkeypatch): """Exercise the actual YAML loader/path resolver and model/train constructors.""" - import embodichain.learning.dynamics.workflow as workflow - from embodichain.learning.dynamics.trainer import DynamicsTrainCfg - from embodichain.learning.models.nerd import NeRDModelCfg + import examples.learning.nerd._support._workflow as workflow + from examples.learning.nerd._support._trainer import DynamicsTrainCfg + from embodichain_models.dynamics.nerd import NeRDModelCfg - host = ModuleType("embodichain.data_pipeline.dynamics.generation") + host = ModuleType("examples.learning.nerd._support._generation") host.generate_dataset = Mock(return_value={"status": "not_run"}) monkeypatch.setitem(sys.modules, host.__name__, host) @@ -406,7 +448,7 @@ def train(config): monkeypatch.setattr(workflow, "run_train", train) monkeypatch.setattr(workflow, "run_eval", lambda config: {"status": "not_run"}) - repo = Path(__file__).resolve().parents[2] + repo = Path(__file__).resolve().parents[4] (tmp_path / "data").mkdir() (tmp_path / "export").mkdir() files = [] @@ -422,15 +464,9 @@ def train(config): } for robot in ("cartpole", "ant"): robot_configs = list( - ( - repo - / "extensions" - / "embodichain_tasks" - / "configs" - / "tasks" - / "classic_control" - / robot - ).glob("dynamics.*.yaml") + (repo / "examples" / "learning" / "nerd" / "configs" / robot).glob( + "dynamics.*.yaml" + ) ) assert {path.name for path in robot_configs} == expected_configs files.extend(robot_configs) @@ -449,61 +485,3 @@ def train(config): copied.write_text(yaml.safe_dump(config, sort_keys=False)) operation = source.name.split(".")[1] main([operation, "--config", str(copied)]) - - -def test_build_py_packages_reference_assets_and_workflow_resources(tmp_path): - """Use actual setup.py package-data rules, without compiling native modules.""" - import ast - import shutil - from setuptools import Distribution - from setuptools.command.build_py import build_py - - repo = Path(__file__).resolve().parents[2] - tree = ast.parse((repo / "setup.py").read_text()) - call = next( - node - for node in ast.walk(tree) - if isinstance(node, ast.Call) - and isinstance(node.func, ast.Name) - and node.func.id == "setup" - ) - package_data = ast.literal_eval( - next(item.value for item in call.keywords if item.arg == "package_data") - ) - source = tmp_path / "source" - assets = source / "nerd" - shutil.copytree(repo / "embodichain" / "data" / "nerd", assets) - configs = source / "configs" - configs.mkdir() - real_configs = repo / "extensions" / "embodichain_tasks" / "configs" - shutil.copyfile(real_configs / "__init__.py", configs / "__init__.py") - expected = [] - for robot in ("cartpole", "ant"): - relative = Path("tasks") / "classic_control" / robot - shutil.copytree(real_configs / relative, configs / relative) - expected.extend( - Path("embodichain_tasks/configs") / relative / file.name - for file in (configs / relative).iterdir() - if file.suffix in {".yaml", ".md"} - ) - expected.extend( - Path("embodichain/data/nerd") / name - for name in ("ant.xml", "cartpole_single.urdf", "ATTRIBUTION.md", "LICENSE.md") - ) - distribution = Distribution( - { - "packages": ["embodichain.data.nerd", "embodichain_tasks.configs"], - "package_dir": { - "embodichain.data.nerd": str(assets), - "embodichain_tasks.configs": str(configs), - }, - "package_data": package_data, - "include_package_data": False, - } - ) - distribution.script_name = str(repo / "setup.py") - command = build_py(distribution) - command.build_lib = str(tmp_path / "build") - command.ensure_finalized() - command.run() - assert all((Path(command.build_lib) / path).is_file() for path in expected) diff --git a/tests/data_pipeline/dynamics/test_generation.py b/tests/examples/learning/nerd/test_generation.py similarity index 94% rename from tests/data_pipeline/dynamics/test_generation.py rename to tests/examples/learning/nerd/test_generation.py index 8912f5f922..001f7e2a6c 100644 --- a/tests/data_pipeline/dynamics/test_generation.py +++ b/tests/examples/learning/nerd/test_generation.py @@ -28,19 +28,19 @@ import pytest import torch -from embodichain.data_pipeline.dynamics.generation import ( +from examples.learning.nerd._support._generation import ( DynamicsGenerationCfg, _sample, generate_dataset, ) -from embodichain.data_pipeline.dynamics._newton import _NewtonTeacher -from embodichain.data_pipeline.dynamics.shards import load_manifest, trajectory_split +from examples.learning.nerd._support._newton import _NewtonTeacher +from embodichain_models.dynamics.nerd.data import load_manifest, trajectory_split @pytest.fixture def deterministic_generation(monkeypatch): """Exercise persistence with reproducible accepted and rejected candidates.""" - from embodichain.data_pipeline.dynamics import generation + from examples.learning.nerd._support import _generation as generation class Teacher: def __init__(self, robot, *, device, num_envs, **kwargs): @@ -80,7 +80,7 @@ def sample(robot, seed, horizon, *args): def test_resume_retains_rejections_when_interrupted_after_shard_manifest( tmp_path, monkeypatch, deterministic_generation ): - from embodichain.data_pipeline.dynamics import shards + from embodichain_models.dynamics.nerd import data as shards config = { "output_dir": str(tmp_path), @@ -258,23 +258,23 @@ def find_spec(self, fullname, path=None, target=None): if fullname == 'dexsim' or fullname.startswith(('dexsim.', 'embodichain.lab.')): raise AssertionError('SDK-free workflow imported ' + fullname) sys.meta_path.insert(0, NoSDK()) -from embodichain.cli.main import main -from embodichain.data_pipeline.dynamics.shards import load_manifest +from examples.learning.nerd.__main__ import main +from embodichain_models.dynamics.nerd.data import load_manifest root, robot = Path(sys.argv[1]), sys.argv[2] def run(operation, config): path = root / (operation + '.json') path.write_text(json.dumps(config)) - main(['dynamics', operation, '--config', str(path)]) + main([operation, '--config', str(path)]) run('generate', dict(robot=robot, output_dir='data', num_envs=2, trajectories=64, horizon=4, num_substeps=2, initial_states='default', effort_scale=0.01)) import yaml -model_cfg = yaml.safe_load((Path.cwd() / 'extensions/embodichain_tasks/configs/tasks/classic_control' / robot / 'dynamics.train.smoke.yaml').read_text())['model'] +model_cfg = yaml.safe_load((Path.cwd() / 'examples/learning/nerd/configs' / robot / 'dynamics.train.smoke.yaml').read_text())['model'] model_cfg.update(n_embd=12, history_length=2, head_hidden_dims=[12]) run('train', dict(dataset='data', output_dir='train', model=model_cfg, train=dict(total_steps=2, batch_size=2, validation_interval=1, checkpoint_interval=1, device='cpu'))) -main(['dynamics', 'export', '--checkpoint', str(root/'train/best.pt'), +main(['export', '--checkpoint', str(root/'train/best.pt'), '--output', str(root/'export')]) run('eval', dict(dataset='data', model='export', split='test', device='cpu', output='reports/evaluation.json')) diff --git a/tests/compute/dynamics/test_dynamics_metrics.py b/tests/examples/learning/nerd/test_metrics.py similarity index 92% rename from tests/compute/dynamics/test_dynamics_metrics.py rename to tests/examples/learning/nerd/test_metrics.py index ba93266d37..ef079b6bc1 100644 --- a/tests/compute/dynamics/test_dynamics_metrics.py +++ b/tests/examples/learning/nerd/test_metrics.py @@ -25,7 +25,7 @@ import numpy as np import torch -from embodichain.compute.dynamics._evaluation import ( +from examples.learning.nerd._support._evaluation import ( _error_values, _safe_number, _teacher_row_sha256, @@ -37,7 +37,7 @@ def test_metrics_import_without_any_simulation_runtime(): [ sys.executable, "-c", - "import sys; from embodichain.compute.dynamics._evaluation import _error_values; assert not any(k == 'newton' or k == 'dexsim' or k.startswith('embodichain.lab.sim') for k in sys.modules)", + "import sys; from examples.learning.nerd._support._evaluation import _error_values; assert not any(k == 'newton' or k == 'dexsim' or k.startswith('embodichain.lab.sim') for k in sys.modules)", ], check=True, capture_output=True, diff --git a/tests/learning/dynamics/test_reference.py b/tests/examples/learning/nerd/test_reference.py similarity index 89% rename from tests/learning/dynamics/test_reference.py rename to tests/examples/learning/nerd/test_reference.py index 470e14e1bc..e1a9d09bf8 100644 --- a/tests/learning/dynamics/test_reference.py +++ b/tests/examples/learning/nerd/test_reference.py @@ -31,11 +31,11 @@ import pytest import torch -from embodichain.data_pipeline.dynamics.shards import ( +from embodichain_models.dynamics.nerd.data import ( DynamicsShardWriter, DynamicsWindowDataset, ) -from embodichain.learning.dynamics.reference import ( +from examples.learning.nerd._support._reference import ( export_legacy_dataset, run_reference, reference_identity, @@ -57,8 +57,8 @@ def test_original_method_adapter_rejects_unsupported_profile_before_source(profi def test_original_method_adapter_rejects_math_model_before_data(tmp_path, monkeypatch): - import embodichain.learning.dynamics.reference as reference - import embodichain.learning.models.nerd as models + import examples.learning.nerd._support._reference as reference + import embodichain_models.dynamics.nerd as models for name in ("algorithms", "utils"): package = ModuleType(name) @@ -99,11 +99,12 @@ def test_original_method_adapter_rejects_math_model_before_data(tmp_path, monkey def test_official_transport_uses_the_exact_native_codec_velocity_order(): - from embodichain.compute.dynamics import NeRDStateCodec - from embodichain.learning.dynamics.reference import _legacy_states + from embodichain_models.dynamics.nerd.codec import NeRDStateCodec + from examples.learning.nerd._support._reference import _legacy_states with np.load( - Path(__file__).parents[2] / "compute/dynamics/data/ant-official-frame.npz" + Path(__file__).resolve().parents[3] + / "models/dynamics/nerd/data/ant-official-frame.npz" ) as fixture: q = torch.from_numpy(fixture["q"].copy()) qd = torch.from_numpy(fixture["qd"].copy()) @@ -119,13 +120,14 @@ def test_official_transport_uses_the_exact_native_codec_velocity_order(): "device", ["cpu", pytest.param("cuda:0", marks=pytest.mark.gpu)] ) def test_training_transport_device_is_independent_of_raw_residency(device): - from embodichain.compute.dynamics import NeRDStateCodec - from embodichain.learning.dynamics.reference import _official_batch + from embodichain_models.dynamics.nerd.codec import NeRDStateCodec + from examples.learning.nerd._support._reference import _official_batch if device.startswith("cuda") and not torch.cuda.is_available(): pytest.skip("CUDA is unavailable") with np.load( - Path(__file__).parents[2] / "compute/dynamics/data/ant-official-frame.npz" + Path(__file__).resolve().parents[3] + / "models/dynamics/nerd/data/ant-official-frame.npz" ) as fixture: q = torch.from_numpy(fixture["q"].copy()) qd = torch.from_numpy(fixture["qd"].copy()) @@ -191,16 +193,19 @@ def test_legacy_export_preserves_lengths_splits_and_native_to_twist(tmp_path): ) def test_actual_official_main_two_updates_match_migrated(tmp_path): """External reference is opt-in; synthetic data tests migration, not accuracy.""" - from embodichain.compute.dynamics import NeRDStateCodec - from embodichain.learning.models.nerd import ( + from embodichain_models.dynamics.nerd.codec import NeRDStateCodec + from embodichain_models.dynamics.nerd import ( NeRDModel, NeRDModelCfg, export_model, load_model, ) - from embodichain.learning.dynamics.trainer import DynamicsTrainer, DynamicsTrainCfg - from embodichain.learning.dynamics.reference import _hash - from embodichain.learning.dynamics.workflow import run_train, run_export + from examples.learning.nerd._support._trainer import ( + DynamicsTrainer, + DynamicsTrainCfg, + ) + from examples.learning.nerd._support._reference import _hash + from examples.learning.nerd._support._workflow import run_train, run_export metadata = {"codec": {"robot": "cartpole"}, "dt": 1 / 60} with DynamicsShardWriter(tmp_path / "data", metadata) as writer: @@ -267,15 +272,22 @@ def test_actual_official_main_two_updates_match_migrated(tmp_path): [ os.environ.get("NERD_REFERENCE_DEPS", ""), source, - str(Path(__file__).resolve().parents[3]), + str(Path(__file__).resolve().parents[4]), + str( + Path(__file__).resolve().parents[4] + / "extensions/embodichain_models" + ), ] ), } process = subprocess.run( [ sys.executable, - "-m", - "embodichain.learning.dynamics.reference", + str( + Path(__file__).resolve().parents[4] + / "examples/learning/nerd/_support/_worker.py" + ), + "_reference", "--spec", str(path), "--trust-external-code", @@ -316,7 +328,7 @@ def test_actual_official_main_two_updates_match_migrated(tmp_path): def test_legacy_roundtrip_preserves_native_qd_and_declared_split(tmp_path): - from embodichain.learning.dynamics.reference import import_legacy_datasets + from examples.learning.nerd._support._reference import import_legacy_datasets metadata = {"codec": {"robot": "cartpole"}, "dt": 1 / 60} with DynamicsShardWriter(tmp_path / "original", metadata) as writer: @@ -385,9 +397,9 @@ def _write_source_ant_velocities(path: Path) -> np.ndarray: def test_original_velocity_import_roundtrip_preserves_fp32_bits_and_defaults(tmp_path): - from embodichain.compute.dynamics import NeRDStateCodec - from embodichain.data_pipeline.dynamics.shards import load_manifest - from embodichain.learning.dynamics.reference import ( + from embodichain_models.dynamics.nerd.codec import NeRDStateCodec + from embodichain_models.dynamics.nerd.data import load_manifest + from examples.learning.nerd._support._reference import ( import_legacy_datasets, _official_batch, ) @@ -463,8 +475,8 @@ def test_original_velocity_import_roundtrip_preserves_fp32_bits_and_defaults(tmp def test_preserved_velocity_follows_prefix_masks_and_cpu_residency(tmp_path): from torch.utils.data._utils.collate import default_collate - from embodichain.compute.dynamics import NeRDStateCodec - from embodichain.learning.dynamics.reference import ( + from embodichain_models.dynamics.nerd.codec import NeRDStateCodec + from examples.learning.nerd._support._reference import ( import_legacy_datasets, _official_batch, ) @@ -506,8 +518,8 @@ def test_preserved_velocity_follows_prefix_masks_and_cpu_residency(tmp_path): def test_disabling_velocity_preservation_clears_only_copied_contract(tmp_path): - from embodichain.data_pipeline.dynamics.shards import load_manifest - from embodichain.learning.dynamics.reference import import_legacy_datasets + from embodichain_models.dynamics.nerd.data import load_manifest + from examples.learning.nerd._support._reference import import_legacy_datasets path = tmp_path / "source.h5" _write_source_ant_velocities(path) @@ -533,7 +545,7 @@ def test_disabling_velocity_preservation_clears_only_copied_contract(tmp_path): @pytest.mark.parametrize("mismatch", ["small_delta", "signed_zero"]) def test_lossless_import_rejects_nonidentical_adjacent_fp32_states(tmp_path, mismatch): - from embodichain.learning.dynamics.reference import import_legacy_datasets + from examples.learning.nerd._support._reference import import_legacy_datasets path = tmp_path / "source.h5" _write_source_ant_velocities(path) @@ -556,7 +568,7 @@ def test_lossless_import_rejects_nonidentical_adjacent_fp32_states(tmp_path, mis def test_lossless_import_rejects_even_explicit_identity_world_rotation(tmp_path): - from embodichain.learning.dynamics.reference import import_legacy_datasets + from examples.learning.nerd._support._reference import import_legacy_datasets with pytest.raises(ValueError, match="source world frame"): import_legacy_datasets( @@ -570,8 +582,8 @@ def test_lossless_import_rejects_even_explicit_identity_world_rotation(tmp_path) def test_official_transport_rejects_malformed_preserved_velocity_shape(): - from embodichain.compute.dynamics import NeRDStateCodec - from embodichain.learning.dynamics.reference import _official_batch + from embodichain_models.dynamics.nerd.codec import NeRDStateCodec + from examples.learning.nerd._support._reference import _official_batch q = torch.zeros(1, 2, 15) q[..., 6] = 1 @@ -588,7 +600,7 @@ def test_official_transport_rejects_malformed_preserved_velocity_shape(): def test_gravity_repair_is_explicit_preserves_raw_and_other_fields(): - from embodichain.learning.dynamics.reference import repair_gravity_transport + from examples.learning.nerd._support._reference import repair_gravity_transport raw = np.array([[[np.nan, -1.0, 8.49e5], [1e-45, -1.0, np.nan]]], np.float32) before = raw.copy() @@ -604,7 +616,7 @@ def test_gravity_repair_is_explicit_preserves_raw_and_other_fields(): def test_import_explicit_world_rotation_and_gravity_transport(tmp_path): - from embodichain.learning.dynamics.reference import import_legacy_datasets + from examples.learning.nerd._support._reference import import_legacy_datasets import math metadata = {"codec": {"robot": "ant", "base_com": [0.0, 0.0, 0.0]}, "dt": 1 / 60} @@ -683,15 +695,22 @@ def test_real_official_collection_streams_batches_and_preserves_raw(tmp_path): [ os.environ.get("NERD_REFERENCE_DEPS", ""), spec["source"], - str(Path(__file__).resolve().parents[3]), + str(Path(__file__).resolve().parents[4]), + str( + Path(__file__).resolve().parents[4] + / "extensions/embodichain_models" + ), ] ), } process = subprocess.run( [ sys.executable, - "-m", - "embodichain.learning.dynamics.reference", + str( + Path(__file__).resolve().parents[4] + / "examples/learning/nerd/_support/_worker.py" + ), + "_reference", "--spec", str(spec_path), "--trust-external-code", @@ -729,7 +748,7 @@ def test_real_official_collection_streams_batches_and_preserves_raw(tmp_path): assert result["teacher_origin"]["gravity_magnitude"] == pytest.approx( 9.81, abs=1e-6 ) - from embodichain.learning.dynamics.reference import _hash + from examples.learning.nerd._support._reference import _hash import hashlib assert physics["source_generator_files"]["generate/trajectory_sampler.py"] == _hash( @@ -742,8 +761,8 @@ def test_real_official_collection_streams_batches_and_preserves_raw(tmp_path): ).hexdigest() assert physics["source_generator_sha256"] == expected_generator_hash assert physics["generation_adapter_sha256"] == _hash( - Path(__file__).resolve().parents[3] - / "embodichain/learning/dynamics/reference.py" + Path(__file__).resolve().parents[4] + / "examples/learning/nerd/_support/_reference.py" ) with h5py.File(root / "train.raw.h5") as file: assert file["data/states"].shape == (4, 5, 4) @@ -772,11 +791,11 @@ def test_real_official_collection_streams_batches_and_preserves_raw(tmp_path): def test_portable_to_official_restoration_preserves_state_dtype_and_ownership(): code = r""" import sys,torch -from embodichain.learning.dynamics.reference import reference_identity,_restore_official_model +from examples.learning.nerd._support._reference import reference_identity,_restore_official_model source=sys.argv[1] reference_identity(source) sys.path.insert(0,source) -from embodichain.learning.models.nerd import NeRDModel,NeRDModelCfg +from embodichain_models.dynamics.nerd import NeRDModel,NeRDModelCfg for dtype in (torch.float32,torch.float64): torch.manual_seed(5) common=NeRDModel(NeRDModelCfg(input_dims={"states_embedding":4,"joint_f":2},output_dim=4,n_embd=12,n_head=3,n_layer=1,block_size=4,history_length=3,head_hidden_dims=(12,),dropout=.25,output_tanh=True)).to(dtype=dtype).eval() @@ -812,16 +831,16 @@ def test_actual_official_trainer_resume_preserves_optimizer_rng_and_best_model( tmp_path, ): """Synthetic CPU data verify lifecycle, not learned-dynamics accuracy.""" - from embodichain.compute.dynamics import NeRDStateCodec - from embodichain.learning.models.nerd import ( + from embodichain_models.dynamics.nerd.codec import NeRDStateCodec + from embodichain_models.dynamics.nerd import ( NeRDModel, NeRDModelCfg, export_model, load_model, ) - from embodichain.learning.dynamics.trainer import compute_statistics - from embodichain.learning.dynamics._sample_order import _write_checkpointed_order - from embodichain.learning.dynamics.reference import _hash + from examples.learning.nerd._support._trainer import compute_statistics + from examples.learning.nerd._support._sample_order import _write_checkpointed_order + from examples.learning.nerd._support._reference import _hash metadata = {"codec": {"robot": "cartpole"}, "dt": 1 / 60} generator = np.random.default_rng(3) @@ -896,7 +915,11 @@ def test_actual_official_trainer_resume_preserves_optimizer_rng_and_best_model( [ os.environ.get("NERD_REFERENCE_DEPS", ""), source, - str(Path(__file__).resolve().parents[3]), + str(Path(__file__).resolve().parents[4]), + str( + Path(__file__).resolve().parents[4] + / "extensions/embodichain_models" + ), ] ), } @@ -907,8 +930,11 @@ def execute(name, extra): result = subprocess.run( [ sys.executable, - "-m", - "embodichain.learning.dynamics.reference", + str( + Path(__file__).resolve().parents[4] + / "examples/learning/nerd/_support/_worker.py" + ), + "_reference", "--spec", str(path), "--trust-external-code", diff --git a/tests/learning/dynamics/test_reference_rollout.py b/tests/examples/learning/nerd/test_reference_rollout.py similarity index 96% rename from tests/learning/dynamics/test_reference_rollout.py rename to tests/examples/learning/nerd/test_reference_rollout.py index 98a62860cd..fa160661a5 100644 --- a/tests/learning/dynamics/test_reference_rollout.py +++ b/tests/examples/learning/nerd/test_reference_rollout.py @@ -25,9 +25,11 @@ import subprocess import sys import pytest -from embodichain.learning.dynamics.reference_rollout import evaluate_reference_rollouts +from examples.learning.nerd._support._reference_rollout import ( + evaluate_reference_rollouts, +) -_ROOT = Path(__file__).resolve().parents[3] +_ROOT = Path(__file__).resolve().parents[4] def test_external_reference_requires_explicit_trust() -> None: @@ -50,9 +52,9 @@ def _check_original_runtime(tmp_path: Path, robot: str, source: str) -> None: import h5py import torch import warp as wp - from embodichain.learning.dynamics.reference import generate_official_dataset - from embodichain.learning.models.nerd import NeRDModel, NeRDModelCfg, export_model - from embodichain.data_pipeline.dynamics.shards import load_manifest + from examples.learning.nerd._support._reference import generate_official_dataset + from embodichain_models.dynamics.nerd import NeRDModel, NeRDModelCfg, export_model + from embodichain_models.dynamics.nerd.data import load_manifest # Generate new physical provenance from the actual original teacher. No # current-Newton manifest is rewritten to pretend it describes old physics. @@ -335,7 +337,9 @@ def test_real_historical_teacher_neural_runtime_and_source_contract( check=True, env={ **os.environ, - "PYTHONPATH": dependencies + os.pathsep + str(_ROOT), + "PYTHONPATH": os.pathsep.join( + (dependencies, str(_ROOT), str(_ROOT / "extensions/embodichain_models")) + ), "OMP_NUM_THREADS": "2", "MKL_NUM_THREADS": "2", "CUDA_VISIBLE_DEVICES": "", @@ -350,7 +354,7 @@ def _check_supported_replay( import h5py import numpy as np import torch - from embodichain.learning.models.nerd import NeRDModel, NeRDModelCfg, export_model + from embodichain_models.dynamics.nerd import NeRDModel, NeRDModelCfg, export_model metadata = json.loads((data / f"{robot}-passive/manifest.json").read_text())[ "metadata" @@ -525,7 +529,9 @@ def test_real_supported_environment_teacher_replay(tmp_path: Path, robot: str) - check=True, env={ **os.environ, - "PYTHONPATH": dependencies + os.pathsep + str(_ROOT), + "PYTHONPATH": os.pathsep.join( + (dependencies, str(_ROOT), str(_ROOT / "extensions/embodichain_models")) + ), "CUDA_VISIBLE_DEVICES": "", "OMP_NUM_THREADS": "2", "MKL_NUM_THREADS": "2", diff --git a/tests/learning/dynamics/test_reference_traces.py b/tests/examples/learning/nerd/test_reference_traces.py similarity index 98% rename from tests/learning/dynamics/test_reference_traces.py rename to tests/examples/learning/nerd/test_reference_traces.py index 7257b0c7bc..69f887c178 100644 --- a/tests/learning/dynamics/test_reference_traces.py +++ b/tests/examples/learning/nerd/test_reference_traces.py @@ -28,11 +28,11 @@ import pytest import torch -from embodichain.learning.dynamics._reference_traces import ( +from examples.learning.nerd._support._reference_traces import ( _load_reference_trace, _reference_trace_batches, ) -from embodichain.learning.dynamics.reference import OFFICIAL_COMMIT +from examples.learning.nerd._support._reference import OFFICIAL_COMMIT def _fixture( diff --git a/tests/learning/dynamics/test_report.py b/tests/examples/learning/nerd/test_report.py similarity index 96% rename from tests/learning/dynamics/test_report.py rename to tests/examples/learning/nerd/test_report.py index 788bca3ab3..9285374181 100644 --- a/tests/learning/dynamics/test_report.py +++ b/tests/examples/learning/nerd/test_report.py @@ -24,7 +24,7 @@ import pytest -from embodichain.learning.dynamics.report import ( +from examples.learning.nerd._support._report import ( _curve, _dataset_evidence, _matched_rollouts, @@ -32,7 +32,7 @@ _rollout_evidence, annotate_rollout_evidence, ) -from embodichain.learning.dynamics.reference import OFFICIAL_COMMIT, _hash +from examples.learning.nerd._support._reference import OFFICIAL_COMMIT, _hash def test_learning_curve_retains_sparse_full_budget_observations(tmp_path: Path) -> None: @@ -65,7 +65,7 @@ def test_learning_curve_keeps_irregular_final_observation(tmp_path: Path) -> Non def test_isolated_supported_source_import_has_verified_dataset_lineage(tmp_path): import h5py import numpy as np - from embodichain.learning.dynamics.reference import import_legacy_datasets + from examples.learning.nerd._support._reference import import_legacy_datasets sources = {} identities = {} @@ -123,7 +123,7 @@ def test_isolated_supported_source_import_has_verified_dataset_lineage(tmp_path) def test_source_split_labels_do_not_replace_supported_runtime_receipts( tmp_path, field, value ): - from embodichain.learning.dynamics.report import _supported_split_sources + from examples.learning.nerd._support._report import _supported_split_sources raw = tmp_path / "source.h5" raw.write_bytes(b"synthetic raw-file identity fixture") @@ -330,7 +330,7 @@ def test_unmatched_teacher_traces_or_fixed_policy_never_form_seed_pair(tmp_path) def test_historical_rollout_uses_pinned_warp_identity_without_newton(tmp_path): - from embodichain.learning.dynamics.reference import OFFICIAL_COMMIT + from examples.learning.nerd._support._reference import OFFICIAL_COMMIT record, cell = rollout_fixture(tmp_path) source = Path(record["source_report"]) @@ -373,9 +373,9 @@ def test_full_profile_checks_actual_feature_order_precision_and_training_control tmp_path, ): import torch - from embodichain.learning.dynamics.report import _full_profile_issues - from embodichain.learning.dynamics.trainer import DynamicsTrainCfg - from embodichain.learning.models.nerd import NeRDModelCfg + from examples.learning.nerd._support._report import _full_profile_issues + from examples.learning.nerd._support._trainer import DynamicsTrainCfg + from embodichain_models.dynamics.nerd import NeRDModelCfg model = NeRDModelCfg(input_dims={"states_embedding": 4, "joint_f": 2}).to_dict() root = tmp_path / "initial" diff --git a/tests/examples/learning/nerd/test_resources.py b/tests/examples/learning/nerd/test_resources.py new file mode 100644 index 0000000000..b42532abd1 --- /dev/null +++ b/tests/examples/learning/nerd/test_resources.py @@ -0,0 +1,212 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Verify the explicitly requested showcase resource-download capability.""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path +import zipfile + +import huggingface_hub +import pytest + +from examples.learning.nerd import _resources as resources + + +def digest(value: bytes) -> str: + return hashlib.sha256(value).hexdigest() + + +@pytest.fixture +def asset_server(tmp_path, monkeypatch): + payloads = { + "ATTRIBUTION.md": b"reference attribution", + "LICENSE.md": b"Apache-2.0 license fixture", + "ant.xml": b"\n", + "cartpole_single.urdf": b"\n", + } + server = tmp_path / "server" + server.mkdir() + archive = server / "reference_assets.zip" + with zipfile.ZipFile(archive, "w") as bundle: + bundle.writestr("nerd_reference_assets/", b"") + for name, value in payloads.items(): + bundle.writestr(f"nerd_reference_assets/{name}", value) + archive_hash = digest(archive.read_bytes()) + member_hashes = {name: digest(value) for name, value in payloads.items()} + manifest = { + "schema_version": 1, + "license": "Apache-2.0", + "upstream": {"revision": resources._UPSTREAM_REVISION}, + "archive": { + "root": "nerd_reference_assets", + "sha256": archive_hash, + "size": archive.stat().st_size, + }, + "files": { + name: { + "path": f"nerd_reference_assets/{name}", + "sha256": digest(value), + "size": len(value), + } + for name, value in payloads.items() + }, + } + metadata = server / "manifest.json" + metadata.write_text(json.dumps(manifest)) + monkeypatch.setattr(resources, "_ASSET_FILES", member_hashes) + monkeypatch.setattr(resources, "_ARCHIVE_SHA256", archive_hash) + monkeypatch.setattr(resources, "_MANIFEST_SHA256", digest(metadata.read_bytes())) + calls = [] + + def download(**options): + assert options["repo_id"] == "DexForceAI/embodichain_data" + assert options["repo_type"] == "dataset" + assert options["revision"] == "97e202e41fe7fb504656a6b163291305e400780d" + calls.append(options) + return str(server / Path(options["filename"]).name) + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", download) + return payloads, calls, archive + + +def test_pinned_asset_cache_repairs_corruption_and_reuses_verified_files( + tmp_path, asset_server +): + payloads, calls, _ = asset_server + cache = tmp_path / "cache" + target = resources.prepare_reference_assets(cache) + assert {path.name for path in target.iterdir()} == set(payloads) + assert all( + (target / name).read_bytes() == value for name, value in payloads.items() + ) + assert len(calls) == 2 + assert resources.prepare_reference_assets(cache, local_files_only=True) == target + assert len(calls) == 2 # A verified cache needs neither Hub nor network access. + (target / "ant.xml").write_bytes(b"damaged local cache") + assert resources.prepare_reference_assets(cache, local_files_only=True) == target + assert len(calls) == 4 and calls[-1]["local_files_only"] is True + assert (target / "ant.xml").read_bytes() == payloads["ant.xml"] + assert not list(target.parent.glob(".staging-*")) + + +def test_invalid_download_archive_is_rejected_before_extraction(tmp_path, asset_server): + _, _, archive = asset_server + archive.write_bytes(archive.read_bytes() + b"changed") + cache = tmp_path / "cache" + with pytest.raises(ValueError, match="SHA256 mismatch"): + resources.prepare_reference_assets(cache) + assert not (cache / "reference_assets").exists() + + +def test_example_resource_environment_and_explicit_cache_precedence( + tmp_path, asset_server, monkeypatch +): + monkeypatch.setenv("EMBODICHAIN_NERD_RESOURCE_ROOT", str(tmp_path / "configured")) + environmental = resources.prepare_reference_assets() + assert environmental.is_relative_to(tmp_path / "configured/assets") + explicit = resources.prepare_reference_assets(tmp_path / "explicit") + assert explicit.is_relative_to(tmp_path / "explicit") + + +def test_assets_only_does_not_start_model_or_dataset_downloads( + tmp_path, asset_server, monkeypatch +): + def unexpected_snapshot(**kwargs): + raise AssertionError("Assets-only preparation must not download large bundles") + + monkeypatch.setattr(huggingface_hub, "snapshot_download", unexpected_snapshot) + prepared = resources.prepare_resources( + tmp_path / "resources", "cartpole", assets_only=True + ) + assert set(prepared) == {"assets"} + assert prepared["assets"].is_dir() + assert not (tmp_path / "resources/datasets").exists() + assert not (tmp_path / "resources/models").exists() + + +def test_published_payload_revisions_paths_and_checksums( + tmp_path, asset_server, monkeypatch +): + calls = [] + manifests = {"dataset": {}, "model": {}} + monkeypatch.setattr(resources, "_PAYLOAD_MANIFEST_SHA256", manifests) + metadata = { + "codec": {"robot": "cartpole"}, + "dt": 1 / 60, + "asset_sha256": resources._ASSET_FILES["cartpole_single.urdf"], + } + + def snapshot(**options): + calls.append(options) + is_dataset = options["repo_type"] == "dataset" + assert options["revision"] == ( + "b9ec58f36da99cc99a93f97adfe74ebc9fd66586" + if is_dataset + else "78e581526570069410bf269b38c0c52c800b24a7" + ) + directory = Path(options["local_dir"]) / "neural_dynamics/nerd/cartpole/v1" + directory.mkdir(parents=True, exist_ok=True) + if is_dataset: + value = b"immutable dataset shard" + (directory / "shard-000001.h5").write_bytes(value) + manifest = { + "schema_version": 1, + "metadata": metadata, + "shards": [{"file": "shard-000001.h5", "sha256": digest(value)}], + } + else: + value = b"immutable portable weights" + (directory / "weights.pt").write_bytes(value) + manifest = { + "format_version": 1, + **metadata, + "weights_sha256": digest(value), + } + (directory / "manifest.json").write_text(json.dumps(manifest)) + manifests["dataset" if is_dataset else "model"]["cartpole"] = digest( + (directory / "manifest.json").read_bytes() + ) + + monkeypatch.setattr(huggingface_hub, "snapshot_download", snapshot) + prepared = resources.prepare_resources(tmp_path / "resources", "cartpole") + assert ( + prepared["dataset"] + == tmp_path / "resources/datasets/neural_dynamics/nerd/cartpole/v1" + ) + assert ( + prepared["model"] + == tmp_path / "resources/models/neural_dynamics/nerd/cartpole/v1" + ) + assert len(calls) == 2 + resources.prepare_resources( + tmp_path / "resources", "cartpole", local_files_only=True + ) + assert len(calls) == 2 + (prepared["model"] / "weights.pt").write_bytes(b"bad weights") + monkeypatch.setattr(huggingface_hub, "snapshot_download", lambda **kwargs: None) + with pytest.raises(ValueError, match="SHA256 mismatch"): + resources.prepare_resources(tmp_path / "resources", "cartpole") + # A self-consistent replacement still cannot impersonate the pinned release. + manifest_path = prepared["model"] / "manifest.json" + manifest = json.loads(manifest_path.read_text()) + manifest["weights_sha256"] = digest(b"bad weights") + manifest_path.write_text(json.dumps(manifest)) + with pytest.raises(ValueError, match="manifest.json"): + resources.prepare_resources(tmp_path / "resources", "cartpole") diff --git a/tests/learning/dynamics/test_sample_order.py b/tests/examples/learning/nerd/test_sample_order.py similarity index 98% rename from tests/learning/dynamics/test_sample_order.py rename to tests/examples/learning/nerd/test_sample_order.py index 7886d24b0b..efca95807f 100644 --- a/tests/learning/dynamics/test_sample_order.py +++ b/tests/examples/learning/nerd/test_sample_order.py @@ -24,7 +24,7 @@ import pytest import torch -from embodichain.learning.dynamics._sample_order import ( +from examples.learning.nerd._support._sample_order import ( _CheckpointedOrder, _OfficialShuffleOrder, _open_order, @@ -204,7 +204,7 @@ def test_trainer_admits_official_order_and_rejects_a_different_dataset_size( tmp_path: Path, ) -> None: from types import SimpleNamespace - from embodichain.learning.dynamics.trainer import DynamicsTrainer + from examples.learning.nerd._support._trainer import DynamicsTrainer path, _, _, expected = _official_order(tmp_path) # Exercise the production order admission boundary without creating a diff --git a/tests/learning/dynamics/test_small_budget_observation.py b/tests/examples/learning/nerd/test_small_budget_observation.py similarity index 97% rename from tests/learning/dynamics/test_small_budget_observation.py rename to tests/examples/learning/nerd/test_small_budget_observation.py index 9f711f3292..c83d7b1159 100644 --- a/tests/learning/dynamics/test_small_budget_observation.py +++ b/tests/examples/learning/nerd/test_small_budget_observation.py @@ -30,11 +30,9 @@ @pytest.fixture(scope="module") def workers() -> tuple[Any, Any]: - directory = ( - Path(__file__).resolve().parents[3] / "scripts/experiments/neural_dynamics" - ) + directory = Path(__file__).resolve().parents[4] / "examples/learning/nerd/_support" modules = [] - for filename in ("official_small_budget", "matched_small_budget"): + for filename in ("_official_small_budget", "_matched_small_budget"): spec = importlib.util.spec_from_file_location( "_test_" + filename, directory / (filename + ".py") ) diff --git a/tests/compute/dynamics/test_dynamics_task_reference.py b/tests/examples/learning/nerd/test_task_values.py similarity index 97% rename from tests/compute/dynamics/test_dynamics_task_reference.py rename to tests/examples/learning/nerd/test_task_values.py index 6f65154d54..e32f28d29c 100644 --- a/tests/compute/dynamics/test_dynamics_task_reference.py +++ b/tests/examples/learning/nerd/test_task_values.py @@ -17,7 +17,7 @@ from __future__ import annotations import pytest import torch -from embodichain.compute.dynamics import reference_task_values +from examples.learning.nerd._support._task_values import reference_task_values def test_cartpole_observation_and_official_cost_definition() -> None: diff --git a/tests/models/__init__.py b/tests/models/__init__.py new file mode 100644 index 0000000000..051f680fe4 --- /dev/null +++ b/tests/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. +# ---------------------------------------------------------------------------- + +"""NeRD model and showcase support.""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/tests/models/dynamics/__init__.py b/tests/models/dynamics/__init__.py new file mode 100644 index 0000000000..051f680fe4 --- /dev/null +++ b/tests/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. +# ---------------------------------------------------------------------------- + +"""NeRD model and showcase support.""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/tests/models/dynamics/nerd/__init__.py b/tests/models/dynamics/nerd/__init__.py new file mode 100644 index 0000000000..051f680fe4 --- /dev/null +++ b/tests/models/dynamics/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/tests/learning/models/_nerd_reference_check.py b/tests/models/dynamics/nerd/_nerd_reference_check.py similarity index 99% rename from tests/learning/models/_nerd_reference_check.py rename to tests/models/dynamics/nerd/_nerd_reference_check.py index 5c31b58f27..d425810097 100644 --- a/tests/learning/models/_nerd_reference_check.py +++ b/tests/models/dynamics/nerd/_nerd_reference_check.py @@ -24,7 +24,7 @@ import sys from pathlib import Path import torch -from embodichain.learning.models.nerd import convert_legacy_checkpoint, load_model +from embodichain_models.dynamics.nerd import convert_legacy_checkpoint, load_model def compare_codec( @@ -50,7 +50,7 @@ def compare_codec( else: from newton import JointType from solvers.solver_neural import NeuralSolver - from embodichain.compute.dynamics import NeRDStateCodec, quat_rotate + from embodichain_models.dynamics.nerd.codec import NeRDStateCodec, quat_rotate is_ant = robot == "Ant" codec = NeRDStateCodec( diff --git a/tests/compute/dynamics/data/ant-official-frame.json b/tests/models/dynamics/nerd/data/ant-official-frame.json similarity index 100% rename from tests/compute/dynamics/data/ant-official-frame.json rename to tests/models/dynamics/nerd/data/ant-official-frame.json diff --git a/tests/compute/dynamics/data/ant-official-frame.npz b/tests/models/dynamics/nerd/data/ant-official-frame.npz similarity index 100% rename from tests/compute/dynamics/data/ant-official-frame.npz rename to tests/models/dynamics/nerd/data/ant-official-frame.npz diff --git a/tests/compute/dynamics/test_codec.py b/tests/models/dynamics/nerd/test_codec.py similarity index 99% rename from tests/compute/dynamics/test_codec.py rename to tests/models/dynamics/nerd/test_codec.py index f2dbb6cdbe..220131d08b 100644 --- a/tests/compute/dynamics/test_codec.py +++ b/tests/models/dynamics/nerd/test_codec.py @@ -20,7 +20,7 @@ import pytest import torch from torch.autograd import gradcheck -from embodichain.compute.dynamics import ( +from embodichain_models.dynamics.nerd.codec import ( NeRDStateCodec, quat_rotate, plane_contact_features, diff --git a/tests/data_pipeline/dynamics/test_shards.py b/tests/models/dynamics/nerd/test_data.py similarity index 99% rename from tests/data_pipeline/dynamics/test_shards.py rename to tests/models/dynamics/nerd/test_data.py index ee82ed0ff1..2576b06117 100644 --- a/tests/data_pipeline/dynamics/test_shards.py +++ b/tests/models/dynamics/nerd/test_data.py @@ -24,7 +24,7 @@ import torch from torch.utils.data import DataLoader -from embodichain.data_pipeline.dynamics.shards import ( +from embodichain_models.dynamics.nerd.data import ( DynamicsShardWriter, DynamicsWindowDataset, load_manifest, @@ -109,7 +109,7 @@ def test_validation_and_integrity(tmp_path): def test_progress_only_publication_preserves_data_identity_and_pending_rows(tmp_path): - from embodichain.data_pipeline.dynamics.shards import dataset_fingerprint + from embodichain_models.dynamics.nerd.data import dataset_fingerprint with DynamicsShardWriter(tmp_path, {}, trajectories_per_shard=4) as writer: writer.add_trajectory("accepted", *trajectory(), split="train") diff --git a/tests/learning/models/test_nerd.py b/tests/models/dynamics/nerd/test_model.py similarity index 97% rename from tests/learning/models/test_nerd.py rename to tests/models/dynamics/nerd/test_model.py index 577b3d6e3e..c14a5c17f3 100644 --- a/tests/learning/models/test_nerd.py +++ b/tests/models/dynamics/nerd/test_model.py @@ -22,8 +22,8 @@ import sys import pytest import torch -from embodichain.compute.dynamics import append_history -from embodichain.learning.models.nerd import ( +from embodichain_models.dynamics.nerd.codec import append_history +from embodichain_models.dynamics.nerd import ( NeRDModel, NeRDModelCfg, export_model, @@ -313,7 +313,11 @@ def test_explicit_official_checkpoint_forward_gradient_and_adam( None, ( os.environ.get("NERD_REFERENCE_DEPS"), - str(Path(__file__).resolve().parents[3]), + str(Path(__file__).resolve().parents[4]), + str( + Path(__file__).resolve().parents[4] + / "extensions/embodichain_models" + ), ), ) ), diff --git a/tests/learning/dynamics/test_trainer.py b/tests/models/dynamics/nerd/test_training.py similarity index 95% rename from tests/learning/dynamics/test_trainer.py rename to tests/models/dynamics/nerd/test_training.py index ad58335428..deab29142b 100644 --- a/tests/learning/dynamics/test_trainer.py +++ b/tests/models/dynamics/nerd/test_training.py @@ -30,15 +30,14 @@ import torch from torch.utils.data import default_collate -from embodichain.compute.dynamics import NeRDStateCodec -from embodichain.data_pipeline.dynamics.shards import ( +from embodichain_models.dynamics.nerd.codec import NeRDStateCodec +from embodichain_models.dynamics.nerd.data import ( DynamicsShardWriter, DynamicsWindowDataset, ) -from embodichain.learning.models.nerd import NeRDModel, NeRDModelCfg, load_model -from embodichain.learning.dynamics.trainer import ( - DynamicsTrainCfg, - DynamicsTrainer, +from embodichain_models.dynamics.nerd import NeRDModel, NeRDModelCfg, load_model +from examples.learning.nerd._support._trainer import DynamicsTrainCfg, DynamicsTrainer +from embodichain_models.dynamics.nerd.integration import ( compute_statistics, encode_batch, evaluate_model, @@ -46,8 +45,8 @@ normalized_mse, save_statistics, ) -from embodichain.learning.dynamics.workflow import run_train, run_export, run_eval -from embodichain.learning.dynamics.qualification import ( +from examples.learning.nerd._support._workflow import run_train, run_export, run_eval +from examples.learning.nerd._support._qualification import ( paired_noninferiority, compare_training_steps, capture_training_step, @@ -281,7 +280,7 @@ def test_shared_recipe_matches_legacy_updates_and_restores_sampler( tmp_path, numerical_profile ): """The new CLI recipe retains the established NeRD schedule and updates.""" - from embodichain.learning.dynamics.sl_recipe import build_supervised_trainer + from embodichain_models.dynamics.nerd.sl_recipe import build_supervised_trainer make_data(tmp_path / "data") legacy = build( @@ -363,7 +362,7 @@ def test_shared_recipe_matches_legacy_updates_and_restores_sampler( report = train_from_config( { - "factory": "embodichain.learning.dynamics.sl_recipe:build_supervised_trainer", + "factory": "embodichain_models.dynamics.nerd.sl_recipe:build_supervised_trainer", **config, }, resume=str(tmp_path / "shared/last.pt"), @@ -535,7 +534,7 @@ def test_resume_new_directory_preserves_best_export_without_new_improvement( make_data(tmp_path / "data") scores = iter([1.0, 2.0, 3.0]) monkeypatch.setattr( - "embodichain.learning.dynamics.trainer.evaluate_model", + "examples.learning.nerd._support._trainer.evaluate_model", lambda *args, **kwargs: {"normalized_mse": next(scores)}, ) original = build(tmp_path / "data", tmp_path / "original", steps=4) @@ -585,7 +584,7 @@ def test_resume_preserves_future_best_after_interrupted_publication( make_data(tmp_path / "data") scores = iter([1.0, 0.5]) monkeypatch.setattr( - "embodichain.learning.dynamics.trainer.evaluate_model", + "examples.learning.nerd._support._trainer.evaluate_model", lambda *args, **kwargs: {"normalized_mse": next(scores)}, ) original = build(tmp_path / "data", tmp_path / "original", steps=4) @@ -848,7 +847,17 @@ def test_training_capture_detects_altered_update(): def test_offline_import_does_not_load_simulation(): - code = "import sys; import embodichain.learning.dynamics.workflow; import embodichain.learning.dynamics.trainer; assert not any(k == 'dexsim' or k.startswith('embodichain.lab.sim') for k in sys.modules)" + code = """ +import sys +from embodichain_models.dynamics.nerd import NeRDModel +assert not any(name in sys.modules for name in ('h5py', 'newton', 'warp')) +from embodichain_models.dynamics.nerd import build_supervised_trainer +assert 'newton' not in sys.modules +assert not any(name.startswith('examples.') for name in sys.modules) +import examples.learning.nerd._support._workflow +import examples.learning.nerd._support._trainer +assert not any(name == 'dexsim' or name.startswith('embodichain.lab.sim') for name in sys.modules) +""" subprocess.run([sys.executable, "-c", code], check=True) @@ -875,7 +884,8 @@ def test_interrupted_workflow_before_validation_is_json_serializable(tmp_path): ) -def test_ant_offline_training_and_metrics(tmp_path): +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_ant_offline_training_and_metrics(tmp_path, dtype): """A floating-base model trains through native-to-legacy targets without sim.""" metadata = {"codec": {"robot": "ant"}, "dt": 1 / 60} with DynamicsShardWriter(tmp_path / "data", metadata) as writer: @@ -904,7 +914,7 @@ def test_ant_offline_training_and_metrics(tmp_path): cfg = architecture() cfg.input_dims = {key: value.shape[-1] for key, value in inputs.items()} cfg.output_dim = 29 - model = NeRDModel(cfg) + model = NeRDModel(cfg).to(dtype=dtype) trainer = DynamicsTrainer( model, codec, @@ -919,6 +929,19 @@ def test_ant_offline_training_and_metrics(tmp_path): assert np.isfinite(report["normalized_mse"]) assert report["base_orientation_error_rad"] >= 0.0 assert report["base_position_error"] >= 0.0 + run_export(tmp_path / "run" / "last.pt", tmp_path / "export") + exported, _ = load_model(tmp_path / "export") + assert next(exported.parameters()).dtype == dtype + offline = run_eval( + {"dataset": str(tmp_path / "data"), "model": str(tmp_path / "export")} + ) + assert offline["split"] == "test" + for metric in ( + "normalized_mse", + "base_orientation_error_rad", + "base_position_error", + ): + assert np.isfinite(offline[metric]) def test_memory_mapped_order_resume_and_integrity(tmp_path): @@ -948,7 +971,7 @@ def test_memory_mapped_order_resume_and_integrity(tmp_path): def test_checkpointed_order_resume_is_identical_to_uninterrupted_training(tmp_path): - from embodichain.learning.dynamics._sample_order import _write_checkpointed_order + from examples.learning.nerd._support._sample_order import _write_checkpointed_order make_data(tmp_path / "data") first = build(tmp_path / "data", tmp_path / "full", steps=10) @@ -1160,7 +1183,7 @@ def _gradient_parameters(): @pytest.mark.parametrize("profile", ["default", "reference"]) def test_numerical_clipping_matches_its_independent_cpu_oracle(profile): - from embodichain.learning.dynamics._numerics import _clip_grad_norm_ + from embodichain_models.dynamics.nerd._numerics import _clip_grad_norm_ parameters = _gradient_parameters() expected = [parameter.grad.clone() for parameter in parameters] @@ -1239,7 +1262,7 @@ def backward(gradient): def test_attention_context_restores_flags_on_failure(): - from embodichain.learning.dynamics._numerics import _attention_context + from embodichain_models.dynamics.nerd._numerics import _attention_context before = ( torch.backends.cuda.flash_sdp_enabled(), @@ -1286,7 +1309,7 @@ def test_unknown_numerical_profile_is_rejected_before_output(tmp_path): def test_reference_clipping_rejects_unsupported_gradient_precision(): - from embodichain.learning.dynamics._numerics import _clip_grad_norm_ + from embodichain_models.dynamics.nerd._numerics import _clip_grad_norm_ parameter = torch.nn.Parameter(torch.ones(2, dtype=torch.float64)) parameter.grad = torch.ones_like(parameter) @@ -1402,7 +1425,7 @@ def test_reference_cuda_resume_preserves_model_adam_and_rng(tmp_path): @pytest.mark.gpu def test_reference_cuda_clipping_stays_on_device_and_preserves_ordered_norm(): - from embodichain.learning.dynamics._numerics import _clip_grad_norm_ + from embodichain_models.dynamics.nerd._numerics import _clip_grad_norm_ if not torch.cuda.is_available(): pytest.skip("CUDA is unavailable") 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_models_package.py b/tests/test_models_package.py new file mode 100644 index 0000000000..cc200a2313 --- /dev/null +++ b/tests/test_models_package.py @@ -0,0 +1,150 @@ +# ---------------------------------------------------------------------------- +# 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 models ship with the core and import without repository fallbacks.""" + +from __future__ import annotations + +import ast +from pathlib import Path +import subprocess +import sys + +import pytest +from setuptools import Distribution +from setuptools.command.build_py import build_py + +from setup import get_package_dir, get_packages + +pytestmark = pytest.mark.no_sim +_ROOT = Path(__file__).resolve().parents[1] + + +@pytest.fixture(scope="module") +def staged_models(tmp_path_factory: pytest.TempPathFactory) -> Path: + """Stage actual runtime sources with the distribution's package-data rules.""" + tree = ast.parse((_ROOT / "setup.py").read_text()) + call = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "setup" + ) + data = ast.literal_eval( + next(item.value for item in call.keywords if item.arg == "package_data") + ) + distribution = Distribution( + { + "packages": [ + name + for name in get_packages() + if name.partition(".")[0] in {"embodichain", "embodichain_models"} + ], + "package_dir": get_package_dir(), + "package_data": data, + "include_package_data": False, + "script_name": str(_ROOT / "setup.py"), + } + ) + output = tmp_path_factory.mktemp("staged_official_models") + command = build_py(distribution) + command.build_lib = str(output) + command.ensure_finalized() + command.run() + return output + + +def test_official_models_and_third_party_license_are_staged( + staged_models: Path, +) -> None: + for name in ( + "embodichain_models/__init__.py", + "embodichain_models/dynamics/__init__.py", + "embodichain_models/dynamics/nerd/__init__.py", + "embodichain_models/dynamics/nerd/model.py", + "embodichain_models/dynamics/nerd/checkpoint.py", + "embodichain_models/dynamics/nerd/codec.py", + "embodichain_models/dynamics/nerd/sl_recipe.py", + ): + assert (staged_models / name).is_file(), name + license_path = Path("embodichain_models/dynamics/nerd/THIRD_PARTY_LICENSES") + source = _ROOT / "extensions/embodichain_models" / license_path + assert (staged_models / license_path).read_bytes() == source.read_bytes() + assert not (staged_models / "embodichain/data/nerd/ant.xml").exists() + assert not (staged_models / "embodichain/data/nerd/cartpole_single.urdf").exists() + + +def test_staged_models_and_sl_factory_work_without_checkout_or_examples( + staged_models: Path, tmp_path: Path +) -> None: + script = """ +import importlib.abc +from pathlib import Path +import sys + +stage, repository = map(Path, sys.argv[1:]) +sys.meta_path = [finder for finder in sys.meta_path if '__editable__' not in str(finder)] +sys.path = [str(stage)] + [entry for entry in sys.path if entry and not Path(entry).is_relative_to(repository)] + +class NoSimulationOrExamples(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname.split('.')[0] in {'dexsim', 'newton', 'examples'} or fullname.startswith('embodichain.lab'): + raise AssertionError('model inference imported ' + fullname) + +sys.meta_path.insert(0, NoSimulationOrExamples()) +from embodichain.learning.models import MLP, EmpiricalNormalizer +import embodichain.learning.sl +assert not any(name == 'embodichain_models' or name.startswith('embodichain_models.') for name in sys.modules) +import embodichain_models +assert not any(name.startswith('embodichain_models.dynamics.') for name in sys.modules) +from embodichain_models.dynamics.nerd import NeRDModel, NeRDModelCfg +from embodichain_models.dynamics.nerd import model +assert Path(model.__file__).is_relative_to(stage) +import embodichain +assert Path(embodichain.__file__).is_relative_to(stage) +cfg = NeRDModelCfg(input_dims={'states_embedding': 4, 'joint_f': 2}, output_dim=4, n_embd=12, n_head=3, n_layer=1, block_size=3, history_length=3, head_hidden_dims=(12,)) +network = NeRDModel(cfg) +import torch +prediction = network({'states_embedding': torch.zeros(1, 3, 4), 'joint_f': torch.zeros(1, 3, 2)}) +assert prediction.shape == (1, 3, 4) +assert torch.isfinite(prediction).all() +from embodichain_models.dynamics.nerd.sl_recipe import build_supervised_trainer +assert callable(build_supervised_trainer) +from embodichain_models.dynamics.nerd.data import DynamicsShardWriter +from embodichain.learning.sl.workflow import train_from_config +import numpy as np +data = Path.cwd() / 'data' +with DynamicsShardWriter(data, {'codec': {'robot': 'cartpole'}, 'dt': 1 / 60}) as writer: + for split in ('train', 'validation'): + writer.add_trajectory(split, {'q': np.zeros((6, 2), np.float32), 'qd': np.zeros((6, 2), np.float32)}, {'joint_f': np.zeros((5, 2), np.float32)}, split=split) +result = train_from_config({ + 'factory': 'embodichain_models.dynamics.nerd.sl_recipe:build_supervised_trainer', + 'model': cfg.to_dict(), + 'data': {'path': str(data)}, + 'trainer': {'total_updates': 2, 'batch_size': 1, 'device': 'cpu', 'output_dir': str(Path.cwd() / 'training'), 'validation_interval': 1}, +}) +assert result['num_updates'] == 2 +assert Path(result['checkpoint']).is_file() +""" + result = subprocess.run( + [sys.executable, "-I", "-c", script, str(staged_models), str(_ROOT)], + cwd=tmp_path, + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/tests/test_release_metadata.py b/tests/test_release_metadata.py index c727e8a147..737f9298cf 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", } From 60f9b070144b66d6a0345d706e0db46496633ec0 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 25 Oct 2026 14:40:21 +0000 Subject: [PATCH 19/21] fix(nerd): retain experiment root for comparison receipts --- .../learning/nerd/_support/_comparison.py | 8 +-- .../examples/learning/nerd/test_comparison.py | 52 +++++++++++++++++++ 2 files changed, 57 insertions(+), 3 deletions(-) diff --git a/examples/learning/nerd/_support/_comparison.py b/examples/learning/nerd/_support/_comparison.py index ac214e1812..dfafe1a70e 100644 --- a/examples/learning/nerd/_support/_comparison.py +++ b/examples/learning/nerd/_support/_comparison.py @@ -439,9 +439,11 @@ def run_comparison( if plan["reference_dependencies"]: pythonpaths.append(plan["reference_dependencies"]) pythonpaths.append(plan["source"]["root"]) - root = Path(__file__).resolve().parents[4] - if (root / "pyproject.toml").is_file(): - pythonpaths.extend((str(root / "extensions/embodichain_models"), str(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) diff --git a/tests/examples/learning/nerd/test_comparison.py b/tests/examples/learning/nerd/test_comparison.py index 6eece29673..b7cbe43d8e 100644 --- a/tests/examples/learning/nerd/test_comparison.py +++ b/tests/examples/learning/nerd/test_comparison.py @@ -121,6 +121,58 @@ def test_changed_common_artifact_fails_before_execution(tmp_path, monkeypatch): _read_plan(path) +def test_comparison_keeps_execution_receipts_and_reuses_completed_cells( + tmp_path, monkeypatch +): + import examples.learning.nerd._support._comparison as comparison + + experiment_root = tmp_path / "experiment" + path = prepare_comparison(configuration(tmp_path, monkeypatch)) + _, plan = _read_plan(path) + prepared = [cell for cell in plan["cells"] if cell["status"] == "prepared"] + assert len(prepared) == 2 + source_root = tmp_path / "checkout" + source_root.mkdir() + (source_root / "pyproject.toml").write_text("") + monkeypatch.setattr( + comparison, + "__file__", + str(source_root / "examples/learning/nerd/_support/_comparison.py"), + ) + commands = [] + + def worker(command, **kwargs): + commands.append(command) + spec = json.loads(Path(command[command.index("--spec") + 1]).read_text()) + (Path(spec["output_dir"]) / "result.json").write_text( + json.dumps( + { + "status": "completed", + "completed_updates": spec["train"]["total_steps"], + } + ) + ) + return subprocess.CompletedProcess(command, 0) + + monkeypatch.setattr(comparison.subprocess, "run", worker) + result = run_comparison(path, trust_external_code=True) + assert result["executed"] == 2 + assert {run["id"] for run in result["runs"] if run["status"] == "completed"} == { + cell["id"] for cell in prepared + } + assert {file.name for file in experiment_root.glob("*.receipt.json")} == { + cell["id"] + ".receipt.json" for cell in prepared + } + assert json.loads((experiment_root / "execution.json").read_text()) == { + "runs": result["runs"] + } + assert not list(source_root.glob("*.receipt.json")) + assert not (source_root / "execution.json").exists() + reused = run_comparison(path, trust_external_code=True) + assert reused == {"runs": result["runs"], "executed": 0} + assert len(commands) == 2 + + def test_candidate_worker_uses_frozen_artifacts(tmp_path, monkeypatch): path = prepare_comparison(configuration(tmp_path, monkeypatch)) _, plan = _read_plan(path) From 2c9ee8875f8617c0d6a33f65907d5be75343af2e Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 25 Oct 2026 15:20:48 +0000 Subject: [PATCH 20/21] fix(nerd): repair asset downloads and remove dedicated unit tests --- agent_context/MAP.yaml | 3 +- .../topics/neural-dynamics/neural-dynamics.md | 6 +- examples/learning/nerd/_resources.py | 2 + tests/examples/learning/__init__.py | 21 - tests/examples/learning/nerd/__init__.py | 21 - .../examples/learning/nerd/test_acceptance.py | 75 - .../examples/learning/nerd/test_comparison.py | 261 --- .../nerd/test_dynamics_runtime_identity.py | 183 -- .../examples/learning/nerd/test_entrypoint.py | 487 ------ .../examples/learning/nerd/test_generation.py | 334 ---- tests/examples/learning/nerd/test_metrics.py | 99 -- .../examples/learning/nerd/test_reference.py | 961 ----------- .../learning/nerd/test_reference_rollout.py | 539 ------ .../learning/nerd/test_reference_traces.py | 247 --- tests/examples/learning/nerd/test_report.py | 414 ----- .../examples/learning/nerd/test_resources.py | 212 --- .../learning/nerd/test_sample_order.py | 328 ---- .../nerd/test_small_budget_observation.py | 167 -- .../learning/nerd/test_task_values.py | 70 - tests/models/__init__.py | 21 - tests/models/dynamics/__init__.py | 21 - tests/models/dynamics/nerd/__init__.py | 21 - .../dynamics/nerd/_nerd_reference_check.py | 348 ---- .../nerd/data/ant-official-frame.json | 64 - .../dynamics/nerd/data/ant-official-frame.npz | Bin 34898 -> 0 bytes tests/models/dynamics/nerd/test_codec.py | 331 ---- tests/models/dynamics/nerd/test_data.py | 369 ----- tests/models/dynamics/nerd/test_model.py | 325 ---- tests/models/dynamics/nerd/test_training.py | 1465 ----------------- tests/test_models_package.py | 150 -- 30 files changed, 6 insertions(+), 7539 deletions(-) delete mode 100644 tests/examples/learning/__init__.py delete mode 100644 tests/examples/learning/nerd/__init__.py delete mode 100644 tests/examples/learning/nerd/test_acceptance.py delete mode 100644 tests/examples/learning/nerd/test_comparison.py delete mode 100644 tests/examples/learning/nerd/test_dynamics_runtime_identity.py delete mode 100644 tests/examples/learning/nerd/test_entrypoint.py delete mode 100644 tests/examples/learning/nerd/test_generation.py delete mode 100644 tests/examples/learning/nerd/test_metrics.py delete mode 100644 tests/examples/learning/nerd/test_reference.py delete mode 100644 tests/examples/learning/nerd/test_reference_rollout.py delete mode 100644 tests/examples/learning/nerd/test_reference_traces.py delete mode 100644 tests/examples/learning/nerd/test_report.py delete mode 100644 tests/examples/learning/nerd/test_resources.py delete mode 100644 tests/examples/learning/nerd/test_sample_order.py delete mode 100644 tests/examples/learning/nerd/test_small_budget_observation.py delete mode 100644 tests/examples/learning/nerd/test_task_values.py delete mode 100644 tests/models/__init__.py delete mode 100644 tests/models/dynamics/__init__.py delete mode 100644 tests/models/dynamics/nerd/__init__.py delete mode 100644 tests/models/dynamics/nerd/_nerd_reference_check.py delete mode 100644 tests/models/dynamics/nerd/data/ant-official-frame.json delete mode 100644 tests/models/dynamics/nerd/data/ant-official-frame.npz delete mode 100644 tests/models/dynamics/nerd/test_codec.py delete mode 100644 tests/models/dynamics/nerd/test_data.py delete mode 100644 tests/models/dynamics/nerd/test_model.py delete mode 100644 tests/models/dynamics/nerd/test_training.py delete mode 100644 tests/test_models_package.py diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 18b44d5967..aaa6071a80 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -283,8 +283,7 @@ topics: - examples/learning/nerd/_support/ - examples/learning/nerd/configs/ watch_paths: [extensions/embodichain_models/, examples/learning/nerd/, - tests/models/dynamics/nerd/, tests/examples/learning/nerd/, - tests/test_models_package.py, docs/source/features/neural_dynamics.md, + 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 diff --git a/agent_context/topics/neural-dynamics/neural-dynamics.md b/agent_context/topics/neural-dynamics/neural-dynamics.md index f397bad326..682240f8fa 100644 --- a/agent_context/topics/neural-dynamics/neural-dynamics.md +++ b/agent_context/topics/neural-dynamics/neural-dynamics.md @@ -81,9 +81,9 @@ 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. Tests are in `tests/models/dynamics/nerd/`, -`tests/examples/learning/nerd/`, `tests/test_models_package.py` and the common SL -suite. The [feature guide](../../../docs/source/features/neural_dynamics.md) +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/examples/learning/nerd/_resources.py b/examples/learning/nerd/_resources.py index 7ca427a998..2cd739b2e3 100644 --- a/examples/learning/nerd/_resources.py +++ b/examples/learning/nerd/_resources.py @@ -200,6 +200,8 @@ def prepare_reference_assets( "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, } diff --git a/tests/examples/learning/__init__.py b/tests/examples/learning/__init__.py deleted file mode 100644 index 051f680fe4..0000000000 --- a/tests/examples/learning/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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/tests/examples/learning/nerd/__init__.py b/tests/examples/learning/nerd/__init__.py deleted file mode 100644 index 051f680fe4..0000000000 --- a/tests/examples/learning/nerd/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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/tests/examples/learning/nerd/test_acceptance.py b/tests/examples/learning/nerd/test_acceptance.py deleted file mode 100644 index 37ecc24df0..0000000000 --- a/tests/examples/learning/nerd/test_acceptance.py +++ /dev/null @@ -1,75 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""CPU-only CLI preflight for independent-seed acceptance evidence.""" - -from __future__ import annotations - -from pathlib import Path -import subprocess -import sys - -import pytest - -from examples.learning.nerd._support import _acceptance as nerd_acceptance - - -@pytest.mark.parametrize( - ("seeds", "message"), - [ - ("0,0,0,0,0", "seeds must be distinct"), - ("0,1,2,3,00", "seeds must be distinct"), - ("0,1,2,3,invalid", "seeds must be comma-separated integers"), - ("", "seeds must be comma-separated integers"), - ], -) -def test_invalid_seed_list_fails_before_output_or_workers( - tmp_path: Path, seeds: str, message: str -) -> None: - runner = ( - Path(__file__).resolve().parents[4] - / "examples/learning/nerd/_support/_acceptance.py" - ) - output = tmp_path / "results" - # Missing artifacts must never be inspected or passed to a worker. - result = subprocess.run( - [ - sys.executable, - str(runner), - "--reference-source", - str(tmp_path / "missing-reference"), - "--datasets", - str(tmp_path / "missing-datasets"), - "--models", - str(tmp_path / "missing-models"), - "--output", - str(output), - "--seeds", - seeds, - ], - capture_output=True, - text=True, - timeout=10, - ) - assert result.returncode == 2 - assert message in result.stderr - assert not output.exists() - - -@pytest.mark.parametrize("seeds", [[0], [0, 1, 2, 3], [0, 1, 2, 3, 4]]) -def test_unique_seed_lists_preserve_short_and_full_runs(seeds: list[int]) -> None: - # Short runs remain available; the existing qualification gate returns pending. - assert nerd_acceptance._parse_seeds(",".join(map(str, seeds))) == seeds diff --git a/tests/examples/learning/nerd/test_comparison.py b/tests/examples/learning/nerd/test_comparison.py deleted file mode 100644 index b7cbe43d8e..0000000000 --- a/tests/examples/learning/nerd/test_comparison.py +++ /dev/null @@ -1,261 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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. -# ---------------------------------------------------------------------------- - -"""Immutable experiment preparation and truthful pending qualification reports.""" - -from __future__ import annotations - -import json -import os -from pathlib import Path -import subprocess - -import numpy as np -import pytest -import torch - -from embodichain_models.dynamics.nerd.data import DynamicsShardWriter -from examples.learning.nerd._support._comparison import ( - prepare_comparison, - run_comparison, - _read_plan, -) -from examples.learning.nerd._support._report import write_comparison_report - - -def configuration(tmp_path, monkeypatch): - import examples.learning.nerd._support._comparison as comparison - from examples.learning.nerd._support._reference import OFFICIAL_COMMIT - - monkeypatch.setattr( - comparison, - "reference_identity", - lambda path: { - "root": str(tmp_path / "source"), - "commit": OFFICIAL_COMMIT, - "files": {}, - "source_patches": [], - }, - ) - with DynamicsShardWriter( - tmp_path / "data", {"codec": {"robot": "cartpole"}, "dt": 1 / 60} - ) as writer: - for split in ("train", "validation", "test"): - q = np.zeros((6, 2), np.float32) - qd = np.zeros_like(q) - writer.add_trajectory( - split, - {"q": q, "qd": qd}, - {"joint_f": np.ones((5, 2), np.float32)}, - split=split, - ) - return { - "output_dir": str(tmp_path / "experiment"), - "source": str(tmp_path / "source"), - "datasets": {"official_warp": None, "current_newton": str(tmp_path / "data")}, - "scope": "smoke", - "seeds": [0], - "model": { - "input_dims": {"states_embedding": 4, "joint_f": 2}, - "output_dim": 4, - "n_embd": 12, - "n_head": 3, - "n_layer": 1, - "history_length": 3, - "block_size": 3, - "head_hidden_dims": [12], - }, - "train": { - "total_steps": 2, - "batch_size": 2, - "validation_interval": 1, - "device": "cpu", - }, - } - - -def test_preparation_shared_artifacts_and_missing_row_are_explicit( - tmp_path, monkeypatch -): - torch.set_num_threads(1) - path = prepare_comparison(configuration(tmp_path, monkeypatch)) - root, plan = _read_plan(path) - assert len(plan["cells"]) == 4 - assert sum(cell["status"] == "pending" for cell in plan["cells"]) == 2 - specs = [ - json.loads(Path(cell["spec"]).read_text()) - for cell in plan["cells"] - if cell["status"] == "prepared" - ] - assert specs[0]["initial_model_sha256"] == specs[1]["initial_model_sha256"] - assert specs[0]["sample_order_sha256"] == specs[1]["sample_order_sha256"] - with pytest.raises(ValueError, match="trust_external_code"): - run_comparison(path) - report = write_comparison_report(path) - assert report["status"] == "pending" - assert any("Five" in reason for reason in report["pending"]) - assert (root / "report" / "training-curves.png").is_file() - assert (root / "report" / "REPORT.md").is_file() - - -def test_changed_common_artifact_fails_before_execution(tmp_path, monkeypatch): - path = prepare_comparison(configuration(tmp_path, monkeypatch)) - _, plan = _read_plan(path) - name = next(iter(plan["artifacts"])) - with (path.parent / name).open("ab") as stream: - stream.write(b"tampered") - with pytest.raises(ValueError, match="Immutable comparison artifact"): - _read_plan(path) - - -def test_comparison_keeps_execution_receipts_and_reuses_completed_cells( - tmp_path, monkeypatch -): - import examples.learning.nerd._support._comparison as comparison - - experiment_root = tmp_path / "experiment" - path = prepare_comparison(configuration(tmp_path, monkeypatch)) - _, plan = _read_plan(path) - prepared = [cell for cell in plan["cells"] if cell["status"] == "prepared"] - assert len(prepared) == 2 - source_root = tmp_path / "checkout" - source_root.mkdir() - (source_root / "pyproject.toml").write_text("") - monkeypatch.setattr( - comparison, - "__file__", - str(source_root / "examples/learning/nerd/_support/_comparison.py"), - ) - commands = [] - - def worker(command, **kwargs): - commands.append(command) - spec = json.loads(Path(command[command.index("--spec") + 1]).read_text()) - (Path(spec["output_dir"]) / "result.json").write_text( - json.dumps( - { - "status": "completed", - "completed_updates": spec["train"]["total_steps"], - } - ) - ) - return subprocess.CompletedProcess(command, 0) - - monkeypatch.setattr(comparison.subprocess, "run", worker) - result = run_comparison(path, trust_external_code=True) - assert result["executed"] == 2 - assert {run["id"] for run in result["runs"] if run["status"] == "completed"} == { - cell["id"] for cell in prepared - } - assert {file.name for file in experiment_root.glob("*.receipt.json")} == { - cell["id"] + ".receipt.json" for cell in prepared - } - assert json.loads((experiment_root / "execution.json").read_text()) == { - "runs": result["runs"] - } - assert not list(source_root.glob("*.receipt.json")) - assert not (source_root / "execution.json").exists() - reused = run_comparison(path, trust_external_code=True) - assert reused == {"runs": result["runs"], "executed": 0} - assert len(commands) == 2 - - -def test_candidate_worker_uses_frozen_artifacts(tmp_path, monkeypatch): - path = prepare_comparison(configuration(tmp_path, monkeypatch)) - _, plan = _read_plan(path) - cell = next( - cell - for cell in plan["cells"] - if cell["status"] == "prepared" and cell["trainer"] == "embodichain" - ) - repository = Path(__file__).resolve().parents[4] - environment = { - **os.environ, - "CUDA_VISIBLE_DEVICES": "", - "OMP_NUM_THREADS": "1", - "PYTHONPATH": os.pathsep.join( - (str(repository / "extensions/embodichain_models"), str(repository)) - ), - } - process = subprocess.run( - cell["command"], - cwd=tmp_path, - env=environment, - capture_output=True, - text=True, - timeout=30, - ) - assert process.returncode == 0, process.stdout + process.stderr - result = json.loads((Path(cell["output_dir"]) / "result.json").read_text()) - assert result["completed_updates"] == 2 - assert result["test"]["split"] == "test" - assert result["qualification"] == "matched_training_run_only" - - -def test_rollout_attachment_preserves_measurements_and_requires_identity(tmp_path): - from examples.learning.nerd._support._report import annotate_rollout_evidence - - source = tmp_path / "source.json" - source.write_text( - json.dumps({"trajectories": 2048, "horizons": {"1": {"q_mae": 0.1}}}) - ) - with pytest.raises(ValueError, match="model_sha256"): - annotate_rollout_evidence(source, "run", tmp_path / "bad.json") - raw = { - "model_sha256": "abc", - "motion_mode": "passive", - "trajectories": 2048, - "horizons": { - "1": { - "q_mae": 0.1, - "qd_mae": 0.2, - "failure_rate": 0.0, - "nonfinite_rate": 0.0, - }, - "10": { - "q_mae": 0.2, - "qd_mae": 0.3, - "failure_rate": 0.1, - "nonfinite_rate": 0.05, - }, - }, - } - source.write_text(json.dumps(raw)) - result = annotate_rollout_evidence(source, "run", tmp_path / "attached.json") - assert result["horizons"] == [1, 10] - assert result["errors"] == raw["horizons"] - assert result["failure_rate"] == 0.1 - assert "policy_return" not in result - - -def test_preparation_keeps_selected_virtualenv_interpreter_symlink( - tmp_path, monkeypatch -): - import sys - - chosen = tmp_path / "selected-env/bin/python" - chosen.parent.mkdir(parents=True) - chosen.symlink_to(sys.executable) - config = configuration(tmp_path, monkeypatch) - config["python"] = str(chosen) - path = prepare_comparison(config) - _, plan = _read_plan(path) - assert plan["python"] == str(chosen) - assert all( - cell["command"][0] == str(chosen) - for cell in plan["cells"] - if cell["status"] == "prepared" - ) diff --git a/tests/examples/learning/nerd/test_dynamics_runtime_identity.py b/tests/examples/learning/nerd/test_dynamics_runtime_identity.py deleted file mode 100644 index a58d1e10a3..0000000000 --- a/tests/examples/learning/nerd/test_dynamics_runtime_identity.py +++ /dev/null @@ -1,183 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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. -# ---------------------------------------------------------------------------- - -"""Execution provenance is read-only, simulation-free and best effort.""" - -from __future__ import annotations - -import json -import subprocess -import sys -from types import SimpleNamespace - -import pytest -import torch - -from examples.learning.nerd._support import _runtime - - -def test_cpu_identity_reflects_actual_parameters_without_side_effects(monkeypatch): - model = torch.nn.Linear(2, 1, dtype=torch.float64) - rng = torch.get_rng_state().clone() - modules = set(sys.modules) - monkeypatch.delitem(sys.modules, "warp", raising=False) - - def unexpected(*args, **kwargs): - pytest.fail("CPU identity must not query CUDA or nvidia-smi") - - monkeypatch.setattr(_runtime, "_driver_inventory", unexpected) - monkeypatch.setattr(torch.cuda, "get_device_name", unexpected) - result = _runtime._training_runtime_identity(model) - assert result["parameter_placements"] == [ - {"device": "cpu", "dtype": "torch.float64"} - ] - assert result["cuda"]["used"] is False - assert result["cuda"]["driver_inventory"] == {"status": "not_used"} - assert result["torch_version"] == str(torch.__version__) - assert result["python"] and result["platform"] - assert result["warp"] == {"status": "not_loaded"} - assert not any( - name.split(".")[0] in {"warp", "dexsim", "newton"} - for name in set(sys.modules) - modules - ) - assert torch.equal(rng, torch.get_rng_state()) - assert json.loads(json.dumps(result)) == result - - -def test_empty_and_mixed_parameter_metadata_and_loaded_warp(monkeypatch): - assert ( - _runtime._training_runtime_identity(torch.nn.Identity())[ - "parameter_placement_status" - ] - == "no_parameters" - ) - model = torch.nn.Module() - model.register_parameter( - "a", torch.nn.Parameter(torch.zeros(1, dtype=torch.float32)) - ) - model.register_parameter( - "b", torch.nn.Parameter(torch.zeros(1, dtype=torch.float64)) - ) - monkeypatch.setitem( - sys.modules, "warp", SimpleNamespace(__version__="test-version") - ) - result = _runtime._training_runtime_identity(model) - assert result["parameter_placements"] == [ - {"device": "cpu", "dtype": "torch.float32"}, - {"device": "cpu", "dtype": "torch.float64"}, - ] - assert result["warp"] == {"status": "already_loaded", "version": "test-version"} - - -def _mock_cuda_model(): - # Metadata only: no CUDA allocations or initialization in these tests. - return SimpleNamespace( - parameters=lambda: iter( - [ - SimpleNamespace(device=torch.device("cuda:1"), dtype=torch.float32), - SimpleNamespace(device=torch.device("cuda:1"), dtype=torch.float32), - ] - ) - ) - - -def test_cuda_inventory_is_bounded_and_keeps_physical_ids_separate(monkeypatch): - monkeypatch.setattr(_runtime.shutil, "which", lambda name: "/usr/bin/nvidia-smi") - calls = [] - - def run(command, **kwargs): - calls.append((command, kwargs)) - return SimpleNamespace( - returncode=0, stdout='0, GPU-physical, "Example, GPU", 595.1\n', stderr="" - ) - - monkeypatch.setattr(_runtime.subprocess, "run", run) - monkeypatch.setattr(torch.cuda, "get_device_name", lambda device: "Logical GPU") - result = _runtime._training_runtime_identity(_mock_cuda_model()) - assert result["cuda"]["devices"] == [ - {"device": "cuda:1", "status": "available", "name": "Logical GPU"} - ] - inventory = result["cuda"]["driver_inventory"] - assert inventory == { - "status": "available", - "physical_gpus": [ - { - "index": "0", - "uuid": "GPU-physical", - "name": "Example, GPU", - "driver_version": "595.1", - } - ], - } - assert calls[0][0][1:] == [ - "--query-gpu=index,uuid,name,driver_version", - "--format=csv,noheader,nounits", - ] - assert calls[0][1]["timeout"] == 2.0 - assert not calls[0][1].get("shell", False) - - -@pytest.mark.parametrize( - "failure, expected", - [ - ("missing", "unavailable"), - ("timeout", "error"), - ("exit", "error"), - ("oserror", "error"), - ("malformed", "error"), - ("oversized", "error"), - ], -) -def test_cuda_inventory_failure_does_not_abort_receipt(monkeypatch, failure, expected): - monkeypatch.setattr( - _runtime.shutil, - "which", - lambda name: None if failure == "missing" else "/mock/nvidia-smi", - ) - - def run(command, **kwargs): - if failure == "timeout": - raise subprocess.TimeoutExpired(command, kwargs["timeout"]) - if failure == "oserror": - raise OSError("driver unavailable") - return SimpleNamespace( - returncode=1 if failure == "exit" else 0, - stdout="x" * 65537 if failure == "oversized" else "bad,row", - stderr="driver error", - ) - - monkeypatch.setattr(_runtime.subprocess, "run", run) - - def device_name(device): - raise RuntimeError("device unavailable") - - monkeypatch.setattr(torch.cuda, "get_device_name", device_name) - result = _runtime._training_runtime_identity(_mock_cuda_model()) - assert result["cuda"]["devices"][0]["status"] == "error" - assert result["cuda"]["driver_inventory"]["status"] == expected - json.dumps(result, allow_nan=False) - - -def test_fresh_process_identity_does_not_import_simulation(): - code = """ -import sys -import torch -from examples.learning.nerd._support._runtime import _training_runtime_identity -_training_runtime_identity(torch.nn.Linear(1, 1)) -assert not any(name.split('.')[0] in {'warp', 'dexsim', 'newton'} for name in sys.modules) -assert not any(name.startswith('embodichain.lab.sim') for name in sys.modules) -""" - subprocess.run([sys.executable, "-c", code], check=True, timeout=30) diff --git a/tests/examples/learning/nerd/test_entrypoint.py b/tests/examples/learning/nerd/test_entrypoint.py deleted file mode 100644 index 1b2e0368da..0000000000 --- a/tests/examples/learning/nerd/test_entrypoint.py +++ /dev/null @@ -1,487 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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 dynamics CLI routing, actual offline workflow and packaged configs.""" - -from __future__ import annotations - -import json -from copy import deepcopy -import os -from pathlib import Path -import subprocess -import sys -from types import ModuleType -from unittest.mock import Mock - -import numpy as np -import pytest -import torch -import yaml - -from examples.learning.nerd.__main__ import main - - -def test_help_and_dispatcher_are_simulation_free(): - code = "from examples.learning.nerd.__main__ import main; import sys\ntry: main(['--help'])\nexcept SystemExit as exc: assert exc.code == 0\nassert not any(k == 'dexsim' or k.startswith('embodichain.lab.sim') for k in sys.modules)" - subprocess.run( - [sys.executable, "-c", code], check=True, capture_output=True, text=True - ) - - -def test_generate_routing_and_path_resolution(tmp_path, monkeypatch): - host = ModuleType("examples.learning.nerd._support._generation") - host.generate_dataset = Mock(return_value={"status": "ok"}) - monkeypatch.setitem(sys.modules, host.__name__, host) - config = tmp_path / "generate.yaml" - config.write_text("robot: cartpole\noutput_dir: data\ndevice: cpu\nresume: false\n") - main(["generate", "--config", str(config), "--seed", "7", "--resume"]) - received = host.generate_dataset.call_args.args[0] - assert received["output_dir"] == str(tmp_path / "data") - assert received["seed"] == 7 and received["resume"] is True - - -def test_default_evaluation_writes_under_current_directory(tmp_path, monkeypatch): - host = ModuleType("examples.learning.nerd._support._workflow") - host.run_eval = Mock(return_value={"status": "ok"}) - monkeypatch.setitem(sys.modules, host.__name__, host) - monkeypatch.chdir(tmp_path) - resources = tmp_path / "resources" - prefix = Path("neural_dynamics/nerd/cartpole/v1") - for kind in ("datasets", "models"): - (resources / kind / prefix).mkdir(parents=True) - main(["eval", "--resource-root", str(resources), "--device", "cpu"]) - received = host.run_eval.call_args.args[0] - assert received["output"] == str(tmp_path / "outputs/nerd/cartpole/offline.json") - assert received["dataset"] == str(resources / "datasets" / prefix) - - -def test_invalid_paths_and_conversion_requires_explicit_trust(tmp_path): - config = tmp_path / "bad.yaml" - config.write_text("dataset: missing\n") - with pytest.raises(SystemExit) as error: - main(["train", "--config", str(config)]) - assert error.value.code == 2 - with pytest.raises(SystemExit) as error: - main( - [ - "convert-checkpoint", - "--checkpoint", - "model.pt", - "--output", - "export", - "--manifest", - "manifest.json", - ] - ) - assert error.value.code == 2 - - -def test_conversion_manifest_and_trust_reach_converter(tmp_path, monkeypatch): - import embodichain_models.dynamics.nerd as nerd - - checkpoint = tmp_path / "old.pt" - checkpoint.write_bytes(b"unused mock fixture") - manifest = tmp_path / "manifest.json" - manifest.write_text(json.dumps({"codec": {"robot": "cartpole"}, "dt": 1 / 60})) - convert = Mock() - monkeypatch.setattr(nerd, "convert_legacy_checkpoint", convert) - main( - [ - "convert-checkpoint", - "--checkpoint", - str(checkpoint), - "--output", - str(tmp_path / "export"), - "--manifest", - str(manifest), - "--trust-pickle", - ] - ) - assert convert.call_args.kwargs == {"trust_pickle": True} - assert convert.call_args.args[2]["codec"]["robot"] == "cartpole" - - -@pytest.mark.parametrize("trainer_kind", ["legacy", "sl"]) -def test_real_offline_cli_train_export_eval(tmp_path, trainer_kind): - from embodichain_models.dynamics.nerd.data import DynamicsShardWriter - - metadata = {"codec": {"robot": "cartpole"}, "dt": 1 / 60} - with DynamicsShardWriter(tmp_path / "data", metadata) as writer: - for split in ("train", "validation", "test"): - q = np.zeros((6, 2), np.float32) - q[:, 0] = np.arange(6) * 0.01 - qd = np.zeros_like(q) - qd[:, 0] = 0.6 - writer.add_trajectory( - split, - {"q": q, "qd": qd}, - {"joint_f": np.zeros((5, 2), np.float32)}, - split=split, - ) - training = { - "dataset": "data", - "output_dir": "train", - "model": { - "input_dims": {"states_embedding": 4, "joint_f": 2}, - "output_dim": 4, - "n_embd": 12, - "n_layer": 1, - "n_head": 3, - "block_size": 3, - "history_length": 3, - "head_hidden_dims": [12], - }, - "train": { - "total_steps": 2, - "batch_size": 2, - "device": "cpu", - "validation_interval": 1, - }, - } - if trainer_kind == "sl": - training = { - "factory": "embodichain_models.dynamics.nerd.sl_recipe:build_supervised_trainer", - "model": training["model"], - "data": {"path": str(tmp_path / "data")}, - "trainer": { - "total_updates": 2, - "batch_size": 2, - "device": "cpu", - "validation_interval": 1, - "checkpoint_interval": 0, - "output_dir": str(tmp_path / "train"), - }, - } - train_cfg = tmp_path / "train.json" - train_cfg.write_text(json.dumps(training)) - evaluation = { - "mode": "offline", - "dataset": "data", - "model": "export", - "output": "reports/evaluation/result.json", - "device": "cpu", - } - eval_cfg = tmp_path / "eval.yaml" - eval_cfg.write_text(yaml.safe_dump(evaluation)) - code = """import sys -from examples.learning.nerd.__main__ import main -main(['train','--config',sys.argv[1]]) -main(['export','--checkpoint',sys.argv[2],'--output',sys.argv[3]]) -main(['eval','--config',sys.argv[4]]) -assert not any(k == 'dexsim' or k.startswith('embodichain.lab.sim') for k in sys.modules) -""" - env = {**os.environ, "OMP_NUM_THREADS": "1", "MKL_NUM_THREADS": "1"} - result = subprocess.run( - [ - sys.executable, - "-c", - code, - str(train_cfg), - str( - tmp_path / "train" / ("last.pt" if trainer_kind == "sl" else "best.pt") - ), - str(tmp_path / "export"), - str(eval_cfg), - ], - capture_output=True, - text=True, - env=env, - ) - assert result.returncode == 0, result.stdout + result.stderr - report = json.loads((tmp_path / "reports/evaluation/result.json").read_text()) - assert report["split"] == "test" and report["windows"] == 3 - if trainer_kind == "sl": - from embodichain_models.dynamics.nerd import load_model - - model, _ = load_model(tmp_path / "export") - assert list(model.cfg.input_dims) == ["states_embedding", "joint_f"] - state = torch.load(tmp_path / "train/last.pt", weights_only=False) - for name, value in model.state_dict().items(): - torch.testing.assert_close( - value, state["learner"]["model"][name], rtol=0, atol=0 - ) - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) -@pytest.mark.parametrize("model_snapshot", [False, True]) -def test_training_checkpoint_cli_export_preserves_dtype_statistics_and_predictions( - tmp_path, dtype, model_snapshot -): - from embodichain_models.dynamics.nerd import NeRDModel, NeRDModelCfg, load_model - - torch.set_num_threads(2) - model = ( - NeRDModel( - NeRDModelCfg( - input_dims={"states_embedding": 4, "joint_f": 2}, - output_dim=4, - n_embd=12, - n_head=3, - n_layer=1, - block_size=3, - history_length=3, - head_hidden_dims=(12,), - ) - ) - .to(dtype=dtype) - .eval() - ) - # Values beyond FP32 precision expose a downcast even if tensors were later - # converted back to FP64. Frozen normalizers must retain their precision too. - with torch.no_grad(): - for parameter in model.parameters(): - parameter.add_(1e-10) - model.set_normalization( - { - key: ( - torch.full((width,), 1 / 7, dtype=dtype), - torch.full((width,), 1 / 3, dtype=dtype), - ) - for key, width in model.cfg.input_dims.items() - }, - torch.full((4,), 1 / 11, dtype=dtype), - torch.full((4,), 1 / 13, dtype=dtype), - ) - checkpoint = tmp_path / "training.pt" - state = { - "format": "embodichain.nerd.training.v1", - "model_cfg": model.cfg.to_dict(), - "model": model.state_dict(), - "train_cfg": {"numerical_profile": "default"}, - "manifest": {"codec": {"robot": "cartpole"}, "dt": 1 / 60}, - "global_step": 7, - "dataset_fingerprint": "fixed-training-corpus", - } - if model_snapshot: - state.pop("train_cfg") - state.update( - resumable=False, numerical_profile="default", attention_backend="default" - ) - torch.save(state, checkpoint) - destination = tmp_path / "export" - main(["export", "--checkpoint", str(checkpoint), "--output", str(destination)]) - loaded, metadata = load_model(destination) - assert metadata["training"]["global_step"] == 7 - assert all(value.dtype == dtype for value in loaded.state_dict().values()) - for name, expected in model.state_dict().items(): - torch.testing.assert_close(loaded.state_dict()[name], expected, atol=0, rtol=0) - inputs = { - name: torch.randn(2, 3, width, dtype=dtype) - for name, width in model.cfg.input_dims.items() - } - torch.testing.assert_close(model(inputs), loaded(inputs), atol=0, rtol=0) - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) -@pytest.mark.parametrize("normalize_input", [False, True]) -@pytest.mark.parametrize("normalize_output", [False, True]) -def test_fine_tune_resume_export_preserves_dtype_updates_and_input_order( - tmp_path, dtype, normalize_input, normalize_output -): - from embodichain_models.dynamics.nerd.data import DynamicsShardWriter - from examples.learning.nerd._support._workflow import run_train - from embodichain_models.dynamics.nerd import ( - NeRDModel, - NeRDModelCfg, - export_model, - load_model, - ) - - torch.set_num_threads(2) - metadata = {"codec": {"robot": "cartpole"}, "dt": 1 / 60} - with DynamicsShardWriter(tmp_path / "data", metadata) as writer: - for split in ("train", "validation", "test"): - q = np.zeros((6, 2), np.float32) - q[:, 0] = np.arange(6) * 0.01 - qd = np.zeros_like(q) - qd[:, 0] = 0.6 - writer.add_trajectory( - split, - {"q": q, "qd": qd}, - {"joint_f": np.zeros((5, 2), np.float32)}, - split=split, - ) - model = NeRDModel( - NeRDModelCfg( - input_dims={"states_embedding": 4, "joint_f": 2}, - output_dim=4, - n_embd=12, - n_head=3, - n_layer=1, - block_size=3, - history_length=3, - head_hidden_dims=(12,), - normalize_input=normalize_input, - normalize_output=normalize_output, - ) - ).to(dtype=dtype) - model.input_rms["states_embedding"].mean.fill_(1 / 7) - export_model(tmp_path / "pretrained", model, metadata) - # Use the same JSON-shaped config as the CLI, including list-valued fields. - config = json.loads( - json.dumps( - { - "dataset": str(tmp_path / "data"), - "output_dir": str(tmp_path / "interrupted"), - "fine_tune": str(tmp_path / "pretrained"), - "model": model.cfg.to_dict(), - "train": { - "total_steps": 2, - "batch_size": 2, - "device": "cpu", - "validation_interval": 1, - "checkpoint_interval": 1, - }, - "stop_after": 1, - } - ) - ) - - def train_cli(value, label): - path = tmp_path / f"{label}.json" - path.write_text(json.dumps(value)) - main(["train", "--config", str(path)]) - - def evaluate_cli(model_path, label): - path = tmp_path / f"{label}.json" - output = tmp_path / f"{label}-result.json" - path.write_text( - json.dumps( - { - "mode": "offline", - "dataset": str(tmp_path / "data"), - "model": str(model_path), - "device": "cpu", - "output": str(output), - } - ) - ) - main(["eval", "--config", str(path)]) - assert json.loads(output.read_text())["windows"] == 3 - - evaluate_cli(tmp_path / "pretrained", "pretrained-eval") - train_cli(config, "fine-tune") - checkpoint = tmp_path / "interrupted/last.pt" - saved = torch.load(checkpoint, weights_only=True) - assert all(value.dtype == dtype for value in saved["model"].values()) - resumed_config = deepcopy(config) - resumed_config.pop("fine_tune") - resumed_config.update( - resume=str(checkpoint), output_dir=str(tmp_path / "resumed"), stop_after=2 - ) - reordered = deepcopy(resumed_config) - reordered["output_dir"] = str(tmp_path / "reordered") - reordered["model"]["input_dims"] = dict( - reversed(list(reordered["model"]["input_dims"].items())) - ) - with pytest.raises(ValueError, match="input_dims order"): - run_train(reordered) - assert list((tmp_path / "reordered").iterdir()) == [] - train_cli(resumed_config, "resume") - uninterrupted = deepcopy(config) - uninterrupted.update(output_dir=str(tmp_path / "full"), stop_after=2) - train_cli(uninterrupted, "uninterrupted") - expected = torch.load(tmp_path / "full/last.pt", weights_only=True) - actual = torch.load(tmp_path / "resumed/last.pt", weights_only=True) - assert actual["global_step"] == 2 - for name, value in expected["model"].items(): - torch.testing.assert_close(actual["model"][name], value, atol=0, rtol=0) - assert actual["best_step"] == expected["best_step"] - main( - [ - "export", - "--checkpoint", - str(tmp_path / "resumed/best.pt"), - "--output", - str(tmp_path / "export"), - ] - ) - loaded, _ = load_model(tmp_path / "export") - assert all(value.dtype == dtype for value in loaded.state_dict().values()) - assert loaded.cfg.normalize_input is normalize_input - assert loaded.cfg.normalize_output is normalize_output - for name, value in actual["best_model"].items(): - torch.testing.assert_close(loaded.state_dict()[name], value, atol=0, rtol=0) - inputs = { - name: torch.randn(2, 3, width) for name, width in model.cfg.input_dims.items() - } - prediction = loaded(inputs) - assert prediction.dtype == dtype - assert torch.isfinite(prediction).all() - evaluate_cli(tmp_path / "export", "resumed-eval") - - -def test_packaged_workflow_configs_decode_through_cli(tmp_path, monkeypatch): - """Exercise the actual YAML loader/path resolver and model/train constructors.""" - import examples.learning.nerd._support._workflow as workflow - from examples.learning.nerd._support._trainer import DynamicsTrainCfg - from embodichain_models.dynamics.nerd import NeRDModelCfg - - host = ModuleType("examples.learning.nerd._support._generation") - host.generate_dataset = Mock(return_value={"status": "not_run"}) - monkeypatch.setitem(sys.modules, host.__name__, host) - - def train(config): - model = NeRDModelCfg(**config["model"]) - model.validate() - cfg = DynamicsTrainCfg(**config["train"]) - cfg.validate() - assert "id" not in config - if cfg.total_steps > 100: - assert cfg.total_steps == 999 * 5000 - assert cfg.lr_update_interval == 5000 and cfg.lr_offset == 1 - return {"status": "not_run"} - - monkeypatch.setattr(workflow, "run_train", train) - monkeypatch.setattr(workflow, "run_eval", lambda config: {"status": "not_run"}) - repo = Path(__file__).resolve().parents[4] - (tmp_path / "data").mkdir() - (tmp_path / "export").mkdir() - files = [] - expected_configs = { - "dynamics.generate.smoke.yaml", - "dynamics.generate.full.yaml", - "dynamics.generate.warp-reference.yaml", - "dynamics.train.smoke.yaml", - "dynamics.train.full.yaml", - "dynamics.eval.smoke.yaml", - "dynamics.eval.full.yaml", - "dynamics.eval.pretrained.yaml", - } - for robot in ("cartpole", "ant"): - robot_configs = list( - (repo / "examples" / "learning" / "nerd" / "configs" / robot).glob( - "dynamics.*.yaml" - ) - ) - assert {path.name for path in robot_configs} == expected_configs - files.extend(robot_configs) - for source in files: - config = yaml.safe_load(source.read_text()) - # Temporary destinations let path validation run without generating data. - if "dataset" in config: - config["dataset"] = "data" - if isinstance(config.get("model"), str): - config["model"] = "export" - if "output_dir" in config: - config["output_dir"] = "out" - if "output" in config: - config["output"] = "report.json" - copied = tmp_path / source.name - copied.write_text(yaml.safe_dump(config, sort_keys=False)) - operation = source.name.split(".")[1] - main([operation, "--config", str(copied)]) diff --git a/tests/examples/learning/nerd/test_generation.py b/tests/examples/learning/nerd/test_generation.py deleted file mode 100644 index 001f7e2a6c..0000000000 --- a/tests/examples/learning/nerd/test_generation.py +++ /dev/null @@ -1,334 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Real analytical collection and an SDK-free generate/train/export/eval path.""" - -from __future__ import annotations - -import os -import json -from pathlib import Path -import subprocess -import sys - -import h5py -import pytest -import torch - -from examples.learning.nerd._support._generation import ( - DynamicsGenerationCfg, - _sample, - generate_dataset, -) -from examples.learning.nerd._support._newton import _NewtonTeacher -from embodichain_models.dynamics.nerd.data import load_manifest, trajectory_split - - -@pytest.fixture -def deterministic_generation(monkeypatch): - """Exercise persistence with reproducible accepted and rejected candidates.""" - from examples.learning.nerd._support import _generation as generation - - class Teacher: - def __init__(self, robot, *, device, num_envs, **kwargs): - self.device = torch.device("cpu") - self.gravity_dir = torch.tensor([0.0, 0.0, -1.0]).expand(num_envs, -1) - self.manifest = { - "physics_signature": "fake-physical-v1", - "teacher_origin": {"collector": "fake"}, - } - - def reset(self, q, qd): - self.q, self.qd = q.clone(), qd.clone() - self.rejected = torch.tensor([int(seed) in {11, 13} for seed in q[:, 0]]) - self.index = 0 - - def snapshot(self): - q = self.q + self.index * 0.01 - q[self.rejected] = float("nan") - return {"q": q, "qd": self.qd.clone()}, {} - - def step(self, force): - self.index += 1 - - def __enter__(self): - return self - - def __exit__(self, *args): - pass - - def sample(robot, seed, horizon, *args): - return torch.tensor([float(seed), 0.0]), torch.zeros(2), torch.zeros(horizon, 2) - - monkeypatch.setattr(generation, "_NewtonTeacher", Teacher) - monkeypatch.setattr(generation, "_sample", sample) - - -def test_resume_retains_rejections_when_interrupted_after_shard_manifest( - tmp_path, monkeypatch, deterministic_generation -): - from embodichain_models.dynamics.nerd import data as shards - - config = { - "output_dir": str(tmp_path), - "seed": 11, - "trajectories": 2, - "num_envs": 2, - "horizon": 2, - "trajectories_per_shard": 1, - } - atomic_json = shards._atomic_json - - def interrupt_after_publish(path, value): - atomic_json(path, value) - if path.name == "manifest.json" and value["shards"]: - raise RuntimeError("interrupted-after-manifest-publication") - - monkeypatch.setattr(shards, "_atomic_json", interrupt_after_publish) - with pytest.raises(RuntimeError, match="interrupted-after-manifest-publication"): - generate_dataset(config) - committed = load_manifest(tmp_path, verify=True) - assert [ - row["id"] for shard in committed["shards"] for row in shard["trajectories"] - ] == ["11:1"] - # This legacy report can lag behind the committed shard; it cannot be the - # recovery authority once progress is included in the shard manifest. - (tmp_path / "generation-progress.json").write_text( - json.dumps({"next_candidate": 0, "rejected": 0, "rejections": []}) - ) - monkeypatch.setattr(shards, "_atomic_json", atomic_json) - result = generate_dataset({**config, "resume": True}) - assert result["trajectories"] == 2 - assert result["rejected"] == 2 - evidence = json.loads(Path(result["rejection_evidence"]).read_text()) - assert evidence["next_candidate"] == 4 - assert [record["id"] for record in evidence["rejections"]] == ["11:0", "11:2"] - final = load_manifest(tmp_path, verify=True) - assert final["metadata"] == committed["metadata"] - assert final["shards"][0] == committed["shards"][0] - baseline = tmp_path / "uninterrupted" - generate_dataset({**config, "output_dir": str(baseline)}) - expected = load_manifest(baseline, verify=True) - assert final["metadata"] == expected["metadata"] - for actual_shard, expected_shard in zip(final["shards"], expected["shards"]): - assert actual_shard["trajectories"] == expected_shard["trajectories"] - with ( - h5py.File(tmp_path / actual_shard["file"]) as actual_file, - h5py.File(baseline / expected_shard["file"]) as expected_file, - ): - for group in ("states", "transitions"): - for field in actual_file[f"0/{group}"]: - assert torch.equal( - torch.from_numpy(actual_file[f"0/{group}/{field}"][()]), - torch.from_numpy(expected_file[f"0/{group}/{field}"][()]), - ) - - -def test_resume_reads_legacy_generation_progress(tmp_path, deterministic_generation): - config = { - "output_dir": str(tmp_path), - "seed": 11, - "trajectories": 1, - "num_envs": 1, - "horizon": 2, - "trajectories_per_shard": 1, - "max_attempts": 1, - } - with pytest.raises(RuntimeError, match="Generation attempt limit"): - generate_dataset(config) - # Even a batch with no accepted trajectories advances a recoverable cursor. - manifest = load_manifest(tmp_path) - assert not manifest["shards"] - assert manifest["generation_progress"]["next_candidate"] == 1 - assert manifest["generation_progress"]["rejected"] == 1 - # Older collectors store progress only in the separate JSON report. - progress = manifest.pop("generation_progress") - (tmp_path / "manifest.json").write_text(json.dumps(manifest)) - (tmp_path / "generation-progress.json").write_text(json.dumps(progress)) - result = generate_dataset({**config, "resume": True, "max_attempts": 2}) - assert result["trajectories"] == result["rejected"] == 1 - evidence = json.loads(Path(result["rejection_evidence"]).read_text()) - assert [record["id"] for record in evidence["rejections"]] == ["11:0"] - - -@pytest.mark.parametrize( - "field,value", - [ - ("num_envs", True), - ("horizon", 2.5), - ("resume", "false"), - ("seed", False), - ("use_cuda_graph", 1), - ("effort_scale", float("nan")), - ("physics_profile", "unregistered"), - ], -) -def test_generation_rejects_invalid_config(field, value): - with pytest.raises((ValueError, TypeError)): - DynamicsGenerationCfg(**{field: value}) - - -@pytest.mark.parametrize("robot", ["cartpole", "ant"]) -def test_real_generation_partial_batch_resume_and_teacher_replay(tmp_path, robot): - config = { - "robot": robot, - "output_dir": str(tmp_path / "data"), - "num_envs": 2, - "trajectories": 3, - "horizon": 3, - "num_substeps": 2, - "seed": 7, - "split_seed": 11, - "trajectories_per_shard": 2, - "initial_states": "default", - "effort_scale": 0.01, - } - assert generate_dataset(config)["transitions"] == 9 - first = load_manifest(tmp_path / "data", verify=True) - original = [row for shard in first["shards"] for row in shard["trajectories"]] - assert ( - generate_dataset({**config, "trajectories": 5, "resume": True})["transitions"] - == 15 - ) - manifest = load_manifest(tmp_path / "data", verify=True) - rows = [row for shard in manifest["shards"] for row in shard["trajectories"]] - assert [row["id"] for row in rows] == [f"7:{i}" for i in range(5)] - assert rows[:3] == original - assert [row["split"] for row in rows] == [ - trajectory_split(row["id"], 11) for row in rows - ] - assert manifest["metadata"]["teacher_origin"]["collector"] == "direct_newton" - shard = manifest["shards"][0] - with h5py.File(tmp_path / "data" / shard["file"]) as file: - row = file[shard["trajectories"][0]["group"]] - q, qd, force = [ - torch.tensor(row[key][()]) - for key in ("states/q", "states/qd", "transitions/joint_f") - ] - contacts = { - key: torch.tensor(value[:1]) - for key, value in row["transitions"].items() - if key.startswith("contact_") - } - assert q.shape[0] == force.shape[0] + 1 == 4 - with _NewtonTeacher(robot, num_substeps=2) as teacher: - assert ( - teacher.manifest["physics_signature"] - == manifest["metadata"]["physics_signature"] - ) - teacher.reset(q[:1], qd[:1]) - previous, actual_contacts = teacher.snapshot() - for key, value in contacts.items(): - torch.testing.assert_close(actual_contacts[key], value) - teacher.step(force[:1]) - actual, _ = teacher.snapshot() - torch.testing.assert_close(actual["q"], q[1:2], atol=1e-5, rtol=1e-4) - torch.testing.assert_close(actual["qd"], qd[1:2], atol=1e-5, rtol=1e-4) - torch.testing.assert_close(previous["q"], q[:1], atol=0, rtol=0) - teacher.reset(q[:1], qd[:1]) - teacher.step(force[:1]) - repeated, _ = teacher.snapshot() - torch.testing.assert_close(repeated["q"], actual["q"], atol=1e-6, rtol=1e-5) - - -@pytest.mark.parametrize("robot", ["cartpole", "ant"]) -def test_actual_cli_pipeline_with_dexsim_imports_forbidden(tmp_path, robot): - """A native SDK or its custom solver extension cannot mask a dependency.""" - script = """ -import importlib.abc -import json -from pathlib import Path -import sys - -class NoSDK(importlib.abc.MetaPathFinder): - def find_spec(self, fullname, path=None, target=None): - if fullname == 'dexsim' or fullname.startswith(('dexsim.', 'embodichain.lab.')): - raise AssertionError('SDK-free workflow imported ' + fullname) -sys.meta_path.insert(0, NoSDK()) -from examples.learning.nerd.__main__ import main -from embodichain_models.dynamics.nerd.data import load_manifest -root, robot = Path(sys.argv[1]), sys.argv[2] -def run(operation, config): - path = root / (operation + '.json') - path.write_text(json.dumps(config)) - main([operation, '--config', str(path)]) -run('generate', dict(robot=robot, output_dir='data', num_envs=2, trajectories=64, - horizon=4, num_substeps=2, initial_states='default', effort_scale=0.01)) -import yaml -model_cfg = yaml.safe_load((Path.cwd() / 'examples/learning/nerd/configs' / robot / 'dynamics.train.smoke.yaml').read_text())['model'] -model_cfg.update(n_embd=12, history_length=2, head_hidden_dims=[12]) -run('train', dict(dataset='data', output_dir='train', - model=model_cfg, - train=dict(total_steps=2, batch_size=2, validation_interval=1, checkpoint_interval=1, - device='cpu'))) -main(['export', '--checkpoint', str(root/'train/best.pt'), - '--output', str(root/'export')]) -run('eval', dict(dataset='data', model='export', split='test', device='cpu', - output='reports/evaluation.json')) -report=json.loads((root/'reports/evaluation.json').read_text()) -assert report['windows'] > 0 -assert not any(name == 'dexsim' or name.startswith('embodichain.lab.') for name in sys.modules) -""" - result = subprocess.run( - [sys.executable, "-c", script, str(tmp_path), robot], - capture_output=True, - text=True, - env={**os.environ, "CUDA_VISIBLE_DEVICES": "", "OMP_NUM_THREADS": "2"}, - ) - assert result.returncode == 0, result.stdout + result.stderr - - -@pytest.mark.gpu -@pytest.mark.parametrize( - "robot,substeps,profile", - [ - ("cartpole", 5, "warp_reference"), - ("ant", 16, "warp_reference"), - ("cartpole", 10, "native"), - ], -) -def test_cuda_graph_and_direct_teacher_replay(robot, substeps, profile): - if not torch.cuda.is_available(): - pytest.skip("CUDA teacher required") - q, qd, effort = _sample(robot, 5, 3, "random_effort", "default", 0.01) - q, qd, effort = q[None].cuda(), qd[None].cuda(), effort.cuda() - trajectories = [] - for capture in (False, True): - with _NewtonTeacher( - robot, - device="cuda:0", - num_substeps=substeps, - physics_profile=profile, - use_cuda_graph=capture, - ) as teacher: - assert (teacher.graph is not None) is capture - teacher.reset(q, qd) - rows = [] - for force in effort: - teacher.step(force[None]) - state, _ = teacher.snapshot() - rows.append(torch.cat((state["q"], state["qd"]), -1)) - trajectories.append(torch.stack(rows)) - teacher.reset(q, qd) - teacher.step(effort[:1]) - repeated, _ = teacher.snapshot() - torch.testing.assert_close( - torch.cat((repeated["q"], repeated["qd"]), -1), - rows[0], - atol=1e-6, - rtol=1e-5, - ) - torch.testing.assert_close(trajectories[0], trajectories[1], atol=1e-5, rtol=1e-4) diff --git a/tests/examples/learning/nerd/test_metrics.py b/tests/examples/learning/nerd/test_metrics.py deleted file mode 100644 index ef079b6bc1..0000000000 --- a/tests/examples/learning/nerd/test_metrics.py +++ /dev/null @@ -1,99 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Independent geometry and simulator-free contracts for rollout metrics.""" - -from __future__ import annotations - -import math -import subprocess -import sys - -import numpy as np -import torch - -from examples.learning.nerd._support._evaluation import ( - _error_values, - _safe_number, - _teacher_row_sha256, -) - - -def test_metrics_import_without_any_simulation_runtime(): - subprocess.run( - [ - sys.executable, - "-c", - "import sys; from examples.learning.nerd._support._evaluation import _error_values; assert not any(k == 'newton' or k == 'dexsim' or k.startswith('embodichain.lab.sim') for k in sys.modules)", - ], - check=True, - capture_output=True, - text=True, - ) - - -def test_continuous_angle_and_bounded_joint_errors_have_distinct_meaning(): - cart = torch.tensor([[0.0, 2 * math.pi]], dtype=torch.float64) - error = _error_values( - "cartpole", - cart, - torch.zeros_like(cart), - torch.zeros_like(cart), - torch.zeros_like(cart), - ) - torch.testing.assert_close(error, torch.zeros_like(error), atol=1e-12, rtol=0) - q = torch.zeros(1, 15, dtype=torch.float64) - q[:, 6] = 1 - target = q.clone() - q[:, 7] = 2 * math.pi - v = torch.zeros(1, 14, dtype=torch.float64) - error = _error_values("ant", q, v, target, v) - assert error[0, 0].item() == math.pi / 4 - - -def test_orientation_metric_is_geodesic_and_quaternion_sign_invariant(): - q = torch.zeros(2, 15, dtype=torch.float64) - q[:, 6] = 1 - target = q.clone() - q[0, 5:7] = math.sqrt(0.5) - q[1, 6] = -1 - v = torch.zeros(2, 14, dtype=torch.float64) - error = _error_values("ant", q, v, target, v) - torch.testing.assert_close( - error[:, 3], - torch.tensor([math.pi / 2, 0], dtype=torch.float64), - atol=1e-12, - rtol=0, - ) - assert _safe_number(float("nan")) is None - assert _safe_number(float("inf")) is None - assert _safe_number(0.0) == 0.0 - - -def test_teacher_content_hash_changes_with_values_and_identity(): - row = {"id": "seed:1", "metadata": {"sampling_seed": 12}} - values = [ - np.zeros((3, 2), np.float32), - np.zeros((3, 2), np.float32), - np.zeros((2, 2), np.float32), - ] - expected = _teacher_row_sha256(row, values) - assert expected == _teacher_row_sha256( - {**row, "file": "different.h5"}, [v.copy() for v in values] - ) - assert expected != _teacher_row_sha256({**row, "id": "seed:2"}, values) - values[-1][0, 0] = 1 - assert expected != _teacher_row_sha256(row, values) diff --git a/tests/examples/learning/nerd/test_reference.py b/tests/examples/learning/nerd/test_reference.py deleted file mode 100644 index e1a9d09bf8..0000000000 --- a/tests/examples/learning/nerd/test_reference.py +++ /dev/null @@ -1,961 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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. -# ---------------------------------------------------------------------------- - -"""Legacy transport and opt-in independent official-trainer execution.""" - -from __future__ import annotations - -import json -import hashlib -import os -from pathlib import Path -import subprocess -import sys -from types import ModuleType, SimpleNamespace - -import h5py -import numpy as np -import pytest -import torch - -from embodichain_models.dynamics.nerd.data import ( - DynamicsShardWriter, - DynamicsWindowDataset, -) -from examples.learning.nerd._support._reference import ( - export_legacy_dataset, - run_reference, - reference_identity, - OFFICIAL_COMMIT, -) - - -def test_external_execution_requires_explicit_trust(): - with pytest.raises(ValueError, match="trust_external_code"): - run_reference({}) - - -@pytest.mark.parametrize("profile", ["reference", "unknown"]) -def test_original_method_adapter_rejects_unsupported_profile_before_source(profile): - with pytest.raises(ValueError, match="numerical_profile"): - run_reference( - {"train": {"numerical_profile": profile}}, trust_external_code=True - ) - - -def test_original_method_adapter_rejects_math_model_before_data(tmp_path, monkeypatch): - import examples.learning.nerd._support._reference as reference - import embodichain_models.dynamics.nerd as models - - for name in ("algorithms", "utils"): - package = ModuleType(name) - package.__path__ = [] - monkeypatch.setitem(sys.modules, name, package) - vanilla = ModuleType("algorithms.vanilla_trainer") - vanilla.VanillaTrainer = object - timers = ModuleType("utils.time_report") - timers.TimeReport = object - monkeypatch.setitem(sys.modules, vanilla.__name__, vanilla) - monkeypatch.setitem(sys.modules, timers.__name__, timers) - monkeypatch.setattr( - reference, "reference_identity", lambda source: {"root": str(tmp_path)} - ) - monkeypatch.setattr( - models, - "load_model", - lambda *args, **kwargs: ( - SimpleNamespace(cfg=SimpleNamespace(attention_backend="math")), - {}, - ), - ) - monkeypatch.setattr(sys, "path", list(sys.path)) - threads, precision = torch.get_num_threads(), torch.get_float32_matmul_precision() - try: - with pytest.raises(ValueError, match="math attention initial_model"): - run_reference( - { - "source": "unused", - "initial_model": "unused", - "train": {"device": "cpu"}, - }, - trust_external_code=True, - ) - finally: - torch.set_num_threads(threads) - torch.set_float32_matmul_precision(precision) - - -def test_official_transport_uses_the_exact_native_codec_velocity_order(): - from embodichain_models.dynamics.nerd.codec import NeRDStateCodec - from examples.learning.nerd._support._reference import _legacy_states - - with np.load( - Path(__file__).resolve().parents[3] - / "models/dynamics/nerd/data/ant-official-frame.npz" - ) as fixture: - q = torch.from_numpy(fixture["q"].copy()) - qd = torch.from_numpy(fixture["qd"].copy()) - codec = NeRDStateCodec("ant") - actual = _legacy_states(q, qd, codec) - torch.testing.assert_close(actual[..., :15], q, atol=0, rtol=0) - torch.testing.assert_close( - actual[..., 15:], codec.native_to_legacy(q, qd), atol=0, rtol=0 - ) - - -@pytest.mark.parametrize( - "device", ["cpu", pytest.param("cuda:0", marks=pytest.mark.gpu)] -) -def test_training_transport_device_is_independent_of_raw_residency(device): - from embodichain_models.dynamics.nerd.codec import NeRDStateCodec - from examples.learning.nerd._support._reference import _official_batch - - if device.startswith("cuda") and not torch.cuda.is_available(): - pytest.skip("CUDA is unavailable") - with np.load( - Path(__file__).resolve().parents[3] - / "models/dynamics/nerd/data/ant-official-frame.npz" - ) as fixture: - q = torch.from_numpy(fixture["q"].copy()) - qd = torch.from_numpy(fixture["qd"].copy()) - shape = (q.shape[0], q.shape[1] - 1) - batch = { - "states": {"q": q, "qd": qd}, - "transitions": { - "joint_f": torch.zeros(*shape, 14), - "contact_depths": torch.zeros(*shape, 25), - }, - } - resident = { - key: {name: tensor.to(device) for name, tensor in fields.items()} - for key, fields in batch.items() - } - codec = NeRDStateCodec("ant") - lazy = _official_batch(batch, codec, device=device) - cached = _official_batch(resident, codec, device=device) - for name, actual in lazy.items(): - assert actual.device == torch.device(device) - torch.testing.assert_close(actual, cached[name], atol=0, rtol=0) - - -def test_legacy_export_preserves_lengths_splits_and_native_to_twist(tmp_path): - metadata = {"codec": {"robot": "ant", "base_com": [0.0, 0.0, 0.0]}, "dt": 1 / 60} - with DynamicsShardWriter(tmp_path / "data", metadata) as writer: - for index, length in enumerate((3, 5)): - q = np.zeros((length + 1, 15), np.float32) - q[:, 0] = 1.0 - q[:, 6] = 1.0 - qd = np.zeros((length + 1, 14), np.float32) - qd[:, 1] = 2.0 - qd[:, 5] = 1.0 - transitions = { - "joint_f": np.zeros((length, 14), np.float32), - "contact_depths": np.zeros((length, 1), np.float32), - "contact_thicknesses_0": np.full((length, 1), 0.02, np.float32), - "contact_thicknesses_1": np.full((length, 1), 0.03, np.float32), - "contact_normals": np.tile([0.0, 0.0, 1.0], (length, 1)).astype( - np.float32 - ), - "contact_points_1": np.zeros((length, 3), np.float32), - } - writer.add_trajectory( - str(index), {"q": q, "qd": qd}, transitions, split="train" - ) - result = export_legacy_dataset(tmp_path / "data", tmp_path / "legacy.h5", "train") - with h5py.File(tmp_path / "legacy.h5") as file: - data = file["data"] - assert data["states"].shape == (5, 2, 29) - assert data["traj_lengths"][()].tolist() == [3, 5] - np.testing.assert_allclose( - data["states"][0, 0, 15:21], [0.0, 0.0, 1.0, 0.0, 1.0, 0.0] - ) - np.testing.assert_allclose(data["contact_thicknesses"][0], 0.05) - assert data["joint_acts"].shape == (5, 2, 8) - assert result["native_to_legacy_qd"] is True - - -@pytest.mark.skipif( - not os.getenv("NERD_REFERENCE_SOURCE"), - reason="Explicit trusted official checkout/dependency environment required", -) -def test_actual_official_main_two_updates_match_migrated(tmp_path): - """External reference is opt-in; synthetic data tests migration, not accuracy.""" - from embodichain_models.dynamics.nerd.codec import NeRDStateCodec - from embodichain_models.dynamics.nerd import ( - NeRDModel, - NeRDModelCfg, - export_model, - load_model, - ) - from examples.learning.nerd._support._trainer import ( - DynamicsTrainer, - DynamicsTrainCfg, - ) - from examples.learning.nerd._support._reference import _hash - from examples.learning.nerd._support._workflow import run_train, run_export - - metadata = {"codec": {"robot": "cartpole"}, "dt": 1 / 60} - with DynamicsShardWriter(tmp_path / "data", metadata) as writer: - for split in ("train", "validation", "test"): - q = np.zeros((7, 2), np.float32) - q[:, 0] = np.arange(7) * 0.02 - qd = np.zeros_like(q) - qd[:, 0] = 1.2 - writer.add_trajectory( - split, - {"q": q, "qd": qd}, - {"joint_f": np.zeros((6, 2), np.float32)}, - split=split, - ) - source = os.environ["NERD_REFERENCE_SOURCE"] - assert reference_identity(source)["commit"] == OFFICIAL_COMMIT - torch.set_num_threads(1) - torch.manual_seed(4) - model = NeRDModel( - NeRDModelCfg( - input_dims={"states_embedding": 4, "joint_f": 2}, - output_dim=4, - n_embd=12, - n_head=3, - n_layer=1, - history_length=3, - block_size=3, - head_hidden_dims=(12,), - ) - ) - train = DynamicsWindowDataset(tmp_path / "data", "train", 3) - validation = DynamicsWindowDataset(tmp_path / "data", "validation", 3) - cfg = DynamicsTrainCfg(total_steps=2, batch_size=2, validation_interval=1, seed=4) - trainer = DynamicsTrainer( - model, - NeRDStateCodec("cartpole"), - train, - validation, - cfg, - tmp_path / "scratch", - manifest=metadata, - ) - export_model(tmp_path / "initial", trainer.model, metadata) - np.save(tmp_path / "order.npy", np.array([[0, 1], [1, 2]], np.int64)) - spec = { - "source": source, - "dataset": str(tmp_path / "data"), - "dataset_fingerprint": train.fingerprint, - "initial_model": str(tmp_path / "initial"), - "initial_model_sha256": _hash(tmp_path / "initial" / "weights.pt"), - "sample_order": str(tmp_path / "order.npy"), - "sample_order_sha256": _hash(tmp_path / "order.npy"), - "train": cfg.to_dict(), - "output_dir": str(tmp_path / "official"), - } - path = tmp_path / "spec.json" - path.write_text(json.dumps(spec)) - env = { - **os.environ, - "OMP_NUM_THREADS": "1", - "MKL_NUM_THREADS": "1", - "CUDA_VISIBLE_DEVICES": "", - "PYTHONPATH": os.pathsep.join( - [ - os.environ.get("NERD_REFERENCE_DEPS", ""), - source, - str(Path(__file__).resolve().parents[4]), - str( - Path(__file__).resolve().parents[4] - / "extensions/embodichain_models" - ), - ] - ), - } - process = subprocess.run( - [ - sys.executable, - str( - Path(__file__).resolve().parents[4] - / "examples/learning/nerd/_support/_worker.py" - ), - "_reference", - "--spec", - str(path), - "--trust-external-code", - ], - env=env, - capture_output=True, - text=True, - ) - assert process.returncode == 0, process.stdout + process.stderr - receipt = json.loads((tmp_path / "official" / "result.json").read_text()) - assert receipt["runtime_identity"]["parameter_placements"] == [ - {"device": "cpu", "dtype": "torch.float32"} - ] - assert ( - receipt["runtime_identity"]["cuda"]["driver_inventory"]["status"] == "not_used" - ) - run_train( - { - "dataset": spec["dataset"], - "output_dir": str(tmp_path / "candidate"), - "fine_tune": spec["initial_model"], - "sample_order": spec["sample_order"], - "train": cfg.to_dict(), - } - ) - run_export(tmp_path / "candidate" / "last.pt", tmp_path / "candidate-export") - reference, _ = load_model(tmp_path / "official" / "last-export") - candidate, _ = load_model(tmp_path / "candidate-export") - for name, value in candidate.state_dict().items(): - torch.testing.assert_close( - value, reference.state_dict()[name], atol=1e-5, rtol=1e-4 - ) - report = json.loads((tmp_path / "official" / "result.json").read_text()) - assert report["completed_updates"] == 2 - assert report["reference"]["source_patches"] == [] - train.close() - validation.close() - - -def test_legacy_roundtrip_preserves_native_qd_and_declared_split(tmp_path): - from examples.learning.nerd._support._reference import import_legacy_datasets - - metadata = {"codec": {"robot": "cartpole"}, "dt": 1 / 60} - with DynamicsShardWriter(tmp_path / "original", metadata) as writer: - q = np.arange(12, dtype=np.float32).reshape(6, 2) * 0.01 - writer.add_trajectory( - "sample", - {"q": q, "qd": q * 0.5}, - {"joint_f": np.zeros((5, 2), np.float32)}, - split="validation", - ) - export_legacy_dataset(tmp_path / "original", tmp_path / "legacy.h5", "validation") - imported = import_legacy_datasets( - {"validation": tmp_path / "legacy.h5"}, tmp_path / "imported", metadata - ) - assert imported["trajectories"] == 1 - data = DynamicsWindowDataset(tmp_path / "imported", "validation", 3) - np.testing.assert_allclose(data[0]["states"]["q"], q[:4]) - np.testing.assert_allclose(data[0]["states"]["qd"], q[:4] * 0.5) - assert len(DynamicsWindowDataset(tmp_path / "imported", "train", 3)) == 0 - data.close() - - -def _write_source_ant_velocities(path: Path) -> np.ndarray: - """Make continuous source states exposing cancellation and signed-zero loss.""" - horizon, trajectories, state_dim = 5, 2, 29 - states = np.zeros((horizon + 1, trajectories, state_dim), np.float32) - states[..., :3] = [10.0, 20.0, 30.0] - states[..., 6] = 1.0 - states[..., 15:18] = [5.0, 6.0, 7.0] - states[..., 18:21] = [0.3, 0.1, 0.2] - states[:, :, 18] += np.arange(horizon + 1, dtype=np.float32)[:, None] * 0.01 - states[0, 1, 19] = -0.0 - with h5py.File(path, "w") as file: - data = file.create_group("data") - data.attrs["mode"] = "trajectory" - data.create_dataset("states", data=states[:-1]) - data.create_dataset("next_states", data=states[1:]) - data.create_dataset("traj_lengths", data=np.array([3, 5], np.int32)) - data.create_dataset("root_body_q", data=states[:-1, :, :7]) - data.create_dataset( - "joint_acts", data=np.zeros((horizon, trajectories, 8), np.float32) - ) - data.create_dataset( - "contact_depths", - data=np.broadcast_to( - np.array([0.01, 100.0], np.float32), (horizon, trajectories, 2) - ), - ) - data.create_dataset( - "contact_thicknesses", - data=np.full((horizon, trajectories, 2), 0.02, np.float32), - ) - data.create_dataset( - "contact_points_1", data=np.zeros((horizon, trajectories, 6), np.float32) - ) - data.create_dataset( - "contact_normals", data=np.zeros((horizon, trajectories, 6), np.float32) - ) - data.create_dataset( - "gravity_dir", - data=np.broadcast_to( - np.array([0.0, -1.0, 0.0], np.float32), (horizon, trajectories, 3) - ), - ) - return states - - -def test_original_velocity_import_roundtrip_preserves_fp32_bits_and_defaults(tmp_path): - from embodichain_models.dynamics.nerd.codec import NeRDStateCodec - from embodichain_models.dynamics.nerd.data import load_manifest - from examples.learning.nerd._support._reference import ( - import_legacy_datasets, - _official_batch, - ) - - path = tmp_path / "source.h5" - original = _write_source_ant_velocities(path) - source_sha = hashlib.sha256(path.read_bytes()).hexdigest() - metadata = {"codec": {"robot": "ant"}, "dt": 1 / 60} - imported = import_legacy_datasets( - {"validation": path}, - tmp_path / "preserved", - metadata, - preserve_legacy_velocity=True, - trajectories_per_shard=1, - ) - import_legacy_datasets({"validation": path}, tmp_path / "default", metadata) - manifest = load_manifest(tmp_path / "preserved", verify=True) - baseline = load_manifest(tmp_path / "default", verify=True) - assert imported["trajectories"] == 2 - assert manifest["fields"]["states"]["legacy_velocity"]["shape"] == [14] - assert "legacy_velocity" not in baseline["fields"]["states"] - assert "legacy_velocity" not in baseline["metadata"] - without_contract = dict(manifest["metadata"]) - contract = without_contract.pop("legacy_velocity") - assert without_contract == baseline["metadata"] - assert contract["world_frame"] == "source" - assert contract["field"] == "states/legacy_velocity" - assert manifest["metadata"]["legacy_sources"]["validation"]["sha256"] == source_sha - assert metadata == {"codec": {"robot": "ant"}, "dt": 1 / 60} - codec = NeRDStateCodec("ant") - for shard in manifest["shards"]: - with h5py.File(tmp_path / "preserved" / shard["file"], "r") as file: - for record in shard["trajectories"]: - index, length = record["metadata"]["legacy_index"], record["length"] - group = file[record["group"]] - stored = group["states/legacy_velocity"][()] - expected = original[: length + 1, index, 15:] - np.testing.assert_array_equal( - stored.view(np.uint32), expected.view(np.uint32) - ) - batch = { - kind: { - name: torch.from_numpy(field[()]).unsqueeze(0) - for name, field in group[kind].items() - } - for kind in ("states", "transitions") - } - raw = _official_batch(batch, codec) - np.testing.assert_array_equal( - raw["states"][0].numpy().view(np.uint32), - original[:length, index].view(np.uint32), - ) - np.testing.assert_array_equal( - raw["next_states"][0].numpy().view(np.uint32), - original[1 : length + 1, index].view(np.uint32), - ) - exported = export_legacy_dataset( - tmp_path / "preserved", tmp_path / "export.h5", "validation" - ) - assert exported["native_to_legacy_qd"] is False - with h5py.File(tmp_path / "export.h5", "r") as file: - for index, length in enumerate((3, 5)): - np.testing.assert_array_equal( - file["data/states"][:length, index].view(np.uint32), - original[:length, index].view(np.uint32), - ) - np.testing.assert_array_equal( - file["data/next_states"][:length, index].view(np.uint32), - original[1 : length + 1, index].view(np.uint32), - ) - assert hashlib.sha256(path.read_bytes()).hexdigest() == source_sha - - -def test_preserved_velocity_follows_prefix_masks_and_cpu_residency(tmp_path): - from torch.utils.data._utils.collate import default_collate - from embodichain_models.dynamics.nerd.codec import NeRDStateCodec - from examples.learning.nerd._support._reference import ( - import_legacy_datasets, - _official_batch, - ) - - source = tmp_path / "source.h5" - original = _write_source_ant_velocities(source) - import_legacy_datasets( - {"train": source}, - tmp_path / "data", - {"codec": {"robot": "ant"}}, - preserve_legacy_velocity=True, - trajectories_per_shard=1, - ) - dataset = DynamicsWindowDataset(tmp_path / "data", "train", 3, include_prefix=True) - indices = [0, 2, 3, 7, 0] - individually = default_collate([dataset[index] for index in indices]) - grouped = dataset.get_batch(indices) - for name, value in grouped["states"].items(): - torch.testing.assert_close(value, individually["states"][name], atol=0, rtol=0) - expected = np.concatenate( - (original[:2, 0, 15:], np.repeat(original[1:2, 0, 15:], 2, axis=0)) - ) - np.testing.assert_array_equal( - grouped["states"]["legacy_velocity"][0].numpy().view(np.uint32), - expected.view(np.uint32), - ) - assert grouped["valid_mask"][0].tolist() == [True, False, False] - codec = NeRDStateCodec("ant") - transported = _official_batch(grouped, codec) - dataset.enable_device_cache("cpu", max_bytes=1024 * 1024) - cached = dataset.get_batch(indices) - for name, value in cached["states"].items(): - torch.testing.assert_close(value, grouped["states"][name], atol=0, rtol=0) - for name, value in _official_batch(cached, codec).items(): - torch.testing.assert_close(value, transported[name], atol=0, rtol=0) - cached["states"]["legacy_velocity"][0, 0, 0] = 999 - assert dataset.get_batch(indices)["states"]["legacy_velocity"][0, 0, 0] != 999 - dataset.close() - - -def test_disabling_velocity_preservation_clears_only_copied_contract(tmp_path): - from embodichain_models.dynamics.nerd.data import load_manifest - from examples.learning.nerd._support._reference import import_legacy_datasets - - path = tmp_path / "source.h5" - _write_source_ant_velocities(path) - import_legacy_datasets( - {"train": path}, - tmp_path / "preserved", - {"codec": {"robot": "ant"}}, - preserve_legacy_velocity=True, - ) - metadata = load_manifest(tmp_path / "preserved")["metadata"] - original = json.loads(json.dumps(metadata)) - import_legacy_datasets( - {"train": path}, - tmp_path / "native", - metadata, - preserve_legacy_velocity=False, - ) - native = load_manifest(tmp_path / "native") - assert "legacy_velocity" not in native["metadata"] - assert "legacy_velocity" not in native["fields"]["states"] - assert metadata == original - - -@pytest.mark.parametrize("mismatch", ["small_delta", "signed_zero"]) -def test_lossless_import_rejects_nonidentical_adjacent_fp32_states(tmp_path, mismatch): - from examples.learning.nerd._support._reference import import_legacy_datasets - - path = tmp_path / "source.h5" - _write_source_ant_velocities(path) - with h5py.File(path, "a") as file: - if mismatch == "small_delta": - file["data/next_states"][0, 0, 18] += np.float32(1e-7) - else: - file["data/next_states"][0, 0, 21] = -0.0 - # The old tolerance remains available to existing callers. - import_legacy_datasets( - {"train": path}, tmp_path / "default", {"codec": {"robot": "ant"}} - ) - with pytest.raises(ValueError, match="bitwise continuous"): - import_legacy_datasets( - {"train": path}, - tmp_path / "preserved", - {"codec": {"robot": "ant"}}, - preserve_legacy_velocity=True, - ) - - -def test_lossless_import_rejects_even_explicit_identity_world_rotation(tmp_path): - from examples.learning.nerd._support._reference import import_legacy_datasets - - with pytest.raises(ValueError, match="source world frame"): - import_legacy_datasets( - {"train": tmp_path / "unused.h5"}, - tmp_path / "preserved", - {"codec": {"robot": "ant"}}, - preserve_legacy_velocity=True, - world_rotation_xyzw=(0.0, 0.0, 0.0, 1.0), - ) - assert not (tmp_path / "preserved").exists() - - -def test_official_transport_rejects_malformed_preserved_velocity_shape(): - from embodichain_models.dynamics.nerd.codec import NeRDStateCodec - from examples.learning.nerd._support._reference import _official_batch - - q = torch.zeros(1, 2, 15) - q[..., 6] = 1 - batch = { - "states": { - "q": q, - "qd": torch.zeros(1, 2, 14), - "legacy_velocity": torch.zeros(1, 2, 13), - }, - "transitions": {"joint_f": torch.zeros(1, 1, 14)}, - } - with pytest.raises(ValueError, match="Preserved legacy velocity shape"): - _official_batch(batch, NeRDStateCodec("ant")) - - -def test_gravity_repair_is_explicit_preserves_raw_and_other_fields(): - from examples.learning.nerd._support._reference import repair_gravity_transport - - raw = np.array([[[np.nan, -1.0, 8.49e5], [1e-45, -1.0, np.nan]]], np.float32) - before = raw.copy() - with pytest.raises(ValueError, match="uninitialized gravity"): - repair_gravity_transport(raw, (0.0, -1.0, 0.0)) - repaired, evidence = repair_gravity_transport(raw, (0.0, -1.0, 0.0), enabled=True) - np.testing.assert_array_equal(raw, before) - np.testing.assert_allclose( - repaired, np.array([[[0.0, -1.0, 0.0], [0.0, -1.0, 0.0]]]) - ) - assert evidence["nonfinite_components"] == 2 - assert evidence["enabled"] is True - - -def test_import_explicit_world_rotation_and_gravity_transport(tmp_path): - from examples.learning.nerd._support._reference import import_legacy_datasets - import math - - metadata = {"codec": {"robot": "ant", "base_com": [0.0, 0.0, 0.0]}, "dt": 1 / 60} - q = np.zeros((4, 15), np.float32) - q[:, 1] = 1.0 - q[:, 3] = -math.sqrt(0.5) - q[:, 6] = math.sqrt(0.5) - qd = np.zeros((4, 14), np.float32) - qd[:, 0] = 1.0 - qd[:, 4] = 1.0 - with DynamicsShardWriter(tmp_path / "source", metadata) as writer: - writer.add_trajectory( - "one", - {"q": q, "qd": qd}, - { - "joint_f": np.zeros((3, 14), np.float32), - "contact_depths": np.zeros((3, 1), np.float32), - "gravity_dir": np.tile([0.0, -1.0, 0.0], (3, 1)).astype(np.float32), - }, - split="train", - ) - export_legacy_dataset(tmp_path / "source", tmp_path / "raw.h5", "train") - with h5py.File(tmp_path / "raw.h5", "a") as file: - file["data/gravity_dir"][:, :, 0] = np.nan - import_legacy_datasets( - {"train": tmp_path / "raw.h5"}, - tmp_path / "canonical", - metadata, - world_rotation_xyzw=(math.sqrt(0.5), 0.0, 0.0, math.sqrt(0.5)), - gravity_direction=(0.0, -1.0, 0.0), - ) - dataset = DynamicsWindowDataset(tmp_path / "canonical", "train", 3) - item = dataset[0] - np.testing.assert_allclose( - item["states"]["q"][0, :7], [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0], atol=1e-6 - ) - np.testing.assert_allclose( - item["states"]["qd"][0, :6], [1.0, 0.0, 0.0, 0.0, 0.0, 1.0], atol=1e-6 - ) - np.testing.assert_allclose( - item["transitions"]["gravity_dir"][0], [0.0, 0.0, -1.0], atol=1e-6 - ) - assert dataset.manifest["metadata"]["gravity_transport_override"] == [ - 0.0, - -1.0, - 0.0, - ] - dataset.close() - - -@pytest.mark.skipif( - not os.getenv("NERD_REFERENCE_SOURCE"), - reason="Explicit trusted official checkout/dependency environment required", -) -def test_real_official_collection_streams_batches_and_preserves_raw(tmp_path): - spec = { - "source": os.environ["NERD_REFERENCE_SOURCE"], - "output_dir": str(tmp_path / "collection"), - "robot": "cartpole", - "warp_version": "1.8.0", - "device": "cpu", - "num_envs": 2, - "horizon": 4, - "splits": {"train": 5, "validation": 2, "test": 2}, - "repair_uninitialized_gravity": True, - "canonical_world": "z_up", - } - spec_path = tmp_path / "generate.json" - spec_path.write_text(json.dumps(spec)) - env = { - **os.environ, - "OMP_NUM_THREADS": "1", - "MKL_NUM_THREADS": "1", - "CUDA_VISIBLE_DEVICES": "", - "PYTHONPATH": os.pathsep.join( - [ - os.environ.get("NERD_REFERENCE_DEPS", ""), - spec["source"], - str(Path(__file__).resolve().parents[4]), - str( - Path(__file__).resolve().parents[4] - / "extensions/embodichain_models" - ), - ] - ), - } - process = subprocess.run( - [ - sys.executable, - str( - Path(__file__).resolve().parents[4] - / "examples/learning/nerd/_support/_worker.py" - ), - "_reference", - "--spec", - str(spec_path), - "--trust-external-code", - "--generate", - ], - env=env, - capture_output=True, - text=True, - ) - assert process.returncode == 0, process.stdout + process.stderr - root = tmp_path / "collection" - result = json.loads((root / "generation-result.json").read_text()) - assert result["source_trajectory_counts"] == { - "train": 5, - "validation": 2, - "test": 2, - } - assert result["gravity_repair"]["train"]["sampler_calls"] == 3 - assert result["source_bug_corrected"] is True - assert result["teacher_origin"]["integrator"] == "featherstone" - physics = json.loads((root / "source-physics.json").read_text()) - assert physics["physics_contract_version"] == "official-warp-physical-v2" - assert ( - result["teacher_origin"]["physics_contract_version"] - == physics["physics_contract_version"] - ) - assert result["teacher_origin"]["up_axis"] == 1 - np.testing.assert_allclose( - physics["source_physical_arrays"]["gravity"], [0.0, -9.81, 0.0], atol=1e-6 - ) - assert ( - physics["source_physical_arrays"]["gravity"] - == result["teacher_origin"]["gravity"] - ) - assert result["teacher_origin"]["gravity_magnitude"] == pytest.approx( - 9.81, abs=1e-6 - ) - from examples.learning.nerd._support._reference import _hash - import hashlib - - assert physics["source_generator_files"]["generate/trajectory_sampler.py"] == _hash( - Path(spec["source"]) / "generate/trajectory_sampler.py" - ) - expected_generator_hash = hashlib.sha256( - json.dumps( - physics["source_generator_files"], sort_keys=True, separators=(",", ":") - ).encode() - ).hexdigest() - assert physics["source_generator_sha256"] == expected_generator_hash - assert physics["generation_adapter_sha256"] == _hash( - Path(__file__).resolve().parents[4] - / "examples/learning/nerd/_support/_reference.py" - ) - with h5py.File(root / "train.raw.h5") as file: - assert file["data/states"].shape == (4, 5, 4) - dataset = DynamicsWindowDataset(root / "canonical", "train", 3) - assert len(dataset) == 10 - assert ( - dataset.manifest["metadata"]["source_physical_arrays"]["gravity"] - == physics["source_physical_arrays"]["gravity"] - ) - assert ( - dataset.manifest["metadata"]["generation_adapter_sha256"] - == physics["generation_adapter_sha256"] - ) - np.testing.assert_allclose( - dataset[0]["transitions"]["gravity_dir"], - np.tile([0.0, 0.0, -1.0], (3, 1)), - atol=1e-6, - ) - dataset.close() - - -@pytest.mark.skipif( - not os.getenv("NERD_REFERENCE_SOURCE"), - reason="Explicit trusted official checkout required", -) -def test_portable_to_official_restoration_preserves_state_dtype_and_ownership(): - code = r""" -import sys,torch -from examples.learning.nerd._support._reference import reference_identity,_restore_official_model -source=sys.argv[1] -reference_identity(source) -sys.path.insert(0,source) -from embodichain_models.dynamics.nerd import NeRDModel,NeRDModelCfg -for dtype in (torch.float32,torch.float64): - torch.manual_seed(5) - common=NeRDModel(NeRDModelCfg(input_dims={"states_embedding":4,"joint_f":2},output_dim=4,n_embd=12,n_head=3,n_layer=1,block_size=4,history_length=3,head_hidden_dims=(12,),dropout=.25,output_tanh=True)).to(dtype=dtype).eval() - common.set_normalization({"states_embedding":(torch.arange(4,dtype=dtype),torch.arange(4,dtype=dtype)+.5),"joint_f":(torch.tensor([.1,-.3],dtype=dtype),torch.tensor([.9,1.2],dtype=dtype))},torch.arange(4,dtype=dtype)*.2,torch.arange(4,dtype=dtype)+.3) - old=_restore_official_model(common,device="cpu") - assert not old.training - for name,value in old.named_parameters(): - assert value.dtype==dtype - torch.testing.assert_close(value,dict(common.named_parameters())[name],atol=0,rtol=0) - inputs={"states_embedding":torch.randn(2,3,4,dtype=dtype),"joint_f":torch.randn(2,3,2,dtype=dtype)} - legacy=lambda:{("joint_acts" if name=="joint_f" else name):value.clone() for name,value in inputs.items()} - torch.testing.assert_close(old(legacy()),common(inputs),atol=1e-6,rtol=1e-5) - torch.testing.assert_close(old.evaluate(legacy()),common(inputs)[:,-1:,:],atol=1e-6,rtol=1e-5) - before=common.input_rms["joint_f"].mean.clone() - old.input_rms["joint_acts"].mean.add_(1.) - torch.testing.assert_close(common.input_rms["joint_f"].mean,before,atol=0,rtol=0) -assert "dexsim" not in sys.modules -""" - env = {**os.environ, "OMP_NUM_THREADS": "1", "MKL_NUM_THREADS": "1"} - result = subprocess.run( - [sys.executable, "-c", code, os.environ["NERD_REFERENCE_SOURCE"]], - env=env, - capture_output=True, - text=True, - ) - assert result.returncode == 0, result.stdout + result.stderr - - -@pytest.mark.skipif( - not os.getenv("NERD_REFERENCE_SOURCE"), reason="Explicit trusted reference required" -) -def test_actual_official_trainer_resume_preserves_optimizer_rng_and_best_model( - tmp_path, -): - """Synthetic CPU data verify lifecycle, not learned-dynamics accuracy.""" - from embodichain_models.dynamics.nerd.codec import NeRDStateCodec - from embodichain_models.dynamics.nerd import ( - NeRDModel, - NeRDModelCfg, - export_model, - load_model, - ) - from examples.learning.nerd._support._trainer import compute_statistics - from examples.learning.nerd._support._sample_order import _write_checkpointed_order - from examples.learning.nerd._support._reference import _hash - - metadata = {"codec": {"robot": "cartpole"}, "dt": 1 / 60} - generator = np.random.default_rng(3) - with DynamicsShardWriter(tmp_path / "data", metadata) as writer: - for split in ["train", "validation", "test"]: - for index in range(2): - q = generator.normal(size=(13, 2)).astype(np.float32) - qd = generator.normal(size=(13, 2)).astype(np.float32) - writer.add_trajectory( - f"{split}:{index}", - {"q": q, "qd": qd}, - {"joint_f": np.zeros((12, 2), np.float32)}, - split=split, - ) - dataset = DynamicsWindowDataset(tmp_path / "data", "train", 3) - codec = NeRDStateCodec("cartpole") - stats = compute_statistics(dataset, codec) - torch.manual_seed(4) - model = NeRDModel( - NeRDModelCfg( - input_dims={"states_embedding": 4, "joint_f": 2}, - output_dim=4, - n_embd=12, - n_head=3, - n_layer=1, - block_size=4, - history_length=3, - dropout=0.1, - ) - ) - model.set_normalization( - {k: (v["mean"], v["var"]) for k, v in stats["inputs"].items()}, - stats["output"]["mean"], - stats["output"]["var"], - ) - export_model(tmp_path / "initial", model, metadata) - order, _ = _write_checkpointed_order( - tmp_path / "order.json", - dataset_size=len(dataset), - steps=6, - batch_size=4, - seed=7, - block_steps=2, - ) - cfg = { - "total_steps": 6, - "batch_size": 4, - "seed": 7, - "validation_interval": 2, - "checkpoint_interval": 2, - "log_interval": 2, - "resident_cache_bytes": 1024 * 1024, - "device": "cpu", - } - source = os.environ["NERD_REFERENCE_SOURCE"] - spec = { - "source": source, - "dataset": str(tmp_path / "data"), - "dataset_fingerprint": dataset.fingerprint, - "initial_model": str(tmp_path / "initial"), - "initial_model_sha256": _hash(tmp_path / "initial/weights.pt"), - "sample_order": str(order), - "sample_order_sha256": _hash(order), - "train": cfg, - } - environment = { - **os.environ, - "CUDA_VISIBLE_DEVICES": "", - "OMP_NUM_THREADS": "1", - "MKL_NUM_THREADS": "1", - "PYTHONPATH": os.pathsep.join( - [ - os.environ.get("NERD_REFERENCE_DEPS", ""), - source, - str(Path(__file__).resolve().parents[4]), - str( - Path(__file__).resolve().parents[4] - / "extensions/embodichain_models" - ), - ] - ), - } - - def execute(name, extra): - path = tmp_path / (name + ".json") - path.write_text(json.dumps({**spec, "output_dir": str(tmp_path / name)})) - result = subprocess.run( - [ - sys.executable, - str( - Path(__file__).resolve().parents[4] - / "examples/learning/nerd/_support/_worker.py" - ), - "_reference", - "--spec", - str(path), - "--trust-external-code", - *extra, - ], - env=environment, - capture_output=True, - text=True, - ) - assert result.returncode == 0, result.stdout + result.stderr - - execute("full", []) - execute("resumed", ["--stop-after", "3"]) - interrupted = json.loads((tmp_path / "resumed/result.json").read_text()) - assert ( - interrupted["status"] == "interrupted" and interrupted["completed_updates"] == 3 - ) - execute("resumed", ["--resume", str(tmp_path / "resumed/reference-last.pt")]) - for artifact in ["best-export", "last-export"]: - first, _ = load_model(tmp_path / "full" / artifact) - second, _ = load_model(tmp_path / "resumed" / artifact) - for key, value in first.state_dict().items(): - torch.testing.assert_close(second.state_dict()[key], value, atol=0, rtol=0) - dataset.close() diff --git a/tests/examples/learning/nerd/test_reference_rollout.py b/tests/examples/learning/nerd/test_reference_rollout.py deleted file mode 100644 index fa160661a5..0000000000 --- a/tests/examples/learning/nerd/test_reference_rollout.py +++ /dev/null @@ -1,539 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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. -# ---------------------------------------------------------------------------- - -"""Historical runtime qualification through explicitly trusted isolated Warp.""" - -from __future__ import annotations -import hashlib -import importlib.util -import json -import os -from pathlib import Path -import subprocess -import sys -import pytest -from examples.learning.nerd._support._reference_rollout import ( - evaluate_reference_rollouts, -) - -_ROOT = Path(__file__).resolve().parents[4] - - -def test_external_reference_requires_explicit_trust() -> None: - with pytest.raises(ValueError, match="trust_external_code"): - evaluate_reference_rollouts({}) - - -@pytest.mark.parametrize("robot,version", [("cartpole", "1.5.1"), ("ant", "1.8.0")]) -def test_wrong_historical_version_rejected_before_source_import( - robot: str, version: str -) -> None: - with pytest.raises(ValueError, match="pinned Warp"): - evaluate_reference_rollouts( - {"robot": robot, "warp_version": version}, trust_external_code=True - ) - - -def _check_original_runtime(tmp_path: Path, robot: str, source: str) -> None: - import numpy as np - import h5py - import torch - import warp as wp - from examples.learning.nerd._support._reference import generate_official_dataset - from embodichain_models.dynamics.nerd import NeRDModel, NeRDModelCfg, export_model - from embodichain_models.dynamics.nerd.data import load_manifest - - # Generate new physical provenance from the actual original teacher. No - # current-Newton manifest is rewritten to pretend it describes old physics. - with wp.ScopedDevice("cpu"): - generation = generate_official_dataset( - { - "source": source, - "robot": robot, - "output_dir": str(tmp_path / "source-data"), - "device": "cpu", - "num_envs": 2, - "horizon": 4, - "splits": {"train": 2, "validation": 2, "test": 2}, - "seed": 3, - "initial_states_source": "env", - "canonical_world": "z_up", - "repair_uninitialized_gravity": True, - "cpu_threads": 2, - }, - trust_external_code=True, - ) - metadata = load_manifest(tmp_path / "source-data/canonical")["metadata"] - assert "gravity" in metadata["source_physical_arrays"] - 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, - } - ) - torch.manual_seed(7) - model = NeRDModel( - NeRDModelCfg( - input_dims=dims, - output_dim=4 if robot == "cartpole" else 29, - n_embd=24, - n_head=3, - n_layer=1, - history_length=3, - block_size=4, - ) - ) - with torch.no_grad(): - model.model.output_net.weight.mul_(0.02) - model.model.output_net.bias.zero_() - if robot == "ant": - model.model.output_net.bias[6] = 1.0 - export = tmp_path / "model" - export_model(export, model, metadata) - reports = {} - for mode in ["passive", "random_effort", "fixed_policy"]: - spec = { - "source": source, - "robot": robot, - "model": str(export), - "output": str(tmp_path / f"{mode}.json"), - "device": "cpu", - "num_envs": 2, - "trajectories": 3, - "horizons": [1, 4], - "seed": 17, - "initial_states": "default", - "motion_mode": mode, - "cpu_threads": 2, - } - if mode == "fixed_policy": - spec["policy"] = {"kind": "pd"} - reports[mode] = evaluate_reference_rollouts(spec, trust_external_code=True) - report = reports[mode] - assert report["teacher_origin"]["kind"] == "official_warp" - assert report["runtime_identity"]["engine"] == "official_warp" - assert report["physical_verification"]["complete"] - assert report["physical_verification"]["missing_recorded_fields"] == [] - assert report["portable_runtime_equivalence"]["status"] == "passed" - assert report["portable_runtime_equivalence"]["max_abs_error"] <= 1e-5 - assert report["trajectories"] == 3 and report["num_envs"] == 2 - assert len(Path(report["trajectory_evidence"]).read_text().splitlines()) == 3 - with h5py.File(report["trace_evidence"], "r") as evidence: - assert evidence["0/q"].shape[0] == 2 and evidence["1/q"].shape[0] == 1 - if mode == "fixed_policy": - assert not np.allclose( - evidence["0/teacher_joint_f"][:, 1:], - evidence["0/joint_f"][:, 1:], - atol=1e-8, - ) - assert report["selected_teacher_sha256"] - assert ( - report["trace_evidence_sha256"] - == hashlib.sha256(Path(report["trace_evidence"]).read_bytes()).hexdigest() - ) - assert ( - report["horizons"]["4"]["valid_trajectories"] - + report["horizons"]["4"]["numerically_invalid_trajectories"] - == 3 - ) - # Run an independently restored equal export: teacher identity and content - # are independent of the neural artifact's path, and results remain exact. - second = tmp_path / "equal-model" - export_model(second, model, metadata) - again = evaluate_reference_rollouts( - { - "source": source, - "robot": robot, - "model": str(second), - "output": str(tmp_path / "equal.json"), - "device": "cpu", - "num_envs": 2, - "trajectories": 3, - "horizons": [1, 4], - "seed": 17, - "initial_states": "default", - "motion_mode": "passive", - "cpu_threads": 2, - }, - trust_external_code=True, - ) - assert ( - again["selected_teacher_sha256"] - == reports["passive"]["selected_teacher_sha256"] - ) - assert again["horizons"] == reports["passive"]["horizons"] - # Original per-trajectory sampling is deterministic across replay batching. - sample_reports = [] - for size in [1, 2]: - sample_reports.append( - evaluate_reference_rollouts( - { - "source": source, - "robot": robot, - "model": str(export), - "output": str(tmp_path / f"sample-{size}.json"), - "device": "cpu", - "num_envs": size, - "trajectories": 3, - "horizons": [1], - "seed": 31, - "initial_states": "sample", - "motion_mode": "passive", - "cpu_threads": 2, - }, - trust_external_code=True, - ) - ) - assert ( - sample_reports[0]["selected_teacher_sha256"] - == sample_reports[1]["selected_teacher_sha256"] - ) - if robot == "cartpole": - with torch.no_grad(): - model.model.output_net.bias[2:] = torch.finfo(torch.float32).max - divergent = tmp_path / "divergent-model" - export_model(divergent, model, metadata) - failure = evaluate_reference_rollouts( - { - "source": source, - "robot": robot, - "model": str(divergent), - "output": str(tmp_path / "nonfinite.json"), - "num_envs": 2, - "trajectories": 3, - "horizons": [1, 3], - "initial_states": "default", - "motion_mode": "passive", - "cpu_threads": 2, - }, - trust_external_code=True, - ) - assert failure["nonfinite_rate"] == 1.0 and failure["failure_rate"] == 1.0 - assert failure["horizons"]["3"]["valid_trajectories"] == 0 - assert len(Path(failure["trajectory_evidence"]).read_text().splitlines()) == 3 - assert failure["portable_runtime_equivalence"]["status"] == "matched_nonfinite" - with torch.no_grad(): - model.model.output_net.bias[2:] = 0.0 - # Missing gravity is a diagnostic run, never complete physical qualification. - legacy = json.loads(json.dumps(metadata)) - legacy["source_physical_arrays"].pop("gravity") - legacy["teacher_origin"].pop("gravity", None) - old = tmp_path / "legacy-model" - export_model(old, model, legacy) - old_report = evaluate_reference_rollouts( - { - "source": source, - "robot": robot, - "model": str(old), - "output": str(tmp_path / "legacy.json"), - "num_envs": 2, - "trajectories": 3, - "horizons": [1], - "initial_states": "default", - "motion_mode": "passive", - "cpu_threads": 2, - }, - trust_external_code=True, - ) - assert old_report["qualification_status"] == "pending_missing_recorded_physics" - assert old_report["physical_verification"]["missing_recorded_fields"] == ["gravity"] - # Independent official-version generation records the physical replica - # count in its source split receipt rather than the old sampling field. - split_metadata = json.loads(json.dumps(metadata)) - batch = split_metadata["sampling"].pop("batch_size") - split_metadata["source_split_identities"] = {"train": {"num_envs": batch}} - split_export = tmp_path / "split-source-model" - export_model(split_export, model, split_metadata) - split_spec = { - "source": source, - "robot": robot, - "model": str(split_export), - "output": str(tmp_path / "split-source.json"), - "num_envs": 2, - "trajectories": 3, - "horizons": [1], - "initial_states": "default", - "motion_mode": "passive", - "cpu_threads": 2, - } - split_report = evaluate_reference_rollouts(split_spec, trust_external_code=True) - assert split_report["physical_verification"]["complete"] - assert ( - split_report["physical_verification"]["sha256"] - == reports["passive"]["physical_verification"]["sha256"] - ) - split_metadata["sampling"]["batch_size"] = batch + 1 - conflicting_export = tmp_path / "conflicting-source-model" - export_model(conflicting_export, model, split_metadata) - with pytest.raises(ValueError, match="Conflicting source sampling"): - evaluate_reference_rollouts( - { - **split_spec, - "model": str(conflicting_export), - "output": str(tmp_path / "conflicting-source.json"), - }, - trust_external_code=True, - ) - # A declared change to gravity must reject, not silently reuse old dynamics. - changed = json.loads(json.dumps(metadata)) - changed["source_physical_arrays"]["gravity"] = [0.0, -1.0, 0.0] - bad = tmp_path / "bad-model" - export_model(bad, model, changed) - with pytest.raises(ValueError, match="gravity"): - evaluate_reference_rollouts( - { - "source": source, - "robot": robot, - "model": str(bad), - "output": str(tmp_path / "bad.json"), - "num_envs": 2, - "trajectories": 3, - "horizons": [1], - }, - trust_external_code=True, - ) - - -@pytest.mark.parametrize("robot", ["cartpole", "ant"]) -def test_real_historical_teacher_neural_runtime_and_source_contract( - tmp_path: Path, robot: str -) -> None: - source = os.environ.get("NERD_MAIN_REFERENCE_ROOT") - dependencies = os.environ.get( - "NERD_WARP18_ROOT" if robot == "cartpole" else "NERD_WARP151_ROOT" - ) - if not source or not dependencies: - pytest.skip("Set explicit trusted reference and legacy Warp dependency paths") - code = "import importlib.util,sys;from pathlib import Path;s=importlib.util.spec_from_file_location('reference_rollout_test',sys.argv[1]);m=importlib.util.module_from_spec(s);s.loader.exec_module(m);m._check_original_runtime(Path(sys.argv[2]),sys.argv[3],sys.argv[4])" - subprocess.run( - [ - sys.executable, - "-c", - code, - str(Path(__file__).resolve()), - str(tmp_path), - robot, - source, - ], - check=True, - env={ - **os.environ, - "PYTHONPATH": os.pathsep.join( - (dependencies, str(_ROOT), str(_ROOT / "extensions/embodichain_models")) - ), - "OMP_NUM_THREADS": "2", - "MKL_NUM_THREADS": "2", - "CUDA_VISIBLE_DEVICES": "", - }, - ) - - -def _check_supported_replay( - tmp_path: Path, robot: str, source: str, data: Path -) -> None: - import shutil - import h5py - import numpy as np - import torch - from embodichain_models.dynamics.nerd import NeRDModel, NeRDModelCfg, export_model - - metadata = json.loads((data / f"{robot}-passive/manifest.json").read_text())[ - "metadata" - ] - 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, - } - ) - torch.manual_seed(7) - model = NeRDModel( - NeRDModelCfg( - input_dims=dims, - output_dim=4 if robot == "cartpole" else 29, - n_embd=24, - n_layer=1, - n_head=3, - history_length=3, - block_size=4, - ) - ) - with torch.no_grad(): - model.model.output_net.weight.mul_(0.02) - model.model.output_net.bias.zero_() - if robot == "ant": - model.model.output_net.bias[6] = 1.0 - export = tmp_path / "export" - export_model(export, model, metadata) - reports = {} - specs = {} - for mode in ["passive", "random_effort", "fixed_policy"]: - spec = { - "source": source, - "robot": robot, - "model": str(export), - "output": str(tmp_path / (mode + ".json")), - "device": "cpu", - "num_envs": 2, - "trajectories": 3, - "horizons": [1, 4, 12], - "seed": 1000000, - "initial_states": "sample", - "motion_mode": mode, - "cpu_threads": 2, - "teacher_dataset": str(data / (robot + "-" + mode)), - } - if mode == "fixed_policy": - spec["policy"] = {"kind": "pd"} - result = evaluate_reference_rollouts(spec, trust_external_code=True) - reports[mode], specs[mode] = result, spec - assert result["teacher_generation_environment"]["python"] == "3.8.20" - assert result["teacher_generation_environment"]["torch"].startswith("2.2.2") - assert ( - result["teacher_generation_environment"]["runtime_imports_embodichain"] - is False - ) - assert result["physical_verification"]["complete"] - assert result["teacher_physical_verification"]["complete"] - assert result["teacher_transport_max_abs_error"] <= 1e-5 - assert result["trajectories"] == 3 and set(result["horizons"]) == { - "1", - "4", - "12", - } - assert result["portable_runtime_equivalence"]["status"] == "passed" - if mode == "fixed_policy": - assert ( - result["teacher_feedback_verification"] - == "all_selected_actions_verified_against_teacher_states" - ) - assert result["teacher_feedback_max_abs_error"] <= 1e-5 - with h5py.File(result["trace_evidence"], "r") as file: - assert np.any( - file["0/joint_f"][:, 1:] != file["0/teacher_joint_f"][:, 1:] - ) - # Different evaluation batching cannot select different teacher contents. - again = evaluate_reference_rollouts( - {**specs["passive"], "num_envs": 1, "output": str(tmp_path / "batch1.json")}, - trust_external_code=True, - ) - assert ( - again["selected_teacher_sha256"] - == reports["passive"]["selected_teacher_sha256"] - ) - # A rehashed but inconsistent raw/canonical frame must not become a new oracle. - changed = tmp_path / "changed-transport" - shutil.copytree(data / (robot + "-passive"), changed) - manifest_path = changed / "manifest.json" - manifest = json.loads(manifest_path.read_text()) - shard = changed / manifest["shards"][0]["file"] - with h5py.File(shard, "a") as file: - file["0/states/q"][1, 0] += 1.0 - manifest["shards"][0]["sha256"] = hashlib.sha256(shard.read_bytes()).hexdigest() - manifest_path.write_text(json.dumps(manifest)) - with pytest.raises(ValueError, match="state transport differs"): - evaluate_reference_rollouts( - { - **specs["passive"], - "teacher_dataset": str(changed), - "output": str(tmp_path / "bad-transport.json"), - }, - trust_external_code=True, - ) - # PD action verification also rejects consistent raw/canonical effort tampering. - changed = tmp_path / "changed-feedback" - shutil.copytree(data / (robot + "-fixed_policy"), changed) - manifest_path = changed / "manifest.json" - manifest = json.loads(manifest_path.read_text()) - shard = changed / manifest["shards"][0]["file"] - with h5py.File(shard, "a") as file: - file["0/transitions/raw_joint_f"][0, 0] += 1.0 - file["0/transitions/joint_f"][0, 6 if robot == "ant" else 0] += 1.0 - manifest["shards"][0]["sha256"] = hashlib.sha256(shard.read_bytes()).hexdigest() - manifest_path.write_text(json.dumps(manifest)) - with pytest.raises(ValueError, match="actions differ from declared feedback"): - evaluate_reference_rollouts( - { - **specs["fixed_policy"], - "teacher_dataset": str(changed), - "output": str(tmp_path / "bad-feedback.json"), - }, - trust_external_code=True, - ) - missing = tmp_path / "missing-teacher-physics" - shutil.copytree(data / (robot + "-passive"), missing) - manifest_path = missing / "manifest.json" - manifest = json.loads(manifest_path.read_text()) - manifest["metadata"]["source_physical_arrays"].pop("gravity") - manifest["metadata"]["teacher_origin"].pop("gravity", None) - manifest_path.write_text(json.dumps(manifest)) - with pytest.raises(ValueError, match="missing measured physical fields"): - evaluate_reference_rollouts( - { - **specs["passive"], - "teacher_dataset": str(missing), - "output": str(tmp_path / "missing-physics.json"), - }, - trust_external_code=True, - ) - - -@pytest.mark.parametrize("robot", ["cartpole", "ant"]) -def test_real_supported_environment_teacher_replay(tmp_path: Path, robot: str) -> None: - source = os.environ.get("NERD_MAIN_REFERENCE_ROOT") - dependencies = os.environ.get( - "NERD_WARP18_ROOT" if robot == "cartpole" else "NERD_WARP151_ROOT" - ) - data = os.environ.get("NERD_SUPPORTED_EVAL_DATA") - if not source or not dependencies or not data: - pytest.skip( - "Set explicit trusted original source/dependencies and supported-environment evaluation corpus" - ) - code = "import importlib.util,sys;from pathlib import Path;s=importlib.util.spec_from_file_location('reference_replay_test',sys.argv[1]);m=importlib.util.module_from_spec(s);s.loader.exec_module(m);m._check_supported_replay(Path(sys.argv[2]),sys.argv[3],sys.argv[4],Path(sys.argv[5]))" - subprocess.run( - [ - sys.executable, - "-c", - code, - str(Path(__file__).resolve()), - str(tmp_path), - robot, - source, - data, - ], - check=True, - env={ - **os.environ, - "PYTHONPATH": os.pathsep.join( - (dependencies, str(_ROOT), str(_ROOT / "extensions/embodichain_models")) - ), - "CUDA_VISIBLE_DEVICES": "", - "OMP_NUM_THREADS": "2", - "MKL_NUM_THREADS": "2", - }, - ) diff --git a/tests/examples/learning/nerd/test_reference_traces.py b/tests/examples/learning/nerd/test_reference_traces.py deleted file mode 100644 index 69f887c178..0000000000 --- a/tests/examples/learning/nerd/test_reference_traces.py +++ /dev/null @@ -1,247 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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. -# ---------------------------------------------------------------------------- - -"""Original trace provenance, ordering, bounded batching and failure retention.""" - -from __future__ import annotations - -import hashlib -import json -from pathlib import Path -from typing import Any - -import h5py -import numpy as np -import pytest -import torch - -from examples.learning.nerd._support._reference_traces import ( - _load_reference_trace, - _reference_trace_batches, -) -from examples.learning.nerd._support._reference import OFFICIAL_COMMIT - - -def _fixture( - root: Path, robot: str -) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: - root.mkdir() - ant = robot == "ant" - q_dim, qd_dim, force_dim = (15, 14, 8) if ant else (2, 2, 2) - codec = { - "robot": robot, - "base_com": [0.0, 0.0, 0.0], - "continuous_q_indices": [] if ant else [1], - } - shared = { - "robot": robot, - "asset_sha256": "fixture-asset", - "dt": 1 / 60, - "codec": codec, - "world_rotation_xyzw": [1.0, 0.0, 0.0, 0.0], - } - source = {"runtime_files": {"fixture-source.py": "fixture-source-hash"}} - version = "1.5.1" if ant else "1.8.0" - meta = { - **shared, - "source_runtime_files": source["runtime_files"], - "teacher_origin": { - "kind": "official_warp", - "source_commit": OFFICIAL_COMMIT, - "warp_version": version, - }, - "generation_environment": { - "python": "3.8.20", - "torch": "2.2.2+cpu", - "numpy": "1.24.4", - "warp": version, - "runtime_imports_embodichain": False, - "runtime_imports_dexsim": False, - "runtime_imports_newton": False, - }, - } - rows = [] - file = root / "teacher.h5" - with h5py.File(file, "w") as stream: - for index in range(3): - group = stream.create_group(str(index)) - q = np.full((5, q_dim), index + 0.25, dtype=np.float32) - qd = np.zeros((5, qd_dim), dtype=np.float32) - if index == 2: - q[-1, 0] = np.nan - raw = np.concatenate((q, qd), -1) - force = np.full((4, force_dim), index, dtype=np.float32) - canonical = ( - np.concatenate((np.zeros((4, 6), dtype=np.float32), force), -1) - if ant - else force - ) - for key, value in { - "states/raw": raw, - "states/q": q, - "states/qd": qd, - "transitions/raw_joint_f": force, - "transitions/joint_f": canonical, - }.items(): - group.create_dataset(key, data=value) - rows.append( - { - "id": f"71:{index}", - "group": str(index), - "length": 4, - "split": "test", - "metadata": {"sampling_seed": 71 + index}, - } - ) - manifest = { - "schema_version": 1, - "trace_type": "evaluation_including_failures", - "metadata": meta, - "evaluation": { - "seed": 71, - "trajectories": 3, - "horizon": 4, - "initial_states": "sample", - "motion_mode": "random_effort", - "policy": None, - }, - "shards": [ - { - "file": file.name, - "sha256": hashlib.sha256(file.read_bytes()).hexdigest(), - "trajectories": rows, - } - ], - } - (root / "manifest.json").write_text(json.dumps(manifest)) - return {"robot": robot, "seed": 71, "motion_mode": "random_effort"}, source, shared - - -@pytest.mark.parametrize("robot", ["cartpole", "ant"]) -def test_original_trace_preserves_order_padding_failure_and_ownership( - tmp_path: Path, robot: str -) -> None: - directory = tmp_path / "teacher" - spec, source, metadata = _fixture(directory, robot) - _, records = _load_reference_trace(directory, spec, source, metadata, None, 3, 4) - batches = list(_reference_trace_batches(directory, records, 2, 4, robot, "cpu")) - assert [[row["id"] for row in selected] for selected, _ in batches] == [ - ["71:0", "71:1"], - ["71:2"], - ] - assert batches[0][1]["q"][0, 0, 0] == 0.25 - assert batches[1][1]["q"][0, 0, 0] == 2.25 - assert torch.isnan(batches[1][1]["q"][:, -1, 0]).all() - assert all(row["reference_sha256"] for row in records) - batches[1][1]["q"][0, 0, 0] = 99 - assert batches[1][1]["q"][1, 0, 0] == 2.25 - assert batches[0][1]["q"][0, 0, 0] == 0.25 - - -@pytest.mark.parametrize( - "field,value", - [ - ("python", "3.11.14"), - ("torch", "2.7.0"), - ("warp", "1.17.0"), - ("runtime_imports_embodichain", True), - ("runtime_imports_newton", True), - ], -) -def test_original_trace_rejects_wrong_generation_environment( - tmp_path: Path, field: str, value: Any -) -> None: - directory = tmp_path / "teacher" - spec, source, metadata = _fixture(directory, "ant") - path = directory / "manifest.json" - manifest = json.loads(path.read_text()) - manifest["metadata"]["generation_environment"][field] = value - path.write_text(json.dumps(manifest)) - with pytest.raises(ValueError, match="supported isolated environment"): - _load_reference_trace(directory, spec, source, metadata, None, 3, 4) - - -@pytest.mark.parametrize("change", ["seed", "order", "count", "asset", "policy"]) -def test_original_trace_rejects_mispaired_evaluation( - tmp_path: Path, change: str -) -> None: - directory = tmp_path / "teacher" - spec, source, metadata = _fixture(directory, "cartpole") - path = directory / "manifest.json" - manifest = json.loads(path.read_text()) - if change == "seed": - manifest["evaluation"]["seed"] += 1 - elif change == "order": - manifest["shards"][0]["trajectories"].reverse() - elif change == "count": - manifest["evaluation"]["trajectories"] += 1 - elif change == "asset": - manifest["metadata"]["asset_sha256"] = "another-asset" - else: - manifest["evaluation"]["policy"] = {"kind": "pd"} - path.write_text(json.dumps(manifest)) - with pytest.raises(ValueError): - _load_reference_trace(directory, spec, source, metadata, None, 3, 4) - - -def test_original_trace_rejects_rehashed_control_mismatch(tmp_path: Path) -> None: - directory = tmp_path / "teacher" - spec, source, metadata = _fixture(directory, "ant") - file = directory / "teacher.h5" - with h5py.File(file, "a") as stream: - stream["0/transitions/joint_f"][1, 0] = 1 - path = directory / "manifest.json" - manifest = json.loads(path.read_text()) - manifest["shards"][0]["sha256"] = hashlib.sha256(file.read_bytes()).hexdigest() - path.write_text(json.dumps(manifest)) - _, records = _load_reference_trace(directory, spec, source, metadata, None, 3, 4) - with pytest.raises(ValueError, match="controls differ"): - list(_reference_trace_batches(directory, records, 2, 4, "ant", "cpu")) - - -def test_original_trace_rejects_raw_shard_tampering(tmp_path: Path) -> None: - directory = tmp_path / "teacher" - spec, source, metadata = _fixture(directory, "cartpole") - with (directory / "teacher.h5").open("ab") as stream: - stream.write(b"tamper") - with pytest.raises(ValueError, match="checksum mismatch"): - _load_reference_trace(directory, spec, source, metadata, None, 3, 4) - - -@pytest.mark.parametrize("change", ["missing", "dtype", "length"]) -def test_original_trace_rejects_rehashed_invalid_state_schema( - tmp_path: Path, change: str -) -> None: - directory = tmp_path / "teacher" - spec, source, metadata = _fixture(directory, "cartpole") - file = directory / "teacher.h5" - with h5py.File(file, "a") as stream: - value = stream["0/states/raw"][()] - del stream["0/states/raw"] - if change != "missing": - stream.create_dataset( - "0/states/raw", - data=value.astype(np.float64) if change == "dtype" else value[:-1], - ) - path = directory / "manifest.json" - manifest = json.loads(path.read_text()) - manifest["shards"][0]["sha256"] = hashlib.sha256(file.read_bytes()).hexdigest() - path.write_text(json.dumps(manifest)) - _, records = _load_reference_trace(directory, spec, source, metadata, None, 3, 4) - with pytest.raises( - ValueError, match="missing raw/canonical fields|dtype, shape or length" - ): - list(_reference_trace_batches(directory, records, 2, 4, "cartpole", "cpu")) diff --git a/tests/examples/learning/nerd/test_report.py b/tests/examples/learning/nerd/test_report.py deleted file mode 100644 index 9285374181..0000000000 --- a/tests/examples/learning/nerd/test_report.py +++ /dev/null @@ -1,414 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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. -# ---------------------------------------------------------------------------- - -"""Qualification cannot substitute pass flags or mismatched rollout populations.""" - -from __future__ import annotations - -import copy -import json -from pathlib import Path - -import pytest - -from examples.learning.nerd._support._report import ( - _curve, - _dataset_evidence, - _matched_rollouts, - _numerical_evidence, - _rollout_evidence, - annotate_rollout_evidence, -) -from examples.learning.nerd._support._reference import OFFICIAL_COMMIT, _hash - - -def test_learning_curve_retains_sparse_full_budget_observations(tmp_path: Path) -> None: - path = tmp_path / "metrics.jsonl" - budget = 4995000 - with path.open("w") as stream: - for step in range(500, budget + 1, 500): - stream.write(json.dumps({"step": step, "loss": 1.0 / step}) + "\n") - steps, losses = _curve(path, budget) - assert steps[0] == 500 and steps[-1] == budget - assert 900 <= len(steps) <= 1000 - assert losses[0] == 1.0 / 500 and losses[-1] == 1.0 / budget - assert all(left < right for left, right in zip(steps, steps[1:])) - - -def test_learning_curve_keeps_irregular_final_observation(tmp_path: Path) -> None: - path = tmp_path / "metrics.jsonl" - observed = [300, 650, 3000, 4500, 12000] - path.write_text( - "".join( - json.dumps({"step": step, "loss": float(step)}) + "\n" for step in observed - ) - ) - steps, losses = _curve(path, 12000, max_points=4) - assert steps[0] == 300 and steps[-1] == 12000 - assert len(steps) <= 4 and all(step in observed for step in steps) - assert losses == [float(step) for step in steps] - - -def test_isolated_supported_source_import_has_verified_dataset_lineage(tmp_path): - import h5py - import numpy as np - from examples.learning.nerd._support._reference import import_legacy_datasets - - sources = {} - identities = {} - for split in ("train", "validation", "test"): - raw = tmp_path / f"{split}.h5" - with h5py.File(raw, "w") as file: - data = file.create_group("data") - data.attrs["mode"] = "trajectory" - data.create_dataset("states", data=np.zeros((2, 1, 4), np.float32)) - data.create_dataset("next_states", data=np.zeros((2, 1, 4), np.float32)) - data.create_dataset("joint_acts", data=np.zeros((2, 1, 2), np.float32)) - sources[split] = raw - identities[split] = { - "source_commit": OFFICIAL_COMMIT, - "robot": "cartpole", - "python": "3.8.20", - "numpy": "1.24.4", - "torch": "2.2.2+cpu", - "warp": "1.8.0", - "runtime_imports_embodichain": False, - "raw_sha256": _hash(raw), - } - metadata = { - "codec": {"robot": "cartpole"}, - "source_split_identities": identities, - "source_physical_arrays": {"body_mass": [50.0, 2.5]}, - "teacher_origin": {"kind": "official_warp", "source_commit": OFFICIAL_COMMIT}, - } - output = tmp_path / "canonical" - result = import_legacy_datasets(sources, output, metadata) - declaration = { - "path": str(output), - "status": "available", - "fingerprint": result["fingerprint"], - } - evidence = _dataset_evidence("official_warp", declaration) - assert evidence["status"] == "available" - assert evidence["transitions"] == {"train": 2, "validation": 2, "test": 2} - with h5py.File(sources["test"], "r+") as file: - file["data/states"][0, 0, 0] = 1.0 - assert _dataset_evidence("official_warp", declaration)["status"] == "pending" - - -@pytest.mark.parametrize( - "field,value", - [ - ("source_commit", "unrelated"), - ("python", "3.11.14"), - ("torch", "2.7.0+cu128"), - ("warp", "1.17.0"), - ("raw_sha256", "0" * 64), - ("runtime_imports_embodichain", True), - ], -) -def test_source_split_labels_do_not_replace_supported_runtime_receipts( - tmp_path, field, value -): - from examples.learning.nerd._support._report import _supported_split_sources - - raw = tmp_path / "source.h5" - raw.write_bytes(b"synthetic raw-file identity fixture") - identity = { - "source_commit": OFFICIAL_COMMIT, - "robot": "cartpole", - "python": "3.8.20", - "numpy": "1.24.4", - "torch": "2.2.2+cpu", - "warp": "1.8.0", - "runtime_imports_embodichain": False, - "raw_sha256": _hash(raw), - } - splits = ("train", "validation", "test") - metadata = { - "codec": {"robot": "cartpole"}, - "legacy_sources": { - split: {"path": str(raw), "sha256": _hash(raw)} for split in splits - }, - "source_split_identities": {split: dict(identity) for split in splits}, - } - assert _supported_split_sources(metadata) - metadata["source_split_identities"]["test"][field] = value - assert not _supported_split_sources(metadata) - - -def numerical_fixture() -> dict: - return { - "robot": "Cartpole", - "reference_commit": "476dc67fccc6a99be293bf611f897b1c06f13e30", - "checkpoint_sha256": "a" * 64, - "codec_fixtures_sha256": "b" * 64, - "history_lengths": [1, 9, 10], - "adam_updates": 3, - "precision": "float32", - "atol": 1e-5, - "rtol": 1e-4, - **{ - key: 1e-7 - for key in ( - "max_codec_feature_error", - "max_codec_target_error", - "max_codec_decode_error", - "max_forward_error", - "max_input_gradient_error", - "max_parameter_gradient_error", - "max_adam_parameter_error", - ) - }, - } - - -def test_numerical_receipt_requires_measurements_and_matching_identity(tmp_path): - path = tmp_path / "numerical.json" - path.write_text(json.dumps(numerical_fixture())) - assert _numerical_evidence(path, "cartpole")["status"] == "passed" - assert _numerical_evidence(path, "ant")["status"] == "failed" - path.write_text( - json.dumps( - { - "checks": { - key: {"passed": True} - for key in ( - "preprocessing", - "forward", - "input_gradients", - "parameter_gradients", - "adam_update", - ) - } - } - ) - ) - assert _numerical_evidence(path, "cartpole")["status"] == "pending" - - -@pytest.mark.parametrize( - "commit,reference_kind", - [ - (OFFICIAL_COMMIT, "official_warp_main"), - ("476dc67fccc6a99be293bf611f897b1c06f13e30", "newton_preview"), - ], -) -def test_pinned_main_and_preview_numerical_sources_remain_identifiable( - tmp_path, commit, reference_kind -): - raw = numerical_fixture() - raw["reference_commit"] = commit - path = tmp_path / "numerical.json" - path.write_text(json.dumps(raw)) - result = _numerical_evidence(path, "cartpole") - assert result["status"] == "passed" - assert result["reference_commit"] == commit - assert result["reference_kind"] == reference_kind - - -@pytest.mark.parametrize( - "field,value", - [ - ("max_forward_error", float("nan")), - ("max_codec_decode_error", -1.0), - ("atol", 1.0), - ("reference_commit", "unrelated"), - ("reference_commit", []), - ("reference_commit", {}), - ], -) -def test_invalid_numerical_measurements_cannot_pass(tmp_path, field, value): - raw = numerical_fixture() - raw[field] = value - path = tmp_path / "numerical.json" - path.write_text(json.dumps(raw)) - assert _numerical_evidence(path, "cartpole")["status"] == "failed" - - -def rollout_fixture(tmp_path: Path, *, motion: str = "passive"): - model = tmp_path / "run" / "best-export" - model.mkdir(parents=True) - (model / "weights.pt").write_bytes(b"synthetic test artifact") - (model / "manifest.json").write_text("{}") - trace = tmp_path / "trace.h5" - trace.write_bytes(b"synthetic test trace") - trajectories = tmp_path / "rows.jsonl" - trajectories.write_bytes(b"synthetic test row index") - raw = { - "robot": "cartpole", - "motion_mode": motion, - "state_dtype": "float32", - "teacher_origin": {"kind": "current_newton"}, - "teacher_split": "test", - "selected_teacher_sha256": "a" * 64, - "selected_teacher_hash_schema": "ordered-teacher-content-v1", - "runtime_identity": { - "device": "cpu", - "engine": "current_newton", - "state_dtype": "float32", - "python": "3.11", - "versions": {"torch": "2.7.1", "newton": "1.6.0", "warp-lang": "1.17.0"}, - }, - "trajectories": 2048, - "horizons": {str(h): {"q_mae": 0.1, "qd_mae": 0.2} for h in (1, 10, 100, 1000)}, - "model_sha256": _hash(model / "weights.pt"), - "model_manifest_sha256": _hash(model / "manifest.json"), - "trace_evidence": str(trace), - "trace_evidence_sha256": _hash(trace), - "trajectory_evidence": str(trajectories), - "trajectory_evidence_sha256": _hash(trajectories), - "failure_rate": 0.0, - "nonfinite_rate": 0.0, - "policy_return": 1.0 if motion == "fixed_policy" else None, - "policy_sha256": "b" * 64 if motion == "fixed_policy" else None, - } - source = tmp_path / "source.json" - source.write_text(json.dumps(raw)) - record = annotate_rollout_evidence( - source, "current_newton-official-0", tmp_path / "attached.json" - ) - cell = {"teacher": "current_newton", "output_dir": str(model.parent)} - return record, cell - - -def test_rollout_source_model_and_trace_identity_are_verified(tmp_path): - record, cell = rollout_fixture(tmp_path) - assert _rollout_evidence(record, cell, "cartpole")[0] == "available" - record["errors"]["1"]["q_mae"] = 0.0 - assert "Annotated" in _rollout_evidence(record, cell, "cartpole")[1] - record["errors"]["1"]["q_mae"] = 0.1 - Path(record["trace_evidence"]).write_bytes(b"changed") - assert "changed" in _rollout_evidence(record, cell, "cartpole")[1] - - -@pytest.mark.parametrize( - "field,value", - [ - ("teacher_origin", {"kind": "official_warp"}), - ("teacher_split", "train"), - ("failure_rate", -1.0), - ("nonfinite_rate", 1.5), - ("robot", "ant"), - ], -) -def test_wrong_domain_split_or_invalid_rates_fail(tmp_path, field, value): - record, cell = rollout_fixture(tmp_path) - record[field] = value - assert _rollout_evidence(record, cell, "cartpole")[0] == "failed" - - -def test_undefined_error_is_failed_instead_of_crashing_statistics(tmp_path): - record, cell = rollout_fixture(tmp_path) - record["errors"]["1000"]["q_mae"] = None - assert _rollout_evidence(record, cell, "cartpole")[0] == "failed" - - -def test_unmatched_teacher_traces_or_fixed_policy_never_form_seed_pair(tmp_path): - record, cell = rollout_fixture(tmp_path, motion="fixed_policy") - assert _rollout_evidence(record, cell, "cartpole")[0] == "available" - other = copy.deepcopy(record) - assert _matched_rollouts(record, other) - other["selected_teacher_sha256"] = "c" * 64 - assert not _matched_rollouts(record, other) - other = copy.deepcopy(record) - other["policy_sha256"] = "d" * 64 - assert not _matched_rollouts(record, other) - - -def test_historical_rollout_uses_pinned_warp_identity_without_newton(tmp_path): - from examples.learning.nerd._support._reference import OFFICIAL_COMMIT - - record, cell = rollout_fixture(tmp_path) - source = Path(record["source_report"]) - raw = json.loads(source.read_text()) - raw["teacher_origin"] = { - "kind": "official_warp", - "source_commit": OFFICIAL_COMMIT, - "warp_version": "1.8.0", - } - raw["runtime_identity"] = { - "device": "cpu", - "engine": "official_warp", - "state_dtype": "float32", - "python": "3.11", - "reference_commit": OFFICIAL_COMMIT, - "versions": {"torch": "2.7.1", "warp-lang": "1.8.0"}, - } - raw["physical_verification"] = { - "schema": "official-warp-physical-arrays-v2", - "complete": True, - "matched_recorded_fields": True, - "missing_recorded_fields": [], - "matched_fields": ["gravity"], - "sha256": "f" * 64, - } - source.write_text(json.dumps(raw)) - record = annotate_rollout_evidence( - source, "official_warp-official-0", tmp_path / "historical.json" - ) - cell["teacher"] = "official_warp" - assert _rollout_evidence(record, cell, "cartpole")[0] == "available" - record["physical_verification"]["missing_recorded_fields"] = ["gravity"] - assert _rollout_evidence(record, cell, "cartpole")[0] == "pending" - record["physical_verification"]["missing_recorded_fields"] = [] - record["runtime_identity"]["versions"]["warp-lang"] = "1.17.0" - assert "pinned original Warp" in _rollout_evidence(record, cell, "cartpole")[1] - - -def test_full_profile_checks_actual_feature_order_precision_and_training_controls( - tmp_path, -): - import torch - from examples.learning.nerd._support._report import _full_profile_issues - from examples.learning.nerd._support._trainer import DynamicsTrainCfg - from embodichain_models.dynamics.nerd import NeRDModelCfg - - model = NeRDModelCfg(input_dims={"states_embedding": 4, "joint_f": 2}).to_dict() - root = tmp_path / "initial" - root.mkdir() - manifest = root / "manifest.json" - manifest.write_text(json.dumps({"model_config": model})) - torch.save({"weight": torch.ones(1)}, root / "weights.pt") - spec = tmp_path / "run.json" - spec.write_text(json.dumps({"initial_model": str(root)})) - train = DynamicsTrainCfg( - total_steps=4995000, - batch_size=512, - validation_interval=5000, - validation_batches=100, - lr_update_interval=5000, - lr_offset=1, - ).to_dict() - # Plan JSON sorts dict keys. This must not reorder the actual network inputs. - plan = { - "model": json.loads(json.dumps(model, sort_keys=True)), - "train": train, - "cells": [{"id": "fixture", "status": "prepared", "spec": str(spec)}], - } - assert _full_profile_issues(plan) == [] - model["input_dims"] = {"joint_f": 2, "states_embedding": 4} - manifest.write_text(json.dumps({"model_config": model})) - assert any("feature order" in issue for issue in _full_profile_issues(plan)) - model["input_dims"] = {"states_embedding": 4, "joint_f": 2} - model["head_hidden_dims"] = [128] - manifest.write_text(json.dumps({"model_config": model})) - torch.save({"weight": torch.ones(1, dtype=torch.float64)}, root / "weights.pt") - plan["train"]["batch_size"] = 16 - issues = _full_profile_issues(plan) - assert any("head_hidden_dims" in issue for issue in issues) - assert any("FP32" in issue for issue in issues) - assert any("batch_size" in issue for issue in issues) diff --git a/tests/examples/learning/nerd/test_resources.py b/tests/examples/learning/nerd/test_resources.py deleted file mode 100644 index b42532abd1..0000000000 --- a/tests/examples/learning/nerd/test_resources.py +++ /dev/null @@ -1,212 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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. -# ---------------------------------------------------------------------------- - -"""Verify the explicitly requested showcase resource-download capability.""" - -from __future__ import annotations - -import hashlib -import json -from pathlib import Path -import zipfile - -import huggingface_hub -import pytest - -from examples.learning.nerd import _resources as resources - - -def digest(value: bytes) -> str: - return hashlib.sha256(value).hexdigest() - - -@pytest.fixture -def asset_server(tmp_path, monkeypatch): - payloads = { - "ATTRIBUTION.md": b"reference attribution", - "LICENSE.md": b"Apache-2.0 license fixture", - "ant.xml": b"\n", - "cartpole_single.urdf": b"\n", - } - server = tmp_path / "server" - server.mkdir() - archive = server / "reference_assets.zip" - with zipfile.ZipFile(archive, "w") as bundle: - bundle.writestr("nerd_reference_assets/", b"") - for name, value in payloads.items(): - bundle.writestr(f"nerd_reference_assets/{name}", value) - archive_hash = digest(archive.read_bytes()) - member_hashes = {name: digest(value) for name, value in payloads.items()} - manifest = { - "schema_version": 1, - "license": "Apache-2.0", - "upstream": {"revision": resources._UPSTREAM_REVISION}, - "archive": { - "root": "nerd_reference_assets", - "sha256": archive_hash, - "size": archive.stat().st_size, - }, - "files": { - name: { - "path": f"nerd_reference_assets/{name}", - "sha256": digest(value), - "size": len(value), - } - for name, value in payloads.items() - }, - } - metadata = server / "manifest.json" - metadata.write_text(json.dumps(manifest)) - monkeypatch.setattr(resources, "_ASSET_FILES", member_hashes) - monkeypatch.setattr(resources, "_ARCHIVE_SHA256", archive_hash) - monkeypatch.setattr(resources, "_MANIFEST_SHA256", digest(metadata.read_bytes())) - calls = [] - - def download(**options): - assert options["repo_id"] == "DexForceAI/embodichain_data" - assert options["repo_type"] == "dataset" - assert options["revision"] == "97e202e41fe7fb504656a6b163291305e400780d" - calls.append(options) - return str(server / Path(options["filename"]).name) - - monkeypatch.setattr(huggingface_hub, "hf_hub_download", download) - return payloads, calls, archive - - -def test_pinned_asset_cache_repairs_corruption_and_reuses_verified_files( - tmp_path, asset_server -): - payloads, calls, _ = asset_server - cache = tmp_path / "cache" - target = resources.prepare_reference_assets(cache) - assert {path.name for path in target.iterdir()} == set(payloads) - assert all( - (target / name).read_bytes() == value for name, value in payloads.items() - ) - assert len(calls) == 2 - assert resources.prepare_reference_assets(cache, local_files_only=True) == target - assert len(calls) == 2 # A verified cache needs neither Hub nor network access. - (target / "ant.xml").write_bytes(b"damaged local cache") - assert resources.prepare_reference_assets(cache, local_files_only=True) == target - assert len(calls) == 4 and calls[-1]["local_files_only"] is True - assert (target / "ant.xml").read_bytes() == payloads["ant.xml"] - assert not list(target.parent.glob(".staging-*")) - - -def test_invalid_download_archive_is_rejected_before_extraction(tmp_path, asset_server): - _, _, archive = asset_server - archive.write_bytes(archive.read_bytes() + b"changed") - cache = tmp_path / "cache" - with pytest.raises(ValueError, match="SHA256 mismatch"): - resources.prepare_reference_assets(cache) - assert not (cache / "reference_assets").exists() - - -def test_example_resource_environment_and_explicit_cache_precedence( - tmp_path, asset_server, monkeypatch -): - monkeypatch.setenv("EMBODICHAIN_NERD_RESOURCE_ROOT", str(tmp_path / "configured")) - environmental = resources.prepare_reference_assets() - assert environmental.is_relative_to(tmp_path / "configured/assets") - explicit = resources.prepare_reference_assets(tmp_path / "explicit") - assert explicit.is_relative_to(tmp_path / "explicit") - - -def test_assets_only_does_not_start_model_or_dataset_downloads( - tmp_path, asset_server, monkeypatch -): - def unexpected_snapshot(**kwargs): - raise AssertionError("Assets-only preparation must not download large bundles") - - monkeypatch.setattr(huggingface_hub, "snapshot_download", unexpected_snapshot) - prepared = resources.prepare_resources( - tmp_path / "resources", "cartpole", assets_only=True - ) - assert set(prepared) == {"assets"} - assert prepared["assets"].is_dir() - assert not (tmp_path / "resources/datasets").exists() - assert not (tmp_path / "resources/models").exists() - - -def test_published_payload_revisions_paths_and_checksums( - tmp_path, asset_server, monkeypatch -): - calls = [] - manifests = {"dataset": {}, "model": {}} - monkeypatch.setattr(resources, "_PAYLOAD_MANIFEST_SHA256", manifests) - metadata = { - "codec": {"robot": "cartpole"}, - "dt": 1 / 60, - "asset_sha256": resources._ASSET_FILES["cartpole_single.urdf"], - } - - def snapshot(**options): - calls.append(options) - is_dataset = options["repo_type"] == "dataset" - assert options["revision"] == ( - "b9ec58f36da99cc99a93f97adfe74ebc9fd66586" - if is_dataset - else "78e581526570069410bf269b38c0c52c800b24a7" - ) - directory = Path(options["local_dir"]) / "neural_dynamics/nerd/cartpole/v1" - directory.mkdir(parents=True, exist_ok=True) - if is_dataset: - value = b"immutable dataset shard" - (directory / "shard-000001.h5").write_bytes(value) - manifest = { - "schema_version": 1, - "metadata": metadata, - "shards": [{"file": "shard-000001.h5", "sha256": digest(value)}], - } - else: - value = b"immutable portable weights" - (directory / "weights.pt").write_bytes(value) - manifest = { - "format_version": 1, - **metadata, - "weights_sha256": digest(value), - } - (directory / "manifest.json").write_text(json.dumps(manifest)) - manifests["dataset" if is_dataset else "model"]["cartpole"] = digest( - (directory / "manifest.json").read_bytes() - ) - - monkeypatch.setattr(huggingface_hub, "snapshot_download", snapshot) - prepared = resources.prepare_resources(tmp_path / "resources", "cartpole") - assert ( - prepared["dataset"] - == tmp_path / "resources/datasets/neural_dynamics/nerd/cartpole/v1" - ) - assert ( - prepared["model"] - == tmp_path / "resources/models/neural_dynamics/nerd/cartpole/v1" - ) - assert len(calls) == 2 - resources.prepare_resources( - tmp_path / "resources", "cartpole", local_files_only=True - ) - assert len(calls) == 2 - (prepared["model"] / "weights.pt").write_bytes(b"bad weights") - monkeypatch.setattr(huggingface_hub, "snapshot_download", lambda **kwargs: None) - with pytest.raises(ValueError, match="SHA256 mismatch"): - resources.prepare_resources(tmp_path / "resources", "cartpole") - # A self-consistent replacement still cannot impersonate the pinned release. - manifest_path = prepared["model"] / "manifest.json" - manifest = json.loads(manifest_path.read_text()) - manifest["weights_sha256"] = digest(b"bad weights") - manifest_path.write_text(json.dumps(manifest)) - with pytest.raises(ValueError, match="manifest.json"): - resources.prepare_resources(tmp_path / "resources", "cartpole") diff --git a/tests/examples/learning/nerd/test_sample_order.py b/tests/examples/learning/nerd/test_sample_order.py deleted file mode 100644 index efca95807f..0000000000 --- a/tests/examples/learning/nerd/test_sample_order.py +++ /dev/null @@ -1,328 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Exact, bounded shared minibatch order and tamper detection.""" - -from __future__ import annotations - -import json -from pathlib import Path -import numpy as np -import pytest -import torch - -from examples.learning.nerd._support._sample_order import ( - _CheckpointedOrder, - _OfficialShuffleOrder, - _open_order, - _permutation_sha256, - _write_checkpointed_order, - _write_official_shuffle_order, -) - - -def test_checkpointed_stream_matches_the_original_matrix_and_random_seeks(tmp_path): - steps, batch, block_steps = 2071, 16, 1024 - generator = torch.Generator().manual_seed(19) - expected = torch.cat( - [ - torch.randint( - 1037, (min(block_steps, steps - start), batch), generator=generator - ) - for start in range(0, steps, block_steps) - ] - ).numpy() - path, rng = _write_checkpointed_order( - tmp_path / "order.json", - dataset_size=1037, - steps=steps, - batch_size=batch, - seed=19, - block_steps=block_steps, - ) - actual = _open_order(path) - assert isinstance(actual, _CheckpointedOrder) - for index in [0, 1023, 1024, 2070, 2, 2048, 1024]: - np.testing.assert_array_equal(actual[index], expected[index]) - row = actual[0] - row[:] = -1 - np.testing.assert_array_equal(actual[0], expected[0]) - assert actual._block.nbytes <= block_steps * batch * 8 - assert rng.stat().st_size < expected.nbytes - with pytest.raises(IndexError): - actual[steps] - with pytest.raises(IndexError): - actual[-1] - - -def test_checkpoint_stream_first_block_is_independent_of_declared_budget(tmp_path): - for name, steps in [("a", 8), ("b", 3000)]: - _write_checkpointed_order( - tmp_path / (name + ".json"), - dataset_size=91, - steps=steps, - batch_size=8, - seed=7, - ) - left, right = [_open_order(tmp_path / (name + ".json")) for name in ["a", "b"]] - for step in range(8): - np.testing.assert_array_equal(left[step], right[step]) - - -def test_changed_rng_artifact_or_torch_version_fails_before_training(tmp_path): - path, rng = _write_checkpointed_order( - tmp_path / "order.json", - dataset_size=97, - steps=40, - batch_size=4, - seed=3, - block_steps=8, - ) - spec = json.loads(path.read_text()) - spec["torch_version"] = "other" - path.write_text(json.dumps(spec)) - with pytest.raises(ValueError, match="Torch version"): - _open_order(path) - spec["torch_version"] = str(torch.__version__) - path.write_text(json.dumps(spec)) - values = np.load(rng).copy() - values[-1, 20] ^= 1 - np.save(rng, values) - with pytest.raises(ValueError, match="checksum"): - _open_order(path) - - -def test_legacy_matrix_still_opens_read_only(tmp_path): - values = np.arange(24).reshape(6, 4) - np.save(tmp_path / "order.npy", values) - actual = _open_order(tmp_path / "order.npy") - assert isinstance(actual, np.memmap) - np.testing.assert_array_equal(actual, values) - - -def _official_order(tmp_path: Path) -> tuple[Path, Path, Path, np.ndarray]: - # Five complete batches/pass; each twelve-update epoch discards the unused - # tail of its third pass. The final three-update prefix uses one more pass. - size, batch, epoch_steps, steps = 23, 4, 12, 27 - seeds = np.arange(101, 108, dtype=np.int64) - hashes = [ - _permutation_sha256( - torch.randperm( - size, generator=torch.Generator().manual_seed(int(seed)) - ).numpy() - ) - for seed in seeds - ] - paths = _write_official_shuffle_order( - tmp_path / "official.json", - dataset_size=size, - steps=steps, - batch_size=batch, - epoch_steps=epoch_steps, - pass_seeds=seeds, - permutation_sha256=hashes, - provenance={ - "source_rng_evidence": "Synthetic conditional RandomSampler fixture", - "producer_torch_version": "2.2.2+cpu", - }, - ) - # Independent oracle: actual Torch RandomSampler/DataLoader, dropping - # incomplete minibatches and abandoning the final iterator each epoch. - from torch.utils.data import DataLoader, RandomSampler - - result = [] - seed_iterator = iter(seeds) - while len(result) < steps: - remaining = min(epoch_steps, steps - len(result)) - while remaining: - sampler = RandomSampler( - range(size), - replacement=False, - generator=torch.Generator().manual_seed(int(next(seed_iterator))), - ) - loader = DataLoader( - range(size), - batch_size=batch, - sampler=sampler, - drop_last=True, - num_workers=0, - generator=torch.Generator().manual_seed(8), - ) - for indices in loader: - result.append(indices.numpy()) - remaining -= 1 - if not remaining: - break - return *paths, np.stack(result) - - -def test_official_order_matches_real_random_sampler_and_epoch_resets( - tmp_path: Path, -) -> None: - path, _, _, expected = _official_order(tmp_path) - actual = _open_order(path) - assert isinstance(actual, _OfficialShuffleOrder) - assert actual.shape == expected.shape - for index in range(len(expected)): - np.testing.assert_array_equal(actual[index], expected[index]) - for index in (26, 0, 5, 11, 12, 9, 25, 1): - np.testing.assert_array_equal(actual[index], expected[index]) - assert actual._permutation.nbytes == actual.dataset_size * 8 - assert not actual.seeds.flags.writeable - for index in (-1, len(expected)): - with pytest.raises(IndexError): - actual[index] - - -def test_official_rows_remain_owned_across_seeks_and_reopen(tmp_path: Path) -> None: - path, _, _, expected = _official_order(tmp_path) - actual = _open_order(path) - row = actual[0] - actual[26] - np.testing.assert_array_equal(row, expected[0]) - row[:] = -1 - np.testing.assert_array_equal(actual[0], expected[0]) - resumed = _open_order(path) - for index in (12, 25, 26): - np.testing.assert_array_equal(resumed[index], expected[index]) - - -def test_trainer_admits_official_order_and_rejects_a_different_dataset_size( - tmp_path: Path, -) -> None: - from types import SimpleNamespace - from examples.learning.nerd._support._trainer import DynamicsTrainer - - path, _, _, expected = _official_order(tmp_path) - # Exercise the production order admission boundary without creating a - # model, optimizer, simulation or allocating CUDA memory. - owner = DynamicsTrainer.__new__(DynamicsTrainer) - owner.cfg = SimpleNamespace(total_steps=27, batch_size=4) - owner.train_dataset = range(23) - order = owner._prepare_order(path) - assert isinstance(order, _OfficialShuffleOrder) - np.testing.assert_array_equal(order[12], expected[12]) - assert owner._sample_order_identity["path"] == str(path.resolve()) - owner.train_dataset = range(24) - with pytest.raises(ValueError, match="dataset size mismatch"): - owner._prepare_order(path) - - -@pytest.mark.parametrize("payload", ["seeds", "hashes"]) -def test_official_payload_tampering_is_rejected(tmp_path: Path, payload: str) -> None: - path, seeds, hashes, _ = _official_order(tmp_path) - target = seeds if payload == "seeds" else hashes - values = np.load(target, allow_pickle=False).copy() - values[0] = int(values[0]) ^ 1 if payload == "seeds" else b"0" * 64 - np.save(target, values, allow_pickle=False) - with pytest.raises(ValueError, match="payload checksum"): - _open_order(path) - - -def test_official_rehashed_wrong_permutation_still_fails_before_yield( - tmp_path: Path, -) -> None: - import hashlib - - path, _, hashes, _ = _official_order(tmp_path) - values = np.load(hashes, allow_pickle=False).copy() - values[0] = b"0" * 64 - np.save(hashes, values, allow_pickle=False) - spec = json.loads(path.read_text()) - spec["permutation_hashes_sha256"] = hashlib.sha256(hashes.read_bytes()).hexdigest() - path.write_text(json.dumps(spec)) - reader = _open_order(path) - with pytest.raises(ValueError, match="permutation checksum"): - reader[0] - - -@pytest.mark.parametrize( - "change", - [ - {"drop_last": False}, - {"discard_epoch_tail": False}, - {"permutation_algorithm": "other"}, - {"seeds_file": "../outside.npy"}, - ], -) -def test_official_unknown_or_escaping_contract_is_rejected( - tmp_path: Path, change: dict -) -> None: - path, _, _, _ = _official_order(tmp_path) - spec = json.loads(path.read_text()) - spec.update(change) - path.write_text(json.dumps(spec)) - with pytest.raises(ValueError, match="contract|sibling file"): - _open_order(path) - - -def test_official_full_budget_storage_is_not_a_dense_minibatch_matrix( - tmp_path: Path, -) -> None: - size, batch, steps, epoch_steps = 931840, 512, 4995000, 5000 - passes = 2997 - # This is a scale fixture, not claimed original RNG capture: repeated seeds - # let the real reader exercise the full formal shape with one hash oracle. - seeds = np.full(passes, 107, dtype=np.int64) - expected = torch.randperm( - size, generator=torch.Generator().manual_seed(107) - ).numpy() - digest = _permutation_sha256(expected) - path, seed_path, hash_path = _write_official_shuffle_order( - tmp_path / "full.json", - dataset_size=size, - steps=steps, - batch_size=batch, - epoch_steps=epoch_steps, - pass_seeds=seeds, - permutation_sha256=[digest] * passes, - provenance={"source_rng_evidence": "Synthetic storage/shape fixture"}, - ) - reader = _open_order(path) - np.testing.assert_array_equal(reader[0], expected[:batch]) - # The final epoch's third pass consumes only 1360 minibatches. - start = (1360 - 1) * batch - np.testing.assert_array_equal(reader[steps - 1], expected[start : start + batch]) - assert reader._permutation.nbytes == size * 8 - assert sum(p.stat().st_size for p in (path, seed_path, hash_path)) < 256 * 1024 - - -@pytest.mark.parametrize( - "change", - [ - {"dataset_size": 0}, - {"dataset_size": 4194305}, - {"epoch_steps": 1}, - {"pass_seeds": [1]}, - {"permutation_sha256": ["wrong"]}, - {"provenance": {}}, - ], -) -def test_official_writer_rejects_invalid_or_unbounded_inputs( - tmp_path: Path, change: dict -) -> None: - options = dict( - dataset_size=23, - steps=27, - batch_size=4, - epoch_steps=12, - pass_seeds=np.arange(7, dtype=np.int64), - permutation_sha256=["0" * 64] * 7, - provenance={"source_rng_evidence": "Synthetic negative fixture"}, - ) - options.update(change) - with pytest.raises(ValueError): - _write_official_shuffle_order(tmp_path / "invalid.json", **options) diff --git a/tests/examples/learning/nerd/test_small_budget_observation.py b/tests/examples/learning/nerd/test_small_budget_observation.py deleted file mode 100644 index c83d7b1159..0000000000 --- a/tests/examples/learning/nerd/test_small_budget_observation.py +++ /dev/null @@ -1,167 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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. -# ---------------------------------------------------------------------------- - -"""Regression coverage for original-model in-place input normalization.""" - -from __future__ import annotations - -import importlib.util -import json -from pathlib import Path -from types import SimpleNamespace -from typing import Any - -import pytest -import torch - - -@pytest.fixture(scope="module") -def workers() -> tuple[Any, Any]: - directory = Path(__file__).resolve().parents[4] / "examples/learning/nerd/_support" - modules = [] - for filename in ("_official_small_budget", "_matched_small_budget"): - spec = importlib.util.spec_from_file_location( - "_test_" + filename, directory / (filename + ".py") - ) - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - modules.append(module) - return tuple(modules) - - -def _inputs() -> tuple[Any, dict[str, torch.Tensor], torch.Tensor]: - # Nonzero mean and nonunit variance expose accidental double normalization. - rms = SimpleNamespace(mean=torch.tensor([3.0]), var=torch.tensor([4.0])) - model = SimpleNamespace(input_rms={"joint_f": rms}) - return model, {"joint_f": torch.tensor([[[7.0], [11.0]]])}, torch.tensor([1.5]) - - -def test_snapshot_retains_physical_features_after_model_mutation( - workers: tuple[Any, Any], tmp_path: Path -) -> None: - source, _ = workers - _, inputs, targets = _inputs() - data = {"joint_acts": inputs["joint_f"].clone(), "target": targets.clone()} - source._save_first_batch(tmp_path, data) - data["joint_acts"].sub_(3.0).div_(torch.sqrt(torch.tensor(4.0 + 1e-5))) - stored = torch.load(tmp_path / "first-preprocessed-batch.pt", weights_only=True) - assert torch.equal(stored["joint_acts"], inputs["joint_f"]) - - -def test_loss_observer_saves_before_original_forward_normalizes_in_place( - workers: tuple[Any, Any], tmp_path: Path -) -> None: - source, _ = workers - _, inputs, targets = _inputs() - data = {"joint_acts": inputs["joint_f"].clone(), "target": targets} - state = {"steps": 0, "last_loss": None} - calls = [] - - def original_loss(trainer: Any, batch: dict[str, Any], train: bool) -> Any: - calls.append(train) - batch["joint_acts"] = (batch["joint_acts"] - 3.0) / (4.0 + 1e-5) ** 0.5 - return torch.tensor(2.0), {} - - source._observed_loss(None, data, True, original_loss, state, tmp_path) - stored = torch.load(tmp_path / "first-preprocessed-batch.pt", weights_only=True) - assert torch.equal(stored["joint_acts"], inputs["joint_f"]) - assert not torch.equal(data["joint_acts"], inputs["joint_f"]) - assert state["last_loss"] == 2.0 and calls == [True] - - -@pytest.mark.parametrize("train,step", [(False, 0), (True, 1)]) -def test_loss_observer_only_captures_the_first_training_batch( - workers: tuple[Any, Any], tmp_path: Path, train: bool, step: int -) -> None: - source, _ = workers - _, inputs, targets = _inputs() - state = {"steps": step, "last_loss": None} - source._observed_loss( - None, - {"joint_acts": inputs["joint_f"], "target": targets}, - train, - lambda *_: (torch.tensor(2.0), {}), - state, - tmp_path, - ) - assert not (tmp_path / "first-preprocessed-batch.pt").exists() - - -@pytest.mark.parametrize("legacy", [False, True]) -def test_feature_comparison_respects_observation_normalization_stage( - workers: tuple[Any, Any], tmp_path: Path, legacy: bool -) -> None: - source, candidate = workers - model, inputs, targets = _inputs() - if legacy: - rms = model.input_rms["joint_f"] - normalized = (inputs["joint_f"] - rms.mean) / torch.sqrt(rms.var + 1e-5) - torch.save( - {"joint_acts": normalized, "target": targets}, - tmp_path / "first-preprocessed-batch.pt", - ) - else: - source._save_first_batch( - tmp_path, {"joint_acts": inputs["joint_f"], "target": targets} - ) - differences, stage = candidate._first_batch_differences( - model, inputs, targets, tmp_path - ) - assert differences == {"joint_f": 0.0, "target": 0.0} - assert stage == ( - "legacy_unlabeled_normalized_after_model_forward" - if legacy - else "physical_before_model_forward" - ) - - -def test_feature_comparison_still_detects_changed_inputs_and_targets( - workers: tuple[Any, Any], tmp_path: Path -) -> None: - source, candidate = workers - model, inputs, targets = _inputs() - source._save_first_batch( - tmp_path, {"joint_acts": inputs["joint_f"], "target": targets} - ) - differences, _ = candidate._first_batch_differences( - model, {"joint_f": inputs["joint_f"] + 2.0}, targets + 0.25, tmp_path - ) - assert differences["joint_f"] == pytest.approx(2.0 / (4.0 + 1e-5) ** 0.5) - assert differences["target"] == 0.25 - - -@pytest.mark.parametrize( - "change", - [ - {"feature_stage": "unknown"}, - {"format_version": 2}, - {"tensor_sha256": "0" * 64}, - ], -) -def test_feature_comparison_rejects_unknown_or_unbound_observations( - workers: tuple[Any, Any], tmp_path: Path, change: dict[str, Any] -) -> None: - source, candidate = workers - model, inputs, targets = _inputs() - source._save_first_batch( - tmp_path, {"joint_acts": inputs["joint_f"], "target": targets} - ) - metadata = tmp_path / "first-preprocessed-batch.json" - value = json.loads(metadata.read_text()) - value.update(change) - metadata.write_text(json.dumps(value)) - with pytest.raises(ValueError, match="observation metadata"): - candidate._first_batch_differences(model, inputs, targets, tmp_path) diff --git a/tests/examples/learning/nerd/test_task_values.py b/tests/examples/learning/nerd/test_task_values.py deleted file mode 100644 index e32f28d29c..0000000000 --- a/tests/examples/learning/nerd/test_task_values.py +++ /dev/null @@ -1,70 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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. -# ---------------------------------------------------------------------------- - -from __future__ import annotations -import pytest -import torch -from examples.learning.nerd._support._task_values import reference_task_values - - -def test_cartpole_observation_and_official_cost_definition() -> None: - q = torch.tensor([[1.0, 0.2], [-0.7, 2 * torch.pi + 0.2]], dtype=torch.double) - qd = torch.tensor([[0.3, -0.4], [0.1, 0.2]], dtype=torch.double) - obs, reward, failed = reference_task_values("cartpole", q, qd) - torch.testing.assert_close(obs[:, 1], torch.full((2,), 0.2, dtype=torch.double)) - expected = -( - torch.full((2,), 0.2**2, dtype=torch.double) - + 0.05 * q[:, 0] ** 2 - + 0.1 * (qd[:, 0] ** 2 + qd[:, 1] ** 2) - ) - torch.testing.assert_close(reward, expected) - assert obs.shape == (2, 4) and not failed.any() - assert torch.autograd.gradcheck( - lambda a, b: reference_task_values("cartpole", a, b)[:2], - (q.requires_grad_(), qd.requires_grad_()), - ) - - -def test_ant_com_velocity_forward_reward_and_translation_invariance() -> None: - q = torch.zeros(2, 15, dtype=torch.double) - q[:, 2] = 1.0 - angle = torch.tensor([0.2, -0.4], dtype=torch.double) - q[:, 5] = torch.sin(angle / 2) - q[:, 6] = torch.cos(angle / 2) - qd = torch.zeros(2, 14, dtype=torch.double) - qd[:, 1] = torch.tensor([0.3, 0.7], dtype=torch.double) - obs, reward, failed = reference_task_values("ant", q, qd) - torch.testing.assert_close(reward, qd[:, 1] + 0.1 + torch.cos(angle)) - moved = q.clone() - moved[:, :2] += torch.tensor([100.0, -100.0], dtype=torch.double) - changed = reference_task_values("ant", moved, qd) - torch.testing.assert_close(changed[0], obs) - torch.testing.assert_close(changed[1], reward) - assert obs.shape == (2, 27) and not failed.any() - - -def test_reference_failures_are_explicit_and_shapes_checked() -> None: - q = torch.zeros(2, 15) - q[:, 2] = torch.tensor([0.2, 1.0]) - q[0, 6] = 1.0 - _, _, failed = reference_task_values("ant", q, torch.zeros(2, 14)) - assert failed.tolist() == [True, True] # low torso / invalid quaternion - _, _, failed = reference_task_values( - "cartpole", torch.zeros(1, 2), torch.tensor([[float("nan"), 0.0]]) - ) - assert failed.item() - with pytest.raises(ValueError, match="dimensions"): - reference_task_values("ant", torch.zeros(2, 2), torch.zeros(2, 2)) diff --git a/tests/models/__init__.py b/tests/models/__init__.py deleted file mode 100644 index 051f680fe4..0000000000 --- a/tests/models/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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/tests/models/dynamics/__init__.py b/tests/models/dynamics/__init__.py deleted file mode 100644 index 051f680fe4..0000000000 --- a/tests/models/dynamics/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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/tests/models/dynamics/nerd/__init__.py b/tests/models/dynamics/nerd/__init__.py deleted file mode 100644 index 051f680fe4..0000000000 --- a/tests/models/dynamics/nerd/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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/tests/models/dynamics/nerd/_nerd_reference_check.py b/tests/models/dynamics/nerd/_nerd_reference_check.py deleted file mode 100644 index d425810097..0000000000 --- a/tests/models/dynamics/nerd/_nerd_reference_check.py +++ /dev/null @@ -1,348 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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 trusted reference qualification helper; never imported by production.""" - -from __future__ import annotations -import json -import hashlib -import platform -import subprocess -import sys -from pathlib import Path -import torch -from embodichain_models.dynamics.nerd import convert_legacy_checkpoint, load_model - - -def compare_codec( - robot: str, destination: str, *, legacy_warp: bool = False -) -> dict[str, float]: - """Compare the actual upstream solver's pure encoding/target/decode methods. - - Construct only its math metadata: no reference simulation/model is initialized. - Native qd is independently converted to the legacy world-origin twist at the - comparison boundary, matching the upstream solver's expected input contract. - """ - if legacy_warp: - from types import SimpleNamespace - import warp as wp - import warp.sim - from integrators.integrator_neural import NeuralIntegrator as NeuralSolver - - JointType = SimpleNamespace( - FREE=wp.sim.JOINT_FREE, - REVOLUTE=wp.sim.JOINT_REVOLUTE, - PRISMATIC=wp.sim.JOINT_PRISMATIC, - ) - else: - from newton import JointType - from solvers.solver_neural import NeuralSolver - from embodichain_models.dynamics.nerd.codec import NeRDStateCodec, quat_rotate - - is_ant = robot == "Ant" - codec = NeRDStateCodec( - robot.lower(), base_com=(0.1, -0.2, 0.3) if is_ant else (0.0, 0.0, 0.0) - ) - reference = NeuralSolver.__new__(NeuralSolver) - reference.torch_device = "cpu" - reference.dof_q_per_env, reference.dof_qd_per_env = codec.q_dim, codec.qd_dim - reference.state_dim = codec.q_dim + codec.qd_dim - reference.prediction_dim = reference.state_dim - reference.num_envs = 2 - reference.num_contacts_per_env = 25 if is_ant else (16 if legacy_warp else 0) - reference.joint_types = ( - [JointType.FREE] + [JointType.REVOLUTE] * 8 - if is_ant - else [JointType.PRISMATIC, JointType.REVOLUTE] - ) - reference.num_joints_per_env = len(reference.joint_types) - reference.joint_q_start = [0] + list(range(7, 15)) if is_ant else [0, 1] - reference.joint_q_end = [7] + list(range(8, 16)) if is_ant else [1, 2] - reference.is_continuous_dof = torch.zeros(reference.state_dim, dtype=torch.bool) - if not is_ant: - reference.is_continuous_dof[1] = True - reference.states_frame = "body" if is_ant else "world" - reference.anchor_frame_step = "every" - reference.states_embedding_type = "identical" - reference.orientation_prediction_parameterization = "quaternion" - reference.prediction_type = "relative" - result = { - "max_codec_feature_error": 0.0, - "max_codec_target_error": 0.0, - "max_codec_decode_error": 0.0, - } - fixtures = {} - for length in [1, 9, 10]: - q = torch.randn(2, length, codec.q_dim) * 0.2 - qd = torch.randn(2, length, codec.qd_dim) * 0.2 - next_q = q + torch.randn_like(q) * 0.03 - next_qd = qd + torch.randn_like(qd) * 0.03 - if is_ant: - q[..., 3:7] += torch.tensor([0.2, -0.1, 0.3, 1.0]) - q[..., 3:7] = q[..., 3:7] / q[..., 3:7].norm(dim=-1, keepdim=True) - next_q[..., 3:7] = q[..., 3:7] + torch.randn_like(q[..., 3:7]) * 0.03 - next_q[..., 3:7] = next_q[..., 3:7] / next_q[..., 3:7].norm( - dim=-1, keepdim=True - ) - effort = torch.randn(2, length, codec.qd_dim) - gravity = torch.tensor([0.0, 0.0, -1.0]).expand(2, length, 3) - contacts = None - if is_ant: - depth = torch.randn(2, length, 25) * 0.1 - contacts = { - "contact_depths": depth, - "contact_normals": torch.tensor([0.0, 0.0, 1.0]) - .expand(2, length, 25, 3) - .flatten(-2), - "contact_points_1": torch.cat( - (torch.randn(2, length, 25, 2), torch.zeros(2, length, 25, 1)), -1 - ).flatten(-2), - "contact_thicknesses_0": torch.full((2, length, 25), 0.05), - "contact_thicknesses_1": torch.zeros(2, length, 25), - } - contacts["contact_masks"] = ( - reference.get_contact_masks( - depth, - contacts["contact_thicknesses_0"] - + contacts["contact_thicknesses_1"], - ) - if legacy_warp - else reference.get_contact_masks( - depth, - contacts["contact_thicknesses_0"], - contacts["contact_thicknesses_1"], - ) - ) - - def legacy(q: torch.Tensor, qd: torch.Tensor) -> torch.Tensor: - if not is_ant: - return torch.cat((q, qd), -1) - # Definition v(x)=omega cross x + nu, evaluated at root COM. - com = q[..., :3] + quat_rotate(q[..., 3:7], q.new_tensor(codec.base_com)) - omega = qd[..., 3:6] - nu = qd[..., :3] - torch.linalg.cross(omega, com) - return torch.cat((q, omega, nu, qd[..., 6:]), -1) - - state, next_state = legacy(q, qd), legacy(next_q, next_qd) - root_q = ( - q[..., :7] - if is_ant - else torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]).expand(2, length, 7) - ) - raw = { - "root_body_q": root_q.clone(), - "states": state.clone(), - "next_states": next_state.clone(), - "joint_f": effort[..., -8:].clone() if is_ant else effort.clone(), - "gravity_dir": gravity.clone(), - } - if contacts: - raw.update({k: v.clone() for k, v in contacts.items()}) - elif legacy_warp: - # The original contact-free Cartpole still authors the two boxes' - # 16 inactive slots; its preprocessing expects their shape metadata. - raw.update( - { - "contact_points_1": torch.zeros(2, length, 48), - "contact_normals": torch.zeros(2, length, 48), - "contact_depths": torch.full((2, length, 16), 1000.0), - "contact_masks": torch.zeros(2, length, 16, dtype=torch.bool), - } - ) - processed = reference.process_neural_model_inputs(raw) - encoded = codec.encode(q, qd, effort, contacts, gravity) - for key, value in encoded.items(): - torch.testing.assert_close(value, processed[key], atol=1e-5, rtol=1e-4) - result["max_codec_feature_error"] = max( - result["max_codec_feature_error"], - (value - processed[key]).abs().max().item(), - ) - target = codec.targets(q, qd, next_q, next_qd) - expected = reference.convert_next_states_to_prediction( - processed["states"], processed["next_states"] - ) - torch.testing.assert_close(target, expected, atol=1e-5, rtol=1e-4) - result["max_codec_target_error"] = max( - result["max_codec_target_error"], (target - expected).abs().max().item() - ) - # Arbitrary predictions also exercise decode independently of target inverse. - prediction = torch.randn_like(target) * 0.03 - if is_ant: - prediction[..., 6] += 1.0 - actual_q, actual_qd = codec.decode(q, qd, prediction) - predicted_body = reference.convert_prediction_to_next_states( - processed["states"], prediction - ) - # The upstream method converts one last step; flatten B*T into batch and T=1. - world = reference.convert_states_back_to_world( - root_q.reshape(-1, 1, 7), predicted_body.reshape(-1, reference.state_dim) - ).reshape_as(state) - reference.wrap2PI(world) - actual = legacy(actual_q, actual_qd) - torch.testing.assert_close(actual, world, atol=1e-5, rtol=1e-4) - result["max_codec_decode_error"] = max( - result["max_codec_decode_error"], (actual - world).abs().max().item() - ) - fixtures[str(length)] = { - "q": q, - "qd": qd, - "next_q": next_q, - "next_qd": next_qd, - "effort": effort, - "contacts": contacts, - "expected_features": {k: processed[k] for k in encoded}, - "expected_target": expected, - "prediction": prediction, - "expected_decoded_legacy": world, - } - torch.save(fixtures, Path(destination) / "codec-fixtures.pt") - return result - - -def main() -> None: - root, robot, destination = sys.argv[1:] - reference_commit = subprocess.check_output( - ["git", "-C", root, "rev-parse", "HEAD"], text=True - ).strip() - if reference_commit not in { - "370e89d2b8ce68aa3578aa7812859c681b7502bc", - "476dc67fccc6a99be293bf611f897b1c06f13e30", - }: - raise ValueError("Use one of the explicitly pinned official references") - sys.path.insert(0, root) - torch.set_num_threads(2) - path = ( - Path(root) - / "pretrained_models" - / "NeRD_models" - / robot - / "model" - / "nn" - / "model.pt" - ) - convert_legacy_checkpoint( - path, - destination, - {"codec": {"robot": robot.lower()}, "dt": 1 / 60}, - trust_pickle=True, - ) - actual, metadata = load_model(destination) - official = torch.load(path, map_location="cpu", weights_only=False)[0].eval() - torch.manual_seed(7) - result = { - "robot": robot, - "checkpoint_sha256": metadata["legacy_checkpoint_sha256"], - "max_forward_error": 0.0, - "max_input_gradient_error": 0.0, - "max_parameter_gradient_error": 0.0, - "max_adam_parameter_error": 0.0, - } - - def reference_inputs(values: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - return {"joint_acts" if k == "joint_f" else k: v for k, v in values.items()} - - for length in [1, 9, 10]: - features = { - k: torch.randn(2, length, d, requires_grad=True) - for k, d in actual.cfg.input_dims.items() - } - reference = { - k: v.detach().clone().requires_grad_() for k, v in features.items() - } - predicted = actual(features) - expected = official(reference_inputs(reference)) - torch.testing.assert_close(predicted, expected, atol=1e-5, rtol=1e-4) - result["max_forward_error"] = max( - result["max_forward_error"], (predicted - expected).abs().max().item() - ) - predicted.square().mean().backward() - expected.square().mean().backward() - for k in features: - torch.testing.assert_close( - features[k].grad, reference[k].grad, atol=1e-5, rtol=1e-4 - ) - result["max_input_gradient_error"] = max( - result["max_input_gradient_error"], - (features[k].grad - reference[k].grad).abs().max().item(), - ) - for (name, p), (ref_name, rp) in zip( - actual.named_parameters(), official.named_parameters() - ): - assert name == ref_name - torch.testing.assert_close(p.grad, rp.grad, atol=1e-5, rtol=1e-4) - result["max_parameter_gradient_error"] = max( - result["max_parameter_gradient_error"], - (p.grad - rp.grad).abs().max().item(), - ) - actual.zero_grad() - official.zero_grad() - # Deterministic three-step normalized MSE/Adam equivalence with shared stats. - opt = torch.optim.Adam(actual.parameters(), lr=1e-3) - ref_opt = torch.optim.Adam(official.parameters(), lr=1e-3) - for _ in range(3): - features = {k: torch.randn(2, 10, d) for k, d in actual.cfg.input_dims.items()} - target = torch.randn(2, 10, actual.cfg.output_dim) - out = actual(features) - ref = official(reference_inputs({k: v.clone() for k, v in features.items()})) - # Same variance normalizer in both runs, isolating trainer update math. - loss = ((out - target).square() / (actual.output_rms.var + 1e-5)).mean() - ref_loss = ((ref - target).square() / (official.output_rms.var + 1e-5)).mean() - torch.testing.assert_close(loss, ref_loss, atol=1e-5, rtol=1e-4) - opt.zero_grad() - ref_opt.zero_grad() - loss.backward() - ref_loss.backward() - norm = torch.nn.utils.clip_grad_norm_(actual.parameters(), 1.0) - ref_norm = torch.nn.utils.clip_grad_norm_(official.parameters(), 1.0) - torch.testing.assert_close(norm, ref_norm, atol=1e-5, rtol=1e-4) - opt.step() - ref_opt.step() - for (name, p), (ref_name, rp) in zip( - actual.named_parameters(), official.named_parameters() - ): - assert name == ref_name - torch.testing.assert_close(p, rp, atol=1e-5, rtol=1e-4) - result["max_adam_parameter_error"] = max( - result["max_adam_parameter_error"], (p - rp).abs().max().item() - ) - result.update( - compare_codec( - robot, destination, legacy_warp=reference_commit.startswith("370e89d") - ) - ) - result["reference_commit"] = subprocess.check_output( - ["git", "-C", root, "rev-parse", "HEAD"], text=True - ).strip() - result["precision"] = "float32" - result["device"] = "cpu" - result["torch_version"] = str(torch.__version__) - result["python_version"] = platform.python_version() - result["history_lengths"] = [1, 9, 10] - result["adam_updates"] = 3 - result["atol"] = 1e-5 - result["rtol"] = 1e-4 - result["qualification_scope"] = ( - "network, gradient, short optimizer and codec parity; not full-budget training quality" - ) - result["codec_fixtures_sha256"] = hashlib.sha256( - (Path(destination) / "codec-fixtures.pt").read_bytes() - ).hexdigest() - (Path(destination) / "parity-result.json").write_text(json.dumps(result, indent=2)) - print(json.dumps(result)) - - -if __name__ == "__main__": - main() diff --git a/tests/models/dynamics/nerd/data/ant-official-frame.json b/tests/models/dynamics/nerd/data/ant-official-frame.json deleted file mode 100644 index 8737a5f7c8..0000000000 --- a/tests/models/dynamics/nerd/data/ant-official-frame.json +++ /dev/null @@ -1,64 +0,0 @@ -{ - "source_commit": "370e89d2b8ce68aa3578aa7812859c681b7502bc", - "robot": "ant", - "python": "3.8.20", - "torch": "2.2.2+cpu", - "warp": "1.5.1", - "numpy": "1.24.4", - "device": "cpu", - "dtype": "float32", - "count": 32, - "horizon": 100, - "checkpoint_sha256": "be7e29f8834ab22d76efd3d101196c77b3b415c3ef91af48aed233c01eb4e950", - "seed": 500, - "runtime_imports_embodichain": false, - "cpu_stream_adapter": "CPU stream request mapped to None; numerical source unchanged", - "gravity_transport": { - "direction": [ - -0.0, - -1.0, - -0.0 - ], - "raw_preserved": true, - "reason": "Official sampler leaves non-up-axis gravity components uninitialized" - }, - "files": { - "official-windows.npz": "2ff996b31e0622e828254bda6703a1bb608dd8e451c20b0b9c4ad0e5fb242cd3", - "test.raw.h5": "2d86248861ace090f1d709f2f65ee96a1ce7a2c7e7a1f19eca5f30929ac940da", - "official-rollout.npz": "b315b2cd360729ddfd175be91bed5d137406c50c60b6e8937e48be98422a1431" - }, - "model_inputs": [ - "states_embedding", - "contact_normals", - "contact_points_1", - "contact_depths", - "joint_acts", - "gravity_dir" - ], - "state_frame": "body", - "body_com": [ - 0.0, - 0.0, - 0.0 - ], - "scope": "Four real source-world windows from official Torch2.2.2/Warp1.5.1; tests only codec preprocessing and targets", - "indices": [ - [ - 30, - 29 - ], - [ - 13, - 0 - ], - [ - 12, - 47 - ], - [ - 21, - 85 - ] - ], - "fixture_sha256": "4d4c4a1db07b35ba06eb94b810621c84d36d2799347864e443adaafb8383c8a2" -} diff --git a/tests/models/dynamics/nerd/data/ant-official-frame.npz b/tests/models/dynamics/nerd/data/ant-official-frame.npz deleted file mode 100644 index 9c39fbd51a0a3f5ba9f8dc2193d74ff1d1d4c9fd..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 34898 zcmb5UV~n6d)9*dlvF({1+qP}nwr$(CZSUB&ZQIt_J;{A??(^pP@cxs&K6F*m)!kQD zeyMIbN#I||0002se;Yo)`^NS+Ed&6-?%xIifDB+yV`JwI2mt@TJN|kA{4@SjcH#%; zqC4L32)=y8sQPrM6=9%J;#+6W0UP^tT%6I7oTh*b@e@@;JvEM*vzvlAadC@vD=G8w zLkDPwT8of+*hwjCp80(gm%bzi> zAKz34Ue}FYndh*>m5JfktvN!i&?XIqs}98dK#Wz>HNmDG1bA%@VdH@#EmwaaosA~N ztaA^%6r~h~`laQ(KFs(@yrlta71clZwiIE_2xti^P`eM8So>{>aGz`lOKdixGI2fg?77b4v;)ZjC5b=qZwzOHKAQO@Wn?F^IxYaCGv~egu{va^IIwDj~6_$4{2x z8*lg-$OX!|3pyYJa1!aRMC*ql#qo zp3DPsV8KaJccSJl1s}|;WFh`|CecTlPU?RZY4tm(1ufI-1ZAy@%Cx)jF^lk7Heq^g z8H3c@jF;W_$7*S?@v}x4l;<4+1c8sIz3pTOj%l_X=CzPq+w2H%#TX?%K5&J(;{lDs zS@ktZAt{kW&vx^sbsqW_OeA*~V#>75mg@{pE@lSaaJL!tl?cCbxU=q!_zKpv-TRTM zLp5XDlA)Fd+Kf6_(Xoq|Fe}^r9#6s4APQ{iObWMb^&19u7*p>v$f}tl_QZn>m!R+q zcLmt2+yTgKR@Rp?#^^pzOc-DAKrd6L>~3TL)wq(tXbw%uJEcff`Yey7eR>EV76Ej3 zGiQ#%HQ9X_W=pHb#Xq@#zDv>bo5ZmfNyKq~0);BrZM}+~Y_$*Nzx^c4u?A(5mpG}g zP!4OJMeRiGA>F$qabvRtBX|l`dt^#>aJd>NzOhM+2C0a{w=}Ath6`LY1Hx(rJV!rt zoPx<5bgyrbOI#Dm(T)$;*VYE_T#E5@yq+a*yjrcJN>$oY3USIL5qZ1BU#+(R!V+mi ze>o>ntqC@$#F6YXD2T3X`E(^qBPF&`_stW?!6i9<(+@ddxlZA(uw!luze~mmdT}L% z&)l3VJZ4`hg&o2;c9~5SJWW)c!EYqUFrvlnMY`y`-DF@nZ+3~O{mzGyKah@kRf8A> zC8@VVKx=KzJvkTkLCoFkL9V%#o$qP;va_hlP=kIVe1Qyf&Nx>anJaDwp_xskG@gI= z398ffHGBqky)2#E++m(NdnRn5vsx70*~SVotxeX8SHQ!awjbUwJa7_=7PWl;%-?QG ze|zBN<#Rk@GJ#S(7<}<<<2}qE4d)>(tlZ}OFd!SN58<^yn7oLe#9Oh{P8#<-RInK-nr7<_Pg;IlhqI zO+*dtH9;y`-gg`Geg1rLAp5!%!4aNubf!7_DF;wJsj+tU^3ywC#=cXe{JDgQ-TLJYb z-ESIOXqN~#J?9anNJU!}eHyID-p3|u_Uoj?$xY}uMqYL~e)`(z-X;dnP9lgXuyis@ zEH>y*zu_r`PIS$|m)hl7U*>fzvHh^_=h9KV#b);(0wVYDF`;uS=;aU6u^4k< zlVxIT8xNfRWg0l94`*;roIv}&CDJ^m>qrk&$z8>E8e6JQ%Pe%*-_PmgMT?`zph#MM zjM+n*<7j@djtDdvCxHxD>@~9p``9{^Ma~|E?$d~Ds-+W>V@~`vKH&8G*Z48oCVEIc zh%L0AVBHcVy)D$evE^P=9N(xTp34O?T2Rm>o_iQS#J5uNv8)8H{Xh%w{T7qNn$=W* zD(wOCt4vvs_Q8HrZ?Df$tf-mY1KHC)Y{uK*6PjG|DmKuvz)x;5&akW9%n@c<>!KN6 zRrse3>{3*j{=MGQ@U3m&C@f$HgQ2dct5?x&Bu{aScA@nhzBx`Ptm4oKY-ujm)LX)W zXQh!YdN_jv23sqGmQ(=H*Di7gNf>vAd+5n$Q~r_kRCPK822H*lVp_P&;!zwpbty

^WYWwH>-iX+nILYD@)gec~4 zUHMKVbk9-vPa7_*;+jME&4$SKRB!!_y7EF-pC98q;jX~j_^BRD!+vTBN}A3I+~p{e zCym)wK;V}+qRZh`Jn zC5Xpb+1V%t#HJtJzZO?Ll!iSeN+iH zDGRp~e~EJw-mwWD-kXHiMj*vNRh{}L!a8+N;_`4q2xiBN9@6;27ixvKH(CpA2>;A1215;N{i@D7GD2K< zmReSI2Av8xx7Yme3hzT?)>>iFZTM6t&RNBvVCNw-(QJ+&!jL|ZTV+g+n!W0Dxv3ol+di4CJoZ3+y%8FqeUBn zp{-eH9~EZWTcB!ETrg32TBiDHgIbz-!yY*UTE1eAhnZ#slT@~BV|rtXXW?{Q#)r>2 zhluseN*mE+^=c~z^BI4OS+|WGq*Wo(cYG&)m3BpW?Hm_St`fRbE)hszMw6*MGWXAh zw@{HyWwqQDf8z46KcGJ4&Wp4s9(MA*fzxsxaX;123S{Zw2#|{ss^E!)ilCH!;Fjfo zm=eFtIFwv;jb~CZC2sWeG(xYX=__+!#|NHjrKcE3TvuZA+nxh&Hw?iJr1cE2n-3FPfI}!0kfvJW z!r=90fC{XMnU9GjpL%>ps7{RtF)u%$YP zZp{Sdn0pOUb$TK?xIky}s+ooKM%a=ash9XMW4N2>=yrE%K%u;JLUngi(%}DM5ZEuG z|l_hxr(`LMlu8qGo z0D0A1$G8`f#XDu1Zu>Hd(j9xtA176#z`*a|u3E&U{TZjlbuG3wu;ZR?PQT(^4o2%q zC4V7;4&Nbuk9x-7lK%x~hC}7&xuQL{Z_W<&Wrfu}$~#hXJznZ-6YrPDh*gEsM)#Ho zi5=fXsBL+pcag<`8pltxwlaA6MQO>xp(%^V@q~lTD3@rqwIx$5sv0=DPg$lhpjTs! zf{M)p(y?D7unFm*yld4R}Kk_e%Z>Do71ks*?i z9TK99rI+!+(gBCJmOzguhh< zi;})G;46SrCy)VzQ@M@w)tFqReM7fV_rsj6K{miwn%Vg*Oi+ze0or%PRxcwLx}2c! z>D~VsP>JVB;u0TknMaZWdqnnk1A>nY|CHu zJJ1FOPW&-|o((F~ZJ&?%1b<;Q_1WjXNPyYp{@suuyIhJsFK(4)OWe4RoGIN%uy?>K z&Ez$?C`t!)tvHW$Jh_+K$riKU6R7gmKdY{n`zwK&w_UWz_CN; zTQ9otFkX;{e9;DfnTBMtoH8@RSH({Oj9cxr8OU+xQ0{D=?tYXyZhG@Mr-n8Xuu-(z zh^|&80Wy)Un>J#^1;@15g#sO^I zsP26O_S(7BFt-wlk#&{`O8L8B+M9r^@~gexiY^^|hBMR3c+H7OiC?k8HgUfF<1;f& zq<${nUJPCR*T;Kt-Qt-_VTdSX(L0R+J6;i45~RWL?!Xe1Pf_Y;?EFSFw^NrmOci!I zd2KXV{#x#|-+i*>-JeKcL{o$&Y`MY=M%ROiKE|S#PksU_`jE%K#It-tz-UfXlm(2l z;J8{keSBZ!x$QK(Q=}nkk5&Oa5{Q@HHQJzO4OhhSQu1@evq5K&oJa)U!X88Q>5v_>2XUh zd1KU7&1X9zy!%;2{x)dG4K>}cOglBQ!QiI#d4&@oDI-7IZppcHG@u>i&2_9asp2?> z`c8qI+(yWNFZ)85dm(~kYRPiV*5j`g#}MJ&xx;JXhHT_z{q>JWOjAd?Ie3rUlR4e6 zO^BaENa3K|lIY4VKQLYdt_6 zvwKD;v+Qnu*Q)ZRCDtGEJBOfQHqi;&$~9bWg#nV>z;7{^6dosCNI!#OkbPTAa=8!O ze|$axMMfD5rvd;5V7_%PqV>i-OW;@pEPRTh{LHCq2yQ-8a>0+A@}@ zPy-KLxzE>4`vON4>AhcRUuplFZt^BmlIj0)%)c5Ektsa;I}cbMf?}xZgZ9#U`fIP7XZu781s|cWz>HW0mU6~K(?CU_ zB)X$K+B6{&W@M3soZHdW#7jfVWWV$)9cm=$;<3YQIYqyg!Jf!hJkNYkzXK z!yhF_lmg4oyE+R)wk1gV08%_MgoY;R`Sx)PqQL4JWcXa(ZdC$8i5Ifh!5B??`%jLO z4t61Dzf;D85i^}nZPr@P7|^#zQQYL_Y`^QCq}b+-F;UF`R*ZIpUxfhk?fId9$cAMG z1r%w>l%Z4B8X>M-tsmn#%H#v0BioY_J_hbWKka+fSh?~}T<&ETz7{v)G8_e<=J~Ib zo6@tjJmiinJp~le)WXfqsWskxDi^foW`_Qthr?1`&SIDp0RE zQ;56kin3=SKP`0@X;#)DMR|3Ipm2(`p&z=yhE&Kg1a0a~Qb~NB5l8W5VV=P*Sl}g3 z7CnW05R02tGuF-+E|l)lX9`_;=~Ay#ZH2x`=1l+H8i59|1DZvH##nEHVH@ zc?a&=JdnJWRl?dXsvJDvlb7>Y`t@lz?+-KE&+2PVn{BoN7_$Z!ou}jTFjEm_WC|aC zqrjzc4fl>rLC~2p^V2NX4r&%4G}!fHM7yFe)v93v<9Ws~4veS3N7WLMaFnFks0*n{ z>3iz|Tt?>|?zUq7N#d5pY#UQs#F_D2N(SgSqMkI4H2iHWYdGP;*M$nHf2K`TO+~$*242^Ct=UOQ-s5?9qFcOnv|JU8GoWZS)SV$*~L0-m@|<} z%FQRSbe6yVsAi_ut39z){|6ufRXp5t+E=c;7a$Q5$d1*Olh(M0Lrc=q_}4w*?fWy6 zK_#ij63_A5QlVRGIyLkm!u#!;;b=_+aodFGZ3hgGn?5rb42Q?i%6B0gy<&CLp9bhd z!kap^%+L*EH209Xo1xO4Z|E9G*#)>xeYS^V26{X!k!@wri{=!)u0ozAT8H3#FY-k8 zN7=B|4?xEo_kFtPBkoy-LeAKYUN+CL?oLj~*l#&N>tuYcl#?hJ+mh2X(A;HM5P?Q< ztcS~;9X`z`s^buk-N$@hZamoc1(7!9KQZgOMtf|7-&1~=z5va?O3Jcc5>ggX<=js} z%X$`ui#pH*ujAP}<9blxe7Xl0?v`Qy)`T!EVuJVY+R*G|SbxES8@sV-U`4(>Ddum$ zaM;m<5(r;OS<542>C%V#Guxc*kXxbVQ0Gq(NRq+A3{O)J zHuz|4A;Db1ZD7dD0ABjiM`m(r&C>BHW}SHZr#PqcpH2Tp#=-NG!@(o7Zj0vr)@HId z%US<^P9fn(UxoBFRNy%V{FqQrBDSliC|CRBTNc{^MHhU5^mQt`2y{x-gEv|-WbH`U z-YkhyMUzrkfa$9vrvm~zfblae9kJ~h)O#LhGiYl|yG&aOnwYsThgZ!G&jt^SS+kXo z!-_1<)wVw zQDr7ZRwKk(57sUXV_4n_ph|rWtfTTfg}Cbdu$CuR`Qm#lSxzA?Nu+4HviJPL)o7)P zf$;_tg4H|5Kr|OE zn`wQTW-!66BTv+Nw)vkQ%7SMB@4slKI z>inZ{CUjjMumNx!f&%yYN)kGu3%0uWtvhHB(V6&EOJG=7tGBzlpkA$@urXHB#HCnp zh=ht*b-E9Jr;sLl0~)HkX1GHz-}-Kq)`)NzZh4oNw0g#igQwr`-1b)%$#tGMb<^|- z5~jh6y5H<|=Rs-m?6QX(`0#Nap%u|IMKUpgPP(&L_ZNoBzj>sics?M*R}v)pJ8*96 zt=V974kT?4t3w0}ejBBKOJ^Mi^pW;t$6%VzpC=)#ceUmo5|{Z&h#)b#|1i))B}zDz zV$Zf|@Oc-2YkOw{JE*hCs7NE~C&P${i$UNXKU?xU68A+WPncuX?{U{~|O~ zm@aNxO|npLoiIeL<@7ZbPk~We{A>U_d%=I|Q=LmB48nICp+C;pQm|yXJ&y#-%_Ty| z&jicFIAn1$HkQE>PPGV88Df1_a3Pb_F$)`!8oZyNgprPe#+Pye&#g6_-#nzqv}=)M z>O5(gKxS@2q;nD`Fh_z=#=e&{nvjy^k{5a8a{I0S>s;C?enY@_l$}WcN%E>J?hEa! z{LHdmcpiMM65qrHHix)k`YHi#T}G|nc!BrAn*2l6bGCXjz)6!LX#z>ze-GP+@J`oq5DsfTz-*;Q?;1{jAiD%oe$571#!EiSQA%i$w|r>!PeGK zzjvq34PML(+=%SR_l(NxJ4hvXX0Xvk*l1Am-7*x>GOqkX^%e25eEnVK{)!cJbaZUk zXRdB12?_pvwyat!3C+EBYq~A1y^F_)m!_w9#NkhPNu4DJMlpG(tDFq%SKe5S&wWA& zvGs6(^YEhP{#edpw_hY&GyUz@4HNGZyzehvy$^N4Vz((>Of=oWHQAHmX4%@gb9CDF zqc3cy#aUw(W*4f5a8Bm+7krsnkg6?L2CNKQkmYk13Hr8(Fx5>qO9XJ)7k@DB8ml`> zS+ACSFxbnF2oCGmzSvaKz2r%#+t6=`Ti=|(b2r3#)nrI<7p*dl2}S`8{LO|CA0J_; z8MEWLhYMT&8wBr7B_k2IXP=3so)47u6qUTM8=BwPA_NxGdjcoL@3CIf{mGe%tIBHx zOb)3_$tOgOHQEkRi%I%c;h?f~KP%YK z*64-Q)WTHMqPZ9&+I}-#&c+{woAU2#=$=nhzZ@-D{fY|1$fl#?dh_h*MlpyhJQdwE zbuo{*s*E*%?)x{58C=-(xD~bKqDnP+^qY_O+*XGHgNo1;_ENyUG!TODKvX1bD{cxV zdT{-pfl!)8*+SA2kEn#y`f5;nhRUV|RJ0$G{-i?v4vV$S#%@s@i$wIgey5HeKhJ+b zww$v}=~g%rhdD1ESAXFWG&zbZ^Qy_7&ZBbpklzYxtma*&1lqymm zYrP>S#jz~SC-kGeOx4@ZOK)J;OPCznCysgQDw6R(Dq?L6iex6)ZN#-)G>-w; zXA@o{^BMN(xWhBUr0}b;#ASQYhk(j$!^u7V>Bm%fIJ=r+fPb{=t6@ZI+Myhhf3uu# zd4^^2mYWJ_46~Y1Cii^UjEP>u9PKad2|*ghWXSFySK{i1ByvpUTeARtz#gyLT~%Lw z8|P!!at?1{Umx2~tvMOL=(ZZ(j~%;QnV_3L2R1QO5gK`F;5Rj-h^nb7EA(WTEuz9S zal26+FK+xz*(bJqOl8U}T_c6dA~cd{Teu^_i7Es;8*%{&@9=QZBbx!4p0c!>{sr@i zbsbe~fZ0^#B^WiGZNv&4Q-2yrd-^EA9`a}VFj=;BetgdXo=X;?gF0hYNhylA=KL3T zB)Qf0r{3mgVOTN^K?gxlV_=JGpOYptERl(z@gnB=$u$3I*<+pNCqmYfh49RU#pyu@ zs?ua<{e}}F7If~w=JlpgB8mf7c@Llguugz(6)(zUJ6{zUY0t~@z*`Prr*{LuCX9@t zVq&q=+?FVEw2yK0+!WRPgQUf1X1SFQ7#a+)*y_2r(AlXh+SO&zio!Xc;hqHBYhFAE z8Yh2SC5)s5uYd_y%siA2>UheWlLvt>W(56}P<>&>XO<>k<_N2NXq;x;@6mj1%X%uC zKDMkt7_SRG_7h8!I$;_A@ge~jm*aI{f@zgvbS*LLl&ie^I;IS%XTIAmuQdGzp_1P) zzg02feBK=4s!rk8RM%j~^OOv=a$m9cedm6B5YpgxXmL`Favtjh|ZyYAQ+hQXo2u8hU;5j@$u-au1;8I?IN9A;isc=g4 z41B3%%{Yf5aF*zJWbj5{VfvH!z^t~Z9gJFrD;COhfv5 zy_j{lyOUMajs*K#u+dAlIvvS?4nma6ZuK&ngeU_;jJ1k@OL`P^t8oUGj3`&RT!Lzz zr7mCWOjU*z#{`;r%Qbs`At&c9lZPaP-xf7*^LK~Fgt*Q(-NPzqZB6xr)h1@(X8TbW z>eTu21sT`0o;RVE@_YE9ZM>)m<*UUpnEKYJ*QM%!bI5R>dz^ASK&sJmXsTl<3%zvm zLP@Ap4Jlqvi)|9kx(Nl!WOpI06=!o<)uQFDdikj^R~587VwwrWQ$j;(Ovo~?c^vX$6nv|Fw z4X;0}(rlEcJt`zGrHJmb9m*A!h8J1AGdb7_`Q1?<-Af9PyEET1{^<4RN>)AmFhaAI+x*m;ryPNRQm|PRki&zfhU?qe?VroS2_3UQq5o2%Mbb{xtJcKIZGj1* zok?`5^C*vO6;U?qqU@z}hZcS@y?Ysjk%h3hCv<$mH1LI&gZB=zis=Kxp&3w=7>>#< zimv?h4g6{7lBcPTntb>r9S3rUo(@9v{v*wm@pv{Ri6Hc>1|XOwxwx=&D%OB`>g^kZ ztF%yrx&l`a@6IS@6CJ;ujLT)^O^W6%@apP0ewxi%^U1A<(7!&BZ?eO>duIctCDY#@ z|M8Ul$i?E}Iji-Wx*H_M;WCWy?R+XoauvEwDYIywYq)fqaw(OD^`i zjmwGvhf(y~>qXw0;#*HS7Rh_FhBAXIszVy-M{UG*!zU`k1{GdI3FPh`wxCg{moxo= zjl3%g7|dt?2jssI2XZFJ2MYuM!0&JShd36t<~B~cCjYrI``7(X|6KklJN|=r$rXDf z+*v;TC@5ojZipXQm`4A^YukZ5YF1_<#32(oMk|HO!3T`Sa@O^X$ou3(7mt>+{*%g2=MSzMl{g4w6mAD#o#h zdL&^gJPHXLB?t#194Scj^JHtl0EzI}+rUZ9Yw(k$vDs)tO|t=Kl6vjz@+TPfhYMJj z=y^ecx$&3L~wq0HDYz{APq7bE{ zOS6ue`j^_uB);2_Y}}=g{C;xmzOPo^pXsnL(<;BOQ;I!Pm{~}FWNeY{VZ7#pIPakk$^$ z*^VZV&ZA(ZxWgz{l9k=vWJ0n^kLic!ZqfDi&{4Z7g}M|Hqy}H59C`&n(=b+aiNk5z==y?lWbQY~&m#Silq||y(`;Md8=qF#~Sta=i5XMs2 zz$Y3r{$?9wPg?1bm<)_=L+xRBJShe71MpqGc5%p4)o<8<1v}JyI`^VfoeUkc-*ZMy;t>4EP~+k**j*_+LbyN2I-4 zaSIER!U_3`rI%krPz9!5_=J3+7!ZE~Ogggj8c0EJz$oipw@k(v#(}Fv@zolQAjEW< zkxH^;eXK16n#ZIGJXM1mA+qhfs>#`CYK6Ne#F2Yi^QWDG;;z^MbB=N8?o-1P?8p4; zi>r<9+XC9xRs#HMab&n>Nhb{6#gCswyf)l0wJ&J(oJL42dh&sTBcbYAonq$Y4XreF zC9WYo$y8lq(EgMMD6Lfw%H^u*RJ@~RnD>xcw2I?u_;pN^SD8PS*wMw`{X7u-R;5R$~0OH@U3~g#Zg~EF)E1qQFwOp@L>fdiCFxMyqvGETf8P;_nfYslg8!m zt0_&~PO96lxk5ydD#dj^A@lmLzpMxfPhZTB<#tEPn?TBBmfxeW%K}4X#(Bgf~{w2G^G3C1o zSBOjrG+?i_h=xtYNZ4#|HJv;%nQ|*a&*7&V(w+5J**HKuC?)U>nJ504$D_T`GsJo( z36$s$2MSEgs#&{G2a~~O^;>smR^4Kcwys+TO>6Y{S@Ab^jsq@zXgs&uPHlFpRrmFk%#lfBvo~kV7w)*y#lSyAX{-=5IBHM z16H0S?1KBSp&>Jv`csC?RI0PVU8i07IB(gaWf8sgzK~=xeY{~jDc(zyv>_g%+ zZ5$9xzK>yWp&WoczcImVq+9^-ygSjO6$3Cjhh)Q9o=XvBy1ne;7DzZt)8lis*>E+} zRnZCYzQ={l*v=p*Qp)=LB9VG6b6oU>H@eMWsd4rX=X}dV*Xv=jxSAhON~K5I@i5$j zN`GyfsFu*(KdBz|L!@HtK5d%+-3P4OxnG=g-H!MfMU$CXn|I};>nOrgSzdjqBCBOM z9O!MFsd^h+^a=4jYbEg4JP!ac;tPDuE(GY-yaJ7h4|AJyTDymm30=V&^r{QK?*MAFQ@|r-k-g0kviu|0v2%;2+sB8fvXQgf zn#YiEDN6c;ILW_$7U*F%lt)hbcFv}ebux?nG=c7Ws|qpYT8MG73b$iS3#gJQzG_0*kV%X$>RT}ttOeb^4>t1}j+CW7fM z<~_}m@7-Wt)Sc0e224mR)O+*F!o>D?i`}yB0DpT>K~Zt214vXY^pbYt6{P2P4UGBv z$ObM?$^zO(R}T&!wtBe@wga!`*knWhd@b_#k=vvFk@z&-r>Y&$(=vs0JJv(a0Gr3k z`UK;JS?_*>blY>)usjlkh;+3sshq!gX2g2ghEPUHB^zqBr9;;ymM5ooaXSOy9q!4MLHxq(gPgjN^ zEdxEEWUe_pal+)J(0fA=_Jv0~Mmo}*{?lw#$I5rBBnXV;5}`$B9U!Du67Z(j*xwV| zC13@PMzGPYu1RI%x5f|c$bY+j@2U+g zda6Y4k%!~#r7p$oJ`{-C$GoIW^oA=&)aExQRMUWedeo;^8WzjU%lFu*=K4;*3)^Uj zF5*|_8LYs7n&YiP49IGm2#AlnFrbX3+f=!_9cOU-4qlxAr-dfM1(J8G5hTlc_zU{m zw~eip=W?2(gLquwt>luMol7L~pW9?gw}&3c+ou?MGm)Vay8Dk8>mYT54dx=#!(ia+ z)yK8JqYfDB2(cUuMVT#^6ba5cpZ`BV%%_Zw|+#*aWlbTZ_I%2rLjJ2A5 zj&AE6{ee_AKw6m}8)Z}|G2B4>Oo&IjbV5AtsRK?`Uk!*2J=&$d^Z88yI+A}DPHrMG zw*%fHva#fwk{-rWy$q6h96QZ;kAF8^WPLhL8r9jqBYw-nh#>XS3?gki){8`$wA(8p z$@JWUDZ=wwlWZquH9(4Ed`MzGy=oS-On*O{TLP&YwqX^ygW>M#uOye_0MM${2m*ze z-RCFAe!5~3ShCd<65qP4mVcF*IWg)S9J9)_e$03;xl>u^-Cdp>-_fe_pg6}Nqr_bW zAh`vzyV%Y#J8Dz}(NSW0xaBx82iiGpa(Exwyc%@1$Z%hg2(!wPL#BVILi+gWmIv55 zs|;b&xDXhxF4{|tD>{2QUTt*flHkzfeFPKrx%y;b#Wa^3#kJjqun_~pA!2i(HgWT^ z9I!I2;vX3Rtv6Z&Ivcm+XR?8UL-;eH9{Fryi9`7ERv6MlZPTxVYjSAY!fbHiSdxN5>% zX?WPC?BOHEMQY=t1cKVU;#Vqo-YX@}oiTq*i_ykZf@rNKhd|gGSvf57V+}^ADsA4| za@065sqExz0h4Q~1mlxe1C&h+YIpk#Zi-he(&#Zx&m#HgXN9vq8aiZ^e!4MbSy1MAlUE+G!fDo*^@zYx#p1UO_o(OP36hbCS{(H`$V{#A!Rr*&Nm+!8 zl|CRZ;$x_Y)51BJ!}a_j@x9oCO(m!7_ky>{CtY;UQwS?3rz!}OGOKuDWmGYeBzY%7>dZ0(WD@P`&a;d^^&X} z^*OE0&4LG=O7Acs&Ackpp*1<0iKcuzT7Y=#Dvd#Qwc>C5rD12{MvvL!?_TQWqPg}7c`Io%(XaN9_{x43mv30Q4xB5Sz z_RsuJ|6Kp+YLp(bpAT7huHG4?=L1ra2p&PGu2E9VP#*=9u}?rM4T1%K7q>ybFc_+? zuDzAxMtX>xb_VK!E^wJgC*WG>MY=?#$y^X`v&%y8YUgKet04Es%{e>wg`zQmM5QIO zIa%Jv+?=vRXr(9OfK`u1! zigwe^mv`piMXd&{@-JFwu;K)b(JUhsmwKvhNZb^@B7T%wBZKL0n|{+n6qADeV2^j7 ztXr_vEx9eWE@^+=J5`TAusF_)Qnbjb%gi-OxQ>4(qYwr{Y7D7J1&RN$aqU=oQ`DkKXk2~NOFms%h&XC#siWJ#9Gu!T^c z2{e=?(-L8uI@QamuX~l+-A6PtzEN?<$CGx^1AY~?%YuuI{EU|=Uf_k`Ey_*#gM^a zoDhFBd6yXF!}9*~iDiN1?9rWdnY&ZOB2!M2WxfDSKJp^j6w8RiTm$?GKSEa`In5g2 zyYoBZOj%K*xAh8A-(z*1C++9}XVOU~*8rB+aKSaqJkbR4Y zCv_b$8{$BHhQlZmh_3$o(NBh-dyE&?f@ssmZ$rEEWeoh4;bvK`fu<)F@>fUCn7Qs@ z7eCudaV}qBrhuJ}v||_EL_gje?x){lXTj<(FF3{=mstkWS~@1qJWsoEN5Vpa6=4diHerzFL)j<{oPw#con%Q{5%&94duzgGp$j|j>@FD1^ zpw^}z>nfw;*n1+}6nIvEuuKJ3>bmSfjIO8tNTlzR=%H$E8np{m!Z1t~H>JWYRd0U! zd>^hU+_o|!nfBMa+|H~nY^L)>@7hWiwLL^@_TESo;eJMW=%>7O5I%DFsEqEJ>R9FB zT8N7K;XfR*ltgA=;WK{F+#xTo0A*$oJ13XR1EvFZl!qQJlh{DTt}Wrmn_K^=j0#cC zah#lMwn3}z$8L~;=+679kky!h{haWc+0$rYN@)Q#aJnFyVO(M zrkmar;Ojv6v|F(N)fl6xpvp`F$^JQ6z`i&Td}yKve*f-&Ch6N1AXU#kh3Vday87sB zma?<#FlrWF7V5q{)yq{H)$YKW#CM#|@O6o3Px<16Sgr0X1nHAOZ~u+C1?nRWrG0V+ zjMaCOF#q(aJ`dELP51;~z6sSh`K+LtXahICXR5W0WAv^r!$?O&Vr{ZM`@(EA{5D%v zo@wD%&8`R@vAB^AU9d+#A3mp zE%8h^>(MW2$x3cDp(y3ZPxv$(@_qV_n|k_Lo9Ji*DJ)P9&Dkt zr6*)EPfTNB>+kGJT7~GUX*}@`O*^T|z?yq9{MkQ$IX~_D{bDfwGTq{nT*K&&Sr(uF zyLVk-Yy;FAy+4JCKC0?%?SP*=C>a-qb3osL^os*g!d-L|G^ymZ@*lU#H>?-q%^5FcseiWAmnk-;Ar=XGsO7GBlld5(5 zCl2|vRdo`>QlH$cF|JD62c7N2e68`$#kFHV5ears3qvO;2*hWJg!+Wu5FfW<3 z#O_1lq6xWRJ|)P`GthsNm{v2d5)@ToOL2Cup!Wb<`e16dYK_24==*`N5YdwimJhW^ZNN=0zNw%a~i$-AB)e3Nd$Fis(iFlBg-GOf8#ZY%)^EV^$Z7G?Z6xelj zp!2ckEXUGx^X?9^GQ-ysD<-yYiU*TLP{r7~3dA zbcYyii!YFsw1@qoqgQ-b?%mI0+}VJ$doT9tSJoK_7D2ZWka4<@?fS$9b~)Kd^8LqS z5?l98iMV5}Md8l}&pq+ZQXfY)NYnJq^QrG+(F&PMzTbzfil0@AJJMOx+&DJx0-|^9 znWeA=cJ3aQ^}4P3X&i03nmS2O(&sW%b1#eL1II!;8gArkdiCG*Q;Xx^_g3)DSQ4!z z_`C5NZ(VY$GPVF>Czbw7_FgaMCe!km?n`hG{#G^^Q8w6;y9((u?wDm{F`$;QIKw+6~4kma2?0-aIiZcPygncNZ^knl#gy#yTy;cS3o*F@q=yy|&X^BH3{C z^ipH?-~THYGkxSymS&*OS4)!X`dPYeI>C~O^<0knb$^&t3U3B`@XnN*vY@76Wz6x% zSZ0)gx{@*Nd6Rng6|UDQsaiGJALJ+HaF}YB`WkcgC}Bu~=JwXYATpEypK~ThA0; zV~^9e89GXR#`>iT*xgr*(oLxnR%peOCO7CtTX|IjTjrQtj4VNJE7_aZ{uF#esOi$0 z_ktn<9}@a-okC5K7eDulcEl1*?7(iL?>OV-8Zwjj1ew@VPVaVB(p-$eyRQ?))5&sk zdB7d&I5!riLhtGu?qZy0)G$@6ecTwVLB_6*OOH{=$e(h*7<({qw)#_li(aOG5E z5t-S1NG>QjI(`}7_+NDrXJ1Kx@r1u(4B-!|8zuuy(!TG+TD=4H80Abb48@w!0@m_| zcjCeOeth|3@7RVFFp`dX|9UN%2@`ca=2Sabrqk#V-otqPC|(?pSb3qdd@w&hAIy{) z)p`rRF1(9h|9h^Ip|r8y>pA|OwF%{Oe2o{vx~-|ankq%!$X5(Q(sm!*AW~pQAq$!X z;8bk?sH1_U0$=ArPv^ZLxQTQfUk&Bam#yazx~~#3oKIojD%18&ez4CVB9sy|bPKA&^=7*SwS%RXL!(rw zp7adzX`>G7zmhyc&ctg)`H#)HMm*3{tEs)l4OK9NwB6rM`ADkmQ6`$+EWn8gtxd$0b?Y`BvT!v z*Wh#A>O(AaKU?gEisTx-UrWK<7G|7rxH2t+^pLucEv&O->{e3DceQk>tGmUqQZ}td znDbWvFeryC3TBqG(EWLlR~_Lwn?k>UIU4uIW4el<5JG`X(_U)YW`rSY8 z@7)XrFj#X^wa51rwbzV&mn$2~#|y?p!9HG}KXI9oMu=J<2f7wucP znp6z5z)?qu?^W)P6j-{n}p7hbn{TBNgq#$?pPAvfVRaz9PW;yT14c4Hm%|(wAIxeo-&ameROZ za}Xg?znXn0j!t2D+|pB3`=sUW=#ehA>K3{I#?M?@pd+70OX>*-2p7Oj=a`@;mU3$# zv7da8bdPQq&>*B+OrQ7pfghnuVi@s1F9Ll$>TpJiGVgG@UX0Ir8yA3gCzRvPabkR- zudqdTW|@=$N*{KtVdS|fWu<4fk~4eutK&6J$_PBl`j`Y-)(&Z7w63xRN}n-l(Lkh? zbM}AgoBVo^8@3cUQEe7^Owd+Bsr)sk|Mubu1mJA+^-X|0W9vgh_(5$DjR(Lse^+~V zNH+(wP-LUB)#b$Ej1kzFm%5n1#_TI-e{_y?aToI)@i%nrk;Zc_Q2o`X zf%K_Q`M}2av^{<7vE#?UH+|p+(iO=^0teA_*|-a!O+i)U13mAhKwgn@Vrf3a9juS< zo9-^sbeplKdmjsfi*)dFAinJ=cSt&=APpei$DIxYbbiul+;{pjMKq0j3~3S#M`sDn zgVUWju4%_FRB-<1&%&IC{b%q*4Gv;(yKB6nh-qvi_@*8`!g8$vsSI?wBR)r{IDd6y(wJ@k68udn;Gjk4chCRSF3E(7}-z`ld% z4!y)uNl)7?07F#lz^&Arb?%t0u;Lkd$dwY;a#-a%Zn!l88+&9FV42lN6VOu(xHDFD z^6M>6j-PpoOfy@^V%WV8-fHj~2p|tXkWkDW21IV`Z+%evh92k34Qz2`mrOU{u+{h4 zUKJlUerwt9t8`KFXill+z-Rn6y=nUVAfwVdlOwQ+Sr}CjIPoTT`>M3031^j|a+r4E zkpqjJ+pAq{J*MQ>)*Qxrd3G!1H-QD}QOwX#;8$@;{$4#UvgKP#CObRn0BBi3ij~X7 zdR?$)ei(Y823+Em5RjAa^;Htv9y94LHaU%)b&2x4+SS5f!Ntm}g$6sv&1Q-C!phGl=a_`{*F} z+_3JjG$^~~tHPh{VGDGd@?ZNVT zDH`BV>H&E@UkWoD8}I^L<`lK{2R8JvO@1S2zC`Y)-W{A@xMRVNYr_)cNmEq?7Bciz zL#RKE$VJx8BOrhX_n-77c2;5JL1g4AD1%%FF6&*h3&M&J+$A>36!xU#(HUW}B6;ua zHUw-^6U*wX!T>IvP-AEQV1$iedqVub(C76hI55;(b_-a5x z%56(4uogH+w7sNwq&uh~Sj~OaXen-pn4Z}0=mA17^gTbtrwtg}cZqKD67CJP)JMF1 z6CP`%seyePG3`U}DuI0CF}zEjHfRW6y)l^8%v+&eDaVgiY}nHf8(_{Y(`YTnf+oVD zPWm+JaIo%cimNY57X--dz`POJ>*V3AdbA6x5g>l;qA%(lQh6NA#FU}O!J;E80;5ar zSGp<~k*rspot_IFkZv6;`-`b+ppO>t0j4)E-UFVm)g{?bUXgmgwj{!|Pq3&G=0B;!l7Ka|FL1EHc+xk=jGjhOs-)KcxO%n`YzG5PW3M6x$kiel=jJb9}Jjy$Ki-3vfr{9~;X_Y=VOz6UeairPwJ|I;eTFXJ+Jcf!-1G-Dapf5F9T*`(_+RI z48y}d4*QzveE%%@(?utPcW#cA8jO!DnZ2yGCiHs_#lY@TGUEG~{Vo2LA20G_ioqRg+-eXCz5@IGf`uEDD?D>VLTdNUL@lwbx+=VrL0-+JFlP~8wG{zVioBHhyx^PyQ zcSpARC*F?4jx5<8%l6dhX-Wa%RnHU7Bd>&i&4AS$ig399%4q$rMPvSvF?9NBZe(d= z;^=7d-&xRqm8*Y;|6c#Ug>=#4e^k=uoZ|*c-&8aut<%_Ic`pMZr5s^A26blR# z+!Z2gJqnHnWHbXKG7dNmv@XF}hueL^mLLCyYqQQ;Sc` zgnVnlt-6mN92NnWu}7!Lrj0sKc11`Gtx>nnk^J|gZ#^f5Mgiib5u$#SJz76X@%fv8Z0V%Z||(!l_Ls?DvUW6O;6Igdg*tb zgfWc+hl5l)T!U>ewy)Oi66DnmK#(6+#cgKn5)4lh)$d9AT{Op^_Hb>B8?@W!%$RP& zpNG9PW;(u6OU{Ly=afJ;eRBCIX9@WXQX>qVSDf_R<35~5-RCF zaSCHl2HIEk#QJghtH{wGh=x~T3QJ*m7=%F(Jbny&{Q8M$h2}F@gA=TAyu!1UrX(D0 z)sIAk4rA$Wn^A$?t76xxwFQGyI_7s+Ld++IYvDhqWK#LiSJJL1m54dC!F{NagxJO( ztbDf7ez28fEX=6XGgg2>NznvGMWieD-EquM;tMoWWSMqYMlY{|%h@m8r|Y-d5EJs5 z(DWdCBiG!ehg@9lxMBKaeb70|hP?hNU!Mj`4#}#%tM%g3iw8NBHO}b58Q1B5CaT z*+b&5Ry|}m=^2n@xi@>`A}aTTnGsP+YvYUqMW~f)gbf_gSvBZ%d3RF(s`-%=j{BmW zg-_N#4+?zYf{rjf`1j#-;n`s1l%}PQq)@DIBjCVc1qS5@KzHdG746g}Cz?>6(!j;erL)jAT@SP*Q1B$PJ-*nR|;QM!D9D=3L6Pm8D z>Si7?Xo20<<)UyOkJI(^^s8jb5vv+4EsE9*CR9z&vhR@SMqax1fq1-mkyF2%KT@ep z`>)HL=4*}P7KBoxoS=uj7Q_`(+SAcgf^J$ftlFQG;bUDBX1WM&$_f^MeadIzO>rOJ zB4uf2`FKEWcPSgEu0u`xH65rc2zBi$KIvGBl=e_li`9g(HNNijRq>JwrQ_5-8=YQF zWnEBrN@Pl>ulfykpg-o726DGslHOrR7quXGZ3~b|+yZd|e1msU_T;p7a4@N)Vo9lZ?3CGj9Jg%LV>S%2gK!NEVPG~7~V zV{%04-8gix%iXRYj@NLTRW}Y>*z@&?>%nK?Lg-`CD)MxamE3wMy?uw%JRq|ypKP5`xXKq;l{;0r!=$WQo@ z&CKN+bUKKGjT`xdG=%b98Fc5FaC6!?)+Tej#S5W89GHQ1(e+^QCfosweN<(rxxqYo zF>RY>es0_BFU+(JfhMbx+f}D@w?^i?_S$0NN42A8edQHxm8KNy*OcTf1K%RlCSoR3 zWK|Qip|#R*X?!*~nytAFR-MP=9waqB8zAOj_D2{K))%{?Uu2sGibcuDv-JqypQ7Tc zzs!cQ%Si+PwuoPwsz7g5IR&j29}JD19@eY ze6ZyX|Javqn_>?2KC`taJNvK;fGe*Sw#ONyrg5HNTrl8_H|l%scfMZR?PS+Y8oOgg z3>?Q5t3~9Rig^9lY7x`U)2$P(ze`DL?qn*K>31ge>Ehj0;6WdKklrZ1iu<|LDVbmF z3-0NChtBPl55|-Hg`QVA2$OHV;lu_dMe_>vGul_7{N=H4tBKrGeKfB?A9?By|C)~F67gP`=Rg*AmE)Cs!n z)&eAl%0iAuBp1!gt*A=8eBhf0W3z6qYqPH?2=@;DcL>2)y9=n_mQolY$DLGNCU8Zy zH4#NCfo`{`FB4OMmGux)^qL2I-X`24nDJmsn&P^#hmUo#18)xp9%iU06Fjxn%s%M3 zY&;-!rlMS2hXOc~^O^)eVl8ZTeuRT0omcCFwxrbe$954Nj%||aH$jl)CbYIolFzU~ z8y~d+%@fRLXT)CZP|1luX~ZXOAU)0wY@3w!9DsZ{4r09KP+m&CTc-@tI0&}jn`~8E zgW4}STwN7V};bmizF$6Y7&`H^@ZM zS-%Iqd?60fW2e8!J7_gZDKmlWQIMw>D>Vt~5*fNNf6$ZGhEz5Lbz0KXBMM1oe{q#i z-7A>Qn%>2OGS}m#z4WOa-zhAeVys@?G1-8ib`B`!HdgFT?hKmX(MnIQ%s|l zxF5=L3M1EP$aASA3{~$yp$4=#&6FXkI(k{m2VfNxX%pxGi%YU~BrCH~7P8kMJ++0t}D$p$P?H!f=!KgJo=${t`dc9H~hL4FGI9 zh3DwPEL5lQa@WisJ}2iL&DgW1uiDE*9kmIB<)%wraR&iS$a1Bg*VH6XlCrdLCZi!4BWilC=*o5O*t?miEMb6q9zg0od1Xi8G3g%5}3I6SwSmMsI8FhABmbV^$OrC+T8Q!`SkX1y%aXwrr^x z$6H{Bl-Y(<(?zN;yoJlfb9QqS^_3Fw%~5wjo7T-n?}U5>^xtr#Cba8@YY_8pi9yew z!J-(~!!aJq|eq_ zjyZS3GzNtegLn|H!Ji2;EGtbSjE8pjF?g@lMUoQM-_Z#rM^bD=t-z(^hA<2#(YVE(TIAzC(rwbDuo*tSLVn_F?8pdlj|O$%j=e6LtNc>YRF7juxj zxLcZ^1CgGm4U>{A$nA7rlI1&U_7cBDqpmLaQ6^J6J%y5W5VCt3*o?&npwlLE7?_4w z$|Jg`ohOy;9XS(bbi3M|_{j)|YqnB9+Q(t|_EyzeJ_~{Ro>Y)H*pKJjVHX9c&0p;i z{~&9=>0YQBrv&{7 z+yLnR#$si$i`Zrsm|$Sp0_S$6pCzEVl=tEidU zJSSB-Kl3w2CJT(o=-)hnvpbIh%P+|jag?xZOO8S%G%NR-Au<$#BJ${a@JyIb*ui%< z%f_3QfCEgUP6y^K=8mqcm!cdT1^ofA47M&1xW_d2Cy+8OZo23lJ?%+B&r|}l9mni&P7xY(KnTlh?-zj20G^4Xkv27zOEIeiRbM!gvAtG|s z5;s+U2J9V^G^?3GoEEu7Q_ATHNRNZd(-URS&tuom6}jfK(}fMNMu)oJc+}E(HtTDf z1fLE2t=7kK8mrARw|t-w@mL5*uKXie)r#q`Ff-!IYzZ7#gE}Hzxx&)7s$vDUadcHMZXkvAECc>`mN^4P-L3<%NPaG&O-`Ev$0n6qvo>-Ud zG!_fc4wZFTgs$?TArqLx@6)JB2)~2>X%NZTY? z#ToHHd%&O-i$U4!-2{Vm`<=n%!@$-)tEo5L(6iNq1lBEWTA@1vkU=}L(n9J^)(&%L zqc^gq9z4-!&KOjja$~Z%!5gb8u4C@iAudbu2jaG=AxNb1>V&qOObdGxl^$M}g1RmsxEB(K?d3bJ5UyRYLTSw>p4204A>Sw=X#c0&SD38$R; zIvh!p=FwZT@I{Y*JLNY}M)#ln(!u}M{owt1M*f#D{r@3&`i=R$&eQb7mF&@m52YfK z)8Y_sBp=SJu{ce>?5$6OClI4~iqEK5ZdX%N>tv>DRY98=e|XVnf;!J8=V->B7Hj{=x`ogtV&87hA1I<@ThP?F3 zya|gZOw-}nW+_O8$2z%xM)1|ce)Sv*pW*K%h%~0E`F8)U*r4(_cBkVXMN(>Ywe+e}yzh$@Sy5o9`ch+^(J3 z!#jtfC%7&%q!Io}KHq{S?ZdG(K5;McUh4 zncM}}$r~T&G5&>e-x0fz4`Z3so*uU?t_lR{sPI4~a)Ste?Pczm?$^7uYz)i~#?S$! zFuR+Z7t$upSM?m%%tQyrbmq4&Etu-N2c;}MNCYUXrF;v#Ij>FCVbx7*y86+X-xgm5 z-Ncvk-j$D~uRXH`b~X;~Ejr2DU^5X%CvFk!2LX!Grj7*BOQf#o4llNU&5xY~rXR8X z>d*Ln=lmtR^shVTzod2lmsIRG=pVOEhP!-{7|sw65*j*W8ypT+n+n)I5_ZT#qR<`= z7PbwtFt7kLQSHg6;jjmP;-XSNcE5-S?E<8vd|X5{OyTelVKp+~-SR7bFk#>ZtA%%l z?DC%{>zWbYcb;rFo)=c-_%@e#nDl%h*E`+3dg_c=V}2G8qTsZUtLFAE9)_eJa z+No~R+h!mU-goY(H8}RDommP0`WV~>XEWkB6gd@H_V`L+ZZ~qvtBBu`HDcUhPx`#w z+nrIWMM6O}SHg95)$_Br=~}zIZL`NP=mGL^r zfEAThz1p-IKB3FrBDjgOxUhL@SLfh?>PZ+a*r6+V=qJ~Z$fP-2dwk-8 zZv`$UIs?V80~uqDqKV3cv=Q&A)>eAU3v@TNjF2{NdAJR4UWX`&LXPE2-1_s~GXJ?) zc_rF~_v8_Fy3t~JEEwJVLxW}xD{xt*zHgnz6(WxSa(q7fYIWe3ilNizsK+dZ6kq!L z$Vp@6Z#WC3qAcYK(H6?zo)rgk6$Z1L$I+b5O5j&7pI)Y^dp9=`@y7PbxUsld`qpv$ zx3(&T2gj?L`d}&}aDF<6D0@9=@O~W$mv@W9?z^<%9AWXO0rg}=3#F-`bj238x(9ue zQ7Ywb?i|%3%Tr6mB!apk|b0H^5uuEsMEE zhX}rJRL?kY0xdS6mvVbyWNztTH|ISienH=C6aenfnc+tTcY09yDXYisQx0X_#7O*z zkm{(7y()oZEx~pnFYXb`Ztz*V*v|PWUdI;~HZ9rrs*@F1mY)E{dyZb%B_IL`4Te*A| zo|i~Tqn#zJfSaH721hfg+fO@C7G^Dni)N9o7;f|Dz=Vb_Qb8xMDf#R&>)W^_`8}M4O1)Ud zPp3n=Z4iQ1H<^JhSw|?JK68VnRRjbzWha#2M)vRAW_s*nB-3WFbb#4YjqbQ*(gKfW zIg-)3m5jtRyiL0>oXd1)`3uElVS`fygF8tz;8S@pIIMK3MQ#82&rSjSm_siDMfH{mSVM zp#r>P8um^^Xel>9855+0M_RKJ6)QBjIG!t*1mllqOFY2$8k8si7xFr`cc?S-x8s=z z6buz)FO{k+_F{1>XO;V?dERk#-BRZoNNH+z22u?bPI7uLbkF5HAO9oY^tn*k_cNGK zio;J5h)p#??5r?EBpj%}sr#|%vYsc9Qf&%vdI~uc(zllU8&%^plGiKCF>EJlBU(K? z#2xY!T9iKR=2=3?1So#*52(wPm$`9}S0lfKlu1E)@MsNG6-7uG5-(RDB~*LY_cN2@ z6k!=I=38X6o!bMPN?)%Osd?#nA4Fs?=)L#j;91~TOB)y8Ebwky~rB1K8SmjCSb5eWvKdmF9-E%gd6 z-kYd2Gqylk?2UCa3C*1wKnb_10@<|NNk0t*WvN5MkWW z$2qn2p+3GVUZ2xB3NAu(6XO2mQvq#NY@^5EH^f@uNzmeyJtU*6SEaODB(AhswN-~kjb}K*`OS*HUW_kHtHsh5YkgtZcH?{FdpWXQ7!E5pj+o8Sq{cn)rxDM4R2Mwt&&YPS14cyC_%Ck=#cgQ<7p zN*&Hth$>+)5)L{ETb4O*z^y8e7m3vKC6&aqs@Qp2Q3Lkt?5mu{+ll_*3VDAr6FRYK zBp0>M3ABS}L%ea1Gt^-Y*(fU46u=s$~; zxc}++@}K)^e?$K97A(b6T1y3Q7#pR4T1wy0zq-G@4t?ZKgj2}wusyA9q8CayOeK>) ze^6jqKbRgxU7dL>M;`niX?<-h z)_CvW-Td+G=EeK+E$8OgeQnN>oW)UZMe7t~4CD)%rjA6lZak^l*@V&;thP@z;(Fyv zcoScKSUTT65y%*XPd#)R_Y|ui&echZzVEfjz@6o%7 zksDJKdAGDIn|rJS!?F(Jqeh(wrO!;_%C_mq3Q%}J{N~wfd4%`Q@EPJZh+L4Kgr<<- zt}LQja3!r4nSpU5)t~nw9Bm@(?U?z*k)aAAcmMLuf3t)SJSuWp3Y^Jh{ShelTak?T zYJm+g&(joyeo&z;icjjO<>2}+T^bfxW2Oqwp)_2wW3J}Q3Xv4*b?)-fL zo)csn!f1-7+WE}uSpS)fhf*r5w4Wv*Z^o(Pmu<`zzNCSoCSCX|!at1Txups2=>+&BCF~r_@{hgU(X4Le~5+UzwS{V@j+t7bdw3uCPahmhjY)G%qByWL~&<#zH(7cmvpNj5g|P2IxCyqc@e-4ONP;(-4(VEJDImj5+i z`CkK;|21Iw|1)6W>byJP_?=n)^@;W83+(^k#PU1-_sSkQmr}CF8oH8*GbNW2m!!yc zHCxC~N>zBt#swdxFYlU}skLe&(UK}z&9dQ$RhK+9P*g3YtR}UCCdNTNKXVl$Hdd6U zCC5lsq(VvH#^B{={s^vC%ROCnFWjr`rj|WO$~KDb?D{cuVWzHFWTFD6Wzxi96_No{ zF*#Cm2O84g!n=4*E<{84B3TtN`lB>03;)Yd|LYRU-0rZs!*<%TqD`OecM?uEH0UzW zh0-qjR2S}=}1kloZ&edrv2N_jhG*PA*MCgPG(v>_tnBenHQc% z;^C+=*e3aBQ?T=Q3*pl7Qg|1RN+PBt2hm_o+Y}yS_Z{{w#DqiVxLdMWuxh&55X0T1 zVVetnDMnR*G`+11f<>dcRmhWdWSg0@-VtlUEas1G5WaJ_^CRapTe*N^Xxy01vyg6U zt-i8~{vv0eEj>m*1FGhCoS?QrJoRQ(BS$y1PCKYQI|wNf7@L}TaC2+~V~o^M<6V=- zw0&XUYGDg89st*hA7s<$S`}R~Qd28^jbSpL8k#@vj_3y>DnrjGSX()?$v5Y|fS=0C z-|q7GsR`gF3^;j;Jr{DjJPn8d5Q%n0Dt_=K2|9~WX>D;p8~+&P2B(eKi=ZEq~ z>?FmDji{Bxna`T;b)D_UorFf5E(H=ej8DinZ=Nu_Ph0NZ_uY9;-#T+28k&P;aWGj2 z!Ki_$vI`EWv$FRJRXO^w{j5@hI z+si^aI?A}cwAYj?ls|HfINnb=kMoC#kGE}6$^1I5OKuvruQ0{LQDsCvNo`N>DL_sA z1k|xamR)}R-f?42FJEz3<(o0~tU(ndsuFjW1?lz8!W9}BHDttAdG?jUw&mKO2XAP4 zl01daD;hBE-?xw06Iw@B7rg&&$3Q}!JFb0hSSUAYjsj@hq1fbwYZ{M{fL2pN-nha) zG3x;KoK|GKGQ|fk~B}v%jR&M@fg--diO3>F?Qq8WEre%Fh2I^s1phIj! z0RG}|Z)8K3ee)vJi}aMbKCFJwEK%FeL!&z3ZAQ@Bd7gJvzJ=e=pJICrTaIYsBi++( zY2+pjewL6MScS%g*)#uF_u_#Gk+cryyD|HlsI6eT08htMpp_AA+@6JTQWuK z^c`A(@{;-I!)WhIMC0cOrLR^!(%Q@}Y&bfbI&~RM0tUdG+@IcwTlQkIlkC4Jq9k70 z4P4Tky*E&4TZxNoJqpR*J$N4ZkVLR9I_GVzVMr*K7g8m{e(rY5Y`n|%pNg{Qm+mI< zf+M6}ZR95T@&vW%>X>543BaTE(5hK`Mj97PUEd$24{?w*QW*8BLcy9_80NMK3xYiz zZu>gQ=DoGW=ZniwAt@oTM$!G_g~-`Qwv6~02(lT=K3Hv_{RoPt6E)z+`Gmq9 zC8hw@53$(;v@7RbCGf2IVDN70p?1qtG@WsT6owtwBunJH`Sk3e6kOI?0@|*T@6nSd zo6{d95OtxgY$UpZIu^bKo`wU`UVYodHSVe{X<zr$!d=T)mwE;C=*I*j^eH~`6Wqsdpq9Gth^(+UL0qSa z;qpy^9ZQ}xkhBt=KtxU_D)QWs6dGZK&wb_OfxO=V7xl~lx{T#F7hXY0Z!(8wO4zJm$;D-0U(n|H` zKsEPjlAA%ghq9c+MOCNEBO3)?gl4V0=_nEDlTF0eFS0 zeJB*hW9%K&BZ{^X7&M2>P`44#L#fLrz;WscJF2^>eJz(#YQ#qCsZhojTe<>i_TU1_4x=og1Q(;mNastf6xe3QrpF|0% z`{NA9dF*r?u4(uFg5$KNHM*Q_N4a1S?rCiEjbuZFa_`ZJf;!Ij>y4 zSU7kKv1lAV-D#b=xg5fPJskACZu-j~or!}&`W~pv%0nqKCf&IR{haB^47xST4NfUn zsDEGZbgD81Lb!>&y-Uv@wEY%g=d6vB!%~Z$-PaEwHK9w|oDPMY-bn1F)2#k6LpUVt zCf}45p%~+mbj4@jwJAs%h5>8f?8yF6YK7{{;6Xp>t~BBqRE_@SvRkC!qeieQCG~Eu zASiN0@9Cs801TOPeZeCBZlumzyyHQ1NMiZ3c^twQY~ISH2kStG)lof26| z;R(HTCJBcBkw!3aqRmUw%meG$u2VN_s&AW#(q_0va|;-jCa2|vgM}ZY?{uOy_U?jy zTpUa1%>}Yn`}u=aO;KUgU!cDA21wyaZp-=j#l4O-r^Xug(N-6Io_?2oz(}UlxgZ@# z+xE-UXdsv>PU%2(Wa1&b-gXD&d9fGQdi&an6-jxfc}#!7;&n?oo_+55x1Vr1g0z~` z9L9r4`aOYMSC!bD;!Rl0?%FGq5>UOU|mzA0zc3 zNKtkTx=_&s!AP<%-!I80=kn9UXYM}tZQ!`y2sN}6E9JnLtjNHM_01IN$MGr~I!3kk zsh3Ra?x5n-Rx84)*D>z94w>DoFuyh9OYxDW<9;ab)uSFu>U8X0wzVU~SYpClow?R7CK?)~l{!o8a) zkjT?Jxv|a16{&z{zt|chrVBoy>GtM3Hy~@o+{41)j^G>b zN2%Y=7-*sMJyc_1!8behPdXYYF+oGf7+yC^l%4C%l9kt8F^sf1OuVrM>&@82Uo4wZKGIz*$mK#X)nX5n%c1w2l)iYZ_{IV8;y|6&;skW;lDd|m z>Sc1vnwl1Kk34@4Ui(mM?mUU=J#!_U+4R?i>!#Qpu8@$6bvXNYj&r#EhN*R7UPGEo zgPEn|W8$KV8}~rY#yu!OXN~NVSb(l-!SEG^n~gYCYI!X^&4=ZliCq}TdlLyrnV0(V zvprqLAln6O{Lc`O?=gGG-sbSReC27Xyr)s`O$aDXS7whgRL{utMA5^S9~kq`Ao)dw z;$`bEK<7VALw75ksFxccsWxqCWqdm#-Sk{NJvcW!G$xyyUW0Ml%%Ew3HVQOI8yCHkG-L5n@3Gx&RGXm+ zMu|ZVT*x8RXr^JJ-4Cyf%WIwm;P#Em5GImPpEbVZnr8eGa#uq+Yh>>4rh3jHV(NIM zYxrWAG0MlIaNC(*)>Yn5U!$d&x}-Su_V3 z`zm0$FQ$IEV}RZm1JNdi?D37{Q_eGy|HAHPyYoJtofep~scgW(kgjB!CPgKmeI(fZ z2n`o>Pwa*buqf0sW-wQbRHZR5BpgPn;+Kp zJReb)g7F5{Gj7T#s=(exkOe+;hQ}SurXDIT`$cTG_8L+#Zd0{>Co+~aMDX#v0Q!m^ z%&~7XI<2dp!+wVggm)2nh8;BG8#gh%Ok%lf`!$Xy>^TPXY<2Z(=U+T7wPLp4^T7a`4q;eYT6*jFiQ({T%zTBH>2(y zm$X9_=aMTq@AzG}k>|7P)ngKHHyGKC)2ljFDwgt^48xV2n_T^R2_#DwUm?^W_Ln`? z|5hLBSvQKW%Ve4;Mz0DneL{+4FtNBvq1d^essC^jRJQIW#(oN-sHq5b=pp6`Iy*X?ehyuC=#(Sg zd&jOtPI1wE4E!}Ug87GK0pr8rAQFDdSpXGX#sl$YrN>P#hQhru?7V1K{2+P(H`hMB zhx!rK9^uJn$YRcs$cXRgUbI=UT)sMGGo7NNJgH3fHf%l#+f9ZZ;v`3rO+`xcJX`Z_ z6aF)tO`*|&8Oj)rv0Bj|_5Su?V9$<3*ItyWqR7%GWJ1Gjq;@r!=nqMfGCo;9n%CazRY1WXv&1Qt299cUdceW=_cV+t_8Y%Z}vj`pcYSxK@kNl1&ib6BOZ$Fl_whr zxUz32>z((gW-9BtD5;mJvLze4l|wpTs6GtVOajFK3-QBY%f*nDIdv&mtrF)z(${I^;cYcVu9pFALD5_bxzTCQ$uGp2ce!8z@D><1iJ&hJ zf{`{h%(ca+=x%dUU-8&ao8jbe%1d2BdnC{eW5?X>rCYi!KVBUxP*$ouFwcV=foLu-EzTQ-gvT|xU z-JXF&3Xtrd09#B#s-C24-ko_AFl3P1L^|lWfx>#Lph*q|HWnU(VAY!e)7vB?a z^^r6k>}x@P%`WO|PBP4t(mt(55`oky64PWE+nNZzrKjUq++|&B=~20-(vK+iE;DqA z5UX=kIcqTsOcG07qO1b-+(v~rYF(rCv6KCnV2nE+D{^kwzV@zmNQ~xI0p}2YjZN)s z0FnLXK)huyiM`xYHF0nNY?0=bv3K2Fx?YB}J@NQc5^&fskG2fm^5q2z8dG=d!TWl) z_G&u5(<|I?`3ksdnJH#5eO(x{^0`hR=CsYvgFsgskcXzMtI~7rx|3iX+GqNWA`AOM zbdUrCOgW0YGA{yka9!BRDFvbxsho0X;Qj>t$ssthQ`Q%#N=6ZyB~^Mq%T3HsL!+|s zwN~X*hh~FFYJL2BCJB?Y0xIKFRIn_2;5(wc6fg)HDB7<)%d&C|ALtQBKhYy>c3nk{5$Ti z6N>)C)#3gN_m6o+e+T|qe}ert4#*EUi4^G1DFA;5{?*v}-$*~;GfJSp+g$%0{8#(w zKf!@tfc}FS_22P-wMG0BPsRLi{2!X(zZ3k`hwx8=mp=&pbS3;d!(W~G{$!wH{r5xo zQ+)Y%hQD?;|H(ij@o$E|^*R5Y;M1& diff --git a/tests/models/dynamics/nerd/test_codec.py b/tests/models/dynamics/nerd/test_codec.py deleted file mode 100644 index 220131d08b..0000000000 --- a/tests/models/dynamics/nerd/test_codec.py +++ /dev/null @@ -1,331 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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. -# ---------------------------------------------------------------------------- - -from __future__ import annotations -from pathlib import Path -import numpy as np -import pytest -import torch -from torch.autograd import gradcheck -from embodichain_models.dynamics.nerd.codec import ( - NeRDStateCodec, - quat_rotate, - plane_contact_features, - append_history, - reset_history, -) - - -def ant_state(batch: int = 2) -> tuple[torch.Tensor, torch.Tensor]: - torch.manual_seed(7) - q = torch.randn(batch, 15, dtype=torch.double) * 0.1 - rotation = q[:, 3:7] + q.new_tensor((0.1, 0.2, 0.3, 1.0)) - q[:, 3:7] = rotation / rotation.norm(dim=-1, keepdim=True) - return q, torch.randn(batch, 14, dtype=torch.double) - - -def test_ant_real_official_frame_features_match_reference_rounding() -> None: - """Frozen actual Warp1.5.1/Torch2.2.2 windows catch FP32 formula drift.""" - with np.load(Path(__file__).parent / "data/ant-official-frame.npz") as data: - values = {key: torch.from_numpy(data[key].copy()) for key in data.files} - contacts = { - key: value for key, value in values.items() if key.startswith("contact_") - } - thickness = contacts.pop("contact_thicknesses") - contacts["contact_thicknesses_0"] = thickness - contacts["contact_thicknesses_1"] = torch.zeros_like(thickness) - codec = NeRDStateCodec("ant") - actual = codec.encode( - values["q"], values["qd"], values["joint_f"], contacts, values["gravity_dir"] - ) - for name, value in actual.items(): - torch.testing.assert_close(value, values["expected_" + name], atol=0, rtol=0) - target = codec.targets( - values["q"], values["qd"], values["next_q"], values["next_qd"] - ) - torch.testing.assert_close(target, values["expected_target"], atol=1e-5, rtol=1e-4) - - -def test_com_reference_shift_and_inverse() -> None: - codec = NeRDStateCodec("ant", base_com=(0.1, -0.2, 0.3)) - q, qd = ant_state() - legacy = codec.native_to_legacy(q, qd) - center = q[:, :3] + quat_rotate(q[:, 3:7], q.new_tensor(codec.base_com)) - # A spatial twist evaluated at the body's COM reproduces its linear velocity. - torch.testing.assert_close( - legacy[:, 3:6] + torch.linalg.cross(legacy[:, :3], center), qd[:, :3] - ) - torch.testing.assert_close(codec.legacy_to_native(q, legacy), qd) - - -def test_original_velocity_override_retains_values_and_input_gradients( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Imported twists need not pass through the lossy COM representation.""" - codec = NeRDStateCodec("ant") - q = torch.zeros(2, 3, 15, dtype=torch.double) - q[..., 6] = 1.0 - qd = torch.zeros(2, 3, 14, dtype=torch.double, requires_grad=True) - original = torch.arange(2 * 6 * 14, dtype=torch.double).reshape(2, 6, 14) / 100 - original = original[:, ::2].requires_grad_() - assert not original.is_contiguous() - following = (original.detach() + 0.25).requires_grad_() - before = original.detach().clone() - contacts = { - "contact_depths": torch.zeros(2, 3, 2, dtype=torch.double), - "contact_normals": torch.ones(2, 3, 6, dtype=torch.double), - "contact_points_1": torch.ones(2, 3, 6, dtype=torch.double), - "contact_masks": torch.tensor([True, False]).expand(2, 3, 2), - } - - def reject_native_conversion(*args: object) -> None: - raise AssertionError("Original velocity must bypass COM conversion") - - monkeypatch.setattr(codec, "native_to_legacy", reject_native_conversion) - inputs = codec.encode( - q, - qd, - torch.zeros(2, 3, 8, dtype=torch.double), - contacts, - legacy_velocity=original, - ) - target = codec.targets( - q, - qd, - q, - qd, - current_embedding=inputs["states_embedding"], - legacy_velocity=original, - next_legacy_velocity=following, - ) - torch.testing.assert_close( - inputs["states_embedding"][..., 15:], original, atol=0, rtol=0 - ) - torch.testing.assert_close(target[..., 15:], following - original, atol=0, rtol=0) - assert torch.count_nonzero(inputs["contact_points_1"][..., 3:]) == 0 - torch.testing.assert_close(original, before, atol=0, rtol=0) - (3 * inputs["states_embedding"][..., 15:].sum() + target[..., 15:].sum()).backward() - torch.testing.assert_close(original.grad, torch.full_like(original, 2)) - torch.testing.assert_close(following.grad, torch.ones_like(following)) - assert qd.grad is None - - -@pytest.mark.parametrize("which", ["encode", "current_target", "next_target"]) -@pytest.mark.parametrize("invalid", ["shape", "dtype", "device"]) -def test_legacy_velocity_override_rejects_mismatched_metadata( - which: str, invalid: str -) -> None: - codec = NeRDStateCodec("cartpole") - q = torch.zeros(2, 3, 2) - qd = torch.zeros_like(q) - value = { - "shape": torch.zeros(2, 3, 1), - "dtype": torch.zeros(2, 3, 2, dtype=torch.float64), - "device": torch.empty(2, 3, 2, device="meta"), - }[invalid] - with pytest.raises(ValueError, match="Legacy velocity override"): - if which == "encode": - codec.encode(q, qd, qd, legacy_velocity=value) - else: - kwargs = { - ( - "legacy_velocity" - if which == "current_target" - else "next_legacy_velocity" - ): value - } - codec.targets(q, qd, q, qd, **kwargs) - - -@pytest.mark.parametrize("robot", ["ant", "cartpole"]) -def test_explicit_absent_legacy_velocity_preserves_native_contract(robot: str) -> None: - codec = NeRDStateCodec(robot) - q = torch.zeros(2, codec.q_dim) - qd = torch.ones(2, codec.qd_dim) - if robot == "ant": - q[:, 6] = 1.0 - contacts = { - "contact_depths": torch.zeros(2, 1), - "contact_normals": torch.zeros(2, 3), - "contact_points_1": torch.zeros(2, 3), - } - default = codec.encode(q, qd, torch.zeros_like(qd), contacts) - optional = codec.encode(q, qd, torch.zeros_like(qd), contacts, legacy_velocity=None) - for name in default: - torch.testing.assert_close(default[name], optional[name], atol=0, rtol=0) - torch.testing.assert_close( - codec.targets(q, qd, q, qd + 0.1), - codec.targets( - q, qd, q, qd + 0.1, legacy_velocity=None, next_legacy_velocity=None - ), - atol=0, - rtol=0, - ) - - -def test_ant_targets_decode_in_frame_at_t() -> None: - codec = NeRDStateCodec("ant", base_com=(0.1, -0.2, 0.3)) - q, qd = ant_state() - next_q, next_qd = ant_state() - next_q = next_q.clone() - next_q[:, :3] += q.new_tensor((0.03, -0.07, 0.02)) - next_q[:, 3:7] = ( - q.new_tensor((0.2, -0.1, 0.3, 0.9)) / q.new_tensor((0.2, -0.1, 0.3, 0.9)).norm() - ) - next_q[:, 7:] += 0.01 - next_qd = next_qd + 0.13 - prediction = codec.targets(q, qd, next_q, next_qd) - actual_q, actual_qd = codec.decode(q, qd, prediction) - torch.testing.assert_close(actual_q, next_q) - torch.testing.assert_close(actual_qd, next_qd) - # Rotation target is relative to orientation at t, not identity at t+1. - assert prediction[:, 3:6].abs().sum() > 0.1 - - -def test_ant_targets_reuse_the_current_encoded_frame() -> None: - codec = NeRDStateCodec("ant", base_com=(0.1, -0.2, 0.3)) - q, qd = ant_state() - next_q, next_qd = q.clone(), qd + 0.13 - next_q[:, :3] += 0.02 - contacts = { - "contact_depths": q.new_full((len(q), 1), 1e4), - "contact_normals": q.new_zeros(len(q), 3), - "contact_points_1": q.new_zeros(len(q), 3), - } - inputs = codec.encode(q, qd, torch.zeros_like(qd), contacts) - target = codec.targets( - q, qd, next_q, next_qd, current_embedding=inputs["states_embedding"] - ) - actual_q, actual_qd = codec.decode( - q, qd, target, current_embedding=inputs["states_embedding"] - ) - torch.testing.assert_close(actual_q, next_q) - torch.testing.assert_close(actual_qd, next_qd) - with pytest.raises(ValueError, match="embedding shape"): - codec.targets(q, qd, next_q, next_qd, current_embedding=q) - - -def test_cartpole_large_angle_wrapping_matches_official_preprocessing() -> None: - codec = NeRDStateCodec("cartpole") - q = torch.tensor([[0.25, 123.4], [0.5, -321.7]]) - next_q = q + torch.tensor([[0.01, 0.007], [-0.02, 0.009]]) - velocity = torch.zeros_like(q) - - def official_wrap(value): - return value - torch.floor((value + torch.pi) / (2 * torch.pi)) * (2 * torch.pi) - - target = codec.targets(q, velocity, next_q, velocity) - expected_angle = official_wrap(official_wrap(next_q[:, 1]) - official_wrap(q[:, 1])) - torch.testing.assert_close(target[:, 1], expected_angle, atol=0, rtol=0) - actual_q, _ = codec.decode(q, velocity, target) - expected_decoded = official_wrap(official_wrap(q[:, 1]) + expected_angle) - torch.testing.assert_close(actual_q[:, 1], expected_decoded, atol=0, rtol=0) - - -def test_cartpole_wrap_and_decode() -> None: - codec = NeRDStateCodec() - q = torch.tensor([[0.4, torch.pi - 0.01]]) - next_q = torch.tensor([[0.5, -torch.pi + 0.02]]) - qd = torch.ones_like(q) - target = codec.targets(q, qd, next_q, qd * 2) - torch.testing.assert_close(target[0, 1], torch.tensor(0.03), atol=1e-6, rtol=0) - actual = codec.decode(q, qd, target) - torch.testing.assert_close(actual[0], next_q) - torch.testing.assert_close(actual[1], qd * 2) - - -def test_static_plane_mask_and_gradient() -> None: - points = torch.tensor( - [[[0.2, 0.1, 0.05], [0.4, -0.1, 0.5]]], dtype=torch.double, requires_grad=True - ) - normal = points.new_tensor((0.0, 0.0, 1.0)) - result = plane_contact_features(points, points.new_tensor([0.02, 0.02]), normal) - torch.testing.assert_close( - result["contact_points_1"], points.new_tensor([[0.2, 0.1, 0.0, 0.4, -0.1, 0.0]]) - ) - assert result["contact_masks"].tolist() == [[True, False]] - assert gradcheck( - lambda x: tuple( - plane_contact_features(x, x.new_tensor([0.02, 0.02]), normal)[k] - for k in ("contact_points_1", "contact_depths") - ), - (points,), - ) - - -def test_ant_encode_mask_and_gradients() -> None: - codec = NeRDStateCodec("ant", base_com=(0.1, -0.2, 0.3)) - q, qd = ant_state(1) - effort = torch.randn(1, 8, dtype=torch.double, requires_grad=True) - point = torch.tensor([[[0.2, 0.1, 0.05], [0.4, -0.1, 0.5]]], dtype=torch.double) - contacts = plane_contact_features( - point, point.new_tensor([0.02, 0.02]), point.new_tensor((0.0, 0.0, 1.0)) - ) - - def function( - q: torch.Tensor, qd: torch.Tensor, effort: torch.Tensor - ) -> torch.Tensor: - # Quaternion constraint stays on the manifold during finite differences. - q = torch.cat( - (q[:, :3], q[:, 3:7] / q[:, 3:7].norm(dim=-1, keepdim=True), q[:, 7:]), -1 - ) - return torch.cat(tuple(codec.encode(q, qd, effort, contacts).values()), -1) - - assert gradcheck(function, (q.requires_grad_(), qd.requires_grad_(), effort)) - output = codec.encode(q, qd, effort, contacts) - assert torch.count_nonzero(output["contact_points_1"][:, 3:]) == 0 - - -def test_decode_gradients_nonzero_com_and_orientation() -> None: - codec = NeRDStateCodec("ant", base_com=(0.1, -0.2, 0.3)) - q, qd = ant_state(1) - pred = torch.randn(1, 29, dtype=torch.double) * 0.01 - pred[:, 6] = 1.0 - - def function( - q: torch.Tensor, qd: torch.Tensor, pred: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: - q = torch.cat( - (q[:, :3], q[:, 3:7] / q[:, 3:7].norm(dim=-1, keepdim=True), q[:, 7:]), -1 - ) - return codec.decode(q, qd, pred) - - assert gradcheck( - function, (q.requires_grad_(), qd.requires_grad_(), pred.requires_grad_()) - ) - - -def test_history_partial_reset_retains_other_graphs() -> None: - history = {"x": torch.zeros(2, 3, 1, dtype=torch.double)} - lengths = torch.zeros(2, dtype=torch.long) - tokens = [ - torch.tensor( - [[float(i)], [float(i + 10)]], dtype=torch.double, requires_grad=True - ) - for i in range(4) - ] - for token in tokens[:3]: - history, lengths = append_history(history, lengths, {"x": token}) - saved = history["x"] - history, lengths = reset_history(history, lengths, torch.tensor([True, False])) - history, lengths = append_history(history, lengths, {"x": tokens[3]}) - assert lengths.tolist() == [1, 3] - assert history["x"].squeeze(-1).tolist() == [[3.0, 0.0, 0.0], [11.0, 12.0, 13.0]] - assert saved.squeeze(-1).tolist() == [[0.0, 1.0, 2.0], [10.0, 11.0, 12.0]] - history["x"].sum().backward() - torch.testing.assert_close( - tokens[1].grad, torch.tensor([[0.0], [1.0]], dtype=torch.double) - ) diff --git a/tests/models/dynamics/nerd/test_data.py b/tests/models/dynamics/nerd/test_data.py deleted file mode 100644 index 2576b06117..0000000000 --- a/tests/models/dynamics/nerd/test_data.py +++ /dev/null @@ -1,369 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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. -# ---------------------------------------------------------------------------- - -"""Persistence, split and bounded-window contracts.""" - -from __future__ import annotations - -import pickle -import numpy as np -import pytest -import torch -from torch.utils.data import DataLoader - -from embodichain_models.dynamics.nerd.data import ( - DynamicsShardWriter, - DynamicsWindowDataset, - load_manifest, - trajectory_split, -) - - -def trajectory(offset=0, length=5): - q = np.arange(length + 1, dtype=np.float32)[:, None] + offset - return {"q": q, "qd": q * 2}, { - "joint_f": q[:-1], - "contact_depths": np.zeros((length, 0), dtype=np.float32), - } - - -def test_commit_resume_and_final_partial_shard(tmp_path): - writer = DynamicsShardWriter(tmp_path, {"dt": 1 / 60}, trajectories_per_shard=2) - writer.add_trajectory("a", *trajectory(), split="train") - assert not load_manifest(tmp_path)["shards"] - writer.add_trajectory("b", *trajectory(20), split="train") - writer.add_trajectory("c", *trajectory(40), split="train") - writer.abort() - resumed = DynamicsShardWriter( - tmp_path, {"dt": 1 / 60}, trajectories_per_shard=2, resume=True - ) - assert resumed.committed_ids == {"a", "b"} - resumed.add_trajectory("c", *trajectory(40), split="test") - resumed.close() - manifest = load_manifest(tmp_path, verify=True) - assert [len(shard["trajectories"]) for shard in manifest["shards"]] == [2, 1] - assert len(DynamicsWindowDataset(tmp_path, "train", 3)) == 6 - assert len(DynamicsWindowDataset(tmp_path, "test", 3)) == 3 - - -def test_windows_do_not_cross_trajectory_and_prefix_masks(tmp_path): - with DynamicsShardWriter(tmp_path, {}) as writer: - writer.add_trajectory("a", *trajectory(), split="train") - writer.add_trajectory("b", *trajectory(100), split="validation") - dataset = DynamicsWindowDataset(tmp_path, "train", 3, include_prefix=True) - assert len(dataset) == 5 - assert dataset[0]["valid_mask"].tolist() == [True, False, False] - assert dataset[0]["states"]["q"].flatten().tolist() == [0.0, 1.0, 1.0, 1.0] - assert dataset[-1]["states"]["q"].flatten().tolist() == [2.0, 3.0, 4.0, 5.0] - assert all(dataset[i]["trajectory_id"] == "a" for i in range(len(dataset))) - with pytest.raises(IndexError): - dataset[5] - - -def test_independent_worker_handles_bounded_lru(tmp_path): - with DynamicsShardWriter(tmp_path, {}, trajectories_per_shard=1) as writer: - for index in range(4): - writer.add_trajectory(str(index), *trajectory(index * 10), split="train") - dataset = DynamicsWindowDataset(tmp_path, "train", 2, max_open_shards=1) - for index in range(len(dataset)): - dataset[index] - assert len(dataset._handles) == 1 - copy = pickle.loads(pickle.dumps(dataset)) - assert not copy._handles - assert copy[0]["states"]["q"].shape == (3, 1) - batches = list(DataLoader(dataset, batch_size=3, num_workers=2)) - assert sum(len(batch["length"]) for batch in batches) == len(dataset) - dataset.close() - copy.close() - - -def test_validation_and_integrity(tmp_path): - with DynamicsShardWriter(tmp_path, {}) as writer: - writer.add_trajectory("a", *trajectory(), split="train") - with pytest.raises(ValueError, match="duplicate"): - writer.add_trajectory("a", *trajectory()) - states, transitions = trajectory() - states["q"][0] = np.nan - with pytest.raises(ValueError, match="Non-finite"): - writer.add_trajectory("invalid", states, transitions) - with pytest.raises(ValueError, match="metadata"): - DynamicsShardWriter(tmp_path, {"changed": True}, resume=True) - shard = tmp_path / load_manifest(tmp_path)["shards"][0]["file"] - with shard.open("ab") as stream: - stream.write(b"changed") - with pytest.raises(ValueError, match="checksum"): - load_manifest(tmp_path, verify=True) - - -def test_progress_only_publication_preserves_data_identity_and_pending_rows(tmp_path): - from embodichain_models.dynamics.nerd.data import dataset_fingerprint - - with DynamicsShardWriter(tmp_path, {}, trajectories_per_shard=4) as writer: - writer.add_trajectory("accepted", *trajectory(), split="train") - progress = {"next_candidate": 1, "accepted": 1, "rejected": 0, "rejections": []} - writer.manifest["generation_progress"] = progress - writer.commit() - first = load_manifest(tmp_path, verify=True) - assert first["shards"][0]["trajectories"][0]["id"] == "accepted" - legacy = { - key: value for key, value in first.items() if key != "generation_progress" - } - assert dataset_fingerprint(first) == dataset_fingerprint(legacy) - progress.update( - next_candidate=3, - rejected=2, - rejections=[{"id": "rejected-1"}, {"id": "rejected-2"}], - ) - writer.commit() - final = load_manifest(tmp_path, verify=True) - assert final["generation_progress"] == progress - assert final["shards"] == first["shards"] - assert dataset_fingerprint(final) == dataset_fingerprint(first) - - -def test_stable_split_not_ingestion_order(): - forward = {str(i): trajectory_split(str(i), 7) for i in range(100)} - backward = {str(i): trajectory_split(str(i), 7) for i in reversed(range(100))} - assert forward == backward - assert set(forward.values()) == {"train", "validation", "test"} - - -@pytest.mark.parametrize("prefix", [False, True]) -def test_resident_windows_preserve_order_padding_ownership_and_bounded_staging( - tmp_path, monkeypatch, prefix -): - with DynamicsShardWriter(tmp_path, {}, trajectories_per_shard=2) as writer: - for index, length in enumerate([3, 5, 11]): - writer.add_trajectory( - str(index), *trajectory(index * 100, length), split="train" - ) - dataset = DynamicsWindowDataset( - tmp_path, "train", 3, include_prefix=prefix, read_page_steps=2 - ) - requests = [-1, 0, 3, 0, len(dataset) // 2] - expected = dataset.get_batch(requests) - original = dataset._read_block - spans = [] - - def tracked(source, name, start, end, shape, dtype): - spans.append(end - start) - return original(source, name, start, end, shape, dtype) - - monkeypatch.setattr(dataset, "_read_block", tracked) - used = dataset.enable_device_cache("cpu", 1024 * 1024) - assert used > 0 and max(spans) <= 2 - - def forbidden(*args, **kwargs): - raise AssertionError("Resident batches must not reopen trajectory data") - - monkeypatch.setattr(dataset, "_read_block", forbidden) - actual = dataset.get_batch(requests) - _assert_same_batch(actual, expected) - actual["states"]["q"][0].zero_() - _assert_same_batch(dataset.get_batch(requests), expected) - monkeypatch.delattr(dataset, "_read_block") - copied = pickle.loads(pickle.dumps(dataset)) - assert copied._resident is None - result = dataset.get_batch(requests) - dataset.close() - _assert_same_batch(result, expected) - - -def test_resident_budget_failure_keeps_lazy_dataset_usable(tmp_path): - with DynamicsShardWriter(tmp_path, {}) as writer: - writer.add_trajectory("a", *trajectory(), split="train") - dataset = DynamicsWindowDataset(tmp_path, "train", 3) - with pytest.raises(ValueError, match="exceeding cache budget"): - dataset.enable_device_cache("cpu", 1) - assert dataset._resident is None - assert dataset[0]["length"] == 3 - - -@pytest.mark.gpu -def test_resident_cuda_windows_match_cpu_and_survive_close(tmp_path): - with DynamicsShardWriter(tmp_path, {}) as writer: - writer.add_trajectory("a", *trajectory(), split="train") - dataset = DynamicsWindowDataset(tmp_path, "train", 3, include_prefix=True) - expected = dataset.get_batch([0, -1, 0]) - dataset.enable_device_cache("cuda:0", 1024 * 1024) - actual = dataset.get_batch([0, -1, 0]) - dataset.close() - for kind in ["states", "transitions"]: - for name, value in actual[kind].items(): - assert value.device.type == "cuda" - torch.testing.assert_close( - value.cpu(), expected[kind][name], atol=0, rtol=0 - ) - - -def test_finalizer_is_best_effort_but_explicit_close_surfaces_errors(): - from collections import OrderedDict - from types import SimpleNamespace - from unittest.mock import Mock - - handle = SimpleNamespace(close=Mock(side_effect=RuntimeError("h5py finalized"))) - dataset = object.__new__(DynamicsWindowDataset) - dataset._handles = OrderedDict([("shard.h5", handle)]) - with pytest.raises(RuntimeError, match="h5py finalized"): - dataset.close() - dataset.__del__() - assert handle.close.call_count == 2 - - -def _assert_same_batch(actual, expected): - import torch - - if isinstance(actual, dict): - assert actual.keys() == expected.keys() - for key in actual: - _assert_same_batch(actual[key], expected[key]) - elif isinstance(actual, torch.Tensor): - assert actual.dtype == expected.dtype - assert torch.equal(actual, expected) - else: - assert actual == expected - - -@pytest.mark.parametrize("history", [1, 3, 10]) -@pytest.mark.parametrize("include_prefix", [False, True]) -def test_grouped_batch_matches_individual_windows_across_shards( - tmp_path, history, include_prefix -): - import torch - from torch.utils.data import default_collate - - with DynamicsShardWriter(tmp_path, {}, trajectories_per_shard=1) as writer: - for index, length in enumerate([1, 3, 9, 25, 30, 11, 7, 41]): - writer.add_trajectory( - str(index), - *trajectory(index * 100, length), - split="train" if index % 3 else "validation", - ) - reader = DynamicsWindowDataset( - tmp_path, - "train", - history, - include_prefix=include_prefix, - max_open_shards=2, - read_page_steps=8, - ) - baseline = DynamicsWindowDataset( - tmp_path, "train", history, include_prefix=include_prefix, max_open_shards=2 - ) - generator = torch.Generator().manual_seed(7) - for size in (1, 16, 512): - indices = torch.randint(len(reader), (size,), generator=generator) - indices[0] = -1 - _assert_same_batch( - reader.get_batch(indices), - default_collate([baseline[int(index)] for index in indices]), - ) - assert len(reader._handles) <= 2 - reader.close() - baseline.close() - - -def test_grouped_batch_preserves_duplicate_and_cross_call_ownership(tmp_path): - import torch - - with DynamicsShardWriter(tmp_path, {}) as writer: - writer.add_trajectory("a", *trajectory(), split="train") - dataset = DynamicsWindowDataset(tmp_path, "train", 3, include_prefix=True) - first = dataset.get_batch([0, 0, 4]) - untouched = first["states"]["q"][1].clone() - first["states"]["q"][0].fill_(-999) - assert torch.equal(first["states"]["q"][1], untouched) - second = dataset.get_batch([0, 4]) - assert torch.equal(second["states"]["q"][0], untouched) - assert torch.all(first["states"]["q"][0] == -999) - dataset.close() - assert torch.equal(second["states"]["q"][0], untouched) - - -def test_grouped_reads_are_bounded_for_sparse_long_trajectory(tmp_path, monkeypatch): - import h5py - from torch.utils.data import default_collate - - with DynamicsShardWriter(tmp_path, {}) as writer: - writer.add_trajectory("long", *trajectory(length=4097), split="train") - dataset = DynamicsWindowDataset(tmp_path, "train", 7, read_page_steps=32) - selections = [] - original = dataset._read_block - - def observed(group, name, start, stop, shape, dtype): - block = original(group, name, start, stop, shape, dtype) - assert block.shape == (stop - start, *shape) - selections.append(slice(start, stop)) - return block - - monkeypatch.setattr(dataset, "_read_block", observed) - indices = [0, 31, 32, 63, 2048, len(dataset) - 1, -1, 1] - batch = dataset.get_batch(indices) - assert selections - assert max(selection.stop - selection.start for selection in selections) <= 39 - _assert_same_batch(batch, default_collate([dataset[index] for index in indices])) - dataset.close() - - -def test_grouped_batch_validates_indices_before_io(tmp_path): - with DynamicsShardWriter(tmp_path, {}) as writer: - writer.add_trajectory("one", *trajectory(), split="train") - dataset = DynamicsWindowDataset(tmp_path, "train", 3) - with pytest.raises(ValueError, match="at least one"): - dataset.get_batch([]) - with pytest.raises(TypeError): - dataset.get_batch([1.5]) - with pytest.raises(IndexError): - dataset.get_batch([len(dataset)]) - assert not dataset._handles - - -def test_grouped_low_level_read_checks_shape_before_buffer_write(tmp_path): - import h5py - - with DynamicsShardWriter(tmp_path, {}) as writer: - writer.add_trajectory("one", *trajectory(), split="train") - manifest = load_manifest(tmp_path) - with h5py.File(tmp_path / manifest["shards"][0]["file"], "r+") as file: - del file["0/states/q"] - file["0/states"].create_dataset("q", data=np.zeros((6, 2), np.float32)) - reader = DynamicsWindowDataset(tmp_path, "train", 3, verify=False) - with pytest.raises(ValueError, match="shape/range mismatch"): - reader.get_batch([0]) - reader.close() - - -def test_grouped_hyperslabs_preserve_dtype_multiaxis_and_utf8_fields(tmp_path): - from torch.utils.data import default_collate - - states = { - "q": np.arange(12, dtype=np.float64).reshape(6, 2), - "qd": np.arange(6, dtype=np.int16).reshape(6, 1), - "额外值": np.arange(36, dtype=np.float32).reshape(6, 2, 3), - } - transitions = { - "joint_f": np.arange(5, dtype=np.uint8).reshape(5, 1), - "contact_masks": np.array([[True], [False], [True], [True], [False]]), - "empty": np.zeros((5, 0), np.float32), - } - with DynamicsShardWriter(tmp_path, {}) as writer: - writer.add_trajectory("typed", states, transitions, split="train") - reader = DynamicsWindowDataset(tmp_path, "train", 3) - indices = [2, 0, 1, 1] - _assert_same_batch( - reader.get_batch(indices), default_collate([reader[index] for index in indices]) - ) - reader.close() diff --git a/tests/models/dynamics/nerd/test_model.py b/tests/models/dynamics/nerd/test_model.py deleted file mode 100644 index c14a5c17f3..0000000000 --- a/tests/models/dynamics/nerd/test_model.py +++ /dev/null @@ -1,325 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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. -# ---------------------------------------------------------------------------- - -from __future__ import annotations -import json -import os -from pathlib import Path -import subprocess -import sys -import pytest -import torch -from embodichain_models.dynamics.nerd.codec import append_history -from embodichain_models.dynamics.nerd import ( - NeRDModel, - NeRDModelCfg, - export_model, - load_model, - convert_legacy_checkpoint, -) - - -def tiny() -> NeRDModel: - return NeRDModel( - NeRDModelCfg( - input_dims={"states_embedding": 4, "joint_f": 2}, - output_dim=4, - n_embd=12, - n_head=3, - n_layer=2, - block_size=10, - history_length=10, - ) - ) - - -def test_normalizers_and_parameters_follow_dtype() -> None: - model = tiny().double() - assert all(v.dtype == torch.double for v in model.state_dict().values()) - model.set_normalization( - { - k: (torch.ones(d), torch.full((d,), 2.0)) - for k, d in model.cfg.input_dims.items() - }, - torch.ones(4), - torch.full((4,), 3.0), - ) - value = torch.randn(2, 3, 4, dtype=torch.double) - torch.testing.assert_close( - model.input_rms["states_embedding"](value), (value - 1) / (2 + 1e-5) ** 0.5 - ) - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) -@pytest.mark.parametrize("normalize_input", [False, True]) -@pytest.mark.parametrize("normalize_output", [False, True]) -def test_float32_features_follow_model_dtype_without_changing_math_or_gradients( - dtype: torch.dtype, normalize_input: bool, normalize_output: bool -) -> None: - torch.set_num_threads(2) - model = tiny().to(dtype=dtype).eval() - model.cfg.normalize_input = normalize_input - model.cfg.normalize_output = normalize_output - model.set_normalization( - { - name: ( - torch.full((width,), 1 / 7, dtype=dtype), - torch.full((width,), 1 / 3, dtype=dtype), - ) - for name, width in model.cfg.input_dims.items() - }, - torch.full((4,), 1 / 11, dtype=dtype), - torch.full((4,), 1 / 13, dtype=dtype), - ) - inputs = { - name: torch.randn(2, 3, width, requires_grad=True) - for name, width in model.cfg.input_dims.items() - } - references = { - name: value.detach().clone().requires_grad_() for name, value in inputs.items() - } - original = {name: value.detach().clone() for name, value in inputs.items()} - buffers = {name: value.clone() for name, value in model.named_buffers()} - # Independent old network/normalizer math with an explicit boundary cast. - values = [] - for name in model.cfg.input_dims: - value = references[name].to(dtype=dtype) - if normalize_input: - normalizer = model.input_rms[name] - value = (value - normalizer.mean) / torch.sqrt(normalizer.var + 1e-5) - values.append(value) - expected = model.model(model.transformer_model(torch.cat(values, -1))) - if model.cfg.normalize_output: - expected = ( - expected * torch.sqrt(model.output_rms.var + 1e-5) + model.output_rms.mean - ) - actual = model(inputs) - assert actual.dtype == dtype - torch.testing.assert_close(actual, expected, atol=0, rtol=0) - actual.square().mean().backward() - gradients = {name: value.grad.clone() for name, value in model.named_parameters()} - model.zero_grad(set_to_none=True) - expected.square().mean().backward() - for name, value in model.named_parameters(): - torch.testing.assert_close(value.grad, gradients[name], atol=0, rtol=0) - for name, value in inputs.items(): - torch.testing.assert_close(value.grad, references[name].grad, atol=0, rtol=0) - torch.testing.assert_close(value, original[name], atol=0, rtol=0) - for name, value in model.named_buffers(): - torch.testing.assert_close(value, buffers[name], atol=0, rtol=0) - - -def test_variable_prefix_matches_individual_unpadded_calls() -> None: - model = tiny().eval() - inputs = {k: torch.randn(3, 10, d) for k, d in model.cfg.input_dims.items()} - copies = {k: v.clone() for k, v in inputs.items()} - lengths = torch.tensor([1, 9, 10]) - outputs = model.evaluate(inputs, lengths) - for row, length in enumerate(lengths): - expected = model.evaluate( - {k: v[row : row + 1, :length] for k, v in inputs.items()} - ) - torch.testing.assert_close( - outputs[row : row + 1], expected, atol=1e-6, rtol=1e-5 - ) - for k in inputs: - torch.testing.assert_close(inputs[k], copies[k]) - - -def test_multistep_history_control_and_parameter_gradients() -> None: - model = tiny().double().eval() - history = { - k: torch.zeros(1, 10, d, dtype=torch.double) - for k, d in model.cfg.input_dims.items() - } - lengths = torch.zeros(1, dtype=torch.long) - initial = torch.randn(1, 4, dtype=torch.double, requires_grad=True) - state = initial - controls = [ - torch.randn(1, 2, dtype=torch.double, requires_grad=True) for _ in range(12) - ] - for control in controls: - history, lengths = append_history( - history, lengths, {"states_embedding": state, "joint_f": control} - ) - state = state + model.evaluate(history, lengths).squeeze(1) - state.square().mean().backward() - assert initial.grad is not None and torch.isfinite(initial.grad).all() - assert all(c.grad is not None and c.grad.abs().sum() > 0 for c in controls) - assert all( - p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters() - ) - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) -def test_export_roundtrip_and_checksum(tmp_path: Path, dtype: torch.dtype) -> None: - model = tiny().to(dtype=dtype).eval() - inputs = { - k: torch.randn(2, 3, d, dtype=dtype) for k, d in model.cfg.input_dims.items() - } - export_model(tmp_path, model, {"codec": {"robot": "cartpole"}, "dt": 1 / 60}) - loaded, manifest = load_model(tmp_path) - assert manifest["codec"]["robot"] == "cartpole" - assert all(value.dtype == dtype for value in loaded.state_dict().values()) - torch.testing.assert_close(model(inputs), loaded(inputs), atol=0, rtol=0) - with pytest.raises(FileExistsError): - export_model(tmp_path, model, {}) - with (tmp_path / "weights.pt").open("ab") as stream: - stream.write(b"changed") - with pytest.raises(ValueError, match="checksum"): - load_model(tmp_path) - - -def test_legacy_pickle_requires_explicit_trust(tmp_path: Path) -> None: - with pytest.raises(ValueError, match="trust_pickle"): - convert_legacy_checkpoint(tmp_path / "missing", tmp_path, {}) - - -def test_attention_backend_preserves_initial_weights_and_rng() -> None: - torch.manual_seed(19) - default = tiny() - expected_rng = torch.get_rng_state().clone() - cfg = default.cfg.copy() - cfg.attention_backend = "math" - torch.manual_seed(19) - math_model = NeRDModel(cfg) - assert torch.equal(expected_rng, torch.get_rng_state()) - assert all( - torch.equal(value, math_model.state_dict()[key]) - for key, value in default.state_dict().items() - ) - - -def test_math_attention_export_load_and_provider_prediction_restore_flags( - tmp_path: Path, monkeypatch -) -> None: - model = tiny().eval() - model.cfg.attention_backend = "math" - export_model( - tmp_path, - model, - { - "codec": {"robot": "cartpole"}, - "dt": 1 / 60, - "numerical_profile": "reference", - }, - ) - loaded, manifest = load_model(tmp_path) - assert loaded.cfg.attention_backend == manifest["attention_backend"] == "math" - assert manifest["numerical_profile"] == "reference" - original = torch.nn.functional.scaled_dot_product_attention - flags = [] - before = ( - torch.backends.cuda.flash_sdp_enabled(), - torch.backends.cuda.math_sdp_enabled(), - torch.backends.cuda.mem_efficient_sdp_enabled(), - ) - - def record(*args, **kwargs): - flags.append( - ( - torch.backends.cuda.flash_sdp_enabled(), - torch.backends.cuda.math_sdp_enabled(), - torch.backends.cuda.mem_efficient_sdp_enabled(), - ) - ) - return original(*args, **kwargs) - - monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", record) - inputs = { - key: torch.randn(2, 3, width) for key, width in loaded.cfg.input_dims.items() - } - assert torch.equal(loaded.evaluate(inputs), model.evaluate(inputs)) - assert flags and all(row == (False, True, False) for row in flags) - assert before == ( - torch.backends.cuda.flash_sdp_enabled(), - torch.backends.cuda.math_sdp_enabled(), - torch.backends.cuda.mem_efficient_sdp_enabled(), - ) - - -def test_legacy_portable_export_without_execution_fields_remains_default( - tmp_path: Path, -) -> None: - model = tiny().eval() - export_model(tmp_path, model, {"codec": {"robot": "cartpole"}, "dt": 1 / 60}) - manifest = json.loads((tmp_path / "manifest.json").read_text()) - manifest.pop("attention_backend") - manifest.pop("numerical_profile") - manifest["model_config"].pop("attention_backend") - (tmp_path / "manifest.json").write_text(json.dumps(manifest)) - loaded, _ = load_model(tmp_path) - assert loaded.cfg.attention_backend == "default" - assert all( - torch.equal(value, loaded.state_dict()[key]) - for key, value in model.state_dict().items() - ) - - -@pytest.mark.parametrize( - "mutation", ["attention_backend", "numerical_profile", "reference_requires_math"] -) -def test_export_rejects_conflicting_execution_metadata( - tmp_path: Path, mutation: str -) -> None: - model = tiny() - model.cfg.attention_backend = "math" - export_model(tmp_path, model, {"numerical_profile": "reference"}) - manifest = json.loads((tmp_path / "manifest.json").read_text()) - if mutation == "attention_backend": - manifest["attention_backend"] = "default" - elif mutation == "numerical_profile": - manifest["numerical_profile"] = "unknown" - else: - manifest["attention_backend"] = "default" - manifest["model_config"]["attention_backend"] = "default" - (tmp_path / "manifest.json").write_text(json.dumps(manifest)) - with pytest.raises(ValueError, match="attention_backend|numerical_profile"): - load_model(tmp_path) - - -@pytest.mark.parametrize("robot", ["Cartpole", "Ant"]) -def test_explicit_official_checkpoint_forward_gradient_and_adam( - robot: str, tmp_path: Path -) -> None: - # Reference code is only loaded when the operator explicitly opts into a local - # trusted checkout. Run it in a subprocess to avoid polluting production imports. - reference = os.environ.get("NERD_REFERENCE_ROOT") - if reference is None: - pytest.skip("Set NERD_REFERENCE_ROOT to the pinned trusted upstream checkout") - script = Path(__file__).with_name("_nerd_reference_check.py") - subprocess.run( - [sys.executable, str(script), reference, robot, str(tmp_path)], - check=True, - env={ - **os.environ, - "OMP_NUM_THREADS": "2", - "MKL_NUM_THREADS": "2", - "PYTHONPATH": os.pathsep.join( - filter( - None, - ( - os.environ.get("NERD_REFERENCE_DEPS"), - str(Path(__file__).resolve().parents[4]), - str( - Path(__file__).resolve().parents[4] - / "extensions/embodichain_models" - ), - ), - ) - ), - }, - ) diff --git a/tests/models/dynamics/nerd/test_training.py b/tests/models/dynamics/nerd/test_training.py deleted file mode 100644 index deab29142b..0000000000 --- a/tests/models/dynamics/nerd/test_training.py +++ /dev/null @@ -1,1465 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""CPU training, exact resume, train-only statistics and standard exports.""" - -from __future__ import annotations - -import json -import math -import random -from copy import deepcopy -import subprocess -import sys - -import numpy as np -import pytest -import torch -from torch.utils.data import default_collate - -from embodichain_models.dynamics.nerd.codec import NeRDStateCodec -from embodichain_models.dynamics.nerd.data import ( - DynamicsShardWriter, - DynamicsWindowDataset, -) -from embodichain_models.dynamics.nerd import NeRDModel, NeRDModelCfg, load_model -from examples.learning.nerd._support._trainer import DynamicsTrainCfg, DynamicsTrainer -from embodichain_models.dynamics.nerd.integration import ( - compute_statistics, - encode_batch, - evaluate_model, - load_statistics, - normalized_mse, - save_statistics, -) -from examples.learning.nerd._support._workflow import run_train, run_export, run_eval -from examples.learning.nerd._support._qualification import ( - paired_noninferiority, - compare_training_steps, - capture_training_step, -) - - -@pytest.fixture(autouse=True) -def bounded_threads(): - previous = torch.get_num_threads() - torch.set_num_threads(1) - yield - torch.set_num_threads(previous) - - -def make_data(root, validation_offset=0.0): - metadata = {"codec": {"robot": "cartpole"}, "dt": 1 / 60} - rng = np.random.default_rng(5) - with DynamicsShardWriter(root, metadata, trajectories_per_shard=3) as writer: - for split in ("train", "validation", "test"): - for index in range(4): - n = 12 - effort = rng.normal(size=(n, 2)).astype(np.float32) - q, qd = np.zeros((n + 1, 2), np.float32), np.zeros( - (n + 1, 2), np.float32 - ) - q[0] = validation_offset if split != "train" else 0 - for step in range(n): - qd[step + 1] = qd[step] + effort[step] * 0.1 - q[step + 1] = q[step] + qd[step] * 0.1 - writer.add_trajectory( - f"{split}-{index}", - {"q": q, "qd": qd}, - {"joint_f": effort}, - split=split, - ) - return metadata - - -def architecture(dropout=0.0): - return NeRDModelCfg( - input_dims={"states_embedding": 4, "joint_f": 2}, - output_dim=4, - n_embd=12, - n_head=3, - n_layer=1, - block_size=4, - history_length=3, - head_hidden_dims=(12,), - dropout=dropout, - ) - - -def build( - root, - out, - *, - steps=30, - dropout=0.0, - log_interval=1, - resident_cache_bytes=0, - numerical_profile="default", - device="cpu", -): - torch.manual_seed(2) - model = NeRDModel(architecture(dropout)) - codec = NeRDStateCodec("cartpole") - train = DynamicsWindowDataset(root, "train", 3) - validation = DynamicsWindowDataset(root, "validation", 3) - cfg = DynamicsTrainCfg( - total_steps=steps, - batch_size=8, - learning_rate=0.005, - final_learning_rate=0.002, - validation_interval=10, - checkpoint_interval=10, - seed=9, - log_interval=log_interval, - resident_cache_bytes=resident_cache_bytes, - numerical_profile=numerical_profile, - device=device, - ) - return DynamicsTrainer( - model, codec, train, validation, cfg, out, manifest=train.manifest["metadata"] - ) - - -def test_statistics_only_training_windows(tmp_path): - make_data(tmp_path / "a", 0.0) - make_data(tmp_path / "b", 1000.0) - codec = NeRDStateCodec("cartpole") - first = compute_statistics(DynamicsWindowDataset(tmp_path / "a", "train", 3), codec) - second = compute_statistics( - DynamicsWindowDataset(tmp_path / "b", "train", 3), codec - ) - for key in first["inputs"]: - torch.testing.assert_close( - first["inputs"][key]["mean"], second["inputs"][key]["mean"] - ) - with pytest.raises(ValueError, match="training split"): - compute_statistics( - DynamicsWindowDataset(tmp_path / "a", "validation", 3), codec - ) - - -@pytest.fixture -def statistics_sources(tmp_path): - """Provide distinct immutable datasets with the same numeric field schema.""" - metadata = make_data(tmp_path / "own") - with DynamicsShardWriter(tmp_path / "foreign", metadata) as writer: - writer.add_trajectory( - "foreign", - { - "q": np.full((6, 2), 17.0, np.float32), - "qd": np.full((6, 2), 3.0, np.float32), - }, - {"joint_f": np.zeros((5, 2), np.float32)}, - split="train", - ) - train = DynamicsWindowDataset(tmp_path / "own", "train", 3) - validation = DynamicsWindowDataset(tmp_path / "own", "validation", 3) - foreign = DynamicsWindowDataset(tmp_path / "foreign", "train", 3) - assert train.manifest["fields"] == foreign.manifest["fields"] - assert train.fingerprint != foreign.fingerprint - try: - yield train, validation, foreign, NeRDStateCodec("cartpole") - finally: - for reader in (train, validation, foreign): - reader.close() - - -@pytest.mark.parametrize("origin", ["foreign", "missing"]) -@pytest.mark.parametrize("initialize_normalization", [True, False]) -def test_supplied_statistics_identity_rejected_before_mutation( - tmp_path, statistics_sources, origin, initialize_normalization -): - train, validation, foreign, codec = statistics_sources - statistics = compute_statistics(foreign if origin == "foreign" else train, codec) - if origin == "missing": - statistics.pop("dataset_fingerprint") - path = tmp_path / "statistics.pt" - save_statistics(path, statistics) - original_file = path.read_bytes() - loaded = load_statistics(path) - original_statistics = deepcopy(loaded) - model = NeRDModel(architecture()) - original_model = {key: value.clone() for key, value in model.state_dict().items()} - original_config = deepcopy(model.cfg.to_dict()) - # A rejected constructor must not switch the attention policy or reset any - # global RNG while admitting an unrelated normalization artifact. - cfg = DynamicsTrainCfg( - total_steps=2, batch_size=4, numerical_profile="reference", seed=97 - ) - random.seed(101) - np.random.seed(102) - torch.manual_seed(103) - rng = (random.getstate(), np.random.get_state(), torch.get_rng_state().clone()) - output = tmp_path / "run" - with pytest.raises(ValueError, match="statistics/dataset fingerprint mismatch"): - DynamicsTrainer( - model, - codec, - train, - validation, - cfg, - output, - manifest=train.manifest["metadata"], - statistics=loaded, - initialize_normalization=initialize_normalization, - ) - assert not output.exists() - assert model.cfg.to_dict() == original_config - for key, value in original_model.items(): - torch.testing.assert_close(model.state_dict()[key], value, atol=0, rtol=0) - assert random.getstate() == rng[0] - numpy_rng = np.random.get_state() - assert numpy_rng[0] == rng[1][0] and numpy_rng[2:] == rng[1][2:] - np.testing.assert_array_equal(numpy_rng[1], rng[1][1]) - torch.testing.assert_close(torch.get_rng_state(), rng[2], atol=0, rtol=0) - assert loaded.keys() == original_statistics.keys() - for group in ("inputs", "output"): - values = loaded[group] if group == "inputs" else {"output": loaded[group]} - expected = ( - original_statistics[group] - if group == "inputs" - else {"output": original_statistics[group]} - ) - for field in values: - for moment in ("mean", "var"): - torch.testing.assert_close( - values[field][moment], expected[field][moment], atol=0, rtol=0 - ) - assert path.read_bytes() == original_file - - -@pytest.mark.parametrize("initialize_normalization", [True, False]) -def test_matching_supplied_statistics_preserve_normalization_and_update( - tmp_path, statistics_sources, initialize_normalization -): - train, validation, _, codec = statistics_sources - baseline = build(tmp_path / "own", tmp_path / "baseline", steps=2) - path = tmp_path / "matching-statistics.pt" - save_statistics(path, baseline.statistics) - loaded = load_statistics(path) - if not initialize_normalization: - # Explicitly skipping initialization must retain existing model buffers, - # even if matching-corpus statistics were supplied for its record. - baseline.model.input_rms["states_embedding"].mean[0] = 0.5 - model = deepcopy(baseline.model) - candidate = DynamicsTrainer( - model, - codec, - train, - validation, - baseline.cfg.copy(), - tmp_path / "matching", - manifest=train.manifest["metadata"], - statistics=loaded, - initialize_normalization=initialize_normalization, - ) - for key, value in baseline.model.state_dict().items(): - torch.testing.assert_close(model.state_dict()[key], value, atol=0, rtol=0) - assert candidate.step() == baseline.step() - for key, value in baseline.model.state_dict().items(): - torch.testing.assert_close(model.state_dict()[key], value, atol=0, rtol=0) - baseline.train_dataset.close() - baseline.validation_dataset.close() - - -@pytest.mark.parametrize("numerical_profile", ["default", "reference"]) -def test_shared_recipe_matches_legacy_updates_and_restores_sampler( - tmp_path, numerical_profile -): - """The new CLI recipe retains the established NeRD schedule and updates.""" - from embodichain_models.dynamics.nerd.sl_recipe import build_supervised_trainer - - make_data(tmp_path / "data") - legacy = build( - tmp_path / "data", - tmp_path / "legacy", - steps=4, - numerical_profile=numerical_profile, - ) - legacy.cfg.lr_update_interval = 3 - legacy.cfg.lr_offset = 1 - config = { - "model": architecture().to_dict(), - "data": {"path": str(tmp_path / "data")}, - "objective": {"numerical_profile": numerical_profile}, - "optimizer": { - "learning_rate": 0.005, - "final_learning_rate": 0.002, - "max_grad_norm": 1.0, - "lr_update_interval": 3, - "lr_offset": 1, - }, - "trainer": { - "total_updates": 4, - "batch_size": 8, - "seed": 9, - "output_dir": str(tmp_path / "shared"), - }, - } - shared = build_supervised_trainer(config) - shared.learner.model.load_state_dict(legacy.model.state_dict()) - legacy_rates = [legacy.step()["learning_rate"] for _ in range(4)] - shared_rates = [] - shared.learner.optimizer.register_step_pre_hook( - lambda optimizer, args, kwargs: shared_rates.append( - optimizer.param_groups[0]["lr"] - ) - ) - shared.fit(stop_after=2) - reordered = build_supervised_trainer( - { - **config, - "model": { - **config["model"], - "input_dims": dict( - reversed(list(config["model"]["input_dims"].items())) - ), - }, - "trainer": {**config["trainer"], "output_dir": str(tmp_path / "reordered")}, - } - ) - before = { - key: value.clone() - for key, value in reordered.learner.model.state_dict().items() - } - with pytest.raises(ValueError, match="identity"): - reordered.resume(tmp_path / "shared/last.pt") - assert reordered.learner.num_updates == 0 - for key, value in before.items(): - torch.testing.assert_close( - reordered.learner.model.state_dict()[key], value, rtol=0, atol=0 - ) - restored = build_supervised_trainer(config) - restored.resume(tmp_path / "shared/last.pt") - restored.learner.optimizer.register_step_pre_hook( - lambda optimizer, args, kwargs: shared_rates.append( - optimizer.param_groups[0]["lr"] - ) - ) - restored.fit() - assert shared_rates == pytest.approx(legacy_rates, rel=1e-15, abs=0) - assert shared_rates[0] == shared_rates[1] == shared_rates[2] - assert shared_rates[3] < shared_rates[2] - for name, value in legacy.model.state_dict().items(): - torch.testing.assert_close( - restored.learner.model.state_dict()[name], value, atol=0, rtol=0 - ) - assert torch.equal(restored.sampler.get_state(), legacy.sampler.get_state()) - from embodichain.learning.sl.workflow import train_from_config - - report = train_from_config( - { - "factory": "embodichain_models.dynamics.nerd.sl_recipe:build_supervised_trainer", - **config, - }, - resume=str(tmp_path / "shared/last.pt"), - ) - assert report["num_updates"] == legacy.global_step - - -@pytest.mark.parametrize("numerical_profile", ["default", "reference"]) -def test_exact_resume_with_dropout_optimizer_sampler_rng(tmp_path, numerical_profile): - make_data(tmp_path / "data") - full = build( - tmp_path / "data", - tmp_path / "full", - steps=20, - dropout=0.1, - numerical_profile=numerical_profile, - ) - full.train() - expected = {key: value.clone() for key, value in full.model.state_dict().items()} - interrupted = build( - tmp_path / "data", - tmp_path / "resume", - steps=20, - dropout=0.1, - numerical_profile=numerical_profile, - ) - assert interrupted.train(stop_after=7)["status"] == "interrupted" - restored = build( - tmp_path / "data", - tmp_path / "resume", - steps=20, - dropout=0.1, - numerical_profile=numerical_profile, - ) - restored.resume(tmp_path / "resume" / "last.pt") - restored.train() - for key, value in expected.items(): - torch.testing.assert_close( - value, restored.model.state_dict()[key], rtol=0, atol=0 - ) - assert torch.equal(full.sampler.get_state(), restored.sampler.get_state()) - full_lines = [ - json.loads(line) - for line in (tmp_path / "full" / "metrics.jsonl").read_text().splitlines() - ] - resumed_lines = [ - json.loads(line) - for line in (tmp_path / "resume" / "metrics.jsonl").read_text().splitlines() - ] - assert full_lines == resumed_lines - - -def test_training_improves_validation_and_export_matches(tmp_path): - make_data(tmp_path / "data") - trainer = build(tmp_path / "data", tmp_path / "run", steps=100) - initial = evaluate_model(trainer.model, trainer.codec, trainer.validation_dataset) - trainer.train() - final = evaluate_model(trainer.model, trainer.codec, trainer.validation_dataset) - assert final["normalized_mse"] < initial["normalized_mse"] * 0.6 - run_export(tmp_path / "run" / "last.pt", tmp_path / "export") - exported, _ = load_model(tmp_path / "export") - inputs, _, _ = encode_batch( - default_collate([trainer.validation_dataset[0]]), trainer.codec, "cpu" - ) - trainer.model.eval() - torch.testing.assert_close(exported(inputs), trainer.model(inputs), rtol=0, atol=0) - report = run_eval( - {"dataset": str(tmp_path / "data"), "model": str(tmp_path / "export")} - ) - assert report["split"] == "test" - assert report["qualification"] == "offline_only" - - -@pytest.mark.parametrize("head_dims", [(12,), [12]]) -def test_workflow_fine_tune_and_resume_contract(tmp_path, head_dims): - metadata = make_data(tmp_path / "data") - model_config = architecture().to_dict() - model_config["head_hidden_dims"] = head_dims - cfg = { - "dataset": str(tmp_path / "data"), - "output_dir": str(tmp_path / "first"), - "manifest": metadata, - "model": model_config, - "train": { - "total_steps": 2, - "batch_size": 4, - "validation_interval": 1, - "checkpoint_interval": 1, - }, - } - run_train(cfg) - run_export(tmp_path / "first" / "last.pt", tmp_path / "export") - existing = { - path.name: path.read_bytes() - for path in (tmp_path / "first").iterdir() - if path.is_file() - } - for options in ({}, {"fine_tune": str(tmp_path / "export")}): - with pytest.raises(FileExistsError, match="another run"): - run_train({**cfg, **options}) - assert existing == { - path.name: path.read_bytes() - for path in (tmp_path / "first").iterdir() - if path.is_file() - } - fine = { - **cfg, - "output_dir": str(tmp_path / "fine"), - "fine_tune": str(tmp_path / "export"), - } - assert run_train(fine)["global_step"] == 2 - assert ( - run_train( - { - **cfg, - "output_dir": str(tmp_path / "fine-resume"), - "resume": str(tmp_path / "fine/last.pt"), - } - )["global_step"] - == 2 - ) - mismatch = { - **cfg, - "output_dir": str(tmp_path / "bad"), - "resume": str(tmp_path / "first" / "last.pt"), - "train": {**cfg["train"], "total_steps": 3}, - } - with pytest.raises(ValueError, match="train_cfg mismatch"): - run_train(mismatch) - - -@pytest.mark.parametrize("operation", ["step", "train", "save"]) -def test_fresh_trainer_rejects_owned_output_before_update(tmp_path, operation): - make_data(tmp_path / "data") - original = build(tmp_path / "data", tmp_path / "run", steps=2) - original.train(stop_after=1) - before = { - path.name: path.read_bytes() - for path in (tmp_path / "run").iterdir() - if path.is_file() - } - candidate = build(tmp_path / "data", tmp_path / "run", steps=2) - parameters = { - key: value.clone() for key, value in candidate.model.state_dict().items() - } - sampler = candidate.sampler.get_state().clone() - with pytest.raises(FileExistsError, match="another run"): - if operation == "save": - candidate.save(tmp_path / "run/new.pt") - else: - getattr(candidate, operation)() - assert candidate.global_step == candidate.learner.num_updates == 0 - assert not candidate.optimizer.state - torch.testing.assert_close(candidate.sampler.get_state(), sampler, rtol=0, atol=0) - for key, value in parameters.items(): - torch.testing.assert_close( - candidate.model.state_dict()[key], value, rtol=0, atol=0 - ) - assert before == { - path.name: path.read_bytes() - for path in (tmp_path / "run").iterdir() - if path.is_file() - } - - -def test_resume_new_directory_preserves_best_export_without_new_improvement( - tmp_path, monkeypatch -): - make_data(tmp_path / "data") - scores = iter([1.0, 2.0, 3.0]) - monkeypatch.setattr( - "examples.learning.nerd._support._trainer.evaluate_model", - lambda *args, **kwargs: {"normalized_mse": next(scores)}, - ) - original = build(tmp_path / "data", tmp_path / "original", steps=4) - original.cfg.validation_interval = 1 - original.train(stop_after=2) - saved_best = torch.load(tmp_path / "original/best.pt", weights_only=True) - assert saved_best["global_step"] == 1 - assert saved_best["resumable"] is False - assert any( - not torch.equal(value, original.model.state_dict()[key]) - for key, value in saved_best["model"].items() - ) - before = { - path.name: path.read_bytes() - for path in (tmp_path / "original").iterdir() - if path.is_file() - } - run_export(tmp_path / "original/best.pt", tmp_path / "original-export") - restored = build(tmp_path / "data", tmp_path / "resumed", steps=4) - restored.cfg.validation_interval = 1 - restored.resume(tmp_path / "original/last.pt") - assert list((tmp_path / "resumed").iterdir()) == [] - restored.train(stop_after=3) - restored_best = torch.load(tmp_path / "resumed/best.pt", weights_only=True) - assert restored_best["global_step"] == 1 - for key, value in saved_best["model"].items(): - torch.testing.assert_close(restored_best["model"][key], value, rtol=0, atol=0) - assert before == { - path.name: path.read_bytes() - for path in (tmp_path / "original").iterdir() - if path.is_file() - } - run_export(tmp_path / "resumed/best.pt", tmp_path / "resumed-export") - expected, _ = load_model(tmp_path / "original-export") - actual, _ = load_model(tmp_path / "resumed-export") - inputs, _, _ = encode_batch( - restored.validation_dataset.get_batch([0]), restored.codec, "cpu" - ) - torch.testing.assert_close(actual(inputs), expected(inputs), rtol=0, atol=0) - with pytest.raises(ValueError, match="export only"): - restored.resume(tmp_path / "resumed/best.pt") - - -def test_resume_preserves_future_best_after_interrupted_publication( - tmp_path, monkeypatch -): - make_data(tmp_path / "data") - scores = iter([1.0, 0.5]) - monkeypatch.setattr( - "examples.learning.nerd._support._trainer.evaluate_model", - lambda *args, **kwargs: {"normalized_mse": next(scores)}, - ) - original = build(tmp_path / "data", tmp_path / "original", steps=4) - original.cfg.validation_interval = 1 - original.train(stop_after=1) - first_best = torch.load(tmp_path / "original/best.pt", weights_only=True) - publish = original._write_best_snapshot - - def interrupt_after_best() -> None: - publish() - raise RuntimeError("interrupted after publishing best") - - monkeypatch.setattr(original, "_write_best_snapshot", interrupt_after_best) - with pytest.raises(RuntimeError, match="publishing best"): - original.train(stop_after=2) - assert ( - torch.load(tmp_path / "original/best.pt", weights_only=True)["global_step"] == 2 - ) - assert ( - torch.load(tmp_path / "original/last.pt", weights_only=True)["global_step"] == 1 - ) - before = { - path.name: path.read_bytes() - for path in (tmp_path / "original").iterdir() - if path.is_file() - } - candidate = build(tmp_path / "data", tmp_path / "original", steps=4) - candidate.cfg.validation_interval = 1 - parameters = { - key: value.clone() for key, value in candidate.model.state_dict().items() - } - sampler, rng = candidate.sampler.get_state().clone(), torch.get_rng_state().clone() - python_rng, numpy_rng = random.getstate(), np.random.get_state() - with pytest.raises(ValueError, match="best.pt.*fresh output"): - candidate.resume(tmp_path / "original/last.pt") - assert candidate.global_step == 0 and not candidate.optimizer.state - 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) - assert random.getstate() == python_rng - actual_numpy = np.random.get_state() - assert actual_numpy[0] == numpy_rng[0] and actual_numpy[2:] == numpy_rng[2:] - np.testing.assert_array_equal(actual_numpy[1], numpy_rng[1]) - for key, value in parameters.items(): - torch.testing.assert_close( - candidate.model.state_dict()[key], value, rtol=0, atol=0 - ) - restored = build(tmp_path / "data", tmp_path / "fresh", steps=4) - restored.cfg.validation_interval = 1 - restored.resume(tmp_path / "original/last.pt") - restored.train(stop_after=1) - restored_best = torch.load(tmp_path / "fresh/best.pt", weights_only=True) - assert restored_best["global_step"] == 1 - for key, value in first_best["model"].items(): - torch.testing.assert_close(restored_best["model"][key], value, rtol=0, atol=0) - assert before == { - path.name: path.read_bytes() - for path in (tmp_path / "original").iterdir() - if path.is_file() - } - - -def test_resume_input_order_rejection_preserves_runtime_state(tmp_path): - make_data(tmp_path / "data") - original = build(tmp_path / "data", tmp_path / "original", steps=4) - original.train(stop_after=1) - candidate = build(tmp_path / "data", tmp_path / "candidate", steps=4) - candidate.model.cfg.input_dims = dict( - reversed(list(candidate.model.cfg.input_dims.items())) - ) - parameters = { - key: value.clone() for key, value in candidate.model.state_dict().items() - } - sampler, rng = candidate.sampler.get_state().clone(), torch.get_rng_state().clone() - with pytest.raises(ValueError, match="input_dims order"): - candidate.resume(tmp_path / "original/last.pt") - assert candidate.global_step == 0 - assert not candidate.optimizer.state - assert list((tmp_path / "candidate").iterdir()) == [] - 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 key, value in parameters.items(): - torch.testing.assert_close( - candidate.model.state_dict()[key], value, rtol=0, atol=0 - ) - - -def test_legacy_best_fallback_requires_matching_source_weights(tmp_path): - make_data(tmp_path / "data") - original = build(tmp_path / "data", tmp_path / "legacy", steps=4) - original.step() - original.best_validation = 1.0 - original.save(tmp_path / "legacy/best.pt") - original.step() - original.save(tmp_path / "legacy/last.pt") - # Construct the historical v1 layout from actual step-one and step-two - # checkpoints, retaining each original optimizer/RNG boundary. - for filename in ("best.pt", "last.pt"): - path = tmp_path / "legacy" / filename - state = torch.load(path, weights_only=True) - for key in ("run_id", "best_model", "best_step"): - state.pop(key) - torch.save(state, path) - expected = torch.load(tmp_path / "legacy/best.pt", weights_only=True) - restored = build(tmp_path / "data", tmp_path / "restored", steps=4) - restored.resume(tmp_path / "legacy/last.pt") - restored.train(stop_after=2) - best = torch.load(tmp_path / "restored/best.pt", weights_only=True) - assert best["global_step"] == 1 - assert best["resumable"] is False - for key, value in expected["model"].items(): - torch.testing.assert_close(best["model"][key], value, rtol=0, atol=0) - (tmp_path / "legacy/best.pt").unlink() - missing = build(tmp_path / "data", tmp_path / "missing", steps=4) - with pytest.raises(ValueError, match="matching best.pt"): - missing.resume(tmp_path / "legacy/last.pt") - assert missing.global_step == 0 and not missing.optimizer.state - assert list((tmp_path / "missing").iterdir()) == [] - - -def test_resume_foreign_directory_or_newer_history_preserves_artifacts(tmp_path): - make_data(tmp_path / "data") - source = build(tmp_path / "data", tmp_path / "source", steps=4) - source.step() - source.save(tmp_path / "source/step_1.pt") - source.train(stop_after=3) - foreign = build(tmp_path / "data", tmp_path / "foreign", steps=4) - foreign.train(stop_after=1) - for output, checkpoint, error, message in ( - ("foreign", "source/last.pt", FileExistsError, "another run"), - ("source", "source/step_1.pt", ValueError, "history is ahead"), - ): - before = { - path.name: path.read_bytes() - for path in (tmp_path / output).iterdir() - if path.is_file() - } - candidate = build(tmp_path / "data", tmp_path / output, steps=4) - parameters = { - key: value.clone() for key, value in candidate.model.state_dict().items() - } - with pytest.raises(error, match=message): - candidate.resume(tmp_path / checkpoint) - assert candidate.global_step == 0 - assert not candidate.optimizer.state - for key, value in parameters.items(): - torch.testing.assert_close( - candidate.model.state_dict()[key], value, rtol=0, atol=0 - ) - assert before == { - path.name: path.read_bytes() - for path in (tmp_path / output).iterdir() - if path.is_file() - } - - -def test_resume_rejects_unknown_preprocessing_execution(tmp_path): - make_data(tmp_path / "data") - trainer = build(tmp_path / "data", tmp_path / "run", steps=2) - trainer.train(stop_after=1) - state = torch.load(tmp_path / "run/last.pt", map_location="cpu", weights_only=True) - state.pop("preprocessing_execution") - path = tmp_path / "old-execution.pt" - torch.save(state, path) - with pytest.raises(ValueError, match="preprocessing execution mismatch"): - trainer.resume(path) - trainer.train_dataset.close() - trainer.validation_dataset.close() - - -def test_official_loss_order_and_mask(): - prediction = torch.tensor([[[1.0, 2.0], [500.0, 500.0]]], requires_grad=True) - target = torch.tensor([[[2.0, 5.0], [0.0, 0.0]]]) - variance = torch.tensor([4.0, 9.0]) - weights = 1.0 / torch.sqrt(variance + 1e-5) - expected = torch.nn.functional.mse_loss( - prediction[:, :1] * weights, target[:, :1] * weights - ) - actual = normalized_mse(prediction, target, variance, torch.tensor([[True, False]])) - torch.testing.assert_close(actual, expected) - actual.backward() - assert torch.count_nonzero(prediction.grad[:, 1]) == 0 - - -@pytest.mark.parametrize("masked", [False, True]) -def test_unnormalized_fp64_loss_backward_matches_analytic_oracle(masked): - prediction = torch.tensor( - [[[0.25, -0.5], [1.25, 0.125]]], dtype=torch.float64, requires_grad=True - ) - target = torch.tensor([[[0.125, -0.25], [0.5, 0.75]]], dtype=torch.float32) - mask = torch.tensor([[True, False]]) if masked else None - actual = normalized_mse(prediction, target, None, mask) - # Binary fractions yield exact independent values and derivatives: - # residuals=(1/8,-1/4,3/4,-5/8), gradient=2*residual/valid scalar count. - expected_loss = 5 / 128 if masked else 33 / 128 - expected_gradient = torch.tensor( - ( - [[[0.125, -0.25], [0.0, 0.0]]] - if masked - else [[[0.0625, -0.125], [0.375, -0.3125]]] - ), - dtype=torch.float64, - ) - assert actual.dtype == torch.float64 - assert actual.item() == expected_loss - actual.backward() - torch.testing.assert_close(prediction.grad, expected_gradient, atol=0, rtol=0) - assert target.dtype == torch.float32 - - -def test_bfloat16_predictions_do_not_quantize_float32_targets(): - prediction = torch.tensor([[1.0, 2.0]], dtype=torch.bfloat16, requires_grad=True) - # Both differences disappear if labels are incorrectly converted to BF16. - target = torch.tensor([[1.0 + 2**-12, 2.0 + 2**-10]], dtype=torch.float32) - oracle = prediction.detach().float().requires_grad_() - expected = (oracle - target).square().mean() - expected.backward() - with torch.autocast("cpu", dtype=torch.bfloat16): - actual = normalized_mse(prediction, target, None) - assert actual.dtype == torch.float32 - assert actual.item() > 0 - torch.testing.assert_close(actual, expected, atol=0, rtol=0) - actual.backward() - torch.testing.assert_close(prediction.grad, oracle.grad.bfloat16(), atol=0, rtol=0) - - -@pytest.mark.parametrize( - "device", ["cpu", pytest.param("cuda:0", marks=pytest.mark.gpu)] -) -def test_heterogeneous_target_statistics_preserve_official_loss_and_gradient(device): - if device.startswith("cuda") and not torch.cuda.is_available(): - pytest.skip("CUDA is unavailable") - generator = torch.Generator().manual_seed(42) - prediction = torch.randn(4, 10, 29, generator=generator).to(device).requires_grad_() - target = torch.randn(4, 10, 29, generator=generator).to(device) - variance = torch.logspace(-6, 3, 29, device=device) - weights = 1.0 / torch.sqrt(variance + 1e-5) - expected = torch.nn.functional.mse_loss(prediction * weights, target * weights) - actual = normalized_mse(prediction, target, variance) - expected_gradient = torch.autograd.grad(expected, prediction, retain_graph=True)[0] - actual_gradient = torch.autograd.grad(actual, prediction)[0] - torch.testing.assert_close(actual, expected, atol=0, rtol=0) - torch.testing.assert_close(actual_gradient, expected_gradient, atol=0, rtol=0) - - -def test_qualification_gate_requires_five_seeds_and_detects_regression(): - assert paired_noninferiority([1.0, 1.0], [1.0, 1.0])["status"] == "pending" - assert paired_noninferiority([1.0] * 5, [1.01] * 5)["status"] == "passed" - assert paired_noninferiority([1.0] * 5, [1.1] * 5)["status"] == "failed" - assert paired_noninferiority([0.0] * 5, [1e-6] * 5)["difference_mode"] == "absolute" - - -def test_training_capture_detects_altered_update(): - model = NeRDModel(architecture()) - inputs = {"states_embedding": torch.randn(2, 3, 4), "joint_f": torch.randn(2, 3, 2)} - reference = capture_training_step(model, inputs, torch.randn(2, 3, 4)) - # Self-comparison only tests the checker, not official parity. - assert compare_training_steps(reference, reference)["passed"] - candidate = dict(reference, loss=reference["loss"] + 1) - assert not compare_training_steps(reference, candidate)["passed"] - - -def test_offline_import_does_not_load_simulation(): - code = """ -import sys -from embodichain_models.dynamics.nerd import NeRDModel -assert not any(name in sys.modules for name in ('h5py', 'newton', 'warp')) -from embodichain_models.dynamics.nerd import build_supervised_trainer -assert 'newton' not in sys.modules -assert not any(name.startswith('examples.') for name in sys.modules) -import examples.learning.nerd._support._workflow -import examples.learning.nerd._support._trainer -assert not any(name == 'dexsim' or name.startswith('embodichain.lab.sim') for name in sys.modules) -""" - subprocess.run([sys.executable, "-c", code], check=True) - - -def test_interrupted_workflow_before_validation_is_json_serializable(tmp_path): - metadata = make_data(tmp_path / "data") - result = run_train( - { - "dataset": str(tmp_path / "data"), - "output_dir": str(tmp_path / "run"), - "manifest": metadata, - "model": architecture().to_dict(), - "train": {"total_steps": 10, "batch_size": 4, "validation_interval": 10}, - "stop_after": 1, - } - ) - assert result["best_validation"] is None - assert result["status"] == "interrupted" - assert result["runtime_identity"]["parameter_placements"] == [ - {"device": "cpu", "dtype": "torch.float32"} - ] - assert ( - json.loads((tmp_path / "run" / "result.json").read_text())["runtime_identity"] - == result["runtime_identity"] - ) - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) -def test_ant_offline_training_and_metrics(tmp_path, dtype): - """A floating-base model trains through native-to-legacy targets without sim.""" - metadata = {"codec": {"robot": "ant"}, "dt": 1 / 60} - with DynamicsShardWriter(tmp_path / "data", metadata) as writer: - for split in ("train", "validation", "test"): - q = np.zeros((7, 15), np.float32) - q[:, 2] = 0.5 - q[:, 6] = 1.0 - q[:, 0] = np.arange(7) * 0.01 - qd = np.zeros((7, 14), np.float32) - qd[:, 0] = 0.6 - normal = np.tile(np.array([[0.0, 0.0, 1.0]], np.float32), (6, 1)) - transitions = { - "joint_f": np.zeros((6, 14), np.float32), - "contact_normals": normal, - "contact_points_1": np.zeros((6, 3), np.float32), - "contact_depths": np.full((6, 1), 0.1, np.float32), - "contact_thicknesses_0": np.zeros((6, 1), np.float32), - "contact_thicknesses_1": np.zeros((6, 1), np.float32), - "gravity_dir": -normal, - } - writer.add_trajectory(split, {"q": q, "qd": qd}, transitions, split=split) - codec = NeRDStateCodec("ant") - train = DynamicsWindowDataset(tmp_path / "data", "train", 3) - validation = DynamicsWindowDataset(tmp_path / "data", "validation", 3) - inputs, _, _ = encode_batch(default_collate([train[0]]), codec, "cpu") - cfg = architecture() - cfg.input_dims = {key: value.shape[-1] for key, value in inputs.items()} - cfg.output_dim = 29 - model = NeRDModel(cfg).to(dtype=dtype) - trainer = DynamicsTrainer( - model, - codec, - train, - validation, - DynamicsTrainCfg(total_steps=2, batch_size=2, validation_interval=1), - tmp_path / "run", - manifest=metadata, - ) - trainer.train() - report = evaluate_model(model, codec, validation) - assert np.isfinite(report["normalized_mse"]) - assert report["base_orientation_error_rad"] >= 0.0 - assert report["base_position_error"] >= 0.0 - run_export(tmp_path / "run" / "last.pt", tmp_path / "export") - exported, _ = load_model(tmp_path / "export") - assert next(exported.parameters()).dtype == dtype - offline = run_eval( - {"dataset": str(tmp_path / "data"), "model": str(tmp_path / "export")} - ) - assert offline["split"] == "test" - for metric in ( - "normalized_mse", - "base_orientation_error_rad", - "base_position_error", - ): - assert np.isfinite(offline[metric]) - - -def test_memory_mapped_order_resume_and_integrity(tmp_path): - make_data(tmp_path / "data") - order = np.zeros((10, 8), np.int64) - path = tmp_path / "order.npy" - np.save(path, order) - trainer = build(tmp_path / "data", tmp_path / "run", steps=10) - trainer.sample_order = trainer._prepare_order(path) - assert isinstance(trainer.sample_order, np.memmap) - trainer.train(stop_after=3) - snapshot = torch.load(tmp_path / "run" / "last.pt", weights_only=True) - assert isinstance(snapshot["sample_order"], dict) - assert "sha256" in snapshot["sample_order"] - restored = build(tmp_path / "data", tmp_path / "restored", steps=10) - restored.resume(tmp_path / "run" / "last.pt") - restored.train(stop_after=4) - order[0, 0] = 1 - np.save(path, order) - invalid = build(tmp_path / "data", tmp_path / "invalid", steps=10) - sampler, rng = invalid.sampler.get_state().clone(), torch.get_rng_state().clone() - with pytest.raises(ValueError, match="checksum"): - invalid.resume(tmp_path / "run" / "last.pt") - assert invalid.sample_order is None and invalid._sample_order_identity is None - torch.testing.assert_close(invalid.sampler.get_state(), sampler, rtol=0, atol=0) - torch.testing.assert_close(torch.get_rng_state(), rng, rtol=0, atol=0) - - -def test_checkpointed_order_resume_is_identical_to_uninterrupted_training(tmp_path): - from examples.learning.nerd._support._sample_order import _write_checkpointed_order - - make_data(tmp_path / "data") - first = build(tmp_path / "data", tmp_path / "full", steps=10) - order, _ = _write_checkpointed_order( - tmp_path / "order.json", - dataset_size=len(first.train_dataset), - steps=10, - batch_size=first.cfg.batch_size, - seed=4, - block_steps=3, - ) - first.sample_order = first._prepare_order(order) - first.train() - interrupted = build(tmp_path / "data", tmp_path / "partial", steps=10) - interrupted.sample_order = interrupted._prepare_order(order) - interrupted.train(stop_after=4) - resumed = build(tmp_path / "data", tmp_path / "resumed", steps=10) - resumed.resume(tmp_path / "partial/last.pt") - resumed.train() - for name, value in first.model.state_dict().items(): - torch.testing.assert_close( - resumed.model.state_dict()[name], value, atol=0, rtol=0 - ) - - -def test_sparse_logging_retains_validation_and_final_committed_step(tmp_path): - make_data(tmp_path / "data") - trainer = build(tmp_path / "data", tmp_path / "run", steps=13, log_interval=4) - result = trainer.train() - rows = [ - json.loads(line) - for line in (tmp_path / "run/metrics.jsonl").read_text().splitlines() - ] - assert [row["step"] for row in rows] == [4, 8, 10, 12, 13] - assert "validation" in rows[2] and "validation" in rows[-1] - assert result["global_step"] == 13 - - -def test_resident_raw_cache_preserves_training_updates_and_validation(tmp_path): - make_data(tmp_path / "data") - lazy = build(tmp_path / "data", tmp_path / "lazy", steps=12) - lazy_result = lazy.train() - cached = build( - tmp_path / "data", - tmp_path / "cached", - steps=12, - resident_cache_bytes=1024 * 1024, - ) - cached_result = cached.train() - assert lazy_result["last"] == cached_result["last"] - for name, value in lazy.model.state_dict().items(): - torch.testing.assert_close( - cached.model.state_dict()[name], value, atol=0, rtol=0 - ) - - -def test_validation_batch_cap_reports_actual_windows(tmp_path): - make_data(tmp_path / "data") - trainer = build(tmp_path / "data", tmp_path / "run") - report = evaluate_model( - trainer.model, - trainer.codec, - trainer.validation_dataset, - batch_size=2, - max_batches=1, - ) - assert report["windows"] == 2 - assert report["valid_positions"] == 6 - - -def test_grouped_input_batches_preserve_exact_legacy_training_updates(tmp_path): - make_data(tmp_path / "data") - old = build(tmp_path / "data", tmp_path / "old", steps=20, dropout=0.1) - old.train_dataset.get_batch = lambda indices: default_collate( - [old.train_dataset[int(index)] for index in indices] - ) - old.validation_dataset.get_batch = lambda indices: default_collate( - [old.validation_dataset[int(index)] for index in indices] - ) - old.train() - expected = {key: value.clone() for key, value in old.model.state_dict().items()} - grouped = build(tmp_path / "data", tmp_path / "grouped", steps=20, dropout=0.1) - grouped.train() - for key, value in expected.items(): - torch.testing.assert_close( - value, grouped.model.state_dict()[key], rtol=0, atol=0 - ) - assert (tmp_path / "old" / "metrics.jsonl").read_text() == ( - tmp_path / "grouped" / "metrics.jsonl" - ).read_text() - old.train_dataset.close() - old.validation_dataset.close() - grouped.train_dataset.close() - grouped.validation_dataset.close() - - -def _add_physical_metadata(root, signature): - path = root / "manifest.json" - data = json.loads(path.read_text()) - data["metadata"].update( - physics_signature=signature, - physics_signature_schema=2, - gravity=[0.0, 0.0, -9.81], - joint_names=["slider", "pole"], - asset_sha256="asset-digest", - contact_schema="fixed-slots", - ) - path.write_text(json.dumps(data)) - return data["metadata"] - - -@pytest.mark.parametrize( - "mutation", - [ - "changed_signature", - "dropped_signature", - "changed_schema", - "dropped_gravity", - "changed_asset", - "empty_override", - ], -) -def test_training_manifest_cannot_relabel_dataset_physics(tmp_path, mutation): - make_data(tmp_path / "data") - metadata = _add_physical_metadata(tmp_path / "data", "old-physics") - override = dict(metadata) - if mutation == "changed_signature": - override["physics_signature"] = "current-physics" - elif mutation == "dropped_signature": - override.pop("physics_signature") - elif mutation == "changed_schema": - override["physics_signature_schema"] = 3 - elif mutation == "dropped_gravity": - override.pop("gravity") - elif mutation == "changed_asset": - override["asset_sha256"] = "other-asset" - else: - override = {} - with pytest.raises(ValueError, match="manifest/data .* mismatch"): - run_train( - { - "dataset": str(tmp_path / "data"), - "output_dir": str(tmp_path / "run"), - "manifest": override, - "model": architecture().to_dict(), - "train": {"total_steps": 1, "batch_size": 2}, - } - ) - assert not (tmp_path / "run").exists() - - -@pytest.mark.parametrize( - ("field", "value"), - [ - ("physics_signature", "invented"), - ("legacy_velocity", {"field": "states/legacy_velocity", "version": 1}), - ], -) -def test_training_annotations_cannot_invent_missing_physics_contract( - tmp_path, field, value -): - metadata = make_data(tmp_path / "data") - with pytest.raises(ValueError, match=f"{field} mismatch"): - run_train( - { - "dataset": str(tmp_path / "data"), - "output_dir": str(tmp_path / "run"), - "manifest": {**metadata, field: value}, - "model": architecture().to_dict(), - "train": {"total_steps": 1, "batch_size": 2}, - } - ) - - -def test_unrelated_annotations_and_fine_tune_to_new_measured_physics(tmp_path): - make_data(tmp_path / "old-data") - make_data(tmp_path / "new-data") - _add_physical_metadata(tmp_path / "old-data", "old-physics") - new_metadata = _add_physical_metadata(tmp_path / "new-data", "new-physics") - config = { - "dataset": str(tmp_path / "old-data"), - "output_dir": str(tmp_path / "old-run"), - "model": architecture().to_dict(), - "train": {"total_steps": 1, "batch_size": 2}, - } - run_train(config) - run_export(tmp_path / "old-run" / "last.pt", tmp_path / "old-export") - run_train( - { - **config, - "dataset": str(tmp_path / "new-data"), - "output_dir": str(tmp_path / "new-run"), - "fine_tune": str(tmp_path / "old-export"), - "manifest": {**new_metadata, "notes": "adapt to new measured teacher"}, - } - ) - state = torch.load(tmp_path / "new-run" / "last.pt", weights_only=True) - assert state["manifest"]["physics_signature"] == "new-physics" - assert state["manifest"]["notes"] == "adapt to new measured teacher" - - -def _gradient_parameters(): - generator = torch.Generator().manual_seed(14) - parameters = [torch.nn.Parameter(torch.zeros(width)) for width in (17, 65, 513)] - for parameter, scale in zip(parameters, (1e-3, 1.0, 10.0), strict=True): - parameter.grad = torch.randn(parameter.shape, generator=generator) * scale - return parameters - - -@pytest.mark.parametrize("profile", ["default", "reference"]) -def test_numerical_clipping_matches_its_independent_cpu_oracle(profile): - from embodichain_models.dynamics.nerd._numerics import _clip_grad_norm_ - - parameters = _gradient_parameters() - expected = [parameter.grad.clone() for parameter in parameters] - if profile == "reference": - # This is the scalar policy validated with the historical CPU trainer; - # production must preserve its left-to-right FP64 rounding order. - norms = [float(torch.linalg.vector_norm(gradient)) for gradient in expected] - total = torch.tensor(math.sqrt(sum(value * value for value in norms))) - coefficient = (1.0 / (total + 1e-6)).clamp(max=1.0) - for gradient in expected: - gradient.mul_(coefficient) - else: - originals = [ - torch.nn.Parameter(torch.zeros_like(gradient)) for gradient in expected - ] - for parameter, gradient in zip(originals, expected, strict=True): - parameter.grad = gradient - total = torch.nn.utils.clip_grad_norm_(originals, 1.0, error_if_nonfinite=True) - actual = _clip_grad_norm_(parameters, 1.0, profile=profile) - assert torch.equal(actual, total) - for parameter, gradient in zip(parameters, expected, strict=True): - assert torch.equal(parameter.grad, gradient) - - -def test_reference_attention_context_covers_backward_eval_and_restores_flags( - tmp_path, monkeypatch -): - make_data(tmp_path / "data") - trainer = build(tmp_path / "data", tmp_path / "run", numerical_profile="reference") - original = torch.nn.functional.scaled_dot_product_attention - before = ( - torch.backends.cuda.flash_sdp_enabled(), - torch.backends.cuda.math_sdp_enabled(), - torch.backends.cuda.mem_efficient_sdp_enabled(), - ) - forward_flags, backward_flags = [], [] - - def record(*args, **kwargs): - forward_flags.append( - ( - torch.backends.cuda.flash_sdp_enabled(), - torch.backends.cuda.math_sdp_enabled(), - torch.backends.cuda.mem_efficient_sdp_enabled(), - ) - ) - return original(*args, **kwargs) - - def backward(gradient): - backward_flags.append( - ( - torch.backends.cuda.flash_sdp_enabled(), - torch.backends.cuda.math_sdp_enabled(), - torch.backends.cuda.mem_efficient_sdp_enabled(), - ) - ) - return gradient - - monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", record) - handle = next(trainer.model.parameters()).register_hook(backward) - try: - trainer.step() - evaluate_model(trainer.model, trainer.codec, trainer.validation_dataset) - finally: - handle.remove() - trainer.train_dataset.close() - trainer.validation_dataset.close() - assert forward_flags and backward_flags - assert all( - flags == (False, True, False) for flags in forward_flags + backward_flags - ) - assert before == ( - torch.backends.cuda.flash_sdp_enabled(), - torch.backends.cuda.math_sdp_enabled(), - torch.backends.cuda.mem_efficient_sdp_enabled(), - ) - - -def test_attention_context_restores_flags_on_failure(): - from embodichain_models.dynamics.nerd._numerics import _attention_context - - before = ( - torch.backends.cuda.flash_sdp_enabled(), - torch.backends.cuda.math_sdp_enabled(), - ) - with pytest.raises(RuntimeError, match="injected"): - with _attention_context("reference"): - raise RuntimeError("injected") - assert before == ( - torch.backends.cuda.flash_sdp_enabled(), - torch.backends.cuda.math_sdp_enabled(), - ) - - -def test_reference_nonfinite_gradient_does_not_commit_optimizer_or_cursor(tmp_path): - make_data(tmp_path / "data") - trainer = build(tmp_path / "data", tmp_path / "run", numerical_profile="reference") - before = {key: value.clone() for key, value in trainer.model.state_dict().items()} - sampler = trainer.sampler.get_state().clone() - handle = next(trainer.model.parameters()).register_hook( - lambda gradient: torch.full_like(gradient, float("nan")) - ) - try: - with pytest.raises(FloatingPointError, match="gradient norm"): - trainer.step() - finally: - handle.remove() - trainer.train_dataset.close() - trainer.validation_dataset.close() - assert trainer.global_step == 0 - assert not trainer.optimizer.state - assert torch.equal(trainer.sampler.get_state(), sampler) - assert all( - torch.equal(value, trainer.model.state_dict()[key]) - for key, value in before.items() - ) - - -def test_unknown_numerical_profile_is_rejected_before_output(tmp_path): - make_data(tmp_path / "data") - with pytest.raises(ValueError, match="numerical_profile"): - build(tmp_path / "data", tmp_path / "run", numerical_profile="unknown") - assert not (tmp_path / "run").exists() - - -def test_reference_clipping_rejects_unsupported_gradient_precision(): - from embodichain_models.dynamics.nerd._numerics import _clip_grad_norm_ - - parameter = torch.nn.Parameter(torch.ones(2, dtype=torch.float64)) - parameter.grad = torch.ones_like(parameter) - with pytest.raises(ValueError, match="FP32 gradients"): - _clip_grad_norm_([parameter], 1.0, profile="reference") - - -@pytest.mark.parametrize( - "old_profile,new_profile", [("default", "reference"), ("reference", "default")] -) -def test_resume_rejects_numerical_profile_changes(tmp_path, old_profile, new_profile): - make_data(tmp_path / "data") - original = build(tmp_path / "data", tmp_path / "old", numerical_profile=old_profile) - original.step() - original.save(tmp_path / "saved.pt") - candidate = build( - tmp_path / "data", tmp_path / "new", numerical_profile=new_profile - ) - try: - with pytest.raises(ValueError, match="numerical_profile mismatch"): - candidate.resume(tmp_path / "saved.pt") - finally: - for trainer in (original, candidate): - trainer.train_dataset.close() - trainer.validation_dataset.close() - - -def test_legacy_default_checkpoint_resume_and_export(tmp_path): - make_data(tmp_path / "data") - trainer = build(tmp_path / "data", tmp_path / "run") - trainer.step() - trainer.save(tmp_path / "new.pt") - state = torch.load(tmp_path / "new.pt", weights_only=True) - state.pop("numerical_profile") - state.pop("attention_backend") - state["train_cfg"].pop("numerical_profile") - state["model_cfg"].pop("attention_backend") - torch.save(state, tmp_path / "legacy.pt") - restored = build(tmp_path / "data", tmp_path / "restored") - try: - restored.resume(tmp_path / "legacy.pt") - assert restored.global_step == 1 - assert restored.step() == trainer.step() - run_export(tmp_path / "legacy.pt", tmp_path / "legacy-export") - model, manifest = load_model(tmp_path / "legacy-export") - assert model.cfg.attention_backend == manifest["attention_backend"] == "default" - assert manifest["numerical_profile"] == "default" - finally: - for instance in (trainer, restored): - instance.train_dataset.close() - instance.validation_dataset.close() - - -def test_reference_profile_export_and_offline_evaluation_preserve_execution(tmp_path): - make_data(tmp_path / "data") - trainer = build( - tmp_path / "data", tmp_path / "run", steps=2, numerical_profile="reference" - ) - try: - trainer.train() - run_export(tmp_path / "run/last.pt", tmp_path / "export") - model, manifest = load_model(tmp_path / "export") - assert model.cfg.attention_backend == manifest["attention_backend"] == "math" - assert ( - manifest["numerical_profile"] - == manifest["training"]["numerical_profile"] - == "reference" - ) - report = run_eval( - {"dataset": str(tmp_path / "data"), "model": str(tmp_path / "export")} - ) - assert ( - report["attention_backend"] == "math" - and report["numerical_profile"] == "reference" - ) - finally: - trainer.train_dataset.close() - trainer.validation_dataset.close() - - -@pytest.mark.gpu -def test_reference_cuda_resume_preserves_model_adam_and_rng(tmp_path): - if not torch.cuda.is_available(): - pytest.skip("CUDA is unavailable") - make_data(tmp_path / "data") - options = { - "steps": 8, - "dropout": 0.1, - "numerical_profile": "reference", - "device": "cuda:0", - } - full = build(tmp_path / "data", tmp_path / "full", **options) - full.train() - partial = build(tmp_path / "data", tmp_path / "partial", **options) - partial.train(stop_after=3) - resumed = build(tmp_path / "data", tmp_path / "resume", **options) - try: - resumed.resume(tmp_path / "partial/last.pt") - resumed.train() - for key, value in full.model.state_dict().items(): - assert torch.equal(value, resumed.model.state_dict()[key]) - left, right = full.optimizer.state_dict(), resumed.optimizer.state_dict() - assert left["param_groups"] == right["param_groups"] - for index, state in left["state"].items(): - for field, tensor in state.items(): - assert torch.equal(tensor, right["state"][index][field]) - assert torch.equal(full.sampler.get_state(), resumed.sampler.get_state()) - finally: - for trainer in (full, partial, resumed): - trainer.train_dataset.close() - trainer.validation_dataset.close() - - -@pytest.mark.gpu -def test_reference_cuda_clipping_stays_on_device_and_preserves_ordered_norm(): - from embodichain_models.dynamics.nerd._numerics import _clip_grad_norm_ - - if not torch.cuda.is_available(): - pytest.skip("CUDA is unavailable") - # The official six-layer networks have 45 parameter tensors. Exercise both - # cold and optimized TorchScript execution at that reduction length. - parameters = [ - torch.nn.Parameter(torch.zeros(17 + index, device="cuda:0")) - for index in range(45) - ] - for index, parameter in enumerate(parameters): - parameter.grad = torch.full_like(parameter, (index + 1) * 1e-3) - original = [parameter.grad.clone() for parameter in parameters] - norms = torch.stack(torch._foreach_norm(original, 2.0)).cpu().tolist() - expected = torch.tensor( - math.sqrt(sum(value * value for value in norms)), device="cuda:0" - ) - for _ in range(10): - for parameter, gradient in zip(parameters, original, strict=True): - parameter.grad.copy_(gradient) - actual = _clip_grad_norm_( - parameters, 0.25, profile="reference", error_if_nonfinite=False - ) - assert torch.equal(actual, expected) - for parameter, gradient in zip(parameters, original, strict=True): - parameter.grad.copy_(gradient) - with torch.profiler.profile( - activities=[ - torch.profiler.ProfilerActivity.CPU, - torch.profiler.ProfilerActivity.CUDA, - ] - ) as profile: - actual = _clip_grad_norm_( - parameters, 0.25, profile="reference", error_if_nonfinite=False - ) - keys = {event.key for event in profile.key_averages()} - assert actual.device == parameters[0].device - assert not {"aten::item", "aten::_local_scalar_dense"} & keys diff --git a/tests/test_models_package.py b/tests/test_models_package.py deleted file mode 100644 index cc200a2313..0000000000 --- a/tests/test_models_package.py +++ /dev/null @@ -1,150 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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 models ship with the core and import without repository fallbacks.""" - -from __future__ import annotations - -import ast -from pathlib import Path -import subprocess -import sys - -import pytest -from setuptools import Distribution -from setuptools.command.build_py import build_py - -from setup import get_package_dir, get_packages - -pytestmark = pytest.mark.no_sim -_ROOT = Path(__file__).resolve().parents[1] - - -@pytest.fixture(scope="module") -def staged_models(tmp_path_factory: pytest.TempPathFactory) -> Path: - """Stage actual runtime sources with the distribution's package-data rules.""" - tree = ast.parse((_ROOT / "setup.py").read_text()) - call = next( - node - for node in ast.walk(tree) - if isinstance(node, ast.Call) - and isinstance(node.func, ast.Name) - and node.func.id == "setup" - ) - data = ast.literal_eval( - next(item.value for item in call.keywords if item.arg == "package_data") - ) - distribution = Distribution( - { - "packages": [ - name - for name in get_packages() - if name.partition(".")[0] in {"embodichain", "embodichain_models"} - ], - "package_dir": get_package_dir(), - "package_data": data, - "include_package_data": False, - "script_name": str(_ROOT / "setup.py"), - } - ) - output = tmp_path_factory.mktemp("staged_official_models") - command = build_py(distribution) - command.build_lib = str(output) - command.ensure_finalized() - command.run() - return output - - -def test_official_models_and_third_party_license_are_staged( - staged_models: Path, -) -> None: - for name in ( - "embodichain_models/__init__.py", - "embodichain_models/dynamics/__init__.py", - "embodichain_models/dynamics/nerd/__init__.py", - "embodichain_models/dynamics/nerd/model.py", - "embodichain_models/dynamics/nerd/checkpoint.py", - "embodichain_models/dynamics/nerd/codec.py", - "embodichain_models/dynamics/nerd/sl_recipe.py", - ): - assert (staged_models / name).is_file(), name - license_path = Path("embodichain_models/dynamics/nerd/THIRD_PARTY_LICENSES") - source = _ROOT / "extensions/embodichain_models" / license_path - assert (staged_models / license_path).read_bytes() == source.read_bytes() - assert not (staged_models / "embodichain/data/nerd/ant.xml").exists() - assert not (staged_models / "embodichain/data/nerd/cartpole_single.urdf").exists() - - -def test_staged_models_and_sl_factory_work_without_checkout_or_examples( - staged_models: Path, tmp_path: Path -) -> None: - script = """ -import importlib.abc -from pathlib import Path -import sys - -stage, repository = map(Path, sys.argv[1:]) -sys.meta_path = [finder for finder in sys.meta_path if '__editable__' not in str(finder)] -sys.path = [str(stage)] + [entry for entry in sys.path if entry and not Path(entry).is_relative_to(repository)] - -class NoSimulationOrExamples(importlib.abc.MetaPathFinder): - def find_spec(self, fullname, path=None, target=None): - if fullname.split('.')[0] in {'dexsim', 'newton', 'examples'} or fullname.startswith('embodichain.lab'): - raise AssertionError('model inference imported ' + fullname) - -sys.meta_path.insert(0, NoSimulationOrExamples()) -from embodichain.learning.models import MLP, EmpiricalNormalizer -import embodichain.learning.sl -assert not any(name == 'embodichain_models' or name.startswith('embodichain_models.') for name in sys.modules) -import embodichain_models -assert not any(name.startswith('embodichain_models.dynamics.') for name in sys.modules) -from embodichain_models.dynamics.nerd import NeRDModel, NeRDModelCfg -from embodichain_models.dynamics.nerd import model -assert Path(model.__file__).is_relative_to(stage) -import embodichain -assert Path(embodichain.__file__).is_relative_to(stage) -cfg = NeRDModelCfg(input_dims={'states_embedding': 4, 'joint_f': 2}, output_dim=4, n_embd=12, n_head=3, n_layer=1, block_size=3, history_length=3, head_hidden_dims=(12,)) -network = NeRDModel(cfg) -import torch -prediction = network({'states_embedding': torch.zeros(1, 3, 4), 'joint_f': torch.zeros(1, 3, 2)}) -assert prediction.shape == (1, 3, 4) -assert torch.isfinite(prediction).all() -from embodichain_models.dynamics.nerd.sl_recipe import build_supervised_trainer -assert callable(build_supervised_trainer) -from embodichain_models.dynamics.nerd.data import DynamicsShardWriter -from embodichain.learning.sl.workflow import train_from_config -import numpy as np -data = Path.cwd() / 'data' -with DynamicsShardWriter(data, {'codec': {'robot': 'cartpole'}, 'dt': 1 / 60}) as writer: - for split in ('train', 'validation'): - writer.add_trajectory(split, {'q': np.zeros((6, 2), np.float32), 'qd': np.zeros((6, 2), np.float32)}, {'joint_f': np.zeros((5, 2), np.float32)}, split=split) -result = train_from_config({ - 'factory': 'embodichain_models.dynamics.nerd.sl_recipe:build_supervised_trainer', - 'model': cfg.to_dict(), - 'data': {'path': str(data)}, - 'trainer': {'total_updates': 2, 'batch_size': 1, 'device': 'cpu', 'output_dir': str(Path.cwd() / 'training'), 'validation_interval': 1}, -}) -assert result['num_updates'] == 2 -assert Path(result['checkpoint']).is_file() -""" - result = subprocess.run( - [sys.executable, "-I", "-c", script, str(staged_models), str(_ROOT)], - cwd=tmp_path, - capture_output=True, - text=True, - timeout=30, - ) - assert result.returncode == 0, result.stdout + result.stderr From f850a5e7e46950b8d03daa39c4a1c73c15c4173b Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sun, 25 Oct 2026 15:32:14 +0000 Subject: [PATCH 21/21] docs(models): define policy and dynamics domain boundaries --- .../topics/neural-dynamics/neural-dynamics.md | 6 +++ extensions/embodichain_models/README.md | 50 +++++++++++++++++-- 2 files changed, 51 insertions(+), 5 deletions(-) diff --git a/agent_context/topics/neural-dynamics/neural-dynamics.md b/agent_context/topics/neural-dynamics/neural-dynamics.md index 682240f8fa..747cb0b625 100644 --- a/agent_context/topics/neural-dynamics/neural-dynamics.md +++ b/agent_context/topics/neural-dynamics/neural-dynamics.md @@ -8,6 +8,12 @@ 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` | diff --git a/extensions/embodichain_models/README.md b/extensions/embodichain_models/README.md index 7016131911..7509d4f05d 100644 --- a/extensions/embodichain_models/README.md +++ b/extensions/embodichain_models/README.md @@ -18,8 +18,48 @@ 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. -Organize implementations by domain and model type. The current package is -`dynamics/nerd`; future contributions can use policy families for BC, -`world_models` for RWM, or `dynamics` for RigidFormer-style models as their -contracts settle. All use the shared SL core. Add a family when its -implementation lands, without empty placeholder directories or a new registry. +## 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.