From 660121cb46185703555665a53a1767ef963d84da Mon Sep 17 00:00:00 2001 From: yuecideng Date: Tue, 20 Oct 2026 10:03:07 +0800 Subject: [PATCH 1/7] feat(physics): add named NeRD dynamics and training lifecycle --- agent_context/MAP.yaml | 24 + .../topics/data-pipeline/data-pipeline.md | 4 + .../differentiable-env/differentiable-env.md | 14 +- .../topics/env-framework/configuration.md | 5 + .../topics/learning-layers/learning-layers.md | 4 + .../topics/neural-dynamics/neural-dynamics.md | 95 ++ .../physics-configuration.md | 7 +- .../simulation-system/simulation-system.md | 5 +- design/neural-dynamics-solvers.md | 335 +++++ design/reports/nerd-ant-parity-20261020.json | 26 + .../nerd-cartpole-parity-20261020.json | 26 + design/reports/neural-dynamics-20261020.md | 290 ++++ ...ral-dynamics-bounded-five-seed-20261020.md | 85 ++ ...ral-dynamics-bounded-pairing-20261020.json | 611 ++++++++ .../neural-dynamics-five-seed-20261020.json | 455 ++++++ ...al-dynamics-historical-rollout-20261020.md | 61 + ...ral-dynamics-implementation-v2-20261020.md | 111 ++ .../reports/neural-dynamics-input-20261020.md | 144 ++ ...ral-dynamics-local-artifacts-20261020.json | 29 + ...s-local-artifacts-historical-20261020.json | 121 ++ ...-dynamics-local-artifacts-v2-20261020.json | 36 + ...namics-official-protocol-audit-20261020.md | 130 ++ .../neural-dynamics-pr-status-20261020.md | 277 ++++ ...eural-dynamics-pr-validation-20261020.json | 182 +++ .../neural-dynamics-v2-smoke-20261020.json | 577 +++++++ .../embodichain/embodichain.data_pipeline.rst | 1 + .../embodichain/embodichain.learning.rst | 1 + docs/source/api_reference/public_api.rst | 316 +++- docs/source/guides/index.rst | 1 + docs/source/guides/neural_dynamics.md | 170 +++ embodichain/cli/dynamics.py | 255 ++++ embodichain/cli/main.py | 5 + embodichain/compute/dynamics/__init__.py | 41 + embodichain/compute/dynamics/_evaluation.py | 103 ++ embodichain/compute/dynamics/codec.py | 348 +++++ 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 | 4 +- .../data_pipeline/dynamics/__init__.py | 21 + embodichain/data_pipeline/dynamics/shards.py | 534 +++++++ .../gym/envs/differentiable_dynamics_env.py | 235 +++ embodichain/lab/gym/utils/gym_utils.py | 8 +- embodichain/lab/sim/cfg/physics.py | 97 +- embodichain/lab/sim/diff/__init__.py | 8 +- embodichain/lab/sim/diff/dynamics.py | 148 ++ embodichain/lab/sim/objects/articulation.py | 8 +- .../lab/sim/physics/solvers/__init__.py | 25 + .../lab/sim/physics/solvers/_kinematics.py | 110 ++ .../lab/sim/physics/solvers/_nerd_binding.py | 476 ++++++ .../sim/physics/solvers/_reference_assets.py | 350 +++++ .../lab/sim/physics/solvers/collection.py | 1327 +++++++++++++++++ embodichain/lab/sim/physics/solvers/nerd.py | 368 +++++ embodichain/lab/sim/sim_manager.py | 164 +- embodichain/learning/__init__.py | 2 +- embodichain/learning/dynamics/__init__.py | 21 + embodichain/learning/dynamics/_runtime.py | 154 ++ embodichain/learning/dynamics/comparison.py | 504 +++++++ .../learning/dynamics/qualification.py | 179 +++ embodichain/learning/dynamics/reference.py | 1072 +++++++++++++ .../learning/dynamics/reference_rollout.py | 995 ++++++++++++ embodichain/learning/dynamics/report.py | 1078 +++++++++++++ embodichain/learning/dynamics/trainer.py | 625 ++++++++ embodichain/learning/dynamics/workflow.py | 245 +++ .../learning/models/nerd/THIRD_PARTY_LICENSES | 30 + embodichain/learning/models/nerd/__init__.py | 30 + .../learning/models/nerd/checkpoint.py | 188 +++ embodichain/learning/models/nerd/model.py | 298 ++++ .../tasks/classic_control/ant/README.md | 93 ++ .../classic_control/ant/agents/apg.nerd.yaml | 44 + .../ant/dynamics.eval.full.yaml | 8 + .../ant/dynamics.eval.smoke.yaml | 8 + .../ant/dynamics.generate.full.yaml | 18 + .../ant/dynamics.generate.smoke.yaml | 19 + .../ant/dynamics.rollout.full.yaml | 21 + .../ant/dynamics.rollout.smoke.yaml | 20 + .../ant/dynamics.train.full.yaml | 36 + .../ant/dynamics.train.smoke.yaml | 36 + .../tasks/classic_control/ant/env.nerd.yaml | 28 + .../tasks/classic_control/cartpole/README.md | 93 ++ .../cartpole/agents/apg.nerd.yaml | 44 + .../cartpole/dynamics.eval.full.yaml | 8 + .../cartpole/dynamics.eval.smoke.yaml | 8 + .../cartpole/dynamics.generate.full.yaml | 18 + .../cartpole/dynamics.generate.smoke.yaml | 19 + .../cartpole/dynamics.rollout.full.yaml | 21 + .../cartpole/dynamics.rollout.smoke.yaml | 20 + .../cartpole/dynamics.train.full.yaml | 32 + .../cartpole/dynamics.train.smoke.yaml | 32 + .../classic_control/cartpole/env.nerd.yaml | 28 + .../embodichain_tasks/classic_control/ant.py | 244 +++ .../classic_control/cartpole.py | 208 +++ pyproject.toml | 3 + scripts/benchmark/learning/dynamics_input.py | 320 ++++ .../benchmark/learning/dynamics_training.py | 309 ++++ setup.py | 2 + tests/cli/test_dynamics.py | 332 +++++ tests/compute/dynamics/test_codec.py | 161 ++ .../compute/dynamics/test_dynamics_metrics.py | 99 ++ .../dynamics/test_dynamics_task_reference.py | 70 + tests/data_pipeline/dynamics/test_shards.py | 274 ++++ .../envs/tasks/test_neural_dynamics_tasks.py | 171 +++ .../envs/test_differentiable_dynamics_env.py | 176 +++ tests/learning/dynamics/test_comparison.py | 170 +++ .../test_dynamics_runtime_identity.py | 183 +++ tests/learning/dynamics/test_reference.py | 454 ++++++ .../dynamics/test_reference_rollout.py | 306 ++++ tests/learning/dynamics/test_report.py | 275 ++++ tests/learning/dynamics/test_trainer.py | 475 ++++++ .../learning/models/_nerd_reference_check.py | 303 ++++ tests/learning/models/test_nerd.py | 149 ++ tests/sim/diff/test_dynamics.py | 218 +++ tests/sim/physics/solvers/test_collection.py | 438 ++++++ tests/sim/physics/solvers/test_kinematics.py | 92 ++ tests/sim/physics/solvers/test_nerd_solver.py | 900 +++++++++++ .../physics/solvers/test_reference_assets.py | 182 +++ tests/sim/test_sim_manager.py | 23 + tests/sim/test_sim_manager_cfg.py | 52 + tests/test_agent_context_map.py | 1 + 121 files changed, 21810 insertions(+), 70 deletions(-) create mode 100644 agent_context/topics/neural-dynamics/neural-dynamics.md create mode 100644 design/neural-dynamics-solvers.md create mode 100644 design/reports/nerd-ant-parity-20261020.json create mode 100644 design/reports/nerd-cartpole-parity-20261020.json create mode 100644 design/reports/neural-dynamics-20261020.md create mode 100644 design/reports/neural-dynamics-bounded-five-seed-20261020.md create mode 100644 design/reports/neural-dynamics-bounded-pairing-20261020.json create mode 100644 design/reports/neural-dynamics-five-seed-20261020.json create mode 100644 design/reports/neural-dynamics-historical-rollout-20261020.md create mode 100644 design/reports/neural-dynamics-implementation-v2-20261020.md create mode 100644 design/reports/neural-dynamics-input-20261020.md create mode 100644 design/reports/neural-dynamics-local-artifacts-20261020.json create mode 100644 design/reports/neural-dynamics-local-artifacts-historical-20261020.json create mode 100644 design/reports/neural-dynamics-local-artifacts-v2-20261020.json create mode 100644 design/reports/neural-dynamics-official-protocol-audit-20261020.md create mode 100644 design/reports/neural-dynamics-pr-status-20261020.md create mode 100644 design/reports/neural-dynamics-pr-validation-20261020.json create mode 100644 design/reports/neural-dynamics-v2-smoke-20261020.json 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/shards.py create mode 100644 embodichain/lab/gym/envs/differentiable_dynamics_env.py create mode 100644 embodichain/lab/sim/diff/dynamics.py create mode 100644 embodichain/lab/sim/physics/solvers/__init__.py create mode 100644 embodichain/lab/sim/physics/solvers/_kinematics.py create mode 100644 embodichain/lab/sim/physics/solvers/_nerd_binding.py create mode 100644 embodichain/lab/sim/physics/solvers/_reference_assets.py create mode 100644 embodichain/lab/sim/physics/solvers/collection.py create mode 100644 embodichain/lab/sim/physics/solvers/nerd.py create mode 100644 embodichain/learning/dynamics/__init__.py create mode 100644 embodichain/learning/dynamics/_runtime.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/agents/apg.nerd.yaml 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.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.rollout.full.yaml create mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.rollout.smoke.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/ant/env.nerd.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/agents/apg.nerd.yaml 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.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.rollout.full.yaml create mode 100644 extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.rollout.smoke.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 extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/env.nerd.yaml create mode 100644 extensions/embodichain_tasks/embodichain_tasks/classic_control/ant.py create mode 100644 extensions/embodichain_tasks/embodichain_tasks/classic_control/cartpole.py create mode 100644 scripts/benchmark/learning/dynamics_input.py create mode 100644 scripts/benchmark/learning/dynamics_training.py create mode 100644 tests/cli/test_dynamics.py 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_shards.py create mode 100644 tests/gym/envs/tasks/test_neural_dynamics_tasks.py create mode 100644 tests/gym/envs/test_differentiable_dynamics_env.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_report.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 create mode 100644 tests/sim/diff/test_dynamics.py create mode 100644 tests/sim/physics/solvers/test_collection.py create mode 100644 tests/sim/physics/solvers/test_kinematics.py create mode 100644 tests/sim/physics/solvers/test_nerd_solver.py create mode 100644 tests/sim/physics/solvers/test_reference_assets.py diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 49db50801..916b615f1 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -255,6 +255,30 @@ 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 and Named Physics Solvers + aliases: [neural dynamics, neural robot dynamics, NeRD, 神经动力学, 动力学模型训练] + keywords: [NeRDProvider, NeuralRobotDynamicSolver, NeRDStateCodec, NeRDModel, + RegisteredNewtonSolverCfg, dexsim.newton_solvers, DifferentiableDynamicsEnv, + FunctionalDynamics, DynamicsShardWriter, DynamicsWindowDataset, DynamicsTrainer, + 官方训练精度对比] + paths: [topics/neural-dynamics/neural-dynamics.md] + source_of_truth: + - embodichain/lab/sim/physics/solvers/ + - embodichain/lab/sim/cfg/physics.py + - embodichain/lab/gym/envs/differentiable_dynamics_env.py + - embodichain/lab/sim/diff/dynamics.py + - embodichain/compute/dynamics/ + - embodichain/learning/models/nerd/ + - embodichain/learning/dynamics/ + - embodichain/data_pipeline/dynamics/ + - embodichain/cli/dynamics.py + watch_paths: [embodichain/lab/sim/physics/solvers/, embodichain/compute/dynamics/, + embodichain/learning/models/nerd/, embodichain/learning/dynamics/, + embodichain/data_pipeline/dynamics/, tests/sim/physics/solvers/, + tests/compute/dynamics/, tests/learning/dynamics/, tests/data_pipeline/dynamics/] + related_topics: [simulation-system, differentiable-env, learning-layers, data-pipeline, rl-learning] + 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 ed7ef9af2..0f5cff9d2 100644 --- a/agent_context/topics/data-pipeline/data-pipeline.md +++ b/agent_context/topics/data-pipeline/data-pipeline.md @@ -7,6 +7,10 @@ belongs to [data-assets](../data-assets/data-assets.md). ## Entry points +Offline dynamics shards and their generation/training contract are owned by +[neural dynamics](../neural-dynamics/neural-dynamics.md); they do not use the +LeRobot expert-action schema described below. + | Request | Owner | |---|---| | Online worker and sample selection | `embodichain/data_pipeline/engine/data.py`: `OnlineDataEngine` | diff --git a/agent_context/topics/differentiable-env/differentiable-env.md b/agent_context/topics/differentiable-env/differentiable-env.md index 7c24ec400..6078c35d5 100644 --- a/agent_context/topics/differentiable-env/differentiable-env.md +++ b/agent_context/topics/differentiable-env/differentiable-env.md @@ -61,8 +61,9 @@ Task authors implement three required hooks and may override one state hook: `NewtonStepFunc`. There is no public `_apply_dynamics_action_kernel` or -`differentiable_step_mode` extension point. Differentiable dynamics are a -future feature, not a `DifferentiableEnv` capability. +`differentiable_step_mode` extension point. Registered solver dynamics use the +separate `DifferentiableDynamicsEnv` and `FunctionalDynamics` interfaces; see +[neural dynamics](../neural-dynamics/neural-dynamics.md). ## Autograd and reset rules @@ -89,9 +90,12 @@ Newton dynamics. ## Dynamics boundary -The public differentiable package exposes no solver stepper, trajectory, or -gradient-rollout API. Add those capabilities as a separate future design when -Newton dynamics are ready for end-to-end validation. +`sim.diff.FunctionalDynamics` exposes registered solver dynamics through the +public manager/Scene path. It does not change `NewtonStepFunc` or this +kinematic environment's callbacks. Functional physical/history tensors own +backward storage independently of detached live publication. Read the +[neural dynamics lifecycle](../neural-dynamics/neural-dynamics.md) before +changing resets, parameter updates or native/functional mode transitions. ## Source of truth diff --git a/agent_context/topics/env-framework/configuration.md b/agent_context/topics/env-framework/configuration.md index 0797b29f7..4a0aa9af8 100644 --- a/agent_context/topics/env-framework/configuration.md +++ b/agent_context/topics/env-framework/configuration.md @@ -69,6 +69,11 @@ without `base_dir` retains its ordinary filesystem/CWD behavior. Component-owned fields and their inline counterparts are mutually exclusive; without a selector, the original inline `robot`, `sensor`, and scene fields continue to parse unchanged. +An explicitly null `robot` selects no configured robot; a Python task can own +procedural articulation declaration through the manager's descriptor interface. +This does not infer a robot or change the omitted-field policy. The manager +creation option `enable_ground_plane` controls physical ground declaration, +independently of ground visibility. Contact-free dynamics disables it explicitly. `build_env_cfg_from_args()` expands `environment.component` before applying launcher arguments so environment-owned run controls such as `max_episodes` remain visible to the run loop while explicit CLI values retain precedence. diff --git a/agent_context/topics/learning-layers/learning-layers.md b/agent_context/topics/learning-layers/learning-layers.md index fd46bc648..26f8cc155 100644 --- a/agent_context/topics/learning-layers/learning-layers.md +++ b/agent_context/topics/learning-layers/learning-layers.md @@ -2,6 +2,10 @@ ## Owners and dependency direction +Robot-specific learned physics, shared codecs and supervised training are routed +through [neural dynamics](../neural-dynamics/neural-dynamics.md). Their model and +offline learning modules preserve the simulation-independent boundary below. + `compute/kinematics/{dls,pose_error,serial_chain,types}.py` owns fixed-step DLS, SE(3) residuals, copied PK geometry, and shared configuration/results. These modules do not import learning or simulation. `SerialChainKinematics` 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 000000000..5c0d89b60 --- /dev/null +++ b/agent_context/topics/neural-dynamics/neural-dynamics.md @@ -0,0 +1,95 @@ +# Neural dynamics + +## Entry points and ownership + +| Request | Owner (relative to `embodichain/`) | +|---|---| +| Named solver selection | `lab/sim/cfg/physics.py`; installed `dexsim.newton_solvers` metadata | +| NeRD native/functional adapter | `lab/sim/physics/solvers/nerd.py` | +| Scene indexing and source assets | `lab/sim/physics/solvers/_nerd_binding.py`, `_reference_assets.py` | +| Tensor state/contact codec | `compute/dynamics/codec.py` | +| Shared rollout error/content calculations | `compute/dynamics/_evaluation.py` | +| Network and portable checkpoints | `learning/models/nerd/` | +| Teacher collection and rollout evaluation | `lab/sim/physics/solvers/collection.py` | +| Shard persistence and sequence windows | `data_pipeline/dynamics/shards.py` | +| Supervised training, resume and export | `learning/dynamics/trainer.py`, `workflow.py` | +| Official comparisons | `learning/dynamics/reference.py`, `comparison.py`, `report.py` | +| Historical original-Warp rollout | `learning/dynamics/reference_rollout.py` | +| Command dispatch | `cli/dynamics.py`: `embodichain dynamics` | + +DexSim owns the generic registry and provider contracts. EmbodiChain registers +`nerd` in `pyproject.toml`; third-party packages use the same entry-point group. +There is no direct provider-path deployment or imperative registration route. +Discovery does not import implementations. Unknown/duplicate names fail instead +of falling back. A DexSim build exposing the extension and Scene snapshot/ +functional APIs is required. + +These physics solvers are separate from FK/IK in `lab/sim/motion/solvers`. +Creation follows [simulation ownership](../simulation-system/simulation-system.md). +Providers receive model-owned context; tasks/collectors use public manager/Scene +APIs, not native backend mutation. + +## Data and state contract + +NeRD validates supported homogeneous articulations, named joint mapping, trained +timestep, gravity, physics signature and contact schema. Reference Cartpole is +contact-free: disable the physical default plane at manager creation, not only +its visibility. Ant uses a free base and fixed plane-contact candidates. Controls +are generalized efforts; unsupported base/external forces are rejected. + +Newton COM velocities and NeRD spatial twists differ in both component ordering +and reference point. The codec owns both conversions. Relative targets use the +frame at t, not t+1. Training and runtime share the codec and frozen normalizers. + +Committed HDF5 shards hold T+1 states and T controls/contact inputs with provenance. +Splits are trajectory-level; windows cannot cross episodes or splits. Workers own +independent file handles and bounded caches. Offline training does not initialize +simulation. Teacher/neural deployments separately own physics configuration. + +Training checkpoints retain optimizer, scheduler, RNG and committed sample +progress. Portable exports contain tensor weights and JSON metadata without +executable model objects. Legacy conversion is explicitly trusted and offline. +Official comparisons share exact statistics and sample order; overlapping-window +statistical weighting must not be confused with unique-transition weighting. + +## Differentiable lifecycle + +`DifferentiableDynamicsEnv` and `sim.diff.FunctionalDynamics` advance functional +trajectories through public SimulationManager APIs. Existing `DifferentiableEnv` +retains its [kinematics contract](../differentiable-env/differentiable-env.md). + +Backward tensors are independent of native ping-pong buffers. Ordinary and +functional NeRD share a transition; publication adopts detached state/history +without stepping twice. World masks clear selected histories. Functional reset +preserves terminal reward graphs and unselected trajectories. `detach_state()` +is an explicit rollout boundary; optimizer updates follow backward. + +Torch and Warp FK use an explicit VJP bridge. Shared storage alone does not +connect autograd. Contact selections, angle wrapping and quaternion signs have +branch-local derivatives. Unsupported native runtime controls fail in functional +execution. First-order gradients do not imply force reports, second-order +derivatives or CUDA Graph support. + +## Validation + +Tests are under `tests/compute/dynamics`, `tests/learning/models`, +`tests/learning/dynamics`, `tests/data_pipeline/dynamics`, `tests/sim/physics/solvers`, +`tests/sim/diff` and `tests/gym/envs/test_differentiable_dynamics_env.py`. +Official task modules/configs use `classic_control/cartpole` and +`classic_control/ant` in the task extension. + +Distinguish feature/forward parity, original-trainer update parity, matched-data +multi-seed accuracy and deployment accuracy. Small synthetic datasets or finite +gradients do not prove the broader claims. Reference identities, budgets and +thresholds are frozen in comparison manifests; missing runs remain pending. +The acceptance contract is `design/neural-dynamics-solvers.md`; user commands +are documented in `docs/source/guides/neural_dynamics.md`. + +Historical rollout runs in a fresh process with the original pinned Warp and +source checkout. It restores either trainer's portable weights into the same +original runtime and checks recorded physical arrays including gravity. This is +training-port evidence, not deployment of historical physics through current +Newton. Shared metrics remain simulation-independent so old Warp never imports +DexSim/Newton. The controlled matrix's shared sampling/statistics/validation +protocol is distinct from literal paper-recipe reproduction; its report labels +must retain that distinction. diff --git a/agent_context/topics/simulation-system/physics-configuration.md b/agent_context/topics/simulation-system/physics-configuration.md index 2d243afe7..5880141f4 100644 --- a/agent_context/topics/simulation-system/physics-configuration.md +++ b/agent_context/topics/simulation-system/physics-configuration.md @@ -24,8 +24,11 @@ and the latter at the top level through `config_to_cfg()`. Newton `solver_cfg=None` preserves DexSim's scene-aware `AutoSolverCfg`. The concrete solver is resolved from the complete Spawn scene, then exposed through `NewtonPhysicsBackend.solver_type`. Do not guess a solver-specific robot -preset while it is unresolved. Gradient mode requires an explicit -`semi_implicit` solver. Read exact solver options/defaults in +preset while it is unresolved. Gradient mode accepts explicit `semi_implicit` +or a registered provider whose differentiable capability is checked at creation. +Installed providers resolve through the DexSim entry-point registry using their +registered `solver_type` and `options`; see +[neural dynamics](../neural-dynamics/neural-dynamics.md). Read exact solver options/defaults in `cfg/simulation.py` and the installed DexSim config types; do not substitute a hardcoded fallback when an expected DexSim API is absent. diff --git a/agent_context/topics/simulation-system/simulation-system.md b/agent_context/topics/simulation-system/simulation-system.md index 6d9c16baa..24f807375 100644 --- a/agent_context/topics/simulation-system/simulation-system.md +++ b/agent_context/topics/simulation-system/simulation-system.md @@ -78,8 +78,9 @@ tests need non-symmetric rotations to expose component-order mistakes. - Configure arenas, devices, renderer and backend before manager construction. - Keep `prepare()` convergent and retryable; a committed scene is not ready until binding, render publication and camera attachment all succeed. -- Physics advances through `update()` only. Rendering, diagnostics and marker - publication must not step or finalize an empty standalone scene. +- Physics advances through `update()` or registered-solver `functional_step()`. + Functional evaluation with `publish=False` does not change live state/time. + Rendering, diagnostics and marker publication must not step or finalize an empty scene. - Keep physics timesteps fixed during trajectory playback; retime planner samples to the command grid. See [motion execution](../motion-planning/motion-planning.md). - Use Scene batch adapters for state selection. Preserve untouched rows/DOFs diff --git a/design/neural-dynamics-solvers.md b/design/neural-dynamics-solvers.md new file mode 100644 index 000000000..1553b6120 --- /dev/null +++ b/design/neural-dynamics-solvers.md @@ -0,0 +1,335 @@ +# Named Newton solvers and the NeRD lifecycle + +Status: accepted design; implementation and qualification in progress. +Date: 2026-10-20. This document is the implementation and acceptance contract, +not a claim that the capabilities or numerical results already exist. + +Executed experiments and remaining qualification gates are recorded separately +in [the evidence report](reports/neural-dynamics-20261020.md). The subsequent +[v2 completion audit](reports/neural-dynamics-implementation-v2-20261020.md) records +contract fixes, refreshed artifacts and remaining requirements. +The latest [bounded five-seed comparison](reports/neural-dynamics-bounded-five-seed-20261020.md) +and [historical rollout evidence](reports/neural-dynamics-historical-rollout-20261020.md) +cover the subsequent implementation work without claiming full qualification. + +## Objective and decisions + +Implement the complete path from analytical simulation data through supervised +training, evaluation, export, named solver deployment, and differentiable control. +The first qualified robots are Cartpole and Ant. Both require CPU/CUDA forward +execution and gradients with respect to controls, initial state, recurrent history, +and neural model parameters. Training is single-device, supports training from +scratch, fine-tuning, and resume. Ant contact derivatives are local derivatives +of continuous features with discrete selections held fixed. + +There is exactly one extension mechanism: installed Python entry points in +`dexsim.newton_solvers`, selected by `solver_type`. Do not add a `custom` solver +name, a provider path in deployment configuration, an instance-injection route, +or a public imperative registration API. Editable installs use the same path. + +```toml +[project.entry-points."dexsim.newton_solvers"] +nerd = "embodichain.lab.sim.physics.solvers.nerd:NeRDProvider" +``` + +```yaml +physics: newton +physics_config: + physics_dt: 0.016666666666666666 + num_substeps: 1 + requires_grad: true + use_cuda_graph: false + collision_cfg: null + solver_cfg: + solver_type: nerd + options: + checkpoint: /path/to/exported-model +``` + +One solver owns the complete scene. Solver-internal composition is permitted; +framework-managed cross-solver contact coupling is not in v1. CUDA Graph, +second-order gradients, physical-parameter identification, and multi-GPU DDP +are outside v1. NeRD-specific restrictions must not restrict other providers. + +## Ownership and public contracts + +DexSim owns discovery, construction, lifecycle, stepping, clock, reset and native +state publication. It must not import EmbodiChain. EmbodiChain owns the NeRD +provider, training and data workflows. Shared tensor mathematics and codecs must +not import simulation. Networks and offline training must work without loading +DexSim or initializing a simulator. + +The DexSim extension protocol exposes: + +- `RegisteredNewtonSolverCfg(solver_type, options)` retaining the actual name. +- `NewtonSolverProvider.register_custom_attributes(builder, options)` before + model finalization, and `create(context, options)` after model, control and + external contacts are available. +- `NewtonSolverExtension` containing the native `SolverBase` implementation, + explicit device/differentiability/contact-report capabilities, optional + functional execution, and a resource-release hook. +- `FunctionalSolverState` containing named physical tensors and named history + tensors, including per-world valid history lengths. Physical tensor layouts + must be explicitly documented; they cannot depend on articulation-count guesses. + +Construction context supplies model, control, contacts, actual solver timestep, +device and scene index metadata. It is not an escape hatch to private backend or +World mutation. Provider objects and mutable histories are per-World instances. + +Native methods retain Newton signatures: `step(state_in, state_out, control, +contacts, dt)`, `reset(state, world_mask=None, flags=None)`, and +`notify_model_changed(flags)`. Functional execution accepts explicit state, +control, parameters and dt, and returns a new functional state. The extension +must make trainable parameters accessible without serializing Python objects. + +Discovery indexes entry-point metadata without importing unused implementations. +Duplicate names, collisions with built-in aliases, missing names and incompatible +interfaces fail with package provenance. Built-ins use the same resolver while +retaining their configuration compatibility and AutoSolver selection semantics. +Adding another extension must not require another hardcoded dispatch branch. + +DexSim continues to schedule collisions, runtime controls and substeps. Reduced +coordinate solvers publish consistent body state through FK; the host must not +unconditionally overwrite maximal-coordinate solver results with FK. Reset is +idempotent, handles local/global world masks, and preserves unselected histories. +Topology replacement closes and rebuilds providers. Unsupported contact impulses +are unavailable, not reported as zero. + +Add a separate `DifferentiableDynamicsEnv`, preserving the existing kinematic +`DifferentiableEnv` contract. Functional trajectories own immutable intermediates; +native ping-pong buffers are detached publication storage, never backward storage. +The native and functional NeRD paths share the same transition function. Explicit +Warp tape/VJP bridges connect FK to Torch; zero-copy views alone are not autograd. +Unsupported differentiable runtime controls fail before stepping. Simulation time +is advanced exactly once for executed substeps. Functional terminal-row reset +preserves old reward graphs and other rows; `detach_state()` is the explicit BPTT +boundary. Optimizer updates occur after backward and before a detached next rollout. + +## NeRD model and state contract + +Pin reference source: main `370e89d2b8ce68aa3578aa7812859c681b7502bc` for +generation/training; Newton preview `476dc67fccc6a99be293bf611f897b1c06f13e30` +for solver integration. Target integration stack is Newton 1.6 / Warp 1.17. +Record the actual DexSim build, Torch, CUDA and driver identities in every report. + +Use upstream-matching Cartpole/Ant assets and explicit joint/shape mapping. Bind +asset hashes, topology, inertia/collision intent, ordered joints, timestep, +coordinate conventions, history window and contact schema in a model manifest. +Validate these at deployment. Both reference models use dt=1/60 and one neural +substep. V1 NeRD consumes effort controls; unconverted position/velocity targets +or unmodeled external forces must not be silently ignored. Additional dynamic +objects or heterogeneous articulations require a separately supported model. + +Asset identity alone does not establish physical equivalence. The isolated +official Warp importers and Newton 1.6 produce different masses, inertia and +friction from these assets. Persist the actual imported physical arrays and their +hashes in teacher provenance. Official weights must not be relabeled with a +current-Newton physics signature; current-Newton deployment uses a model trained +and validated for that domain. An exact historical-physics reconstruction would +require separate qualification. + +Physical signatures have an explicit schema version. The current v2 contract +includes fixed-root orientation, parent/child topology, relevant body/joint flags, +robot and plane contact parameters, and the effective collision relationships. +Changing an omitted field must not evade validation. Contact-free Cartpole root +translations may vary under uniform gravity; fixed-root rotation may not. A new +signature schema requires regenerated/verified data and exports, not relabeling +an older checkpoint. + +Implement a simulation-independent `NeRDStateCodec` shared by data, training, +evaluation and deployment. It handles Newton coordinates versus legacy twist, +body frames, relative targets, angle wrapping and quaternion increments. Targets +are expressed in the frame at t, not the frame at t+1. Ant conversion includes +the COM/reference-point shift, not only reordering velocity components. + +Ant contact inputs use stable robot points projected onto a static plane. The +continuous geometry is differentiable. Mask decisions, wrap cuts and quaternion +sign choices have branch-local derivatives; tests exclude discontinuities when +checking finite differences. Neural contact features are not force reports. + +Histories are per-world, with explicit valid lengths. Fixed-capacity eager buffers +must preserve upstream variable-prefix positional encoding and causal-attention +semantics after partial reset. Functional histories must not overwrite retained +graphs. Model parameters and normalizer statistics use standard Torch parameters +and buffers so dtype/device migration and serialization are complete. + +## Data generation, storage and training + +Generate trajectories through EmbodiChain scene ownership and DexSim analytical +solvers, with MuJoCo-Warp as the initial teacher default. Teacher and neural +deployments separately own physics configuration; share embodiment/assets rather +than overriding an environment component's physics settings at runtime. Record +teacher substeps and use a common macro timestep of 1/60. + +Support sampled/default initial state, random effort and passive trajectories. +Translate upstream robot-specific sampling ranges to current state conventions. +Record state/contact features at t, apply the recorded effort through the macro +interval, then record state at t+1. Invalid numerical trajectories have bounded +retry and explicit rejection counts, not silent quality-based filtering. + +Store chunked HDF5 shards with a manifest. Physical states have T+1 samples; +controls and contact inputs have T samples. Include q, qd, base poses, gravity, +actual generalized forces, contact data, trajectory IDs and valid lengths. +Record schema/codec versions, units, asset/config hashes, teacher versions and +sampling seeds. Use bounded staging/writer buffers and commit complete shards; +resume only committed work. Dataset workers open independent HDF5 handles and +read windows lazily. Never retain the complete dataset on GPU or in host memory. + +Default dataset splits are trajectory-level 80/10/10. Windows cannot cross resets, +trajectories or splits. Normalization uses training data only and is frozen for +validation, test and inference. Official-parity experiments share exact official +statistics rather than comparing different statistical weighting schemes. + +Provide `embodichain dynamics generate|train|eval|export|convert-checkpoint` with +lazy subcommand imports. Only generation/online evaluation start simulation. +Models belong to learning/models, supervised training/export to learning/dynamics, +storage to data_pipeline/dynamics, simulation sampling/adaptation to lab/sim. +Do not route supervised dynamics training through an RL trainer. + +Port the causal Transformer and relative-state normalized MSE objective. Defaults: +history 10, six layers, twelve heads, embedding 192 (Cartpole) / 384 (Ant), Adam, +FP32, gradient norm clip 1.0, linear lr 1e-3 to 1e-4. Keep normalization +`sqrt(var + 1e-5)`. Training budget/batch are configurable. Upstream +`truncate_grad` means gradient clipping, not truncated temporal backpropagation. + +Training checkpoints include network, weights, statistics, optimizer, scheduler, +global step, RNG state, committed sampler progress and dataset fingerprints. +Prefetch cannot advance the durable cursor. Resume checks the feature/data +contract. Export writes configuration, state_dict, buffers, manifest and checksums, +without optimizer or serialized Python objects. Newly trained exports load directly +through the named solver. A separate tool converts trusted upstream checkpoints. + +## Official correctness and accuracy qualification + +Qualification is a required deliverable. Passing smoke tests, finite gradients or +small-dataset overfitting does not establish parity with official implementation +or training results. Never mark missing runs as passed or silently change limits. + +### Reference identity and fixtures + +Freeze source commits, asset/checkpoint/dataset hashes, precision, software and +hardware. Preserve an unmodified reference checkout outside the production +package. Reference adapters may change data transport or compatibility boundaries, +not the algorithm. Record their patches. Experimental initialization, +normalization, sampling and checkpoint-selection controls must be recorded +separately; a controlled comparison with changed controls is not proof of literal +reference-recipe equivalence. No test imports downloaded code implicitly. + +The official sampler allocates gravity features with `torch.empty` and writes only +the up-axis component. Preserve raw output and fail on invalid components unless +an explicit gravity-transport repair is enabled. Record repaired component counts, +the physical gravity direction and hashes of raw and transported data. Label this +reference as corrected transport, not untouched official data. Likewise record +the Y-up to Z-up rotation and origin-twist to COM-velocity conversion explicitly. + +Export deterministic fixtures for preprocessing, history lengths 1/9/10, partial +reset, normalization, network output, relative targets and decoded next state. +Use identical weights, state, contacts, effort and valid history. In FP32 on the +same device start with atol=1e-5 and rtol=1e-4; every discrepancy must be localized +before changing a threshold. Do not require bitwise CPU/CUDA identity. + +### Training-step equivalence + +Use identical initialized weights, minibatches, normalizer statistics and optimizer +state. Compare encoded inputs/targets, loss reduction, parameter gradients, clipping, +one Adam update and a deterministic short training run. Freeze learning-rate update +timing and checkpoint selection. In particular, official statistics computed over +overlapping sampled windows differ from unique-transition statistics: share the +same statistics artifact for parity experiments and label other profiles separately. + +### Matched-data training matrix + +Run both trainers against both immutable datasets: + +| Teacher dataset | Official trainer | EmbodiChain trainer | +|---|---|---| +| Official Warp-generated data | Reference | Training-port qualification | +| Current DexSim/Newton-generated data | Format-adapted reference | Deployment qualification | + +Within each row fix train/validation/test partitions, initialization seeds, +normalization, sample order, optimizer update budget, network architecture and +selection metric. Run five training seeds. Keep actual configured update counts +in reports; a shortened run must be labeled and cannot stand in for full-budget +qualification. Paper reproduction and current-teacher qualification are distinct: +the official Ant paper baseline requires Warp 1.5.1, not current Newton physics. + +The current matrix is a controlled training-port experiment, not a literal +reproduction of the paper's full data recipe. It shares frozen statistics and +sampling indices and uses one common validation-selection protocol. Its default +replacement sampling, window-population statistics and single ordered validation +split differ from upstream shuffled/drop-last loaders, RunningMeanStd accumulation, +epoch-zero eligibility and separate experiment/passive validation selections. +Those differences must remain visible even if the matched experiment passes. +Matching architecture and optimizer budget does not by itself establish paper +reproduction. Inspect the ordered feature list in the actual immutable initial +model artifact; sorted experiment JSON mappings cannot prove feature order. + +Historical-domain rollout qualification runs in a fresh process with the pinned +original Warp environment and Transformer neural integrator. Restore portable +weights from either trainer into the same original network implementation and +frozen normalizers. This controls the physics/runtime domain while measuring the +training port. Validate recorded original physical arrays, source and asset +identity before running. Use explicit canonical coordinate/velocity transport +only for common policies and metrics. This does not imply that the historical +model can be deployed through the current Newton reference assets. + +The historical runner is an explicitly trusted reference tool, separate from +the installed solver entry-point mechanism. It must not replace the installed +Warp package or import DexSim/Newton into the old-Warp process. Each simulator +recomputes fixed-policy actions from its own state. Record which reference +runtime executed the network; only current-domain tests establish the native +Newton/functional adapter behavior. + +### Evaluation and acceptance + +For Cartpole and Ant evaluate passive motion, random effort and fixed-policy +closed-loop trajectories. Use 2048 common test trajectories at horizons 1, 10, +100 and 1000. Report timestep error curves, q/qd errors, Ant base-position error, +quaternion geodesic error, non-finite/failure rates and policy return. Periodic +angles use wrapped error; orientations are not evaluated only by component MSE. + +Single-step implementation/gradient/update parity is a hard gate. For matched-data +training, the one-sided 95% confidence upper bound of paired relative degradation +in predeclared primary errors must not exceed 5%. Use seed-level paired bootstrap +(10,000 resamples, fixed seed 0), not independent per-timestep samples. If the +official metric is effectively zero, report the absolute difference and use the +predeclared numerical-parity tolerance rather than an unstable ratio. Failures, +non-finite states and policy returns are separately reported; an improved average +cannot hide instability. Validation chooses checkpoints; test never selects them. + +Before computing paired statistics, verify source/model/trace hashes, the robot, +teacher domain and test split. Paired runs must identify the same ordered teacher +trajectory content and fixed policy, independent of output paths or batching. +Numerical qualification requires measured residuals and pinned reference identity; +an unbound pass flag is insufficient. Missing or undefined errors remain explicit +and must not be replaced with zero or excluded to improve acceptance rates. + +Artifacts include machine-readable results and manifests, reproduction commands, +fixtures, training curves, rollout error plots and a human-readable comparison +report. Report unsupported/unavailable reference environments as pending. Do not +assert equivalence to paper results from a different teacher or training budget. + +## Implementation work packages and completion ledger + +1. DexSim named discovery, public contracts, lifecycle/functional execution and + focused regression tests; package metadata must work in the installed artifact. +2. Shared codec, causal network, manifests, checkpoint conversion and reference + fixtures with feature/forward/gradient tests. +3. Sharded datasets, trainer/resume/export, CLI and data/training tests. +4. EmbodiChain configuration, NeRD provider, simulation data generation, + differentiable environment and Cartpole/Ant deployments. +5. End-to-end data→train→eval→export→native/functional simulation validation, + official matched-data qualification and benchmark/quality reports. +6. Public API documentation, solver/training guides and routed project context. + +All items remain pending until backed by authoritative code and executed results. +The implementation log must distinguish smoke, numerical-parity, full-budget +quality, performance and regression evidence. The goal is not complete while any +required capability or qualification artifact is unimplemented or unverified. + +Current evidence status (2026-10-20): packages 1–4 and 6 have implementation, +focused runtime/tests and local-artifact evidence in the linked reports. Package 5 +has executed end-to-end flows, numerical parity, 40 bounded training cells and +both-domain short rollouts. Its full-size data, declared full training budget, +2048-trajectory long-rollout and final acceptance gates remain pending. The +controlled comparison is not substituted for literal paper-recipe reproduction. diff --git a/design/reports/nerd-ant-parity-20261020.json b/design/reports/nerd-ant-parity-20261020.json new file mode 100644 index 000000000..daaa93d53 --- /dev/null +++ b/design/reports/nerd-ant-parity-20261020.json @@ -0,0 +1,26 @@ +{ + "robot": "Ant", + "checkpoint_sha256": "be7e29f8834ab22d76efd3d101196c77b3b415c3ef91af48aed233c01eb4e950", + "max_forward_error": 0.0, + "max_input_gradient_error": 0.0, + "max_parameter_gradient_error": 0.0, + "max_adam_parameter_error": 0.0, + "max_codec_feature_error": 1.430511474609375e-06, + "max_codec_target_error": 2.5192275643348694e-07, + "max_codec_decode_error": 2.980232238769531e-07, + "reference_commit": "476dc67fccc6a99be293bf611f897b1c06f13e30", + "precision": "float32", + "device": "cpu", + "torch_version": "2.7.1+cu128", + "python_version": "3.11.14", + "history_lengths": [ + 1, + 9, + 10 + ], + "adam_updates": 3, + "atol": 1e-05, + "rtol": 0.0001, + "qualification_scope": "network, gradient, short optimizer and codec parity; not full-budget training quality", + "codec_fixtures_sha256": "760282a0081d74f85e106ac826c287d81f1be4371a77ab47dfc02b2b94a323e6" +} diff --git a/design/reports/nerd-cartpole-parity-20261020.json b/design/reports/nerd-cartpole-parity-20261020.json new file mode 100644 index 000000000..84365cbcf --- /dev/null +++ b/design/reports/nerd-cartpole-parity-20261020.json @@ -0,0 +1,26 @@ +{ + "robot": "Cartpole", + "checkpoint_sha256": "287dde2134eb8f5277c17a8f7cc7abbe64063e97cb481bdb6e8d76f9536e5366", + "max_forward_error": 0.0, + "max_input_gradient_error": 0.0, + "max_parameter_gradient_error": 0.0, + "max_adam_parameter_error": 0.0, + "max_codec_feature_error": 1.1920928955078125e-07, + "max_codec_target_error": 1.1920928955078125e-07, + "max_codec_decode_error": 1.1920928955078125e-07, + "reference_commit": "476dc67fccc6a99be293bf611f897b1c06f13e30", + "precision": "float32", + "device": "cpu", + "torch_version": "2.7.1+cu128", + "python_version": "3.11.14", + "history_lengths": [ + 1, + 9, + 10 + ], + "adam_updates": 3, + "atol": 1e-05, + "rtol": 0.0001, + "qualification_scope": "network, gradient, short optimizer and codec parity; not full-budget training quality", + "codec_fixtures_sha256": "53574a5cd3a47f4e7acf4eb72aedf64976206057c4688799bdfbf5eaed03d07c" +} diff --git a/design/reports/neural-dynamics-20261020.md b/design/reports/neural-dynamics-20261020.md new file mode 100644 index 000000000..4934ca8bc --- /dev/null +++ b/design/reports/neural-dynamics-20261020.md @@ -0,0 +1,290 @@ +# NeRD implementation and training evidence — 2026-10-20 + +**Status: implementation/short-run evidence exists; full official-training +qualification remains pending.** No 100-update run below replaces the declared +4,995,000-update budget, five paired seeds, or full rollout evaluation. + +## Reference and execution identity + +- Official source: `NVlabs/neural-robot-dynamics@370e89d2b8ce68aa3578aa7812859c681b7502bc`, + untouched checkout `/tmp/embodichain-nerd-main-reference`. +- Official data: original `NeuralEnvironment(default_env_mode="ground-truth")`, + `TrajectorySampler` / `TrajectorySamplerAnt`, sampling ranges and Featherstone + integration. Cartpole uses isolated Warp **1.8.0**, five substeps; Ant uses the + paper baseline's Warp **1.5.1**, sixteen substeps. Both macro timesteps are 1/60 s. +- The original `VanillaTrainer.preprocess_data_batch`, `compute_loss`, `one_epoch` + and `get_scheduled_learning_rate`, original `NeuralIntegrator` math and original + `ModelMixedInput` execute directly. No official source methods are replaced. + The wrapper binds common initialization, lazy data, frozen statistics, logging + and validation evaluation; it bypasses the constructor's eager data allocation. +- Current teacher: DexSim 0.5.1 development build, Newton 1.6.0, Warp 1.17.0, + MuJoCo CPU with separate worlds, ten substeps. Installed artifact identity is + recorded in each dataset manifest; this is not a claim about the rc1 wheel. +- Training/benchmark hardware: NVIDIA GeForce RTX 5090; driver **595.58.03**, + PyTorch **2.7.1+cu128**, CUDA runtime **12.8**, Python **3.11.14**. + +## Official checkpoint and numerical parity + +The separate Newton-preview reference commit is +`476dc67fccc6a99be293bf611f897b1c06f13e30`. Both official pretrained checkpoints +were checked on CPU in FP32 with history lengths 1, 9 and 10. Maximum forward, +input-gradient, parameter-gradient and three-Adam-update parameter errors were +**zero** for both robots. Maximum codec errors: + +| Robot | Features | Targets | Decoded state | +|---|---:|---:|---:| +| Cartpole | 1.19e-7 | 1.19e-7 | 1.19e-7 | +| Ant | 1.43e-6 | 2.52e-7 | 2.98e-7 | + +These passed the declared `atol=1e-5`, `rtol=1e-4` tolerances. They verify the +checkpoint/network/codec port for these fixtures, not general training quality. +Checked-in receipts retain checkpoint and fixture hashes: +[Cartpole](nerd-cartpole-parity-20261020.json), +[Ant](nerd-ant-parity-20261020.json). +The full fixtures and portable exports remain under +`/tmp/embodichain-nerd-parity-20261020/{cartpole,ant}`. +Reproduce each with: + +```bash +OMP_NUM_THREADS=2 MKL_NUM_THREADS=2 python tests/learning/models/_nerd_reference_check.py /path/to/pinned-newton-preview Cartpole /path/to/new-output +``` + +Use `Ant` for the other checkpoint. This runner executes explicitly selected +trusted reference code; it does not download or import an arbitrary repository. + +## Genuine teacher data and explicit compatibility transport + +Both official robots produced **48 real trajectories × 100 transitions**: +32 train, 8 validation and 8 test. Every raw original HDF5 artifact is retained. +Canonical conversion preserves the declared split and checks within-trajectory +continuity. Dataset paths: + +- `/tmp/nerd-official-real-cartpole-20261020/canonical` +- `/tmp/nerd-official-real-ant-20261020/canonical` + +Each parent directory contains `train.raw.h5`, `validation.raw.h5`, `test.raw.h5`, +`source-physics.json` and `generation-result.json`, including source/file hashes. +The production wrapper now streams one sampler batch at a time; a separately +executed regression verifies multiple batches and a partial final batch. + +**Confirmed official-source bug:** `TrajectorySampler` allocates `gravity_dir` +with `torch.empty` and initializes only the up-axis component. The actual CPU +probe produced NaNs and values around 849,000 in the other components. The +collection explicitly enables `repair_uninitialized_gravity: true`; raw HDF5 +values remain unchanged. Only the canonical gravity field is replaced by the +known physical direction, with changed/nonfinite counts recorded per split. +These datasets are therefore labeled **source-bug-corrected official teacher +inputs**, not byte-for-byte reproductions of an unspecified published dataset. + +Both original environments are Y-up. Explicit world-frame transport rotates Ant +by +90° about X; Cartpole additionally rotates +90° about Z to match the current +reference asset placement. Base position/orientation, COM linear/angular velocity, +world contact points/normals and gravity rotate together. Body-local contact +points remain local. Angular-first world-origin twist is converted to native +linear-first COM velocity. Source-coordinate metadata and rotation are retained. + +### Identical asset bytes do not imply identical physical systems + +Actual initialized model arrays, not XML declarations alone, give these values: + +| Property | Official Warp teacher | Current Newton asset import | +|---|---:|---:| +| Cartpole cart/pole mass | 50 / 2.5 | 1 / 10 | +| Ant torso mass | 31.6534519 | 0.483880247 | +| Ant upper/lower leg mass | 7.83155060 / 13.5184412 | 0.0391577530 / 0.0675922039 | +| Ant robot-shape friction | 0.75 | 1.5 | +| Ant actuated armature | 0.01 | 0.01 | +| Official ground friction | 0.5 | Not independently measured in this comparison | + +The old importers' inertia/density handling differs. Both versions use the same +Ant asset SHA256 `23e9a720465a71fde3fd9fd2d5a18f8c11ec642f3dca8921c66cb4e169421fa7` +and Cartpole SHA256 +`f31276531c03050a3ac3a3e24443899f38a4424586db17140978383433dec237`. +Original Ant contact thickness is radius (0.25 / 0.08) plus a ground thickness +around 1e-5; current Newton contact feature construction includes its configured +margins. Full masses, inertia, COM, joint parameters and materials remain in +`source-physics.json`. + +**Deployment implication:** an official checkpoint cannot be declared physically +compatible using only asset SHA or joint count. Preserve its old physical profile +or retrain against the current teacher. Do not replace an official-data manifest +with current-Newton physics identity merely to bypass deployment validation. + +## Executed real-data matched training matrix + +All **8 cells completed**: 2 robots × 2 teacher datasets × 2 trainers, **one seed +(0), 100 Adam updates per cell**, batch 512, history 10, six layers, twelve heads, +embedding 192 for Cartpole / 384 for Ant. Both trainers within a row share exact +initial weights, train-only statistics, explicit batch-index order, split and +validation selection. These runs use a deliberately short learning-rate schedule; +they are not the first 100 steps of the full epoch schedule. + +Official-data rows use the genuine 48×100 datasets above. Current-Newton rows use +64×32-transition smoke datasets, with observed split counts **57/1/6**. Their +initial states are defaults; Cartpole uses full force scale, Ant uses **0.05** of +the full force scale. The small validation/test sets limit any accuracy inference. + +| Robot | Teacher | Trainer | Test q error norm | Test qd error norm | Worker wall s | +|---|---|---|---:|---:|---:| +| Cartpole | Official Warp | Official | 0.031869880 | 0.307753670 | 29.403 | +| Cartpole | Official Warp | EmbodiChain | 0.031835163 | 0.307730865 | 29.094 | +| Cartpole | Current Newton | Official | 0.028611711 | 0.708493175 | 30.270 | +| Cartpole | Current Newton | EmbodiChain | 0.028664216 | 0.700447635 | 29.686 | +| Ant | Official Warp | Official | 0.077358087 | 7.257187259 | 33.837 | +| Ant | Official Warp | EmbodiChain | 0.077455009 | 7.250118046 | 32.056 | +| Ant | Current Newton | Official | 0.090978993 | 7.981286798 | 34.236 | +| Ant | Current Newton | EmbodiChain | 0.091046900 | 7.983307150 | 33.434 | + +Ant base position error (m) / orientation error (rad): official-Warp row, +0.005254870 / 0.018305524 (official) versus 0.005307779 / 0.018332163 (migrated); +current-Newton row, 0.004343362 / 0.015835387 versus 0.004330985 / 0.015851749. + +Errors are within-row matched comparisons. Cartpole q and generalized qd norms +combine coordinates with different units and must not be described as pure +Cartesian distances. A one-seed result has no five-seed confidence interval and +cannot establish the predeclared 5% noninferiority gate. Rounded values above +are presentation only; machine reports retain full precision. + +Evidence directories (each includes immutable experiment/run specs, common +weights/statistics/order hashes, subprocess commands, stdout/stderr, result +receipts, exports and generated plots): + +- `/tmp/nerd-real-matched-cartpole-20261020/report/REPORT.md` +- `/tmp/nerd-real-matched-ant-20261020/report/REPORT.md` +- Compact precise metrics: `/tmp/nerd-real-matched-summary-20261020.json` + +The report states `training_matrix_status: completed` and overall qualification +`pending`. Earlier synthetic-fixture tests exercised infrastructure only and are +not included in the table above. + +## Measured training-budget feasibility + +Device-resident FP32, TF32 disabled, batch 512, history 10; 20 warmup updates and +100 synchronized timed updates per implementation/robot. The operation includes +forward, normalized MSE, backward, norm clipping and Adam. Initialization and +first-update latency are reported separately. Final parameters match exactly +between models for this seeded compute-only workload. + +| Robot | Implementation | Median ms/update | One 4,995,000-update run, hours | +|---|---|---:|---:| +| Cartpole | Official | 7.539 | 10.460 | +| Cartpole | EmbodiChain | 7.495 | 10.399 | +| Ant | Official | 14.488 | 20.103 | +| Ant | EmbodiChain | 14.263 | 19.790 | + +The declared 40-run matrix (2 robots × 2 datasets × 2 trainers × 5 seeds) has an +**optimistic compute-only lower bound of 607.52 GPU-hours / 25.31 days on one +RTX 5090**. Peak allocated GPU memory was at most 1.08 GB for the harness. +This excludes data loading, encoding, validation, checkpoints, teacher generation +and rollout evaluation. No 40-run full-budget job was launched. + +The current input pipeline is a material bottleneck: independent CPU timing of +HDF5 read + collate + encode, batch 512/history 10, measured medians **259 ms for +Cartpole and 295 ms for Ant**. cProfile attributes most time to individual HDF5 +window/group access. This agrees with the 29–34 s observed worker times for 100 +updates, including startup and evaluation. Full-data timing is unmeasured; +batched reads/caching/prefetch need representative validation before a practical +full-budget duration can be promised. + +Artifacts: `/tmp/nerd-full-budget-estimate-20261020.json` and +`/tmp/nerd-input-pipeline-profile-20261020.json`. + +## Reproduction and remaining gates + +Use a separate pinned reference checkout and isolated dependencies; never replace +EmbodiChain's installed Warp: + +```bash +GIT_LFS_SKIP_SMUDGE=1 git -C /path/to/neural-robot-dynamics worktree add --detach /tmp/nerd-main-reference 370e89d2b8ce68aa3578aa7812859c681b7502bc +python -m pip install --target /tmp/nerd-warp18 --no-deps warp-lang==1.8.0 +python -m pip install --target /tmp/nerd-warp151 --no-deps warp-lang==1.5.1 +PYTHONPATH=/tmp/nerd-warp18:/tmp/nerd-main-reference:$PWD python -m embodichain.learning.dynamics.reference --spec generation.json --trust-external-code --generate +python -m embodichain.learning.dynamics.comparison --prepare comparison.json +python -m embodichain.learning.dynamics.comparison --run /path/to/experiment.json --trust-external-code +python -m embodichain.learning.dynamics.report --experiment /path/to/experiment.json +python -m scripts.benchmark.learning.dynamics_training --source /tmp/nerd-main-reference --trust-external-code --output budget.json +``` + +For Ant generation, use the Warp 1.5.1 PYTHONPATH and `warp_version: "1.5.1"`. +Generation specs used here are `/tmp/nerd-official-generate-{cartpole,ant}.json`; +exact training specs and commands are stored under the matrix directories. +Reference code execution and gravity repair require explicit opt-ins. The source +checkout remains unchanged; its hashes and the adapter identity are recorded. + +Still pending: full-size data distributions, five paired seeds at the declared +full budget, numerical qualification reports attached to final runs, all matched +2048-trajectory horizons/scenarios, fixed-policy/stability gates, physical-profile +compatible official-model deployment, and representative end-to-end throughput. +User budget guidance is needed before launching a multiweek full matrix. Until +those gates are backed by executed artifacts, the original qualification objective +remains open. + +## Integration and packaging validation + +Focused EmbodiChain regression groups completed after integration: + +- CLI, simulation manager/configuration and Gym construction: **398 passed, + 3 skipped**. +- Articulation drive/state compatibility, official task layout and package data: + **23 passed**. +- Context schema/routing: **40 passed**. The changed contracts have a new neural + dynamics owner and links from simulation, environment and learning topics. + Broad watches also match Task Programs, Gen Sim and existing asset/RL context; + their unchanged contracts do not require unrelated rewrites. +- API coverage: **2747/2747** public exports documented; checker regressions: + **9 passed**. +- Black **26.3.1** applied to the repository; changed Python headers, future + imports and public-module export declarations passed structural checks. +- Built an EmbodiChain wheel and inspected the NeRD entry point, modules, + third-party licenses, robot assets and all 16 workflow YAML files. + +The full Sphinx dummy build is **not verified**: the existing architecture +extension refuses generation while relevant source files are uncommitted. Its +failure occurred before page rendering. No source was committed merely to bypass +that precondition. + +An isolated Sphinx dummy build of the new neural-dynamics guide and API sections +passed with `-W` and no warnings. This focused check excludes the architecture +extension and does not replace the full repository build. + +The normal dependency pin still names `dexsim_engine==0.5.1rc1`; it must be +updated with the eventual extension-enabled release. Until then this is a paired +local-development deployment, not an install-ready claim about that published +pin. The tested local DexSim artifact and source/build hashes are supplied by +its adjacent implementation workstream. Native teardown warnings also occurred +in the clean DexSim baseline and are not claimed fixed by this work. + +A dependency-consistent local wheel pair is available. The EmbodiChain wheel was +repacked with only its `Requires-Dist` DexSim pin changed to the exact tested local +version and its `RECORD` regenerated; all runtime files match the source-built +wheel. This overlay is not a source dependency release update. Paths, hashes and +the exact metadata change are recorded in +[the local artifact manifest](neural-dynamics-local-artifacts-20261020.json). +Install both listed paired wheels together in an environment with the ordinary +project dependencies. Installing the unmodified editable checkout instead still +leaves its rc1 requirement unsatisfied; `--no-deps` does not make that conflict +disappear. + +## Subsequent completion audit (implementation still in progress) + +The previous local wheel pair was independently installed and exercised: +**156 tests passed, 21 deselected**, with import paths and dependency matching +recorded in DexSim's `validation-content/manifest.json`. This verifies that +specific snapshot, not later source edits. + +A subsequent runtime audit reproduced missing fixed-root-frame and Ant contact/ +collision fields in the physical signature. Those omissions can accept a model +for a changed physical system. Their corrections and regenerated simulation +fixtures are required before the current worktree or old portable exports can +be considered qualified. Old fingerprints must not be relabeled to bypass the +new checks. Data-input batching and report provenance checks are also undergoing +focused implementation and validation; rebuild the paired artifact after those +changes settle. + +The report reader now accepts the measured numerical receipts directly. Both +historical 8-cell experiments were re-evaluated under `report-audit/`: numerical +parity **passed**, recorded teacher provenance verified, overall qualification +still **pending**. Rollout attachments must prove model/manifest/source/trace +hashes, the test split, teacher domain and identical selected teacher content; +fixed-policy pairs must identify the same policy. Bare `passed: true` summaries +and mismatched rollout populations cannot establish qualification. diff --git a/design/reports/neural-dynamics-bounded-five-seed-20261020.md b/design/reports/neural-dynamics-bounded-five-seed-20261020.md new file mode 100644 index 000000000..bbdcaa829 --- /dev/null +++ b/design/reports/neural-dynamics-bounded-five-seed-20261020.md @@ -0,0 +1,85 @@ +# Five-seed bounded NeRD comparison — 2026-10-20 + +**Status: all 40 short training cells completed; full qualification remains pending.** Each cell ran 100 updates with batch 512/history 10. This does not replace the 4,995,000-update budget or full-size data/evaluation. + +## Scope and reproducibility + +- Two robots × two teacher domains × two trainers × seeds 0–4 = 40 cells, 4,000 optimizer updates total. +- Networks use the official 192/384 embedding, six layers, twelve heads and 64-unit output head in FP32. Actual initial artifacts preserve the required feature order; sorted plan mappings are not used to infer it. +- Current Newton datasets: 64 trajectories × 32 transitions, splits 57/1/6, default initial states; Ant collection uses effort scale 0.05. Historical datasets: 48 × 100, splits 32/8/8, measured gravity and historical physical-v2 provenance. +- Each paired row shares exact initial weights, statistics, indices and checkpoint-selection controls. Both teacher origins and source hashes are preserved. +- This is controlled matched-data comparison, not literal paper-recipe reproduction. See [the protocol audit](neural-dynamics-official-protocol-audit-20261020.md) for sampling/statistics/validation differences. +- Observed sequential matrix time was about 98 s for Cartpole and 139 s for Ant on this workstation, excluding preparation. This small-data timing must not be extrapolated to the full corpus. + +## Paired offline results + +The table shows the one-sided 95% paired bootstrap upper bound on relative degradation, using five seeds and 10,000 resamples. Every short-run offline metric is below the predeclared 5% margin. These conditional results do not establish final-budget quality. q/qd norms can combine coordinates with different units; they are not Cartesian distances. + +| Robot | Teacher | Metric | Upper bound | Short-run criterion | +|---|---|---|---:|---| +| cartpole | official_warp | q_error_norm | 0.7563% | passed | +| cartpole | official_warp | qd_error_norm | -0.0300% | passed | +| cartpole | current_newton | q_error_norm | 0.0609% | passed | +| cartpole | current_newton | qd_error_norm | 2.1676% | passed | +| ant | official_warp | q_error_norm | 0.0786% | passed | +| ant | official_warp | qd_error_norm | 0.0514% | passed | +| ant | official_warp | base_position_error | 0.9660% | passed | +| ant | official_warp | base_orientation_error_rad | 0.6761% | passed | +| ant | current_newton | q_error_norm | 0.0042% | passed | +| ant | current_newton | qd_error_norm | 0.0515% | passed | +| ant | current_newton | base_position_error | -0.1524% | passed | +| ant | current_newton | base_orientation_error_rad | 0.1560% | passed | + +The largest upper bound is about **2.17%**, for current-Newton Cartpole velocity error. This is evidence of short-run port agreement, not a claim that the learned dynamics are accurate enough for deployment. + +## Bounded closed-loop and rollout evidence + +Seed 0 was additionally evaluated in both domains: 24 measured runs, three trajectories each, horizons 1/10/12, two environments with a partial final batch. All 12 paired teacher-content and policy identities match. Source, weights, manifest, HDF5 and JSONL hashes were verified for this bounded scope. The formal report still requires 2048 trajectories, horizon 1000 and five rollout seeds. + +Historical execution uses the actual original Warp environment/integrator; current-domain execution uses the registered Newton provider. See [historical rollout results](neural-dynamics-historical-rollout-20261020.md). The models are not numerically identical and neither domain is declared physically qualified. + +| Current-Newton scenario | Learned failure rate, both trainers | Teacher failure rate | +|---|---:|---:| +| Cartpole passive | 0.6667 | 0.3333 | +| Cartpole random_effort | 0.6667 | 0.6667 | +| Cartpole fixed_policy | 0.3333 | 0.0000 | + +For historical Ant fixed PD, teacher return was 9.7618 versus 1.4588/1.5726 for the official/migrated trained models. Current-Newton Ant returned 4.3053 versus 2.2850/2.3616. These are material teacher errors despite close trainer-to-trainer results. Cartpole negative returns also stop at failure, so a less negative return must not be interpreted as superior control while ignoring the failure count. + +## Evidence and remaining requirements + +- [Machine-readable five-seed summary](neural-dynamics-five-seed-20261020.json). +- [Verified bounded attachments and paired identities](neural-dynamics-bounded-pairing-20261020.json). +- Specs, command/stdout/stderr receipts, common weights/statistics/order and training/paired-error plots are retained under `/tmp/nerd-five-seed-{cartpole,ant}-20261020`. +- Raw current-domain rollout traces are under `/tmp/nerd-current-paired-{cartpole,ant}-20261020`; historical traces are under the matching `nerd-historical-paired-*` directories. +- The exact pre-runtime-metadata reference/workflow sources used by these runs are preserved with hashes in `/tmp/nerd-five-seed-source-snapshot-20261020`. +- Future runs now capture per-worker runtime/driver identity. Existing 40 receipts were not backfilled; they retain their original Torch/device/source records and the documented workstation context. + +Remaining: full-size data, declared full training budget, complete long-rollout and stability/return gates, and an explicitly identified paper data/validation recipe before any literal reproduction claim. No full-budget job was launched. + + +## Latest installation and code checks + +The latest paired local EC/SDK wheel is identified by +[its artifact manifest](neural-dynamics-local-artifacts-historical-20261020.json). +Actual imports, named-provider metadata and exact dependency matching were checked. +The installed EC package ran both historical CPU cases without importing or +attempting to import DexSim/Newton in those isolated processes. + +One tiny installed CUDA probe per historical robot also completed on RTX 5090: +B1/count1, passive/default initialization, horizons 1/3. Cartpole Warp 1.8.0 +portable/original maximum absolute error was zero; Ant Warp 1.5.1 was 1.4782e-5, +within atol=1e-5/rtol=1e-4. Both physical contracts matched. Old Warp emitted a +cuDeviceGetUuid initialization diagnostic, retained in the logs. This is a +bounded version-compatibility check, not a broad GPU or long-horizon qualification. + +Focused reference/report/comparison/shared-metric tests: 36 passed with actual +legacy references enabled. Runtime-provenance tests: 31 passed; an actual CUDA +identity-query smoke also passed without changing Torch RNG state. Context tests: +40 passed. Public API coverage: 2748/2748. Updated neural-dynamics Sphinx pages +built with warnings treated as errors. Black 26.3.1 and diff checks passed. + +No formal-budget or long-evaluation job is running. Full acceptance remains +pending the full experiment scope and compute-budget decision; the existing +budget question has not been answered. Code and bounded evidence are available +for review while that decision is outstanding. diff --git a/design/reports/neural-dynamics-bounded-pairing-20261020.json b/design/reports/neural-dynamics-bounded-pairing-20261020.json new file mode 100644 index 000000000..0634cb2de --- /dev/null +++ b/design/reports/neural-dynamics-bounded-pairing-20261020.json @@ -0,0 +1,611 @@ +{ + "scope": "bounded integrity and pairing only; not formal rollout qualification", + "robots": { + "cartpole": { + "checks": [ + { + "run_id": "official_warp-official-0", + "mode": "passive", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-cartpole-20261020/bounded-rollout-attachments/official_warp-official-0-passive.json" + }, + { + "run_id": "official_warp-official-0", + "mode": "random_effort", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-cartpole-20261020/bounded-rollout-attachments/official_warp-official-0-random_effort.json" + }, + { + "run_id": "official_warp-official-0", + "mode": "fixed_policy", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-cartpole-20261020/bounded-rollout-attachments/official_warp-official-0-fixed_policy.json" + }, + { + "run_id": "official_warp-embodichain-0", + "mode": "passive", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-cartpole-20261020/bounded-rollout-attachments/official_warp-embodichain-0-passive.json" + }, + { + "run_id": "official_warp-embodichain-0", + "mode": "random_effort", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-cartpole-20261020/bounded-rollout-attachments/official_warp-embodichain-0-random_effort.json" + }, + { + "run_id": "official_warp-embodichain-0", + "mode": "fixed_policy", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-cartpole-20261020/bounded-rollout-attachments/official_warp-embodichain-0-fixed_policy.json" + }, + { + "run_id": "current_newton-official-0", + "mode": "passive", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-cartpole-20261020/bounded-rollout-attachments/current_newton-official-0-passive.json" + }, + { + "run_id": "current_newton-official-0", + "mode": "random_effort", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-cartpole-20261020/bounded-rollout-attachments/current_newton-official-0-random_effort.json" + }, + { + "run_id": "current_newton-official-0", + "mode": "fixed_policy", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-cartpole-20261020/bounded-rollout-attachments/current_newton-official-0-fixed_policy.json" + }, + { + "run_id": "current_newton-embodichain-0", + "mode": "passive", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-cartpole-20261020/bounded-rollout-attachments/current_newton-embodichain-0-passive.json" + }, + { + "run_id": "current_newton-embodichain-0", + "mode": "random_effort", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-cartpole-20261020/bounded-rollout-attachments/current_newton-embodichain-0-random_effort.json" + }, + { + "run_id": "current_newton-embodichain-0", + "mode": "fixed_policy", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-cartpole-20261020/bounded-rollout-attachments/current_newton-embodichain-0-fixed_policy.json" + } + ], + "all_pairs_match": true, + "full_report_status": "pending", + "report": "/tmp/nerd-five-seed-cartpole-20261020/report-with-bounded-rollouts/REPORT.md", + "offline_comparisons": [ + { + "teacher": "official_warp", + "metric": "q_error_norm", + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "official": [ + 0.03186987992171403, + 0.02983515760400793, + 0.027172292981828963, + 0.030564142583490727, + 0.028380315382402023 + ], + "embodichain": [ + 0.03183516303261558, + 0.02994343474671081, + 0.027537767179719695, + 0.030569061866173378, + 0.028339288522908977 + ], + "gate": { + "status": "passed", + "seeds": 5, + "mean_difference": 0.0029410878709215105, + "upper_95": 0.0075631630189805675, + "threshold": 0.05, + "difference_mode": "relative", + "resamples": 10000, + "seed": 0 + } + }, + { + "teacher": "official_warp", + "metric": "qd_error_norm", + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "official": [ + 0.30775366982260904, + 0.3291599650959392, + 0.31195038177154877, + 0.33893180260291467, + 0.33426121303013395 + ], + "embodichain": [ + 0.30773086547851564, + 0.32897301139412344, + 0.30997772216796876, + 0.33886656708769747, + 0.33389833890474757 + ], + "gate": { + "status": "passed", + "seeds": 5, + "mean_difference": -0.0016487556883492814, + "upper_95": -0.0003000744695058809, + "threshold": 0.05, + "difference_mode": "relative", + "resamples": 10000, + "seed": 0 + } + }, + { + "teacher": "current_newton", + "metric": "q_error_norm", + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "official": [ + 0.028611711142719657, + 0.03116420524707739, + 0.029795721302861754, + 0.04210370381673177, + 0.03285226683685745 + ], + "embodichain": [ + 0.028664215751316235, + 0.030975601638572804, + 0.029802200759666553, + 0.036095138218091884, + 0.03286314701688462 + ], + "gate": { + "status": "passed", + "seeds": 5, + "mean_difference": -0.029275385391873097, + "upper_95": 0.0006092183787881106, + "threshold": 0.05, + "difference_mode": "relative", + "resamples": 10000, + "seed": 0 + } + }, + { + "teacher": "current_newton", + "metric": "qd_error_norm", + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "official": [ + 0.7084931746773098, + 0.7104982680168704, + 0.7293366805366848, + 0.7345860135728034, + 0.7366227301998415 + ], + "embodichain": [ + 0.7004476353742074, + 0.6937628704568614, + 0.7296680395153986, + 0.7696718962296195, + 0.7347368323284647 + ], + "gate": { + "status": "passed", + "seeds": 5, + "mean_difference": 0.002149326603464712, + "upper_95": 0.021675618573048703, + "threshold": 0.05, + "difference_mode": "relative", + "resamples": 10000, + "seed": 0 + } + } + ] + }, + "ant": { + "checks": [ + { + "run_id": "official_warp-official-0", + "mode": "passive", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-ant-20261020/bounded-rollout-attachments/official_warp-official-0-passive.json" + }, + { + "run_id": "official_warp-official-0", + "mode": "random_effort", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-ant-20261020/bounded-rollout-attachments/official_warp-official-0-random_effort.json" + }, + { + "run_id": "official_warp-official-0", + "mode": "fixed_policy", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-ant-20261020/bounded-rollout-attachments/official_warp-official-0-fixed_policy.json" + }, + { + "run_id": "official_warp-embodichain-0", + "mode": "passive", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-ant-20261020/bounded-rollout-attachments/official_warp-embodichain-0-passive.json" + }, + { + "run_id": "official_warp-embodichain-0", + "mode": "random_effort", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-ant-20261020/bounded-rollout-attachments/official_warp-embodichain-0-random_effort.json" + }, + { + "run_id": "official_warp-embodichain-0", + "mode": "fixed_policy", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-ant-20261020/bounded-rollout-attachments/official_warp-embodichain-0-fixed_policy.json" + }, + { + "run_id": "current_newton-official-0", + "mode": "passive", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-ant-20261020/bounded-rollout-attachments/current_newton-official-0-passive.json" + }, + { + "run_id": "current_newton-official-0", + "mode": "random_effort", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-ant-20261020/bounded-rollout-attachments/current_newton-official-0-random_effort.json" + }, + { + "run_id": "current_newton-official-0", + "mode": "fixed_policy", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-ant-20261020/bounded-rollout-attachments/current_newton-official-0-fixed_policy.json" + }, + { + "run_id": "current_newton-embodichain-0", + "mode": "passive", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-ant-20261020/bounded-rollout-attachments/current_newton-embodichain-0-passive.json" + }, + { + "run_id": "current_newton-embodichain-0", + "mode": "random_effort", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-ant-20261020/bounded-rollout-attachments/current_newton-embodichain-0-random_effort.json" + }, + { + "run_id": "current_newton-embodichain-0", + "mode": "fixed_policy", + "bounded_integrity": "available", + "formal_scope": "pending", + "path": "/tmp/nerd-five-seed-ant-20261020/bounded-rollout-attachments/current_newton-embodichain-0-fixed_policy.json" + } + ], + "all_pairs_match": true, + "full_report_status": "pending", + "report": "/tmp/nerd-five-seed-ant-20261020/report-with-bounded-rollouts/REPORT.md", + "offline_comparisons": [ + { + "teacher": "official_warp", + "metric": "q_error_norm", + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "official": [ + 0.07735808655455872, + 0.07836143787090595, + 0.08000357910826966, + 0.08015401232373584, + 0.07747088841029576 + ], + "embodichain": [ + 0.07745500920893071, + 0.07813525776286702, + 0.08003018976567866, + 0.08014878493088942, + 0.07752861819424473 + ], + "gate": { + "status": "passed", + "seeds": 5, + "mean_difference": -0.0001241758157730334, + "upper_95": 0.0007861924289500945, + "threshold": 0.05, + "difference_mode": "relative", + "resamples": 10000, + "seed": 0 + } + }, + { + "teacher": "official_warp", + "metric": "qd_error_norm", + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "official": [ + 7.257187258542239, + 7.171238222226992, + 7.083392132769574, + 7.319544664320055, + 7.1057311341002745 + ], + "embodichain": [ + 7.250118046016484, + 7.1802550598815245, + 7.0804086538461535, + 7.3203793033138735, + 7.1045006385216345 + ], + "gate": { + "status": "passed", + "seeds": 5, + "mean_difference": -3.941419260966793e-05, + "upper_95": 0.0005139221681135962, + "threshold": 0.05, + "difference_mode": "relative", + "resamples": 10000, + "seed": 0 + } + }, + { + "teacher": "official_warp", + "metric": "base_position_error", + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "official": [ + 0.005254870456653637, + 0.005087500090127463, + 0.0054122172869168796, + 0.005360990304213303, + 0.005156655625982599 + ], + "embodichain": [ + 0.005307778945335975, + 0.005091159946315891, + 0.005486480482331999, + 0.0053636068826193335, + 0.005174731422256638 + ], + "gate": { + "status": "passed", + "seeds": 5, + "mean_difference": 0.005700531791566215, + "upper_95": 0.00965982291230781, + "threshold": 0.05, + "difference_mode": "relative", + "resamples": 10000, + "seed": 0 + } + }, + { + "teacher": "official_warp", + "metric": "base_orientation_error_rad", + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "official": [ + 0.018305523840935676, + 0.01898016038831774, + 0.0195915101648687, + 0.01981754826975393, + 0.017857071593567565 + ], + "embodichain": [ + 0.01833216279417604, + 0.018993363013634315, + 0.01980317451141693, + 0.01985698741870922, + 0.017856536068759123 + ], + "gate": { + "status": "passed", + "seeds": 5, + "mean_difference": 0.002982969323553847, + "upper_95": 0.006760569207506567, + "threshold": 0.05, + "difference_mode": "relative", + "resamples": 10000, + "seed": 0 + } + }, + { + "teacher": "current_newton", + "metric": "q_error_norm", + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "official": [ + 0.09097899284915648, + 0.08960340748662533, + 0.09382180476534194, + 0.09231608294058537, + 0.09260715539904608 + ], + "embodichain": [ + 0.09104689998903136, + 0.08955901878467505, + 0.09358879863352015, + 0.092288915661798, + 0.09246751536493716 + ], + "gate": { + "status": "passed", + "seeds": 5, + "mean_difference": -0.0008069286806190477, + "upper_95": 4.154844541429713e-05, + "threshold": 0.05, + "difference_mode": "relative", + "resamples": 10000, + "seed": 0 + } + }, + { + "teacher": "current_newton", + "metric": "qd_error_norm", + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "official": [ + 7.981286798007246, + 7.992876754981884, + 8.10761081861413, + 7.940719825634058, + 8.12097698709239 + ], + "embodichain": [ + 7.98330715013587, + 7.990900985054348, + 8.114658344655798, + 7.943337437726449, + 8.113467362998188 + ], + "gate": { + "status": "passed", + "seeds": 5, + "mean_difference": 5.6023575830166544e-05, + "upper_95": 0.0005148825688513269, + "threshold": 0.05, + "difference_mode": "relative", + "resamples": 10000, + "seed": 0 + } + }, + { + "teacher": "current_newton", + "metric": "base_position_error", + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "official": [ + 0.004343361785446388, + 0.0048682841701784, + 0.0042658142421556555, + 0.004624356739762901, + 0.0050422865411509634 + ], + "embodichain": [ + 0.004330985442451809, + 0.00485701422760452, + 0.0042299236076465555, + 0.004623724757761195, + 0.005029605091481969 + ], + "gate": { + "status": "passed", + "seeds": 5, + "mean_difference": -0.00324593774083294, + "upper_95": -0.0015236678033083008, + "threshold": 0.05, + "difference_mode": "relative", + "resamples": 10000, + "seed": 0 + } + }, + { + "teacher": "current_newton", + "metric": "base_orientation_error_rad", + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "official": [ + 0.01583538746488267, + 0.016103356126425925, + 0.015208294080651325, + 0.01632652559142182, + 0.01621722622194152 + ], + "embodichain": [ + 0.015851749199024146, + 0.016095023224319237, + 0.015240774292876754, + 0.01631476360818614, + 0.01624640450961348 + ], + "gate": { + "status": "passed", + "seeds": 5, + "mean_difference": 0.0007460519254374892, + "upper_95": 0.0015601198654396253, + "threshold": 0.05, + "difference_mode": "relative", + "resamples": 10000, + "seed": 0 + } + } + ] + } + } +} diff --git a/design/reports/neural-dynamics-five-seed-20261020.json b/design/reports/neural-dynamics-five-seed-20261020.json new file mode 100644 index 000000000..b294eda21 --- /dev/null +++ b/design/reports/neural-dynamics-five-seed-20261020.json @@ -0,0 +1,455 @@ +{ + "scope": "five-seed, 100-update controlled matched-data experiment; full qualification pending", + "runs": 40, + "updates_per_run": 100, + "robots": { + "cartpole": { + "experiment": "/tmp/nerd-five-seed-cartpole-20261020/experiment.json", + "report": "/tmp/nerd-five-seed-cartpole-20261020/report-with-bounded-rollouts/REPORT.md", + "status": "pending", + "comparisons": [ + { + "embodichain": [ + 0.03183516303261558, + 0.02994343474671081, + 0.027537767179719695, + 0.030569061866173378, + 0.028339288522908977 + ], + "gate": { + "difference_mode": "relative", + "mean_difference": 0.0029410878709215105, + "resamples": 10000, + "seed": 0, + "seeds": 5, + "status": "passed", + "threshold": 0.05, + "upper_95": 0.0075631630189805675 + }, + "metric": "q_error_norm", + "official": [ + 0.03186987992171403, + 0.02983515760400793, + 0.027172292981828963, + 0.030564142583490727, + 0.028380315382402023 + ], + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "teacher": "official_warp" + }, + { + "embodichain": [ + 0.30773086547851564, + 0.32897301139412344, + 0.30997772216796876, + 0.33886656708769747, + 0.33389833890474757 + ], + "gate": { + "difference_mode": "relative", + "mean_difference": -0.0016487556883492814, + "resamples": 10000, + "seed": 0, + "seeds": 5, + "status": "passed", + "threshold": 0.05, + "upper_95": -0.0003000744695058809 + }, + "metric": "qd_error_norm", + "official": [ + 0.30775366982260904, + 0.3291599650959392, + 0.31195038177154877, + 0.33893180260291467, + 0.33426121303013395 + ], + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "teacher": "official_warp" + }, + { + "embodichain": [ + 0.028664215751316235, + 0.030975601638572804, + 0.029802200759666553, + 0.036095138218091884, + 0.03286314701688462 + ], + "gate": { + "difference_mode": "relative", + "mean_difference": -0.029275385391873097, + "resamples": 10000, + "seed": 0, + "seeds": 5, + "status": "passed", + "threshold": 0.05, + "upper_95": 0.0006092183787881106 + }, + "metric": "q_error_norm", + "official": [ + 0.028611711142719657, + 0.03116420524707739, + 0.029795721302861754, + 0.04210370381673177, + 0.03285226683685745 + ], + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "teacher": "current_newton" + }, + { + "embodichain": [ + 0.7004476353742074, + 0.6937628704568614, + 0.7296680395153986, + 0.7696718962296195, + 0.7347368323284647 + ], + "gate": { + "difference_mode": "relative", + "mean_difference": 0.002149326603464712, + "resamples": 10000, + "seed": 0, + "seeds": 5, + "status": "passed", + "threshold": 0.05, + "upper_95": 0.021675618573048703 + }, + "metric": "qd_error_norm", + "official": [ + 0.7084931746773098, + 0.7104982680168704, + 0.7293366805366848, + 0.7345860135728034, + 0.7366227301998415 + ], + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "teacher": "current_newton" + } + ], + "full_profile_differences": [ + "train.total_steps differs from the matched official profile", + "train.lr_update_interval differs from the matched official profile", + "train.lr_offset differs from the matched official profile", + "train.validation_interval differs from the matched official profile", + "train.validation_batches differs from the matched official profile" + ] + }, + "ant": { + "experiment": "/tmp/nerd-five-seed-ant-20261020/experiment.json", + "report": "/tmp/nerd-five-seed-ant-20261020/report-with-bounded-rollouts/REPORT.md", + "status": "pending", + "comparisons": [ + { + "embodichain": [ + 0.07745500920893071, + 0.07813525776286702, + 0.08003018976567866, + 0.08014878493088942, + 0.07752861819424473 + ], + "gate": { + "difference_mode": "relative", + "mean_difference": -0.0001241758157730334, + "resamples": 10000, + "seed": 0, + "seeds": 5, + "status": "passed", + "threshold": 0.05, + "upper_95": 0.0007861924289500945 + }, + "metric": "q_error_norm", + "official": [ + 0.07735808655455872, + 0.07836143787090595, + 0.08000357910826966, + 0.08015401232373584, + 0.07747088841029576 + ], + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "teacher": "official_warp" + }, + { + "embodichain": [ + 7.250118046016484, + 7.1802550598815245, + 7.0804086538461535, + 7.3203793033138735, + 7.1045006385216345 + ], + "gate": { + "difference_mode": "relative", + "mean_difference": -3.941419260966793e-05, + "resamples": 10000, + "seed": 0, + "seeds": 5, + "status": "passed", + "threshold": 0.05, + "upper_95": 0.0005139221681135962 + }, + "metric": "qd_error_norm", + "official": [ + 7.257187258542239, + 7.171238222226992, + 7.083392132769574, + 7.319544664320055, + 7.1057311341002745 + ], + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "teacher": "official_warp" + }, + { + "embodichain": [ + 0.005307778945335975, + 0.005091159946315891, + 0.005486480482331999, + 0.0053636068826193335, + 0.005174731422256638 + ], + "gate": { + "difference_mode": "relative", + "mean_difference": 0.005700531791566215, + "resamples": 10000, + "seed": 0, + "seeds": 5, + "status": "passed", + "threshold": 0.05, + "upper_95": 0.00965982291230781 + }, + "metric": "base_position_error", + "official": [ + 0.005254870456653637, + 0.005087500090127463, + 0.0054122172869168796, + 0.005360990304213303, + 0.005156655625982599 + ], + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "teacher": "official_warp" + }, + { + "embodichain": [ + 0.01833216279417604, + 0.018993363013634315, + 0.01980317451141693, + 0.01985698741870922, + 0.017856536068759123 + ], + "gate": { + "difference_mode": "relative", + "mean_difference": 0.002982969323553847, + "resamples": 10000, + "seed": 0, + "seeds": 5, + "status": "passed", + "threshold": 0.05, + "upper_95": 0.006760569207506567 + }, + "metric": "base_orientation_error_rad", + "official": [ + 0.018305523840935676, + 0.01898016038831774, + 0.0195915101648687, + 0.01981754826975393, + 0.017857071593567565 + ], + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "teacher": "official_warp" + }, + { + "embodichain": [ + 0.09104689998903136, + 0.08955901878467505, + 0.09358879863352015, + 0.092288915661798, + 0.09246751536493716 + ], + "gate": { + "difference_mode": "relative", + "mean_difference": -0.0008069286806190477, + "resamples": 10000, + "seed": 0, + "seeds": 5, + "status": "passed", + "threshold": 0.05, + "upper_95": 4.154844541429713e-05 + }, + "metric": "q_error_norm", + "official": [ + 0.09097899284915648, + 0.08960340748662533, + 0.09382180476534194, + 0.09231608294058537, + 0.09260715539904608 + ], + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "teacher": "current_newton" + }, + { + "embodichain": [ + 7.98330715013587, + 7.990900985054348, + 8.114658344655798, + 7.943337437726449, + 8.113467362998188 + ], + "gate": { + "difference_mode": "relative", + "mean_difference": 5.6023575830166544e-05, + "resamples": 10000, + "seed": 0, + "seeds": 5, + "status": "passed", + "threshold": 0.05, + "upper_95": 0.0005148825688513269 + }, + "metric": "qd_error_norm", + "official": [ + 7.981286798007246, + 7.992876754981884, + 8.10761081861413, + 7.940719825634058, + 8.12097698709239 + ], + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "teacher": "current_newton" + }, + { + "embodichain": [ + 0.004330985442451809, + 0.00485701422760452, + 0.0042299236076465555, + 0.004623724757761195, + 0.005029605091481969 + ], + "gate": { + "difference_mode": "relative", + "mean_difference": -0.00324593774083294, + "resamples": 10000, + "seed": 0, + "seeds": 5, + "status": "passed", + "threshold": 0.05, + "upper_95": -0.0015236678033083008 + }, + "metric": "base_position_error", + "official": [ + 0.004343361785446388, + 0.0048682841701784, + 0.0042658142421556555, + 0.004624356739762901, + 0.0050422865411509634 + ], + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "teacher": "current_newton" + }, + { + "embodichain": [ + 0.015851749199024146, + 0.016095023224319237, + 0.015240774292876754, + 0.01631476360818614, + 0.01624640450961348 + ], + "gate": { + "difference_mode": "relative", + "mean_difference": 0.0007460519254374892, + "resamples": 10000, + "seed": 0, + "seeds": 5, + "status": "passed", + "threshold": 0.05, + "upper_95": 0.0015601198654396253 + }, + "metric": "base_orientation_error_rad", + "official": [ + 0.01583538746488267, + 0.016103356126425925, + 0.015208294080651325, + 0.01632652559142182, + 0.01621722622194152 + ], + "seeds": [ + 0, + 1, + 2, + 3, + 4 + ], + "teacher": "current_newton" + } + ], + "full_profile_differences": [ + "train.total_steps differs from the matched official profile", + "train.lr_update_interval differs from the matched official profile", + "train.lr_offset differs from the matched official profile", + "train.validation_interval differs from the matched official profile", + "train.validation_batches differs from the matched official profile" + ] + } + } +} diff --git a/design/reports/neural-dynamics-historical-rollout-20261020.md b/design/reports/neural-dynamics-historical-rollout-20261020.md new file mode 100644 index 000000000..cab40d684 --- /dev/null +++ b/design/reports/neural-dynamics-historical-rollout-20261020.md @@ -0,0 +1,61 @@ +# Historical NeRD rollout evidence — 2026-10-20 + +Executed 12 bounded CPU runs using the newly trained seed-0 exports: two robots × two trainers × three motion modes. All runs completed, all six pairs used identical teacher trajectories and policy identities, and all original-model restoration probes passed with maximum absolute state error **0**. This establishes executable historical-domain comparison; it does not establish full-budget accuracy, paper reproduction, or deployment of historical physics into current Newton. + +## Protocol + +- Original untouched NeRD main commit `370e89d2b8ce68aa3578aa7812859c681b7502bc`; original `NeuralEnvironment`, Featherstone teacher and `TransformerNeuralIntegrator` student. +- Cartpole uses Warp 1.8.0 / 5 teacher substeps; Ant uses Warp 1.5.1 / 16 teacher substeps. Each run is a separate CPU process with CUDA hidden. The scoped CPU stream adapter maps the original CUDA stream request to `None`; source math remains unchanged. +- Both trainers completed 100 updates on the historical physical-v2 data. These are bounded training checkpoints, not converged models. Best-validation export from each trainer is evaluated without changing weights or manifests. +- Batch size 2, 3 trajectories (actual batches 2 + 1), original sampled initial states with seeds 500–502; horizons 1, 10 and 12. Passive and random effort share commands within a pair. Fixed PD recomputes teacher and student controls from their respective states; trace inspection confirms their commands diverge. +- Each run additionally compares reconstructed original `ModelMixedInput` with the portable model inside the same original integrator for 11 steps, covering history length 10 and its first rollover. All 12 probes were finite and exactly equal. +- Live source physical arrays, gravity `[0, -9.81, 0]`, axis, timestep, substeps, asset hash and codec metadata matched the recorded historical contract. No recorded physical field was missing. +- Metrics use shared canonical Z-up / COM velocity transport and shared task reward definitions. Ant q MAE covers its scalar joints; base position and quaternion angular errors are reported separately. Horizon state errors average steps 1 through H; returns sum until the first failure. All requested trajectories remain in evidence. + +## Measured results at horizon 12 + +All 12 runs had 3/3 numerically valid trajectories and **0/3 teacher failures, student failures, teacher nonfinite states and student nonfinite states**. No rows were rejected or resampled. The two trained models are not numerically identical; the table reports both. + +| Robot | Motion | Trainer | q MAE | qd MAE | Base position | Base angle (rad) | Fixed-PD return | +|---|---|---|---:|---:|---:|---:|---:| +| cartpole | passive | official | 0.06185331 | 0.5009031 | 0 | 0 | — | +| cartpole | passive | embodichain | 0.06182037 | 0.5007337 | 0 | 0 | — | +| cartpole | random_effort | official | 0.08175079 | 0.5582347 | 0 | 0 | — | +| cartpole | random_effort | embodichain | 0.08169176 | 0.555877 | 0 | 0 | — | +| cartpole | fixed_policy | official | 0.06957492 | 0.5151157 | 0 | 0 | -170.8377 | +| cartpole | fixed_policy | embodichain | 0.06958081 | 0.5154237 | 0 | 0 | -170.8884 | +| ant | passive | official | 0.3156452 | 2.652962 | 0.1061554 | 0.247806 | — | +| ant | passive | embodichain | 0.3172099 | 2.657473 | 0.1043622 | 0.2473033 | — | +| ant | random_effort | official | 0.2701225 | 2.734384 | 0.08245032 | 0.1912716 | — | +| ant | random_effort | embodichain | 0.2708459 | 2.725167 | 0.08210292 | 0.192527 | — | +| ant | fixed_policy | official | 0.2945366 | 2.536861 | 0.1055715 | 0.2423552 | 1.458784 | +| ant | fixed_policy | embodichain | 0.2932671 | 2.54465 | 0.1017988 | 0.2401732 | 1.572607 | + +Fixed-PD teacher returns: Cartpole **−163.68427**, Ant **9.7618113**. These fixed baselines are explicitly not the paper policies. The short Ant learned rollouts have materially lower return than their teacher; zero failures over 12 steps must not be read as accuracy qualification. + +## Pairing and evidence integrity + +Recomputed every trajectory content digest from the saved teacher state/control arrays and logical seed/ID, then recomputed the ordered selected-teacher digest. All six pairs match exactly despite distinct model artifacts. Verified trace, JSONL, model weights and manifest hashes against the result receipts. Each pair also matches teacher origin, dtype, count and policy hash. + +| Robot | Motion | Ordered teacher SHA-256 | Policy SHA-256 | +|---|---|---|---| +| cartpole | passive | `1754aa11691ee45e4f285796f5323b859b14a747250af4a2efca1053048118e5` | none | +| cartpole | random_effort | `b28a63ba9822549a910ac640953ef581c0f0fe5a9cd0ef029f2221341c192d71` | none | +| cartpole | fixed_policy | `427bcb6f013d43eb31c00e0f071ce085f2b079b5c134263a8b0bbe1d91549b86` | `5f373a5c1aab23985abc978b15ed04891beed8a47e5350e6064f2d9d5ab5930a` | +| ant | passive | `4522f97dbaf9841a2df00630161927836633ac3f4639097e8f2689b6ac423956` | none | +| ant | random_effort | `16e537821f84443b3ec7328bc820049ce3b198b96fc5ad5b08fb83965447b113` | none | +| ant | fixed_policy | `c87a0b7ed58a157ec0cf6818d0e01f3ee1847f6083768b4a4bb807c6ad32ec83` | `2dfde4898331fc07566d72fdfd90e78cbf60eeb6f006e9d151c5d636992f937a` | + +Artifacts are local verification receipts; they must be retained when sharing this report: + +- [Cartpole specs, logs, JSON metrics, HDF5 traces and JSONL trajectories](/tmp/nerd-historical-paired-cartpole-20261020/). +- [Ant specs, logs, JSON metrics, HDF5 traces and JSONL trajectories](/tmp/nerd-historical-paired-ant-20261020/). +- [Independent pairing/integrity receipt](/tmp/nerd-historical-paired-validation-20261020.json); SHA-256 `e8b6bcbba7034a6c4b7eeb33c92fd8d2bc7c573ea50398fdc80f4728f7deb4bb`. +- Model roots: `/tmp/nerd-five-seed-{cartpole,ant}-20261020/official_warp/seed-0/{official,embodichain}/best-export`. Adjacent `result.json` files record each completed 100-update training run. +- [Harness](/home/dex/.codex/worktrees/e142/EmbodiChain/embodichain/learning/dynamics/reference_rollout.py); [focused real-runtime tests](/home/dex/.codex/worktrees/e142/EmbodiChain/tests/learning/dynamics/test_reference_rollout.py), **5 passed** with both pinned legacy environments enabled. + +Reproduce one receipt from its `*-spec.json` using `python -m embodichain.learning.dynamics.reference_rollout --spec --trust-external-code`, with the matching legacy dependency directory first on `PYTHONPATH` and a fresh output destination. + +## Remaining qualification + +Only seed 0 and three trajectories per mode were evaluated here. Five-seed aggregate rollout gates, 2,048 trajectories through horizon 1,000, full training budgets and paper-policy comparisons remain unexecuted in this report. Historical runtime measurements must be compared within the historical physical domain; they do not authorize loading these old-physics models into the current-Newton provider. diff --git a/design/reports/neural-dynamics-implementation-v2-20261020.md b/design/reports/neural-dynamics-implementation-v2-20261020.md new file mode 100644 index 000000000..a3e12ec3d --- /dev/null +++ b/design/reports/neural-dynamics-implementation-v2-20261020.md @@ -0,0 +1,111 @@ +# NeRD implementation audit and remediation — physics signature v2 + +Status: implementation progressed and bounded evidence refreshed; complete +qualification is **not achieved**. The acceptance contract remains +`design/neural-dynamics-solvers.md`. + +## Correctness fixes + +The audit reproduced acceptance of a changed physical system by the original +signature: a 180-degree fixed Cartpole root rotation, changed Ant robot contact +properties, and changed collision/topology relationships could escape validation. +The v2 signature now binds fixed-root orientation, parent/child topology, body/ +joint flags, relevant material properties and effective collision relationships. +Direct model and public metadata construction agree. Replica count and harmless +contact-free Cartpole translation do not change the signature. + +The provider explicitly requires the v2 schema and FP32 model weights/buffers. +Old manifests must not be relabeled. The offline model API still supports FP64; +Newton execution rejects that dtype before stepping. All-false functional reset +masks now preserve the supplied graph and avoid native reset/publication. + +Training manifest overrides may add annotations, but cannot change/drop measured +dataset provenance or introduce missing reserved physical fields. Legitimate +fine-tuning on a new measured dataset remains supported. + +## Verification coverage + +These are separate focused groups, not a deduplicated aggregate test count: + +- Runtime audit: **24 passed**, including **6 CUDA cases** across Cartpole/Ant: + five-step initial q/qd and control gradients, all named network parameters, + explicit incoming history gradients, partial reset and native-storage isolation. +- Follow-up runtime/collection/schema tests: **13 passed**, including regenerated + real data and content-hash invariance across replay batch sizes/output paths. +- Input pipeline, trainer and reference group: **70 passed**, with the actual + official trainer and CPU reference collection enabled. +- Root manager/CLI/comparison regression: **231 passed, 2 skipped**. +- Report qualification tests: **17 passed** after the final provenance changes. +- FP64 fail-fast and actual collector runtime-identity checks: **2 passed**. +- The original official-main two-update training comparison was independently + rerun after batch changes: **1 passed**. + +FP32 numerical receipts still pass for both robots. The two historical matched +training experiments were reread under their `report-audit/` directories: +recorded teacher provenance and numerical parity pass; overall qualification is +pending. Historical experiments retain their original code/data scope. + +## Refreshed current-physics pipeline + +Both robots ran through the actual CLI with newly generated v2 data: +64 trajectories × 32 steps, zero rejected trajectories; 100 updates of the small +one-layer smoke network; export; offline evaluation; then passive, random-effort +and fixed-policy rollouts. Each rollout uses four trajectories, two environments +and horizons 1/10/32. All six runs completed with zero nonfinite trajectories. +This proves executable flow, not low prediction error or official quality. + +The final reports separately identify teacher physics and actual neural +execution device, Python/Torch/CUDA/software versions and adapter source hashes. +They retain selected teacher-content digests, per-trajectory indices, trace files, +model/manifest hashes and policy identity. Reporter gates verify those artifacts +and reject mismatched teacher domains, test populations or fixed policies. +Undefined/nonfinite errors cannot become zero-valued passing evidence. Bare +Boolean numerical pass flags are insufficient. + +Compact exact results and artifact paths: +[refreshed smoke receipts](neural-dynamics-v2-smoke-20261020.json). +Full artifacts: `/tmp/embodichain-nerd-v2-qualification/{cartpole,ant}`. + +## Performance and delivery + +The synchronous grouped batch reader preserves order, duplicates, prefix masks, +owned outputs and exact resume. It groups I/O without prefetch or a full-data +cache. Genuine B512 read/collate/encode timings improved about 11–20× on the +current small Newton datasets. The distinct-trajectory, many-shard stress case +improved only 1.86× and still took 181.72 ms per batch. Full-scale training +throughput remains unmeasured; the former 25.31-day estimate is compute-only. +See [the bounded input benchmark](neural-dynamics-input-20261020.md). + +The replacement DexSim artifact is +`0.5.1+g4f093c4a88ff.d2bf6eb9a5b93fd36`; its independent installation passed +**157 tests, 21 deselected**. The refreshed EmbodiChain wheel contains v2 code +and an explicit dependency metadata overlay for this exact local SDK version. +[Artifact hashes and overlay](neural-dynamics-local-artifacts-v2-20261020.json) +identify the pair. Fresh paired installation verified exact dependency matching, +installed package/provider import paths and the FP32 fail-fast case. Proof is +`output/neural-solver-wheel/validation-final-pair/manifest.json` in the DexSim +checkout. The source dependency pin remains rc1, so these are paired +local development artifacts; a published release/pin update is a separate handoff. + +Black 26.3.1, diff whitespace, context validation and **2747/2747** public API +coverage passed. A complete Sphinx dummy build succeeded on an isolated committed +source snapshot, avoiding the architecture generator's clean-tree precondition +without committing the user's worktree. It emitted 785 repository warnings; +none mentioned the new NeRD/dynamics pages or module symbols. This is not a +warning-clean documentation claim. The previously isolated NeRD page build also +passed with warnings treated as errors. + +## Remaining requirements + +- Full-size matched data distributions, five paired seeds and the declared + 4,995,000-update training budget have not been executed. +- The 2048-trajectory, 1/10/100/1000-step complete evaluation and final + noninferiority/stability/policy-return gates remain pending. +- Official historical Warp data has no compatible v2 current-Newton deployment + profile. Official-domain closed-loop evaluation needs a separately faithful + reference execution path; changing metadata is not an implementation. +- Representative full-scale end-to-end throughput is unverified. The bounded + optimization does not remove the resource-budget question for full training. + +No full-budget training job has been launched. The original goal remains active; +unit tests, local wheels and smoke results are not substituted for these gates. diff --git a/design/reports/neural-dynamics-input-20261020.md b/design/reports/neural-dynamics-input-20261020.md new file mode 100644 index 000000000..3bf75843c --- /dev/null +++ b/design/reports/neural-dynamics-input-20261020.md @@ -0,0 +1,144 @@ +# NeRD bounded input pipeline — 2026-10-20 + +**The synchronous input path is faster without changing sampled data, optimizer +order or checkpoint semantics.** The final benchmark observes 11–23× speedups for +batch 512 on the genuine small datasets. A distinct-trajectory, many-shard stress +case improves 1.86×; it remains substantially slower than GPU model computation. +These measurements do not qualify a full training run or reduce the previously +measured 607.52 GPU-hour compute-only lower bound. + +## Implementation and resource model + +`DynamicsWindowDataset.get_batch(indices)` replaces individual window reads plus +`default_collate` in statistics computation, training, offline evaluation and the +original-trainer transport adapter. The original `__getitem__` path remains intact +as the independent comparison baseline and ordinary Dataset API. + +The method groups its I/O schedule by shard, trajectory and fixed time page, +then writes results directly into owned batch arrays in the **original requested +order**, including duplicate indices. Requests never cross a trajectory or split. +Startup prefixes retain the same right padding, masks and valid lengths. + +Each field read spans at most `read_page_steps + history_length` frames: the +benchmark uses 256 + 10 = **266 frames**. Sparse requests in a long trajectory +therefore cannot cause the whole trajectory to be loaded. Public HDF5 hyperslab +reads avoid repeatedly constructing high-level Dataset/File wrappers. The memory +space is created from the actual owned contiguous destination array, and file +shape/range is checked before reading. Empty fields retain their shape and dtype. + +Storage is the requested output batch, O(batch-size) indexing, one bounded source +field block and one per-field gather. The file-handle LRU remains capped at four +per reader/worker. There is **no retained trajectory cache, whole-dataset load, +background worker or speculative prefetch**. Native HDF5 caches remain separate +from the explicitly measured Python/NumPy allocations. The existing O(trajectory +count) metadata index is unchanged. + +The sampler advances only after an optimizer update, exactly as before. Returned +rows do not alias one another when indices repeat, and later calls or file closure +do not overwrite previously returned tensors. + +Changed production owners: + +- `embodichain/data_pipeline/dynamics/shards.py`: bounded grouped/hyperslab reader. +- `embodichain/learning/dynamics/trainer.py`: batched statistics/train/evaluation. +- `embodichain/learning/dynamics/reference.py`: identical batch transport into the + untouched official trainer methods. + +## Measurements + +CPU read + collation + shared codec encoding, one CPU thread, history 10. Each +implementation uses the same seeded index sequence, two warmup batches and ten +timed batches. Source verification, index generation and correctness checks are +outside timing. Filesystem pages are warm; this is not a cold-storage benchmark. +Python 3.11.14, Torch 2.7.1+cu128 and h5py 3.15.1 were used. No GPU was used. + +| Dataset / batch | Individual windows, ms | Grouped reads, ms | Speedup | +|---|---:|---:|---:| +| Current Newton Cartpole / 1 | 0.742 | 0.443 | 1.68× | +| Current Newton Cartpole / 16 | 8.425 | 3.720 | 2.26× | +| Current Newton Cartpole / 512 | 257.874 | 12.749 | 20.23× | +| Current Newton Ant / 1 | 1.142 | 0.909 | 1.26× | +| Current Newton Ant / 16 | 9.336 | 5.249 | 1.78× | +| Current Newton Ant / 512 | 285.034 | 25.883 | 11.01× | +| Official Warp Cartpole / 1 | 0.618 | 0.358 | 1.73× | +| Official Warp Cartpole / 16 | 7.592 | 3.229 | 2.35× | +| Official Warp Cartpole / 512 | 241.613 | 10.622 | 22.75× | +| Official Warp Ant / 1 | 1.097 | 0.816 | 1.34× | +| Official Warp Ant / 16 | 8.795 | 4.504 | 1.95× | +| Official Warp Ant / 512 | 271.967 | 20.767 | 13.10× | +| 128-shard distinct-trajectory stress / 1 | 1.304 | 1.066 | 1.22× | +| 128-shard distinct-trajectory stress / 16 | 11.353 | 7.237 | 1.57× | +| 128-shard distinct-trajectory stress / 512 | 337.134 | 181.723 | 1.86× | + +The current Newton datasets have 57 training trajectories across 16 shards; the +official datasets have 32 training trajectories in one shard. Their batch-512 +requests consequently reuse many trajectories, explaining the larger gains. + +The stress dataset repeats an actual Ant trajectory under 512 independent IDs +across 128 shards. Every batch-512 sample chooses a **distinct trajectory** and +the file cache holds only four shards. This tests a weak-reuse case beyond cache +capacity, not physics accuracy. Cold-batch file opens fall from **503 to 128**; +retained handles never exceed four. Output/encoding tensors match bitwise. + +For stress batch 512, returned tensor data is 7,316,480 bytes. Tracemalloc's +Python/NumPy peak falls from **10,646,618 to 8,207,876 bytes**. These numbers exclude +native HDF5/Torch allocator memory and must not be presented as total process RSS. +For the small genuine Ant datasets, grouped peaks are about 7.48 MB. No retained +data cache scales with the dataset's total sample count. + +The stress result still spends roughly 182 ms preparing a batch, compared with +about 14 ms of Ant GPU training computation in the earlier separate benchmark. +The small-dataset 11–23× gains must therefore **not** be extrapolated to full-size +random sampling. Packed layouts or controlled prefetch would need a separate, +representative experiment before changing the full-budget duration estimate. + +## Correctness and provenance checks + +Focused tests cover: + +- Exact tensor/metadata equality with individual reads for B=1/16/512, histories + 1/3/10, partial prefixes, repeated/negative indices and more shards than cache. +- Independent duplicate rows, cross-call ownership and validity after close. +- Sparse windows across a 4,097-step trajectory with an observed 39-frame bound + (page 32 + history 7). +- Float64, float32, integer, Boolean, multidimensional, zero-width and UTF-8 fields. +- Shape validation before a low-level buffer write, invalid indices before I/O, + worker-local file handles and existing teardown behavior. +- Bitwise-identical 20-step training weights and logs between individual and + grouped reads with dropout, plus the existing exact-resume and independent + official-trainer tests. + +A separate provenance fix in `workflow.run_train` closes a relabeling hole: +manifest annotations may add unrelated keys, but cannot replace/drop existing +dataset metadata or invent a missing deployment-contract field. Protected fields +include physics signature/schema, gravity, ordered joints, asset identity, codec, +timestep and contact schema. An explicitly empty override is rejected. Fine-tuning +may initialize from another physical model with compatible codec/timestep; its +result retains the **new measured dataset's** identity. Negative tests cover +changed, omitted and invented signatures, schema/gravity/assets, while positive +tests cover annotations and legitimate fine-tuning. + +The combined scoped suite passed **70 tests** with the explicitly enabled real +reference checkout: storage, training/resume, original trainer and original CPU +collection, experiment/report helpers, CLI and packaging checks. The changed +Python files pass the scoped Black check. No GPU tests were needed for this +storage/transport change. + +## Reproduction and artifacts + +Benchmark entry point: `scripts/benchmark/learning/dynamics_input.py`. + +```bash +OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 python -m scripts.benchmark.learning.dynamics_input \ + --dataset /tmp/embodichain-nerd-v2-qualification/cartpole/data \ + --dataset /tmp/embodichain-nerd-v2-qualification/ant/data \ + --dataset /tmp/nerd-official-real-cartpole-20261020/canonical \ + --dataset /tmp/nerd-official-real-ant-20261020/canonical \ + --stress-dir /tmp/nerd-input-new-stress-run \ + --output /tmp/nerd-input-new-benchmark.json +``` + +The stress destination must be new. Full measured samples, dataset fingerprints, +allocation scope and handle counts are saved in +`/tmp/nerd-input-final-benchmark-20261020.json`; console output is in the adjacent +`.log`. No multiweek training job or new GPU workload was launched for this change. diff --git a/design/reports/neural-dynamics-local-artifacts-20261020.json b/design/reports/neural-dynamics-local-artifacts-20261020.json new file mode 100644 index 000000000..bb597f8a3 --- /dev/null +++ b/design/reports/neural-dynamics-local-artifacts-20261020.json @@ -0,0 +1,29 @@ +{ + "scope": "unpublished paired local development wheels; not a source-free release qualification", + "embodichain_source_wheel": { + "path": "/tmp/embodichain-nerd-wheel-final/embodichain-0.3.1-py3-none-any.whl", + "sha256": "69f39e855564b4d683d309e35c8a4ca55a182e39bee1a26afed7556fc1916c4d", + "bytes": 3105173 + }, + "embodichain_paired_wheel": { + "path": "/tmp/embodichain-nerd-paired-wheels/embodichain-0.3.1-py3-none-any.whl", + "sha256": "6a3ad03535f8b911a0d0152db0188dfe8bcdd52404a13dae55b20e5b9a457401", + "bytes": 3105202 + }, + "dexsim_wheel": { + "path": "/home/dex/workspace/sources/dexsim/output/neural-solver-wheel/wheels/dexsim_engine-0.5.1+g4f093c4a88ff.dd70636e2d1f028ad-cp311-cp311-manylinux_2_35_x86_64.whl", + "sha256": "0ab7a36b8a3f80d9f6afa1f0ff86827fba112150f053a4853aae8c9a2ca3741f", + "bytes": 374974063 + }, + "metadata_overlay": { + "from": "dexsim_engine==0.5.1rc1", + "to": "dexsim_engine==0.5.1+g4f093c4a88ff.dd70636e2d1f028ad", + "only_changed_members": [ + "embodichain-0.3.1.dist-info/METADATA", + "embodichain-0.3.1.dist-info/RECORD" + ], + "source_pyproject_modified": false, + "record_regenerated": true + }, + "qualification": "pending" +} diff --git a/design/reports/neural-dynamics-local-artifacts-historical-20261020.json b/design/reports/neural-dynamics-local-artifacts-historical-20261020.json new file mode 100644 index 000000000..e91e07a14 --- /dev/null +++ b/design/reports/neural-dynamics-local-artifacts-historical-20261020.json @@ -0,0 +1,121 @@ +{ + "scope": "unpublished paired local artifacts including isolated historical reference harness", + "embodichain_source_wheel": { + "path": "/tmp/embodichain-nerd-historical-wheel/embodichain-0.3.1-py3-none-any.whl", + "sha256": "f40e498637ad6dbfe7ea75a1c4701d606bd4e28c3626174c8b9678ed1d5a2d16", + "bytes": 3129248 + }, + "embodichain_paired_wheel": { + "path": "/tmp/embodichain-nerd-historical-paired-wheels/embodichain-0.3.1-py3-none-any.whl", + "sha256": "722046aab36844669862b9feff3165800576316872858d1c0719f1b094926c38", + "bytes": 3129277 + }, + "dexsim_wheel": { + "path": "/home/dex/workspace/sources/dexsim/output/neural-solver-wheel/wheels/dexsim_engine-0.5.1+g4f093c4a88ff.d2bf6eb9a5b93fd36-cp311-cp311-manylinux_2_35_x86_64.whl", + "sha256": "2b0a37872ddb690392450b4b9764f128d27b8e53138d7973e87459e722504b3f", + "bytes": 374974433 + }, + "metadata_overlay": { + "from": "dexsim_engine==0.5.1rc1", + "to": "dexsim_engine==0.5.1+g4f093c4a88ff.d2bf6eb9a5b93fd36", + "only_changed_members": [ + "embodichain-0.3.1.dist-info/METADATA", + "embodichain-0.3.1.dist-info/RECORD" + ], + "source_dependency_pin_modified": false, + "record_regenerated": true + }, + "formal_accuracy_qualification": "pending", + "durable_paired_wheel": "/home/dex/workspace/sources/dexsim/output/neural-solver-wheel/paired-ec/722046aab368/embodichain-0.3.1-py3-none-any.whl", + "installed_validation": { + "artifact_manifest": "/home/dex/workspace/sources/dexsim/output/neural-solver-wheel/artifact-manifest.json", + "sha256": "c8e27323f1cbc4151832ea46e9b685338cdec9cbdf69200be0b7db17f2982bb0", + "pair_imports": "/home/dex/workspace/sources/dexsim/output/neural-solver-wheel/validation-historical/pair-imports.json", + "historical_cpu": { + "cartpole": { + "receipt": "/home/dex/workspace/sources/dexsim/output/neural-solver-wheel/validation-historical/cartpole/installation-receipt.json", + "warp_version": "1.8.0", + "runtime_engine": "official_warp", + "physical_complete": true, + "portable_max_abs_error": 0.0, + "scope": "tiny installed historical CPU execution; not accuracy qualification" + }, + "ant": { + "receipt": "/home/dex/workspace/sources/dexsim/output/neural-solver-wheel/validation-historical/ant/installation-receipt.json", + "warp_version": "1.5.1", + "runtime_engine": "official_warp", + "physical_complete": true, + "portable_max_abs_error": 0.0, + "scope": "tiny installed historical CPU execution; not accuracy qualification" + } + }, + "historical_cuda": { + "cartpole": { + "status": "passed", + "receipt": "/home/dex/workspace/sources/dexsim/output/neural-solver-wheel/validation-historical-cuda/cartpole/installation-receipt.json", + "receipt_sha256": "5a84ed60c94f911e178a384ae3ba60413a12c5d809f3cac831874e1f0e913111", + "warp_version": "1.8.0", + "device": "cuda:0", + "gpu": "NVIDIA GeForce RTX 5090", + "runtime_engine": "official_warp", + "physical_complete": true, + "portable_runtime_equivalence": { + "steps": 3, + "batch_size": 1, + "atol": 1e-05, + "rtol": 0.0001, + "max_abs_error": 0.0, + "matched_nonfinite": false, + "status": "passed" + }, + "num_envs": 1, + "trajectories": 1, + "horizons": [ + 1, + 3 + ], + "seed": 500, + "motion_mode": "passive", + "initial_states": "default", + "attempts": 1, + "current_simulator_import_attempts": [], + "log": "/home/dex/workspace/sources/dexsim/output/neural-solver-wheel/validation-historical-cuda/cartpole.log", + "warning": "Old Warp reports unavailable cuDeviceGetUuid driver entry point/API during initialization; CUDA execution and this probe completed successfully.", + "scope": "Tiny installed historical CUDA execution only, not full training, long-horizon accuracy, or general old-Warp GPU compatibility qualification." + }, + "ant": { + "status": "passed", + "receipt": "/home/dex/workspace/sources/dexsim/output/neural-solver-wheel/validation-historical-cuda/ant/installation-receipt.json", + "receipt_sha256": "b63eb5ebe25b20ed284192fdc925f80824f9337673fd12b6aa25d5d769313835", + "warp_version": "1.5.1", + "device": "cuda:0", + "gpu": "NVIDIA GeForce RTX 5090", + "runtime_engine": "official_warp", + "physical_complete": true, + "portable_runtime_equivalence": { + "steps": 3, + "batch_size": 1, + "atol": 1e-05, + "rtol": 0.0001, + "max_abs_error": 1.4781951904296875e-05, + "matched_nonfinite": false, + "status": "passed" + }, + "num_envs": 1, + "trajectories": 1, + "horizons": [ + 1, + 3 + ], + "seed": 500, + "motion_mode": "passive", + "initial_states": "default", + "attempts": 1, + "current_simulator_import_attempts": [], + "log": "/home/dex/workspace/sources/dexsim/output/neural-solver-wheel/validation-historical-cuda/ant.log", + "warning": "Old Warp reports unavailable cuDeviceGetUuid driver entry point/API during initialization; CUDA execution and this probe completed successfully.", + "scope": "Tiny installed historical CUDA execution only, not full training, long-horizon accuracy, or general old-Warp GPU compatibility qualification." + } + } + } +} diff --git a/design/reports/neural-dynamics-local-artifacts-v2-20261020.json b/design/reports/neural-dynamics-local-artifacts-v2-20261020.json new file mode 100644 index 000000000..fd34f11c1 --- /dev/null +++ b/design/reports/neural-dynamics-local-artifacts-v2-20261020.json @@ -0,0 +1,36 @@ +{ + "scope": "unpublished paired local development wheels; physics signature v2", + "embodichain_source_wheel": { + "path": "/tmp/embodichain-nerd-v2-final-wheel/embodichain-0.3.1-py3-none-any.whl", + "sha256": "950be3b1f7a1929249ef09fbe3f5963fd633be329b76335c0f88055d3747823a", + "bytes": 3112657 + }, + "embodichain_paired_wheel": { + "path": "/tmp/embodichain-nerd-v2-final-paired-wheels/embodichain-0.3.1-py3-none-any.whl", + "sha256": "57a37b22f74ba272625c076f7419101061952c6087bd970480121667f55eb81f", + "bytes": 3112686 + }, + "dexsim_wheel": { + "path": "/home/dex/workspace/sources/dexsim/output/neural-solver-wheel/wheels/dexsim_engine-0.5.1+g4f093c4a88ff.d2bf6eb9a5b93fd36-cp311-cp311-manylinux_2_35_x86_64.whl", + "sha256": "2b0a37872ddb690392450b4b9764f128d27b8e53138d7973e87459e722504b3f", + "bytes": 374974433 + }, + "metadata_overlay": { + "from": "dexsim_engine==0.5.1rc1", + "to": "dexsim_engine==0.5.1+g4f093c4a88ff.d2bf6eb9a5b93fd36", + "only_changed_members": [ + "embodichain-0.3.1.dist-info/METADATA", + "embodichain-0.3.1.dist-info/RECORD" + ], + "source_dependency_pin_modified": false, + "record_regenerated": true + }, + "formal_accuracy_qualification": "pending", + "installed_pair_validation": { + "manifest": "/home/dex/workspace/sources/dexsim/output/neural-solver-wheel/validation-final-pair/manifest.json", + "sha256": "0d9fa38a2d574b33f20b463b5791eb443b9d7e82a22047c0f90a2887362f91d1", + "status": "final_pair_imports_dependency_and_focused_runtime_verified", + "exit_code": 0, + "scope": "No SDK changes since the previous installed 157-test suite; final EC-only refresh verified by exact dependency/import and dtype fail-fast test. Full numerical qualification is outside this packaging check." + } +} diff --git a/design/reports/neural-dynamics-official-protocol-audit-20261020.md b/design/reports/neural-dynamics-official-protocol-audit-20261020.md new file mode 100644 index 000000000..67caaec00 --- /dev/null +++ b/design/reports/neural-dynamics-official-protocol-audit-20261020.md @@ -0,0 +1,130 @@ +# NeRD official protocol audit — 2026-10-20 + +Reference: unmodified `NVlabs/neural-robot-dynamics` commit +`370e89d2b8ce68aa3578aa7812859c681b7502bc`. This audit distinguishes the official +architecture/optimizer-budget profile from controlled matched-training experiments +and literal published-result reproduction. Disclosed differences do not waive the +required official comparison, full budget, five seeds or rollout gates. + +## Effective official model profile + +`train/cfg/{Cartpole,Ant}/transformer.yaml` and the saved pretrained-model configs +agree on these numerical architecture fields: + +| Field | Cartpole | Ant | +|---|---|---| +| Ordered inputs (canonical names) | `states_embedding:4`, `joint_f:2` | `states_embedding:29`, `contact_normals:75`, `contact_points_1:75`, `contact_depths:25`, `joint_f:8`, `gravity_dir:3` | +| Output dimension | 4 | 29 | +| Transformer embedding | 192 | 384 | +| Layers / heads / block size | 6 / 12 / 32 | 6 / 12 / 32 | +| History and training window | 10 / 10 | 10 / 10 | +| Prediction-head hidden dimensions | `[64]` | `[64]` | +| Transformer bias / dropout | false / 0.0 | false / 0.0 | +| Normalize input / output | true / true | true / true | +| Effective output tanh | false | false | +| State frame | world scalar joint coordinates | body, anchor at every step | +| Embedding / prediction / orientation | identical / relative / quaternion | identical / relative / quaternion | + +The original input name is `joint_acts`; `joint_f` is the explicit transport rename. +Feature **order** is part of the network contract, not just names and widths. +Input encoders are identity (no hidden layers); the head is Linear/ReLU with no +head LayerNorm. Transformer MLP uses GELU and a 4× expansion. These choices are +fixed in the implementation rather than separately configurable in the migrated +model. Transformer `bias:false` does not mean the prediction head has no bias. +The source YAML nests `output_tanh:false` under `network.model`, while the source +model reads a top-level network flag with a false default; effective behavior is +false in either case. + +A full-profile gate should inspect **hash-verified actual model artifacts** for +all these fields. `comparison._json(sort_keys=True)` sorts the plan's nested +`input_dims` map, so plan dictionary iteration is not authoritative feature order. +The initial-model export's `model_config.input_dims` preserves actual order. +Actual floating weights/normalizers must be FP32 for the official numeric profile; +JSON architecture fields alone do not prove dtype. + +## Effective official optimizer and budget + +| Contract | Expected | +|---|---| +| Algorithm / optimizer | SequenceModelTrainer / Adam | +| Batch size | 512 full windows | +| Adam defaults | betas `(0.9,0.999)`, eps `1e-8`, weight decay `0`, amsgrad false | +| Gradient clipping | enabled, norm limit `1.0` | +| Updates per minibatch | one backward and one Adam step; no accumulation/AMP | +| Learning-rate endpoints | `1e-3`, `1e-4`; linear epoch schedule | +| Source epochs / updates per trained epoch | 1000 / 5000 | +| Actual optimizer updates | 4,995,000 (epoch zero evaluates without training) | +| Migrated schedule interval / offset | 5000 / 1 | +| Validation timing / maximum batches | every 5000 updates / 100 batches | +| Loss | mean MSE after dividing prediction and target by `sqrt(output_var + 1e-5)` | + +For zero-based optimizer step `s`, the matching rate is +`r=(floor(s/5000)+1)/1000`, `lr=1e-3*(1-r)+1e-4*r`. The first actual update uses +**0.0009991** and the final update uses **0.0001009**; the loop never trains at the +nominal epoch-1000 endpoint `1e-4`. + +Current full workflow YAMLs match the translated architecture and numerical +optimizer/budget fields above. `output_tanh:false` is supplied by the model config +default. Enforce the fields, rather than accepting an arbitrary architecture or +zero learning rate merely because it runs 4,995,000 updates. Checkpoint frequency, +worker count and device name are execution choices, not architecture gates. Five +predeclared paired seeds are a project qualification requirement; the saved +supervised source checkpoint config itself records seed zero. + +## Controlled shared-protocol differences + +These are currently explicit adaptations, not equivalent implementations of every +outer-loop behavior in `VanillaTrainer.train()`: + +| Aspect | Pinned original implementation | Current matched adapter | +|---|---|---| +| Training sample order | DataLoader shuffle, `drop_last=True`; iterator behavior governs permutations | Shared explicit batch-index artifact, generated with replacement by default | +| Iterator exhaustion | `one_epoch` rebinds an exhausted iterator locally; outer iterator is not replaced | Durable completed-update cursor and exact shared index rows | +| Normalizer weighting | Overlapping full windows, batches `max(512,B)`, final partial batch dropped | Overlapping full windows including final partial batch unless exact shared stats supplied | +| Normalizer estimator | RunningMeanStd FP32 with initial count `1e-4`, mean zero, variance one | Streaming FP64 population moments, then FP32 frozen buffers | +| Validation order | Shuffled loader, full batches only, up to 100 batches, persistent/exhaustion behavior | Same first up-to-100 validation batches for both trainers, partial final batch included | +| Initial checkpoint candidacy | Epoch-zero validation may select untrained initialization as best | First selection after a trained validation interval | +| Validation collections | Saved checkpoint configs contain separate experimental and passive datasets, each with its own best checkpoint | One canonical validation split and one minimum normalized-MSE selector | + +For strict sample-order reproduction, capture the original loader's actual emitted +indices; an approximate "epoch shuffle" is insufficient because of its exhaustion +behavior. Shared mean/variance artifacts can make both trainers use exactly the +same normalization, but their origin and estimator must be recorded. The normalizer +count prior `1e-4` is distinct from the normalization denominator epsilon `1e-5`. + +These controlled experiments can establish migration correctness/noninferiority +under identical data and optimizer conditions. Matching architecture and update +budget alone does **not** establish literal original validation/statistics/sampling +reproduction or equivalence to published paper numbers. Further protocol work or +explicit evidence remains necessary for those claims; no data-size floor or +alternate acceptance threshold is introduced by this audit. + +Source recipe identity must also be explicit: current `train/cfg` paths name 1M +training transitions, while saved pretrained configs name 5M for Cartpole and 10M +for Ant, plus experimental/passive validation files. These labels are separate +from the required immutable dataset fingerprints and do not justify assuming a +single universal "official dataset". + +## New source-generation provenance + +New calls to `generate_official_dataset` emit +`physics_contract_version: official-warp-physical-v2` and preserve: + +- `source_physical_arrays.gravity`: the actual original model gravity vector in + the source world frame. +- `teacher_origin.gravity`, `gravity_magnitude`, `up_axis` and contract version. +- `source_generator_files`: SHA256 identities for the original sampler, common + sampling definitions, environment/contact code and robot-specific environment. +- `source_generator_sha256`: SHA256 of the canonical sorted generator-file map. +- `generation_adapter_sha256`: the exact EmbodiChain reference-wrapper file hash. + +The source vector remains Y-up even when canonical state transport rotates to Z-up. +The measured original environment magnitude is 9.81: its explicit +`Environment.gravity=-9.81` overrides the Warp Model fallback of 9.80665. This is +a measured original-runtime value, not an inferred Newton setting. Actual +CPU-generation tests inspect the vector, magnitude, up-axis, +file hashes and their preservation in canonical metadata. + +Existing smoke datasets/checkpoints are **not rewritten or relabeled**. Old +artifacts lacking measured gravity retain that provenance gap; they cannot acquire +a v2 physical contract by changing a model manifest annotation. diff --git a/design/reports/neural-dynamics-pr-status-20261020.md b/design/reports/neural-dynamics-pr-status-20261020.md new file mode 100644 index 000000000..061befcb1 --- /dev/null +++ b/design/reports/neural-dynamics-pr-status-20261020.md @@ -0,0 +1,277 @@ +# NeRD draft PR status and evidence index — 2026-10-20 + +**Development implementation and bounded experiments are available for review. +Full accuracy qualification and release integration remain pending.** This is the +current status index; the dated reports below retain their original experimental +scope and may describe limitations subsequently addressed by later reports. + +The authoritative acceptance contract is [the complete design](../neural-dynamics-solvers.md). +There is no claim of full-budget training, production dynamics accuracy, or +literal reproduction of published paper results. + +## Implemented scope + +| Work package | Current progress | Evidence boundary | +|---|---|---| +| Named Newton solver extensions | DexSim discovery, capability/lifecycle and public functional/metadata APIs implemented in the adjacent SDK workstream | Tested local SDK artifact; SDK source/release is a separate prerequisite | +| NeRD integration | Installed `dexsim.newton_solvers` entry point named `nerd`, native and functional stepping, state/history/reset, supported controls and physical-profile checks | Homogeneous Cartpole/Ant scenes, CPU/CUDA, first-order gradients | +| Shared mathematics and network | Simulation-independent codec, contact inputs, upstream causal Transformer, normalizers and portable checkpoint conversion | Official checkpoint/fixture and update parity; no full-accuracy claim | +| Data lifecycle | EmbodiChain-owned analytical teacher generation, HDF5 committed shards, trajectory splits and bounded lazy sequence windows | Genuine bounded historical/current datasets; full-size distributions pending | +| Training lifecycle | Supervised train/fine-tune/resume, train-only statistics, offline evaluation and export, CLI and task workflow YAMLs | Exact-resume and original-trainer checks, 40 bounded training runs | +| Differentiable control | Public manager/Scene adapter and separate `DifferentiableDynamicsEnv`, detached publication, partial reset and explicit BPTT boundary | State/control/history/model gradients and scoped CPU/CUDA integration | +| Comparison tooling | Frozen two-teacher/two-trainer plans, shared initialization/statistics/order, evidence validation, paired bootstrap and historical runtime | Formal reports retain `pending` when required evidence is absent | +| Documentation and project context | Full design, guide/API entries, task/configuration examples and routed neural-dynamics context | API coverage and focused documentation checks; broader build caveats below | + +The only deployment mechanism is installed entry-point metadata selected by +`solver_type: nerd`. There is no `custom + provider` configuration, direct +provider-path deployment or public imperative registration API. One solver owns +the scene. CUDA Graph, second-order derivatives, physical-parameter gradients, +heterogeneous scene coupling and DDP are outside v1. + +## Reference and execution identity + +| Component | Identity | +|---|---| +| Official generation/training source | `NVlabs/neural-robot-dynamics` main `370e89d2b8ce68aa3578aa7812859c681b7502bc` | +| Official Newton-preview/checkpoint reference | `476dc67fccc6a99be293bf611f897b1c06f13e30` | +| Current runtime | Python 3.11.14, Torch 2.7.1+cu128, Newton 1.6.0, Warp 1.17.0 | +| Historical runtime | Cartpole Warp 1.8.0; Ant Warp 1.5.1; isolated fresh processes | +| Workstation | NVIDIA GeForce RTX 5090, driver 595.58.03, CUDA runtime 12.8 | +| EmbodiChain starting revision | `66114674668b258c15feae43c8939a0ef1ec4ac4` | +| Adjacent DexSim starting revision | `4f093c4a88ffd869570faf800341a1ad496b7fec` | + +The historical and current teacher domains are physically different. Measured +examples include Cartpole masses 50/2.5 versus 1/10 and Ant torso mass 31.65345 +versus approximately 0.48388. Historical/current Ant friction also differs. +Asset byte identity therefore cannot qualify a physical model. Compare trainers +within each measured domain; historical weights cannot acquire current-Newton +deployment compatibility by editing manifest annotations. + +## Executed experiments + +### Official implementation correctness + +Official pretrained checkpoints were compared on CPU FP32 at history lengths +1, 9 and 10. Maximum forward, input-gradient, parameter-gradient and +three-Adam-update parameter differences were **zero** for both robots. + +| Robot | Maximum feature error | Maximum target error | Maximum decoded-state error | +|---|---:|---:|---:| +| Cartpole | 1.19e-7 | 1.19e-7 | 1.19e-7 | +| Ant | 1.43e-6 | 2.52e-7 | 2.98e-7 | + +Receipts: [Cartpole](nerd-cartpole-parity-20261020.json), +[Ant](nerd-ant-parity-20261020.json). These are fixture-level correctness results +under `atol=1e-5`, `rtol=1e-4`, rather than general prediction accuracy. + +### Controlled five-seed training + +Two robots × two teacher datasets × two trainers × five seeds = **40 completed +runs**, each with **100 optimizer updates**, batch 512 and history 10. Official +network widths are retained: embedding 192/384, six layers, twelve heads, FP32. +Paired runs share exact initialization, statistics, indices and selection rules. + +Current data: 64 trajectories × 32 transitions, splits 57/1/6; default initial +states and Ant effort scale 0.05. Fresh historical data: 48 × 100 transitions, +splits 32/8/8, with measured historical physical-v2 gravity provenance. + +Every short-run offline metric has a one-sided 95% paired-bootstrap upper bound +below the predeclared 5% margin. The largest is **2.1676%**, current-Newton +Cartpole velocity error. This conditional comparison concerns agreement between +training implementations at 100 updates; it does not qualify the learned model. + +The [five-seed report](neural-dynamics-bounded-five-seed-20261020.md) contains +all twelve metrics. The [machine-readable summary](neural-dynamics-five-seed-20261020.json) +retains experiment identities. Observed sequential matrix execution was about +98 s for Cartpole and 139 s for Ant, excluding preparation; do not extrapolate +this small-data timing to the full corpus. + +### Historical/current rollout comparison + +Seed 0 was evaluated in 24 runs: two robots × two domains × two trainers × +passive/random-effort/fixed-policy modes. Each run used three trajectories, +horizons 1/10/12 and batches 2 + 1. All twelve pairs have matching teacher-content +and policy identities, with source/model/manifest/HDF5/JSONL hashes checked. + +The historical harness uses the original environment, Featherstone teacher and +Transformer integrator. All twelve original-versus-portable restoration probes +had maximum absolute error zero through eleven steps, including history rollover. +Historical runs had no nonfinite states or failures within this bounded horizon. + +Material model errors remain: + +| Scenario | Teacher | Official-trained model | EmbodiChain-trained model | +|---|---:|---:|---:| +| Historical Ant fixed-PD return | 9.7618 | 1.4588 | 1.5726 | +| Current-Newton Ant fixed-PD return | 4.3053 | 2.2850 | 2.3616 | +| Current Cartpole passive failure rate | 1/3 | 2/3 | 2/3 | +| Current Cartpole random-effort failure rate | 2/3 | 2/3 | 2/3 | +| Current Cartpole fixed-PD failure rate | 0/3 | 1/3 | 1/3 | + +The fixed PD is a diagnostic baseline, not the paper policy. Cartpole returns +stop at failure; a less negative return cannot hide an earlier failure. +See [historical measured errors](neural-dynamics-historical-rollout-20261020.md) +and [bounded pairing evidence](neural-dynamics-bounded-pairing-20261020.json). + +### Input and training-compute measurements + +Bounded grouped HDF5 reads preserve request order, repeated-index ownership, +prefixes, split boundaries and exact training/resume behavior. Reads span at most +266 frames in the measured configuration and keep at most four file handles; +there is no whole-dataset cache or speculative prefetch. + +| CPU read/collate/encode workload, B512/H10 | Individual reads | Grouped reads | Speedup | +|---|---:|---:|---:| +| Current Cartpole | 257.87 ms | 12.75 ms | 20.23× | +| Current Ant | 285.03 ms | 25.88 ms | 11.01× | +| Historical Cartpole | 241.61 ms | 10.62 ms | 22.75× | +| Historical Ant | 271.97 ms | 20.77 ms | 13.10× | +| 128-shard, 512-distinct-trajectory stress | 337.13 ms | 181.72 ms | 1.86× | + +Warm-page, single-CPU-thread measurements; the stress corpus tests poor I/O +reuse and is not physics data qualification. Full scope and reproduction commands +are in [the input-pipeline report](neural-dynamics-input-20261020.md). + +Device-resident official-width FP32 training, TF32 disabled, measured medians: +Cartpole 7.539/7.495 ms and Ant 14.488/14.263 ms per update for official/migrated +implementations. The full 40-run, 4,995,000-update matrix has an optimistic +**607.52 GPU-hour / 25.31-day compute-only lower bound** on this RTX 5090. +Data loading, validation, generation, checkpoints and rollout evaluation are extra. +Full-corpus end-to-end throughput is unmeasured. + +## Correctness audit and protocol limits + +- Physical signature `nerd-reference-physics-v2` binds fixed-root rotation, + topology/flags, contact materials and effective collision relationships. + Datasets/exports were regenerated; older artifacts were not relabeled. +- Training metadata may add unrelated annotations but cannot change/drop + measured physics fields or invent missing reserved deployment provenance. +- Deployment fails early for unsupported controls, incompatible scenes/physics + and non-FP32 NeRD weights/buffers. Empty functional reset masks are true no-ops. +- Official generation has an uninitialized gravity-feature defect. Raw output is + preserved; opt-in transport repair and Y-up/twist-to-Z-up/COM conversion are + recorded. Actual original gravity is `[0, -9.81, 0]`. +- Acceptance checks measured residuals and immutable artifacts, actual ordered + inputs and FP32 weights, common teacher content/policy, completed update budget + and selected test split. Missing/undefined evidence remains pending or fails. +- The controlled comparison changes outer-loop sampling/statistics/validation + relative to the literal original recipe. Shared replacement indices, + population moments and one ordered validation split are disclosed in + [the official protocol audit](neural-dynamics-official-protocol-audit-20261020.md). + Even a full-budget controlled pass would not by itself prove paper reproduction. + +The translated official budget is 4,995,000 updates because epoch zero evaluates +without training. LR starts at 0.0009991 and ends at 0.0001009, with interval 5000 +and offset 1. The source and saved pretrained configs also name different dataset +sizes/validation collections; an exact paper recipe must be identified explicitly. + +## Local installation and dependency handoff + +The adjacent SDK implements the required generic extension APIs. The tested +unpublished local SDK version is +`0.5.1+g4f093c4a88ff.d2bf6eb9a5b93fd36`; wheel SHA-256: +`2b0a37872ddb690392450b4b9764f128d27b8e53138d7973e87459e722504b3f`. +The paired EmbodiChain wheel SHA-256 is +`722046aab36844669862b9feff3165800576316872858d1c0719f1b094926c38`. + +Only the paired EC wheel's dependency metadata and RECORD were overlaid to +require that exact SDK. Repository `pyproject.toml` still requires +`dexsim_engine==0.5.1rc1`, which does not establish compatibility with the new +APIs. SDK source/release integration and the final source dependency pin remain +merge prerequisites. This PR does not publish SDK source or wheels. + +Independently installed SDK regressions passed 157 tests (21 deselected). +Installed EC historical CPU probes passed for both robots without loading or +attempting to load DexSim/Newton. Tiny installed historical CUDA probes also +passed at B1/count1/horizons1/3: Cartpole max difference zero, Ant 1.4782e-5 +within `atol=1e-5`, `rtol=1e-4`. Old Warp emits a cuDeviceGetUuid initialization +diagnostic; successful tiny execution is not general old-Warp GPU qualification. +Artifacts and validation receipts are indexed in +[the latest local artifact manifest](neural-dynamics-local-artifacts-historical-20261020.json). + +## Validation scope and remaining work + +Earlier focused groups cover CLI/manager/configuration, task packaging, +CPU/CUDA gradients and stream behavior, partial resets, exact training resume, +reference transport and report integrity. Counts in different dated reports +overlap and must not be added into a single total. + +The submission-time combined regression selected 479 cases with the actual +official reference, isolated historical runtimes and CUDA tests enabled: +**476 passed, 2 skipped, 1 failed** in 191.49 s. The failure is +`test_robot_mounted_camera_backend_lifecycle[newton]`: `Robot.link_names` indexes +an empty entity list while the mounted camera is declared before articulation +binding. The same test and failure were reproduced from an unchanged archive of +the starting EmbodiChain commit against the same local SDK (8.16 s). This is +baseline evidence under this SDK environment, not a passing test or a claim about +all published SDK combinations. It remains recorded rather than being hidden or +fixed as unrelated scope. + +Initial combined collection also exposed three new test-file names colliding +with modules elsewhere under pytest's default import mode. The new files were +renamed to `test_dynamics_task_reference.py`, `test_dynamics_metrics.py` and +`test_dynamics_runtime_identity.py`; production behavior is unchanged. Default +collection with the previously conflicting existing modules now succeeds (106 +cases). The full behavioral run used `--import-mode=importlib`. Submission +commands, final counts, skips and log hashes are retained in +[the validation receipt](neural-dynamics-pr-validation-20261020.json). +The three renamed modules subsequently passed all 17 cases with GPU tests +enabled. The two checkpoint cases skipped by the combined command were also +rerun with the explicit pinned Newton-preview checkout: both passed (7.25 s). +These follow-ups overlap the combined selection and do not increase its size. + +Public API coverage is 2748/2748. Black 26.3.1, changed-source structural checks +and context validation pass. New guide/API pages built with warnings treated as +errors. A broader Sphinx dummy build succeeded from a clean source snapshot with +785 existing warnings; that snapshot predates the final historical harness phase. +Do not claim a warning-free full current documentation build. + +The two vendored official assets retain their exact upstream bytes and six +upstream trailing-whitespace lines. SHA-256 identities match the pinned main +checkout (`ant.xml`: `23e9a720465a71fde3fd9fd2d5a18f8c11ec642f3dca8921c66cb4e169421fa7`; +`cartpole_single.urdf`: `f31276531c03050a3ac3a3e24443899f38a4424586db17140978383433dec237`). +Staged `git diff --check` passes with only these byte-preserved asset files +excluded; the unrestricted check reports their inherited whitespace. + +Affected simulation, data, environment and learning context has been revised, +and the new neural-dynamics owner added to MAP. Broader routing matches do not +require unrelated edits where their contracts remain accurate. + +- [ ] Release/integrate the extension-enabled SDK and update the source dependency. +- [ ] Run full-size data distributions and 40 × 4,995,000-update training cells. +- [ ] Run five-seed, 2048-trajectory evaluation at horizons 1/10/100/1000 for + passive, random-effort and fixed-policy scenarios. +- [ ] Pass final accuracy, stability/nonfinite/failure and policy-return gates. +- [ ] Measure representative full-corpus end-to-end throughput. +- [ ] Establish the exact data/statistics/sampling/validation recipe before + asserting literal published-result reproduction. + +No full-budget training job or long evaluation is running. The original +qualification goal is blocked pending the experiment scope/compute-budget +decision. Creating this Draft makes current code, design and evidence reviewable; +it does not close the qualification goal. + +## Complete checked-in evidence index + +| Record | Scope | +|---|---| +| [Complete design](../neural-dynamics-solvers.md) | Architecture, state/data/training lifecycle, comparisons and formal gates | +| [Initial evidence](neural-dynamics-20261020.md) | Reference fixtures, first real-data experiments, compute budget and initial installation | +| [v2 implementation audit](neural-dynamics-implementation-v2-20261020.md) | Physics-signature fixes, regenerated current-domain flow and integration checks | +| [v2 smoke receipt](neural-dynamics-v2-smoke-20261020.json) | Measured current-domain end-to-end artifacts | +| [Input-pipeline experiment](neural-dynamics-input-20261020.md) | Bounded reader correctness, throughput and stress limits | +| [Official protocol audit](neural-dynamics-official-protocol-audit-20261020.md) | Architecture, optimizer schedule and reproduction differences | +| [Five-seed comparison](neural-dynamics-bounded-five-seed-20261020.md) | Latest 40-cell bounded matrix and both-domain rollout limitations | +| [Five-seed metrics](neural-dynamics-five-seed-20261020.json) | Machine-readable aggregate results | +| [Historical rollout](neural-dynamics-historical-rollout-20261020.md) | Original-runtime protocol, measured state errors and returns | +| [Pairing receipt](neural-dynamics-bounded-pairing-20261020.json) | Both-domain evidence/teacher/policy integrity | +| [Cartpole parity](nerd-cartpole-parity-20261020.json), [Ant parity](nerd-ant-parity-20261020.json) | Official checkpoint and numerical residual identities | +| [Initial artifacts](neural-dynamics-local-artifacts-20261020.json), [v2 artifacts](neural-dynamics-local-artifacts-v2-20261020.json) | Historical artifact snapshots; superseded for latest installation | +| [Latest artifacts](neural-dynamics-local-artifacts-historical-20261020.json) | Current paired wheel identities and installed CPU/CUDA probes | +| [Draft submission validation](neural-dynamics-pr-validation-20261020.json) | Combined regression, baseline failure reproduction and test-discovery fix | + +The committed reports/aggregate receipts are reviewable in GitHub. Large raw +HDF5 datasets, weights, traces, plots, source snapshots and complete process logs +remain local at the paths in those reports; hashes identify them, but a local +`/tmp` path is not a remotely downloadable artifact. No binary or durable public +experiment archive is implied by this PR. diff --git a/design/reports/neural-dynamics-pr-validation-20261020.json b/design/reports/neural-dynamics-pr-validation-20261020.json new file mode 100644 index 000000000..2e3220dfd --- /dev/null +++ b/design/reports/neural-dynamics-pr-validation-20261020.json @@ -0,0 +1,182 @@ +{ + "date": "2026-10-20", + "timezone": "Asia/Shanghai", + "scope": "Draft submission verification; not formal accuracy qualification", + "tests": { + "counts": { + "tests": 479, + "failures": 1, + "errors": 0, + "skipped": 2, + "time_seconds": 191.45 + }, + "skips": [ + { + "test": "tests.learning.models.test_nerd::test_explicit_official_checkpoint_forward_gradient_and_adam[Cartpole]", + "reason": "Set NERD_REFERENCE_ROOT to the pinned trusted upstream checkout" + }, + { + "test": "tests.learning.models.test_nerd::test_explicit_official_checkpoint_forward_gradient_and_adam[Ant]", + "reason": "Set NERD_REFERENCE_ROOT to the pinned trusted upstream checkout" + } + ], + "failures": [ + { + "test": "tests.sim.test_sim_manager::test_robot_mounted_camera_backend_lifecycle[newton]", + "message": "IndexError: list index out of range", + "trace": "tmp_path = PosixPath('/tmp/pytest-of-dex/pytest-752/test_robot_mounted_camera_back1')\nbackend = 'newton'\n\n @pytest.mark.gpu\n @pytest.mark.requires_sim\n @pytest.mark.parametrize(\"backend\", [\"default\", \"newton\"])\n def test_robot_mounted_camera_backend_lifecycle(tmp_path: Path, backend: str) -> None:\n \"\"\"Mounted cameras follow state writes and survive physical topology rebuilds.\"\"\"\n urdf_path = tmp_path / \"camera_robot.urdf\"\n urdf_path.write_text(\"\"\"\n \n \n \n \n \n \n \n \n \n \n \n \"\"\")\n physics_cfg = (\n DefaultPhysicsCfg(device=\"cpu\")\n if backend == \"default\"\n else NewtonPhysicsCfg(\n device=\"cuda:0\",\n use_cuda_graph=False,\n num_substeps=1,\n solver_cfg={\"solver_type\": \"mujoco_warp\"},\n )\n )\n sim = SimulationManager(\n SimulationManagerCfg(\n headless=True,\n enable_entity_gizmo=False,\n physics_cfg=physics_cfg,\n render_cfg=RenderCfg(renderer=\"hybrid\"),\n )\n )\n try:\n camera_cfg = CameraCfg(\n uid=\"mounted\",\n width=32,\n height=24,\n extrinsics=CameraCfg.ExtrinsicsCfg(parent=\"tool\", pos=(0.1, 0.0, 0.1)),\n )\n> robot = sim.add_robot(\n RobotCfg(\n uid=\"robot\",\n fpath=str(urdf_path),\n init_pos=(0.0, 0.0, 0.5),\n enable_sensor=True,\n sensors=[camera_cfg],\n )\n )\n\ntests/sim/test_sim_manager.py:2013: \n_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ \nembodichain/lab/sim/sim_manager.py:3453: in add_robot\n self.add_sensor(sensor_cfg)\nembodichain/lab/sim/sim_manager.py:3978: in add_sensor\n parent_nodes = self._resolve_spawn_sensor_parent_nodes(\nembodichain/lab/sim/sim_manager.py:4042: in _resolve_spawn_sensor_parent_nodes\n return resolve_parent_nodes(parent, assets, self.num_envs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\nembodichain/lab/sim/sensors/attachment.py:62: in resolve_parent_nodes\n if link_name not in asset.link_names:\n ^^^^^^^^^^^^^^^^\n/home/dex/miniconda3/envs/open/lib/python3.11/functools.py:1001: in __get__\n val = self.func(instance)\n ^^^^^^^^^^^^^^^^^^^\n_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ \n\nself = : declared 1 Spawn articulations | uid: robot | device: cuda:0 | control_parts: None, solvers: {}\n\n @cached_property\n def link_names(self) -> List[str]:\n \"\"\"Get the names of the links in the articulation.\n \n Returns:\n List[str]: The names of the links in the articulation.\n \"\"\"\n if self._data is not None:\n return self._data.link_names\n> return self._entities[0].get_link_names()\n ^^^^^^^^^^^^^^^^^\nE IndexError: list index out of range\n\nembodichain/lab/sim/objects/articulation.py:1021: IndexError" + } + ], + "junit": { + "path": "/tmp/neural-dynamics-pr-tests-20261020.xml", + "sha256": "dbef6c757b88c57bc0bef60b17d4fe389a0e6e434c33c3a038166a561ea846dc" + }, + "log": { + "path": "/tmp/neural-dynamics-pr-tests-20261020.log", + "sha256": "a47cabc8206fedefd7fd70dcd4b6fba466363caa0fa80cfec4f6899a5c6cd4b8" + }, + "command": "OMP_NUM_THREADS=2 MKL_NUM_THREADS=2 NERD_REFERENCE_SOURCE=/tmp/embodichain-nerd-main-reference NERD_REFERENCE_DEPS=/tmp/nerd-reference-deps NERD_MAIN_REFERENCE_ROOT=/tmp/embodichain-nerd-main-reference NERD_WARP18_ROOT=/tmp/nerd-reference-deps NERD_WARP151_ROOT=/tmp/nerd-reference-warp151 python -m pytest -q --import-mode=importlib --run-gpu tests/compute/dynamics tests/data_pipeline/dynamics tests/learning/models/test_nerd.py tests/learning/dynamics tests/sim/physics/solvers tests/sim/diff tests/gym/envs/test_differentiable_dynamics_env.py tests/gym/envs/tasks/test_neural_dynamics_tasks.py tests/cli/test_dynamics.py tests/sim/test_sim_manager.py tests/sim/test_sim_manager_cfg.py tests/test_agent_context_map.py tests/test_agent_context_tools.py", + "source_filename_scope": "Before collision-avoidance test renames; renamed code rerun separately" + }, + "baseline_failure_reproduction": { + "counts": { + "tests": 1, + "failures": 1, + "errors": 0, + "skipped": 0, + "time_seconds": 8.152 + }, + "skips": [], + "failures": [ + { + "test": "tests.sim.test_sim_manager::test_robot_mounted_camera_backend_lifecycle[newton]", + "message": "IndexError: list index out of range", + "trace": "tmp_path = PosixPath('/tmp/pytest-of-dex/pytest-753/test_robot_mounted_camera_back0')\nbackend = 'newton'\n\n @pytest.mark.gpu\n @pytest.mark.requires_sim\n @pytest.mark.parametrize(\"backend\", [\"default\", \"newton\"])\n def test_robot_mounted_camera_backend_lifecycle(tmp_path: Path, backend: str) -> None:\n \"\"\"Mounted cameras follow state writes and survive physical topology rebuilds.\"\"\"\n urdf_path = tmp_path / \"camera_robot.urdf\"\n urdf_path.write_text(\"\"\"\n \n \n \n \n \n \n \n \n \n \n \n \"\"\")\n physics_cfg = (\n DefaultPhysicsCfg(device=\"cpu\")\n if backend == \"default\"\n else NewtonPhysicsCfg(\n device=\"cuda:0\",\n use_cuda_graph=False,\n num_substeps=1,\n solver_cfg={\"solver_type\": \"mujoco_warp\"},\n )\n )\n sim = SimulationManager(\n SimulationManagerCfg(\n headless=True,\n enable_entity_gizmo=False,\n physics_cfg=physics_cfg,\n render_cfg=RenderCfg(renderer=\"hybrid\"),\n )\n )\n try:\n camera_cfg = CameraCfg(\n uid=\"mounted\",\n width=32,\n height=24,\n extrinsics=CameraCfg.ExtrinsicsCfg(parent=\"tool\", pos=(0.1, 0.0, 0.1)),\n )\n> robot = sim.add_robot(\n RobotCfg(\n uid=\"robot\",\n fpath=str(urdf_path),\n init_pos=(0.0, 0.0, 0.5),\n enable_sensor=True,\n sensors=[camera_cfg],\n )\n )\n\ntests/sim/test_sim_manager.py:1991: \n_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ \nembodichain/lab/sim/sim_manager.py:3293: in add_robot\n self.add_sensor(sensor_cfg)\nembodichain/lab/sim/sim_manager.py:3818: in add_sensor\n parent_nodes = self._resolve_spawn_sensor_parent_nodes(\nembodichain/lab/sim/sim_manager.py:3882: in _resolve_spawn_sensor_parent_nodes\n return resolve_parent_nodes(parent, assets, self.num_envs)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\nembodichain/lab/sim/sensors/attachment.py:62: in resolve_parent_nodes\n if link_name not in asset.link_names:\n ^^^^^^^^^^^^^^^^\n/home/dex/miniconda3/envs/open/lib/python3.11/functools.py:1001: in __get__\n val = self.func(instance)\n ^^^^^^^^^^^^^^^^^^^\n_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ \n\nself = : declared 1 Spawn articulations | uid: robot | device: cuda:0 | control_parts: None, solvers: {}\n\n @cached_property\n def link_names(self) -> List[str]:\n \"\"\"Get the names of the links in the articulation.\n \n Returns:\n List[str]: The names of the links in the articulation.\n \"\"\"\n if self._data is not None:\n return self._data.link_names\n> return self._entities[0].get_link_names()\n ^^^^^^^^^^^^^^^^^\nE IndexError: list index out of range\n\nembodichain/lab/sim/objects/articulation.py:1021: IndexError" + } + ], + "junit": { + "path": "/tmp/neural-dynamics-pr-baseline-camera-20261020.xml", + "sha256": "5d8fe605afd57dfa3a33bf0cb494e5e9a1d25b15b5138a1dee1f09895ef309c1" + }, + "log": { + "path": "/tmp/neural-dynamics-pr-baseline-camera-20261020.log", + "sha256": "c73694613f15d20f53fd6de4b06204623cf8dd3bc4f201dfb841a96a66c36c58" + }, + "source_commit": "66114674668b258c15feae43c8939a0ef1ec4ac4", + "source_snapshot": "/tmp/neural-dynamics-pr-baseline-20261020-um6vpq8e", + "same_local_sdk_environment": true, + "same_failure_signature": true + }, + "renamed_test_regression": { + "counts": { + "tests": 17, + "failures": 0, + "errors": 0, + "skipped": 0, + "time_seconds": 1.864 + }, + "skips": [], + "failures": [], + "junit": { + "path": "/tmp/neural-dynamics-pr-renamed-tests-20261020.xml", + "sha256": "924ce321b65f775ec7847e67519a8fe3b105d6a3eb65cb05c61cac8db56e7359" + }, + "log": { + "path": "/tmp/neural-dynamics-pr-renamed-tests-20261020.log", + "sha256": "4e7b24fcf877bf8a085a5251afcc7e8c5b2ff1972ec87b6893692c6d6d2bed11" + }, + "command": "python -m pytest -q --run-gpu tests/compute/dynamics/test_dynamics_task_reference.py tests/compute/dynamics/test_dynamics_metrics.py tests/learning/dynamics/test_dynamics_runtime_identity.py" + }, + "default_collection": { + "exit_code": 0, + "tests_collected": 106, + "log": { + "path": "/tmp/neural-dynamics-pr-default-collection-20261020.log", + "sha256": "0d8f30c10f8b0b051b5d070c4868ab83681191a894abb360180e9654389fa914" + }, + "command": "python -m pytest --collect-only -q tests/compute/dynamics tests/learning/dynamics tests/learning/test_runtime.py tests/learning/test_evaluation.py" + }, + "initial_collection_conflict": { + "exit_code": 2, + "reason": "Default pytest prepend import mode loaded two distinct test_reference.py files as the same module; additional runtime/evaluation names were also made unique", + "log": { + "path": "/tmp/neural-dynamics-pr-tests-collection-20261020.log", + "sha256": "12579ae7d616708224ec2af4c36f7e601e09fde1413dfd9da4bc33e03e768d49" + } + }, + "checks": { + "black": { + "version": "26.3.1", + "command": "python -m black .", + "unchanged_files": 1398, + "status": "passed" + }, + "api_docs": { + "command": "python docs/scripts/check_api_docs.py", + "documented_exports": 2748, + "total_exports": 2748, + "status": "passed" + }, + "context": { + "commands": [ + "python .agents/skills/project-dev-context/scripts/context.py check", + "python .agents/skills/project-dev-context/scripts/context.py affected --base origin/main --explain" + ], + "status": "passed", + "review": "Changed simulation/data/environment/learning contracts revised; neural-dynamics topic added. Unchanged contracts matched by broad watch scopes need no unrelated rewrite." + }, + "python_structure": { + "changed_or_new_files": 63, + "checks": [ + "Apache header", + "future annotations", + "public module __all__" + ], + "status": "passed" + }, + "diff_whitespace": { + "status": "passed_except_preserved_upstream_assets", + "inherited_whitespace_lines": 6, + "byte_preserved_exceptions": { + "embodichain/data/nerd/ant.xml": "23e9a720465a71fde3fd9fd2d5a18f8c11ec642f3dca8921c66cb4e169421fa7", + "embodichain/data/nerd/cartpole_single.urdf": "f31276531c03050a3ac3a3e24443899f38a4424586db17140978383433dec237" + } + } + }, + "runtime_scope": { + "python": "3.11.14", + "torch": "2.7.1+cu128", + "newton": "1.6.0", + "warp": "1.17.0", + "gpu": "NVIDIA GeForce RTX 5090", + "sdk_source_package": "/home/dex/workspace/sources/dexsim/build_Release/lib/python_package", + "installed_metadata_version": "0.5.1", + "native_version_in_test_log": "0.5.1+00106a1db10e330c62cb62d8015b8474e997701b", + "note": "Source-backed SDK checkout, not a test of the unchanged published rc1 dependency." + }, + "overall_status": "baseline_failure_recorded; formal_accuracy_pending", + "limitations": [ + "Counts from overlapping test groups must not be added.", + "Large raw experiment artifacts and full local logs are not uploaded by this receipt.", + "Full training budgets, long-horizon qualification and a release-consistent source dependency remain pending." + ], + "official_checkpoint_followup": { + "command": "OMP_NUM_THREADS=2 MKL_NUM_THREADS=2 NERD_REFERENCE_ROOT=/tmp/embodichain-nerd-analysis-20261019 python -m pytest -q tests/learning/models/test_nerd.py::test_explicit_official_checkpoint_forward_gradient_and_adam", + "counts": { + "tests": 2, + "failures": 0, + "errors": 0, + "skipped": 0 + }, + "source_commit": "476dc67fccc6a99be293bf611f897b1c06f13e30", + "log": { + "path": "/tmp/neural-dynamics-pr-official-checkpoints-20261020.log", + "sha256": "c412967b03ecd435d39a6db77259bebf50685116e714f8b90dfa679ecccbb408" + }, + "junit": { + "path": "/tmp/neural-dynamics-pr-official-checkpoints-20261020.xml", + "sha256": "47d56dec913c874c2d4f3b4c34f48bb3354da8d20cfbefc1d3da1c378487250d" + }, + "scope": "Explicitly execute the two reference-checkpoint tests skipped by the combined command; not an independent enlarged population." + } +} diff --git a/design/reports/neural-dynamics-v2-smoke-20261020.json b/design/reports/neural-dynamics-v2-smoke-20261020.json new file mode 100644 index 000000000..2b6aa4f8a --- /dev/null +++ b/design/reports/neural-dynamics-v2-smoke-20261020.json @@ -0,0 +1,577 @@ +{ + "scope": "end-to-end smoke, not full official accuracy qualification", + "robots": { + "cartpole": { + "dataset": "/tmp/embodichain-nerd-v2-qualification/cartpole/data", + "model": "/tmp/embodichain-nerd-v2-qualification/cartpole/export", + "training_updates": 100, + "network_scope": "one-layer smoke architecture", + "trajectories": 64, + "training_horizon": 32, + "offline": { + "normalized_mse": 0.8590003512192356, + "target_mse": 1.3739163670486405, + "q_error_norm": 0.047975492563800536, + "qd_error_norm": 1.218408178937608, + "valid_positions": 1380, + "windows": 138, + "split": "test", + "dataset_fingerprint": "3c4dc35103ad0be3fcee6357d3e50002bb7515abfa4e5de0dcb16d427b1332eb", + "qualification": "offline_only" + }, + "rollouts": { + "passive": { + "source": "/tmp/embodichain-nerd-v2-qualification/cartpole/audited-passive.json", + "source_sha256": "d780c6ad75764acd1ad2c354f2aac939f1207e3e040b3049e4b16e6d3dcc8975", + "trajectories": 4, + "horizons": { + "1": { + "step": 1, + "q_mae": 0.007476515136659145, + "qd_mae": 0.08783073537051678, + "base_position_error": 0.0, + "base_orientation_error": 0.0, + "valid_trajectories": 4, + "teacher_nonfinite_count": 0, + "neural_nonfinite_count": 0, + "teacher_failed_count": 0, + "failed_trajectories": 0, + "teacher_nonfinite_rate": 0.0, + "nonfinite_rate": 0.0, + "teacher_failure_rate": 0.0, + "failure_rate": 0.0, + "teacher_policy_return": 0.0, + "policy_return": -0.0024109298828989267, + "numerically_invalid_trajectories": 0 + }, + "10": { + "step": 10, + "q_mae": 0.04696366265416145, + "qd_mae": 0.3728337919339538, + "base_position_error": 0.0, + "base_orientation_error": 0.0, + "valid_trajectories": 4, + "teacher_nonfinite_count": 0, + "neural_nonfinite_count": 0, + "teacher_failed_count": 0, + "failed_trajectories": 0, + "teacher_nonfinite_rate": 0.0, + "nonfinite_rate": 0.0, + "teacher_failure_rate": 0.0, + "failure_rate": 0.0, + "teacher_policy_return": 0.0, + "policy_return": -0.7320659160614014, + "numerically_invalid_trajectories": 0 + }, + "32": { + "step": 32, + "q_mae": 0.15463752974756062, + "qd_mae": 1.1777281721006148, + "base_position_error": 0.0, + "base_orientation_error": 0.0, + "valid_trajectories": 4, + "teacher_nonfinite_count": 0, + "neural_nonfinite_count": 0, + "teacher_failed_count": 0, + "failed_trajectories": 0, + "teacher_nonfinite_rate": 0.0, + "nonfinite_rate": 0.0, + "teacher_failure_rate": 0.0, + "failure_rate": 0.0, + "teacher_policy_return": 0.0, + "policy_return": -20.58209991455078, + "numerically_invalid_trajectories": 0 + } + }, + "failure_rate": 0.0, + "nonfinite_rate": 0.0, + "selected_teacher_sha256": "00c1d144aae2b3e76cd809733f4305e14d1c52a36a39faa701c3a4ac46459174", + "runtime_identity": { + "device": "cpu", + "device_name": "x86_64", + "python": "3.11.14", + "platform": "Linux-6.8.0-138-generic-x86_64-with-glibc2.35", + "state_dtype": "float32", + "cuda_runtime": "12.8", + "versions": { + "embodichain": "0.3.1", + "dexsim-engine": "0.5.1", + "newton": "1.6.0", + "warp-lang": "1.17.0", + "torch": "2.7.1+cu128" + }, + "source_sha256": { + "collection.py": "e9e5e8cd8eb80a23b831b96116be388ff55b1a874608b62379d564a8bd24c8ab", + "nerd.py": "6731d308e76093e75e530e3cde9df6bf094fad07a999a8f6b9929f4261f242ba", + "_nerd_binding.py": "0240a6eeb2ff23cdafbae4c09da08a2665bc8318fdd29c43199e2ee72a1467ae", + "_reference_assets.py": "ece8077d9e51597efff1925abc96fbeb9bb3998d52031356d8f993db7bc271c0" + } + } + }, + "random_effort": { + "source": "/tmp/embodichain-nerd-v2-qualification/cartpole/audited-random_effort.json", + "source_sha256": "763de0cf46058247b4ee9e428b62a07fc4d64deaeb7be6b557a026aa453b7400", + "trajectories": 4, + "horizons": { + "1": { + "step": 1, + "q_mae": 0.015570371644571424, + "qd_mae": 1.4385360640008003, + "base_position_error": 0.0, + "base_orientation_error": 0.0, + "valid_trajectories": 4, + "teacher_nonfinite_count": 0, + "neural_nonfinite_count": 0, + "teacher_failed_count": 0, + "failed_trajectories": 0, + "teacher_nonfinite_rate": 0.0, + "nonfinite_rate": 0.0, + "teacher_failure_rate": 0.0, + "failure_rate": 0.0, + "teacher_policy_return": -1.1786372736096382, + "policy_return": -0.0032918811484705657, + "numerically_invalid_trajectories": 0 + }, + "10": { + "step": 10, + "q_mae": 0.13329091079649516, + "qd_mae": 2.09402509664651, + "base_position_error": 0.0, + "base_orientation_error": 0.0, + "valid_trajectories": 4, + "teacher_nonfinite_count": 0, + "neural_nonfinite_count": 0, + "teacher_failed_count": 1, + "failed_trajectories": 0, + "teacher_nonfinite_rate": 0.0, + "nonfinite_rate": 0.0, + "teacher_failure_rate": 0.25, + "failure_rate": 0.0, + "teacher_policy_return": -19.488471031188965, + "policy_return": -0.9762343615293503, + "numerically_invalid_trajectories": 0 + }, + "32": { + "step": 32, + "q_mae": 0.46816649168795266, + "qd_mae": 2.3633901313951355, + "base_position_error": 0.0, + "base_orientation_error": 0.0, + "valid_trajectories": 4, + "teacher_nonfinite_count": 0, + "neural_nonfinite_count": 0, + "teacher_failed_count": 1, + "failed_trajectories": 0, + "teacher_nonfinite_rate": 0.0, + "nonfinite_rate": 0.0, + "teacher_failure_rate": 0.25, + "failure_rate": 0.0, + "teacher_policy_return": -61.668283462524414, + "policy_return": -41.33479881286621, + "numerically_invalid_trajectories": 0 + } + }, + "failure_rate": 0.0, + "nonfinite_rate": 0.0, + "selected_teacher_sha256": "a06c266833ae45d4a93f1862d77f588aa99bf068a34a441634ae0b34e8ee6c11", + "runtime_identity": { + "device": "cpu", + "device_name": "x86_64", + "python": "3.11.14", + "platform": "Linux-6.8.0-138-generic-x86_64-with-glibc2.35", + "state_dtype": "float32", + "cuda_runtime": "12.8", + "versions": { + "embodichain": "0.3.1", + "dexsim-engine": "0.5.1", + "newton": "1.6.0", + "warp-lang": "1.17.0", + "torch": "2.7.1+cu128" + }, + "source_sha256": { + "collection.py": "e9e5e8cd8eb80a23b831b96116be388ff55b1a874608b62379d564a8bd24c8ab", + "nerd.py": "6731d308e76093e75e530e3cde9df6bf094fad07a999a8f6b9929f4261f242ba", + "_nerd_binding.py": "0240a6eeb2ff23cdafbae4c09da08a2665bc8318fdd29c43199e2ee72a1467ae", + "_reference_assets.py": "ece8077d9e51597efff1925abc96fbeb9bb3998d52031356d8f993db7bc271c0" + } + } + }, + "fixed_policy": { + "source": "/tmp/embodichain-nerd-v2-qualification/cartpole/audited-fixed_policy.json", + "source_sha256": "5435cad41ffbbf8a6084a777a5b76ad6289f2a24b3fbdc54ca2aab7a64394622", + "trajectories": 4, + "horizons": { + "1": { + "step": 1, + "q_mae": 0.007476515136659145, + "qd_mae": 0.08783073537051678, + "base_position_error": 0.0, + "base_orientation_error": 0.0, + "valid_trajectories": 4, + "teacher_nonfinite_count": 0, + "neural_nonfinite_count": 0, + "teacher_failed_count": 0, + "failed_trajectories": 0, + "teacher_nonfinite_rate": 0.0, + "nonfinite_rate": 0.0, + "teacher_failure_rate": 0.0, + "failure_rate": 0.0, + "teacher_policy_return": 0.0, + "policy_return": -0.0024109298828989267, + "numerically_invalid_trajectories": 0 + }, + "10": { + "step": 10, + "q_mae": 0.04694647546857596, + "qd_mae": 0.37376401368528606, + "base_position_error": 0.0, + "base_orientation_error": 0.0, + "valid_trajectories": 4, + "teacher_nonfinite_count": 0, + "neural_nonfinite_count": 0, + "teacher_failed_count": 0, + "failed_trajectories": 0, + "teacher_nonfinite_rate": 0.0, + "nonfinite_rate": 0.0, + "teacher_failure_rate": 0.0, + "failure_rate": 0.0, + "teacher_policy_return": 0.0, + "policy_return": -0.7368431687355042, + "numerically_invalid_trajectories": 0 + }, + "32": { + "step": 32, + "q_mae": 0.15562747859803494, + "qd_mae": 1.1873545465641655, + "base_position_error": 0.0, + "base_orientation_error": 0.0, + "valid_trajectories": 4, + "teacher_nonfinite_count": 0, + "neural_nonfinite_count": 0, + "teacher_failed_count": 0, + "failed_trajectories": 0, + "teacher_nonfinite_rate": 0.0, + "nonfinite_rate": 0.0, + "teacher_failure_rate": 0.0, + "failure_rate": 0.0, + "teacher_policy_return": 0.0, + "policy_return": -20.996171951293945, + "numerically_invalid_trajectories": 0 + } + }, + "failure_rate": 0.0, + "nonfinite_rate": 0.0, + "selected_teacher_sha256": "aae1a54e5d57b9f4eff7700f920d92be19cc8eaa74e3ce0bcc1c24702494d4e3", + "runtime_identity": { + "device": "cpu", + "device_name": "x86_64", + "python": "3.11.14", + "platform": "Linux-6.8.0-138-generic-x86_64-with-glibc2.35", + "state_dtype": "float32", + "cuda_runtime": "12.8", + "versions": { + "embodichain": "0.3.1", + "dexsim-engine": "0.5.1", + "newton": "1.6.0", + "warp-lang": "1.17.0", + "torch": "2.7.1+cu128" + }, + "source_sha256": { + "collection.py": "e9e5e8cd8eb80a23b831b96116be388ff55b1a874608b62379d564a8bd24c8ab", + "nerd.py": "6731d308e76093e75e530e3cde9df6bf094fad07a999a8f6b9929f4261f242ba", + "_nerd_binding.py": "0240a6eeb2ff23cdafbae4c09da08a2665bc8318fdd29c43199e2ee72a1467ae", + "_reference_assets.py": "ece8077d9e51597efff1925abc96fbeb9bb3998d52031356d8f993db7bc271c0" + } + } + } + } + }, + "ant": { + "dataset": "/tmp/embodichain-nerd-v2-qualification/ant/data", + "model": "/tmp/embodichain-nerd-v2-qualification/ant/export", + "training_updates": 100, + "network_scope": "one-layer smoke architecture", + "trajectories": 64, + "training_horizon": 32, + "offline": { + "normalized_mse": 1.2797744299075904, + "target_mse": 14.06827224018997, + "q_error_norm": 0.16857597309610117, + "qd_error_norm": 11.861796492424563, + "base_position_error": 0.009656980768709944, + "base_orientation_error_rad": 0.024108630159626835, + "valid_positions": 1380, + "windows": 138, + "split": "test", + "dataset_fingerprint": "3d2cfcef1f2b8d931ee19386c0a65af0e90927d4caf277f698b20a3d2105a22d", + "qualification": "offline_only" + }, + "rollouts": { + "passive": { + "source": "/tmp/embodichain-nerd-v2-qualification/ant/audited-passive.json", + "source_sha256": "4a6358444482466ee7caa6e5989a2567c9d9a0f136bbb7699849cbb0ed661f5e", + "trajectories": 4, + "horizons": { + "1": { + "step": 1, + "q_mae": 0.003830164307146333, + "qd_mae": 0.12577928000661942, + "base_position_error": 0.0008632538739382386, + "base_orientation_error": 0.0016251316831310493, + "valid_trajectories": 4, + "teacher_nonfinite_count": 0, + "neural_nonfinite_count": 0, + "teacher_failed_count": 0, + "failed_trajectories": 0, + "teacher_nonfinite_rate": 0.0, + "nonfinite_rate": 0.0, + "teacher_failure_rate": 0.0, + "failure_rate": 0.0, + "teacher_policy_return": 1.100000023841858, + "policy_return": 1.0901442766189575, + "numerically_invalid_trajectories": 0 + }, + "10": { + "step": 10, + "q_mae": 0.020724892207363156, + "qd_mae": 0.7133112022777953, + "base_position_error": 0.026907689162675486, + "base_orientation_error": 0.008084428001299037, + "valid_trajectories": 4, + "teacher_nonfinite_count": 0, + "neural_nonfinite_count": 0, + "teacher_failed_count": 0, + "failed_trajectories": 0, + "teacher_nonfinite_rate": 0.0, + "nonfinite_rate": 0.0, + "teacher_failure_rate": 0.0, + "failure_rate": 0.0, + "teacher_policy_return": 11.0, + "policy_return": 10.46364688873291, + "numerically_invalid_trajectories": 0 + }, + "32": { + "step": 32, + "q_mae": 0.06304648406386215, + "qd_mae": 2.3249791183746704, + "base_position_error": 0.1405739385777644, + "base_orientation_error": 0.023205568469436236, + "valid_trajectories": 4, + "teacher_nonfinite_count": 0, + "neural_nonfinite_count": 0, + "teacher_failed_count": 0, + "failed_trajectories": 0, + "teacher_nonfinite_rate": 0.0, + "nonfinite_rate": 0.0, + "teacher_failure_rate": 0.0, + "failure_rate": 0.0, + "teacher_policy_return": 35.20000076293945, + "policy_return": 27.262149810791016, + "numerically_invalid_trajectories": 0 + } + }, + "failure_rate": 0.0, + "nonfinite_rate": 0.0, + "selected_teacher_sha256": "8a55ddb52203f9399e0a8f44ee594f7225e3d1acefca031b9ae06ec118b163cd", + "runtime_identity": { + "device": "cpu", + "device_name": "x86_64", + "python": "3.11.14", + "platform": "Linux-6.8.0-138-generic-x86_64-with-glibc2.35", + "state_dtype": "float32", + "cuda_runtime": "12.8", + "versions": { + "embodichain": "0.3.1", + "dexsim-engine": "0.5.1", + "newton": "1.6.0", + "warp-lang": "1.17.0", + "torch": "2.7.1+cu128" + }, + "source_sha256": { + "collection.py": "e9e5e8cd8eb80a23b831b96116be388ff55b1a874608b62379d564a8bd24c8ab", + "nerd.py": "6731d308e76093e75e530e3cde9df6bf094fad07a999a8f6b9929f4261f242ba", + "_nerd_binding.py": "0240a6eeb2ff23cdafbae4c09da08a2665bc8318fdd29c43199e2ee72a1467ae", + "_reference_assets.py": "ece8077d9e51597efff1925abc96fbeb9bb3998d52031356d8f993db7bc271c0" + } + } + }, + "random_effort": { + "source": "/tmp/embodichain-nerd-v2-qualification/ant/audited-random_effort.json", + "source_sha256": "f7697a43f4e25975bcc38fd33e0dc1162caa7835e23598853cbf09a959cbb56a", + "trajectories": 4, + "horizons": { + "1": { + "step": 1, + "q_mae": 0.4316115162873757, + "qd_mae": 9.574224930988358, + "base_position_error": 0.016332467047095917, + "base_orientation_error": 0.12790930832119574, + "valid_trajectories": 4, + "teacher_nonfinite_count": 0, + "neural_nonfinite_count": 0, + "teacher_failed_count": 0, + "failed_trajectories": 0, + "teacher_nonfinite_rate": 0.0, + "nonfinite_rate": 0.0, + "teacher_failure_rate": 0.0, + "failure_rate": 0.0, + "teacher_policy_return": 0.933454766869545, + "policy_return": 1.0897283256053925, + "numerically_invalid_trajectories": 0 + }, + "10": { + "step": 10, + "q_mae": 0.44887487458545366, + "qd_mae": 18.719124242423067, + "base_position_error": 0.1171342632746346, + "base_orientation_error": 0.4802983186673072, + "valid_trajectories": 4, + "teacher_nonfinite_count": 0, + "neural_nonfinite_count": 0, + "teacher_failed_count": 0, + "failed_trajectories": 0, + "teacher_nonfinite_rate": 0.0, + "nonfinite_rate": 0.0, + "teacher_failure_rate": 0.0, + "failure_rate": 0.0, + "teacher_policy_return": 9.253726541996002, + "policy_return": 11.031569480895996, + "numerically_invalid_trajectories": 0 + }, + "32": { + "step": 32, + "q_mae": 0.5748123849493254, + "qd_mae": 61.47184968087277, + "base_position_error": 0.5353049895188364, + "base_orientation_error": 0.8997124689806466, + "valid_trajectories": 4, + "teacher_nonfinite_count": 0, + "neural_nonfinite_count": 0, + "teacher_failed_count": 0, + "failed_trajectories": 0, + "teacher_nonfinite_rate": 0.0, + "nonfinite_rate": 0.0, + "teacher_failure_rate": 0.0, + "failure_rate": 0.0, + "teacher_policy_return": -8.692098617553711, + "policy_return": 31.04864501953125, + "numerically_invalid_trajectories": 0 + } + }, + "failure_rate": 0.0, + "nonfinite_rate": 0.0, + "selected_teacher_sha256": "8b043ed7695c90dccc2c3fa0dc7f24e9d903ebaef9a8f5c87f392e78fa60b8b3", + "runtime_identity": { + "device": "cpu", + "device_name": "x86_64", + "python": "3.11.14", + "platform": "Linux-6.8.0-138-generic-x86_64-with-glibc2.35", + "state_dtype": "float32", + "cuda_runtime": "12.8", + "versions": { + "embodichain": "0.3.1", + "dexsim-engine": "0.5.1", + "newton": "1.6.0", + "warp-lang": "1.17.0", + "torch": "2.7.1+cu128" + }, + "source_sha256": { + "collection.py": "e9e5e8cd8eb80a23b831b96116be388ff55b1a874608b62379d564a8bd24c8ab", + "nerd.py": "6731d308e76093e75e530e3cde9df6bf094fad07a999a8f6b9929f4261f242ba", + "_nerd_binding.py": "0240a6eeb2ff23cdafbae4c09da08a2665bc8318fdd29c43199e2ee72a1467ae", + "_reference_assets.py": "ece8077d9e51597efff1925abc96fbeb9bb3998d52031356d8f993db7bc271c0" + } + } + }, + "fixed_policy": { + "source": "/tmp/embodichain-nerd-v2-qualification/ant/audited-fixed_policy.json", + "source_sha256": "cdb5b61269b337a03fdbedf672de94559732f49d7ef6807e7222cf106d5f043a", + "trajectories": 4, + "horizons": { + "1": { + "step": 1, + "q_mae": 0.003830164307146333, + "qd_mae": 0.12577928000661942, + "base_position_error": 0.0008632538739382386, + "base_orientation_error": 0.0016251316831310493, + "valid_trajectories": 4, + "teacher_nonfinite_count": 0, + "neural_nonfinite_count": 0, + "teacher_failed_count": 0, + "failed_trajectories": 0, + "teacher_nonfinite_rate": 0.0, + "nonfinite_rate": 0.0, + "teacher_failure_rate": 0.0, + "failure_rate": 0.0, + "teacher_policy_return": 1.100000023841858, + "policy_return": 1.0901442766189575, + "numerically_invalid_trajectories": 0 + }, + "10": { + "step": 10, + "q_mae": 0.02086211291170912, + "qd_mae": 0.6958701942962111, + "base_position_error": 0.026912807493379723, + "base_orientation_error": 0.008242092204885535, + "valid_trajectories": 4, + "teacher_nonfinite_count": 0, + "neural_nonfinite_count": 0, + "teacher_failed_count": 0, + "failed_trajectories": 0, + "teacher_nonfinite_rate": 0.0, + "nonfinite_rate": 0.0, + "teacher_failure_rate": 0.0, + "failure_rate": 0.0, + "teacher_policy_return": 11.0, + "policy_return": 10.458333969116211, + "numerically_invalid_trajectories": 0 + }, + "32": { + "step": 32, + "q_mae": 0.06248866482036029, + "qd_mae": 2.0772392001703865, + "base_position_error": 0.14014308332201178, + "base_orientation_error": 0.023002099837859766, + "valid_trajectories": 4, + "teacher_nonfinite_count": 0, + "neural_nonfinite_count": 0, + "teacher_failed_count": 0, + "failed_trajectories": 0, + "teacher_nonfinite_rate": 0.0, + "nonfinite_rate": 0.0, + "teacher_failure_rate": 0.0, + "failure_rate": 0.0, + "teacher_policy_return": 35.20000076293945, + "policy_return": 27.674280166625977, + "numerically_invalid_trajectories": 0 + } + }, + "failure_rate": 0.0, + "nonfinite_rate": 0.0, + "selected_teacher_sha256": "65a3f29137ea7419c45b56896964cedbcd0e53a91a6ae156451de46de99e4c57", + "runtime_identity": { + "device": "cpu", + "device_name": "x86_64", + "python": "3.11.14", + "platform": "Linux-6.8.0-138-generic-x86_64-with-glibc2.35", + "state_dtype": "float32", + "cuda_runtime": "12.8", + "versions": { + "embodichain": "0.3.1", + "dexsim-engine": "0.5.1", + "newton": "1.6.0", + "warp-lang": "1.17.0", + "torch": "2.7.1+cu128" + }, + "source_sha256": { + "collection.py": "e9e5e8cd8eb80a23b831b96116be388ff55b1a874608b62379d564a8bd24c8ab", + "nerd.py": "6731d308e76093e75e530e3cde9df6bf094fad07a999a8f6b9929f4261f242ba", + "_nerd_binding.py": "0240a6eeb2ff23cdafbae4c09da08a2665bc8318fdd29c43199e2ee72a1467ae", + "_reference_assets.py": "ece8077d9e51597efff1925abc96fbeb9bb3998d52031356d8f993db7bc271c0" + } + } + } + } + } + } +} diff --git a/docs/source/api_reference/embodichain/embodichain.data_pipeline.rst b/docs/source/api_reference/embodichain/embodichain.data_pipeline.rst index b8955372e..df938f1af 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 Datasets -------- diff --git a/docs/source/api_reference/embodichain/embodichain.learning.rst b/docs/source/api_reference/embodichain/embodichain.learning.rst index 10f9c879a..f8c60d7ed 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 b24eda12a..3a5758fa1 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 ----------------------------------- @@ -88,6 +138,23 @@ embodichain.data_pipeline.depth_video DEPTH_MILLIMETER_UNIT DEPTH_QMAX +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 ----------------------- @@ -332,6 +399,19 @@ embodichain.lab.gym.envs.demo execute_demo_episode resolve_demo_segments +embodichain.lab.gym.envs.differentiable_dynamics_env +---------------------------------------------------- + +Graph-connected dynamics environments use explicit action/control and output +hooks, preserve terminal reward graphs during partial reset, and expose an +explicit detach boundary for rollout optimization. + +.. currentmodule:: embodichain.lab.gym.envs.differentiable_dynamics_env + +.. autosummary:: + + DifferentiableDynamicsEnv + embodichain.lab.gym.envs.embodied_env ------------------------------------- @@ -342,6 +422,49 @@ embodichain.lab.gym.envs.embodied_env EmbodiedEnvCfg EmbodiedEnv +embodichain.lab.sim.diff.dynamics +--------------------------------- + +Advance immutable functional state through SimulationManager public methods. +The host owns substep scheduling, live publication and simulation time. + +.. currentmodule:: embodichain.lab.sim.diff.dynamics + +.. autosummary:: + + FunctionalDynamics + detach_dynamics_state + +embodichain.lab.sim.physics.solvers.collection +---------------------------------------------- + +Collect reference trajectories through SimulationManager and the public Scene. +Analytical teacher and neural deployments use separate scene instances; rollout +accuracy reports are distinct from official training qualification. + +.. currentmodule:: embodichain.lab.sim.physics.solvers.collection + +.. autosummary:: + + DynamicsGenerationCfg + DynamicsSimulation + generate_dataset + evaluate_rollouts + +embodichain.lab.sim.physics.solvers.nerd +---------------------------------------- + +The nerd entry-point provider validates its trained physical contract and shares +one transition between native and functional execution. Cartpole is contact-free; +Ant uses fixed plane features and exposes no contact-force report. + +.. currentmodule:: embodichain.lab.sim.physics.solvers.nerd + +.. autosummary:: + + NeRDProvider + NeuralRobotDynamicSolver + embodichain.lab.task_program.integrations.simulation.handover ------------------------------------------------------------- @@ -979,8 +1102,9 @@ embodichain.lab.sim.atomic_actions.transports embodichain.lab.sim.diff ------------------------ -Public bridge from task-defined Newton kinematics and Warp tapes into PyTorch -autograd. It does not advance the Newton dynamics solver. +Task-defined Newton kinematics and registered solver dynamics have separate +autograd boundaries. Functional dynamics own physical/history tensors and publish +detached state through the Scene; the kinematics bridge never steps a solver. .. currentmodule:: embodichain.lab.sim.diff @@ -989,6 +1113,12 @@ autograd. It does not advance the Newton dynamics solver. NewtonStepFunc tape_context + +.. autosummary:: + + FunctionalDynamics + detach_dynamics_state + embodichain.lab.sim.diff.bridge ------------------------------- @@ -1228,6 +1358,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 -------------------------------------------------- @@ -2208,6 +2494,19 @@ embodichain.toolkits.urdf_assembly.urdf_assembly_manager URDFAssemblyManager +embodichain_tasks.classic_control.ant +------------------------------------- + +Free-base Ant dynamics task with explicit effort controls and differentiable +state ownership, registered with Gym and the APG learning-environment catalog. + +.. currentmodule:: embodichain_tasks.classic_control.ant + +.. autosummary:: + + AntDynamicsEnv + make_ant_dynamics_env + embodichain_tasks.classic_control.cart_pole -------------------------------------------- @@ -2219,6 +2518,19 @@ Cart-pole environment registration under the task-first import path. CartPoleEnv +embodichain_tasks.classic_control.cartpole +------------------------------------------ + +Contact-free Cartpole dynamics task sharing the native and functional NeRD +transition, with Gym and APG learning-environment factories. + +.. currentmodule:: embodichain_tasks.classic_control.cartpole + +.. autosummary:: + + CartpoleDynamicsEnv + make_cartpole_dynamics_env + embodichain_tasks.classic_control.point_mass --------------------------------------------- diff --git a/docs/source/guides/index.rst b/docs/source/guides/index.rst index ba6a9518c..b892a7a30 100644 --- a/docs/source/guides/index.rst +++ b/docs/source/guides/index.rst @@ -19,5 +19,6 @@ Practical guides for common tasks in EmbodiChain. task_catalog robot_catalog policy_evaluation + neural_dynamics rlinf cli diff --git a/docs/source/guides/neural_dynamics.md b/docs/source/guides/neural_dynamics.md new file mode 100644 index 000000000..1c950722a --- /dev/null +++ b/docs/source/guides/neural_dynamics.md @@ -0,0 +1,170 @@ +# Named dynamics solvers and NeRD + +NeRD integrates a robot-specific learned state transition into the DexSim Newton +scene lifecycle. Cartpole uses contact-free dynamics; Ant uses fixed candidate +points projected onto a static plane. A model is bound to its ordered joints, +physical signature, gravity, feature schema and time interval. + +This feature requires a DexSim build exporting `RegisteredNewtonSolverCfg`, the +`dexsim.newton_solvers` registry and the Scene functional/snapshot APIs. The +development JIT package and a released source-free package are different artifact +types; use the actual tested build identity rather than assuming every 0.5.x wheel +contains these interfaces. + +The repository's existing `dexsim_engine==0.5.1rc1` dependency pin does not identify +the extension-enabled development artifact. Until a matching release and pin are +available, deployment requires the tested local DexSim wheel and this EmbodiChain +checkout installed as a pair. Install the ordinary dependencies first, then install +the local wheel and editable checkout with `--no-deps` so the older pin does not +replace it. Record the wheel hash and check the required APIs before running. +This development procedure does not qualify the existing rc1 distribution. + +## Register an implementation + +Declare one installed entry point whose name becomes the configuration value: + +```toml +[project.entry-points."dexsim.newton_solvers"] +my_solver = "my_package.solver:MySolverProvider" +``` + +The provider registers builder attributes before model finalization and creates +its extension after model, control and contacts exist. It returns a Newton +`SolverBase` implementation, explicit capabilities, and optional functional +services. Provider instances and mutable histories are per World. Install local +development packages with an editable install; a source file alone does not +register an entry point. + +```yaml +physics: newton +physics_config: + solver_cfg: + solver_type: my_solver + options: {} +``` + +There is no separate direct-provider configuration or imperative registration +route. Names must be unique and cannot override built-in aliases. AutoSolver +retains its scene-aware selection; it does not guess which neural model applies. + +## Generate, train and export + +The Cartpole and Ant workflow configurations live under +`extensions/embodichain_tasks/configs/tasks/classic_control//`. +These `dynamics.*.yaml` files configure commands, not Gym 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 +``` + +Use `embodichain dynamics --help` for path/device/seed overrides. +Configuration-file paths resolve relative to the configuration file; explicit CLI +path overrides resolve relative to the working directory. The `full` profiles +carry a much larger training budget. Smoke profiles establish executable flow, +not model accuracy or official benchmark qualification. + +Teacher collection owns an ordinary `SimulationManager` and uses its public +articulation/Scene APIs. Cartpole sets `enable_ground_plane: false`: hiding a +ground mesh would leave collisions active. CPU MuJoCo batches use +`separate_worlds: true`. Record the actual teacher configuration, including +ignored solver-specific material settings, with every dataset. + +HDF5 shards contain complete trajectories: T+1 states and T controls/contact +features. Only committed shards appear in the manifest. Sampling windows cannot +cross trajectory/reset/split boundaries. Collection and reading use bounded +buffers, and each data-loader worker owns its file handles. Training, export and +offline evaluation do not initialize DexSim. + +Training checkpoints retain optimizer, scheduler, RNG and committed sample cursor +for resume. Exported models contain tensor state, normalization buffers and JSON +metadata with checksums. They can be loaded directly by `solver_type: nerd`. +Legacy official Python-object checkpoints require the separate +`convert-checkpoint --trust-pickle` operation and an explicit deployment manifest. + +## Functional dynamics and model gradients + +`DifferentiableDynamicsEnv` uses graph-connected state and output hooks. Its +functional state contains physical tensors, per-world history and valid history +lengths. Native state publication is detached and does not own autograd storage. +`FunctionalDynamics` advances these states through SimulationManager; do not also +call `World.update` for the same interval. + +Torch network operations preserve gradients in evaluation mode. Newton FK uses +an explicit first-order VJP bridge. Tensor sharing between Warp and Torch alone +does not establish a gradient path. Ant's continuous contact geometry has local +derivatives; contact masks, wrap cuts and quaternion signs remain discrete branch +choices. Physical parameters and contact topology are not differentiable in v1. + +Resetting selected environments clears their histories without breaking other +rows' graphs or already-computed terminal rewards. Call `detach_state()` at an +intentional rollout boundary, and update model weights only after backward has +completed. Resident model weights are fixed by default; explicitly request +trainable weights or supply parameter tensors for a functional optimization. + +NeRD consumes effort controls and rejects unsupported free-base/external forces. +It provides no physical contact impulse/friction-force report. Its trained time +interval is validated; changing `dt` cannot be implemented by scaling a predicted +state delta. Both reference exports use 1/60 s and one neural substep. + +## Correctness and official accuracy evidence + +Implementation parity, training-step parity, matched-budget training accuracy, +and deployment rollout accuracy are separate gates. A completed short training +run or a finite gradient is not an official accuracy result. + +Comparisons pin source commits, assets, data, initialization, batch order, +normalization statistics, optimizer settings and checkpoint-selection criteria. +The official trainer adapter executes original methods in a separate process. +Current-Newton teacher data and official-Warp teacher data occupy separate rows +of the matched-data matrix. Do not attribute teacher changes to the model port. + +The historical importers and current Newton produce different masses, inertias +and materials even from identical asset bytes. Official checkpoint conversion +does not prove deployment compatibility. Preserve the historical physical profile +or retrain for the current teacher. The reference data adapter also records an +explicit repair of the official sampler's uninitialized gravity components; +raw arrays remain available and corrected inputs are labeled in the report. + +Formal reports require five training seeds, matched declared update budgets, +fixed evaluation trajectories and error-versus-time metrics. Test data never +selects checkpoints. Missing experiments remain explicitly pending. The complete +acceptance contract is maintained in the repository design document +`design/neural-dynamics-solvers.md`. + +## Historical Warp comparison + +Use `embodichain.learning.dynamics.reference_rollout` in a fresh process to +evaluate either trainer's portable weights with the original environment and +Transformer integrator. Cartpole requires Warp 1.8.0; Ant requires Warp 1.5.1. +Keep these dependencies in separate directories rather than replacing the +installed Newton/Warp runtime. The source checkout must match the pinned main +commit checked by the reference tools. + +```bash +PYTHONPATH=/path/to/isolated-warp:/path/to/pinned-nerd-source:$PWD \ + python -m embodichain.learning.dynamics.reference_rollout \ + --spec historical-rollout.json --trust-external-code +``` + +The JSON spec provides `source`, `model`, `robot`, `output`, `device`, `num_envs`, +`trajectories`, `horizons`, `seed`, `motion_mode` and `initial_states`. +`fixed_policy` currently accepts `policy: {"kind": "pd"}`. Teacher and learned +simulators evaluate the policy on their own states. CPU execution uses a recorded +stream compatibility adapter while retaining original simulation/network math. + +Historical physical parameters must match the model's recorded training source. +Older manifests without measured gravity remain pending; do not add a new label +to make them pass. Original/portable network equivalence, teacher-content hashes, +source/trace/model identities, error curves and every failed trajectory are +retained. This path does not make historical weights compatible with the current +Newton asset import. + +The training comparison is explicitly a controlled matched-data experiment. +Shared initialization, statistics, sample order and validation selection permit +port comparisons; they do not reproduce every upstream paper-recipe detail. +Full-profile gates inspect the actual immutable model's feature order, FP32 +weights, architecture and optimizer settings. Short runs remain short-run +evidence even when all paired metrics pass their numerical thresholds. diff --git a/embodichain/cli/dynamics.py b/embodichain/cli/dynamics.py new file mode 100644 index 000000000..c5dfac159 --- /dev/null +++ b/embodichain/cli/dynamics.py @@ -0,0 +1,255 @@ +# ---------------------------------------------------------------------------- +# 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" + inputs["teacher_dataset"] = "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)) + if operation == "eval" and isinstance(config.get("policy"), dict): + policy = dict(config["policy"]) + for name in ("config", "checkpoint"): + if name in policy: + path = _path(policy[name], base) + if not path.is_file(): + raise ValueError(f"policy.{name} file does not exist: {path}") + policy[name] = str(path) + config["policy"] = policy + 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 offline data or matched simulation rollouts."), + ): + 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 a training checkpoint for the named solver." + ) + 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.lab.sim.physics.solvers.collection 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) + elif mode == "rollout": + from embodichain.lab.sim.physics.solvers.collection import ( + evaluate_rollouts, + ) + + result = evaluate_rollouts(config) + else: + raise ValueError("Evaluation mode must be offline or rollout") + 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 75a6f8f68..37211bae6 100644 --- a/embodichain/cli/main.py +++ b/embodichain/cli/main.py @@ -97,6 +97,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 solvers.", + ), 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 000000000..ffc3c550f --- /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 000000000..e0aceb53b --- /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 000000000..6eb353a02 --- /dev/null +++ b/embodichain/compute/dynamics/codec.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. +# ---------------------------------------------------------------------------- + +"""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) + + +def quat_multiply(a: Tensor, b: Tensor) -> Tensor: + """Multiply broadcastable xyzw quaternions.""" + av, bv = torch.broadcast_tensors(a[..., :3], b[..., :3]) + return torch.cat( + ( + a[..., 3:] * bv + b[..., 3:] * av + torch.linalg.cross(av, bv), + a[..., 3:] * b[..., 3:] - (av * bv).sum(-1, keepdim=True), + ), + -1, + ) + + +def quat_rotate(q: Tensor, v: Tensor) -> Tensor: + """Rotate vectors by broadcastable unit xyzw quaternions.""" + xyz, v = torch.broadcast_tensors(q[..., :3], v) + uv = torch.linalg.cross(xyz, v) + return v + 2 * (q[..., 3:] * uv + torch.linalg.cross(xyz, uv)) + + +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 torch.remainder(x + torch.pi, 2 * torch.pi) - 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 _body( + self, anchor: Tensor, q: Tensor, velocity: Tensor + ) -> tuple[Tensor, Tensor]: + inv = quat_inverse(anchor[..., 3:7]) + pos = quat_rotate(inv, q[..., :3] - anchor[..., :3]) + rot = quat_multiply(inv, q[..., 3:7]) + omega = quat_rotate(inv, velocity[..., :3]) + nu = quat_rotate( + inv, + velocity[..., 3:6] - torch.linalg.cross(anchor[..., :3], velocity[..., :3]), + ) + return torch.cat((pos, rot, q[..., 7:]), -1), torch.cat( + (omega, nu, velocity[..., 6:]), -1 + ) + + def encode( + self, + q: Tensor, + qd: Tensor, + effort: Tensor, + contacts: Mapping[str, Tensor] | None = None, + gravity_dir: Tensor | None = None, + ) -> dict[str, Tensor]: + """Encode inputs at t; contact points/normals are world-frame flattened slots. + + Returns: + Named unnormalized model features, retaining their input graphs. + """ + velocity = self.native_to_legacy(q, qd) + 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) -> Tensor: + """Encode relative next state in the frame at t, including quaternion delta.""" + velocity = self.native_to_legacy(q, qd) + next_velocity = self.native_to_legacy(next_q, next_qd) + if self.robot == "ant": + anchor = q + q, velocity = self._body(anchor, q, velocity) + next_q, next_velocity = self._body(anchor, next_q, next_velocity) + 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(next_q - q), next_velocity - velocity), -1) + + def decode( + self, q: Tensor, qd: Tensor, prediction: Tensor + ) -> tuple[Tensor, Tensor]: + """Decode a relative prediction into native Newton q/qd without detaching.""" + if prediction.shape != (*q.shape[:-1], self.q_dim + self.qd_dim): + raise ValueError("Prediction shape does not match state") + velocity = self.native_to_legacy(q, qd) + if self.robot == "cartpole": + return self._wrap_q(q + prediction[..., :2]), qd + prediction[..., 2:] + bq, bv = self._body(q, q, velocity) + 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 000000000..0b332e2b6 --- /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 000000000..e198e2787 --- /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 000000000..2bb9ad240 --- /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 000000000..04c237b0d --- /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 000000000..3f5981a49 --- /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 744419060..f5fe86d95 100644 --- a/embodichain/data_pipeline/__init__.py +++ b/embodichain/data_pipeline/__init__.py @@ -27,9 +27,9 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from . import datasets, depth_video, engine, recording + from . import datasets, depth_video, dynamics, engine, recording -__all__ = ["datasets", "depth_video", "engine", "recording"] +__all__ = ["datasets", "depth_video", "engine", "recording", "dynamics"] 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 000000000..b1f8e18a5 --- /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/shards.py b/embodichain/data_pipeline/dynamics/shards.py new file mode 100644 index 000000000..aa3d1748a --- /dev/null +++ b/embodichain/data_pipeline/dynamics/shards.py @@ -0,0 +1,534 @@ +# ---------------------------------------------------------------------------- +# 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() + + 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") + 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: + """Close this worker's bounded file-handle cache.""" + for handle in self._handles.values(): + handle.close() + self._handles.clear() + + def __getstate__(self) -> dict[str, Any]: + state = self.__dict__.copy() + state["_handles"] = OrderedDict() + 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/lab/gym/envs/differentiable_dynamics_env.py b/embodichain/lab/gym/envs/differentiable_dynamics_env.py new file mode 100644 index 000000000..6f1aa74fb --- /dev/null +++ b/embodichain/lab/gym/envs/differentiable_dynamics_env.py @@ -0,0 +1,235 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Newton solver dynamics with immutable functional state for analytic policy gradients.""" + +from __future__ import annotations +from collections.abc import Mapping +from typing import Any, TYPE_CHECKING +import torch +from embodichain.lab.gym.envs.embodied_env import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.sim.cfg import NewtonPhysicsCfg +from embodichain.lab.sim.diff.dynamics import FunctionalDynamics, detach_dynamics_state + +if TYPE_CHECKING: + from dexsim.engine.newton_physics import FunctionalSolverState + +__all__ = ["DifferentiableDynamicsEnv"] + + +class DifferentiableDynamicsEnv(EmbodiedEnv): + """Environment for a registered solver's functional dynamics implementation. + + Task subclasses implement ``_action_to_control(action, state)`` and + ``_read_outputs(state)``. Observations and rewards must be derived from the + supplied functional tensors, never read back from detached native buffers. + Normal EmbodiedEnv scene, robot and initialization hooks still apply. + + Finished rows auto-reset before the next observation, preserving terminal + reward graphs and unselected state/history graphs. ``detach_state()`` is the + explicit rollout boundary. Retained history costs O(rollout length); callers + own that boundary and must backward before updating model parameters. + """ + + def __init__(self, cfg: EmbodiedEnvCfg, *args: Any, **kwargs: Any) -> None: + self._validate_dynamics_cfg(cfg) + self._dynamics: FunctionalDynamics | None = None + self._functional_state: FunctionalSolverState | None = None + self._resetting_functional = False + super().__init__(cfg, *args, **kwargs) + try: + if ( + getattr(self, "event_manager", None) is not None + and "interval" in self.event_manager.available_modes + ): + raise ValueError( + "DifferentiableDynamicsEnv does not support native interval events; implement changes in the functional transition" + ) + self._ensure_dynamics() + except Exception: + self._cleanup_failed_initialization() + raise + + @staticmethod + def _validate_dynamics_cfg(cfg: EmbodiedEnvCfg) -> None: + if cfg.sim_cfg is None or not isinstance( + cfg.sim_cfg.physics_cfg, NewtonPhysicsCfg + ): + raise ValueError("DifferentiableDynamicsEnv requires NewtonPhysicsCfg") + for name in ("actions", "observations", "rewards"): + configured = getattr(cfg, name, None) + terms = ( + configured + if isinstance(configured, Mapping) + else vars(configured) if configured is not None else {} + ) + if any( + value is not None + for key, value in terms.items() + if not key.startswith("_") + ): + raise ValueError( + f"DifferentiableDynamicsEnv uses functional hooks instead of configured {name} functors" + ) + if not cfg.sim_cfg.physics_cfg.requires_grad: + raise ValueError("DifferentiableDynamicsEnv requires requires_grad=True") + + def _ensure_dynamics(self) -> FunctionalDynamics: + if self._dynamics is None: + self._dynamics = FunctionalDynamics(self.sim) + if self._functional_state is None: + self._functional_state = self._dynamics.create_state() + return self._dynamics + + @property + def functional_state(self) -> FunctionalSolverState: + """Return current graph-connected state; callers must treat it as immutable.""" + self._ensure_dynamics() + return self._functional_state + + def dynamics_parameters(self) -> Mapping[str, torch.Tensor]: + """Return named solver parameters for explicit optimization or freezing.""" + return self._ensure_dynamics().parameters() + + def _action_to_control( + self, action: torch.Tensor, state: FunctionalSolverState + ) -> Any: + """Map actions to differentiable controls held for one control interval. + + Args: + action: Batched task action on the environment device. + state: Current graph-connected physical state and history. + + Returns: + The selected provider's control structure, retaining action gradients. + """ + raise NotImplementedError("Implement _action_to_control(action, state)") + + def _read_outputs(self, state: FunctionalSolverState) -> dict[str, Any]: + """Compute obs, reward, terminated, truncated and optional info from tensors. + + The observation's leading axis is environment. Reward and boolean done + fields have shape (num_envs,). This hook must be pure: it is also called + for reset, get_obs and detach_state without a physics transition. + """ + raise NotImplementedError("Implement _read_outputs(state)") + + def _outputs(self, state: FunctionalSolverState) -> dict[str, Any]: + outputs = self._read_outputs(state) + for name in ("obs", "reward", "terminated", "truncated"): + if name not in outputs or not isinstance(outputs[name], torch.Tensor): + raise TypeError(f"_read_outputs must return tensor {name!r}") + if outputs["obs"].ndim < 1 or outputs["obs"].shape[0] != self.num_envs: + raise ValueError("Observation must have one leading row per environment") + for name in ("reward", "terminated", "truncated"): + if outputs[name].shape != (self.num_envs,): + raise ValueError(f"{name} must have shape (num_envs,)") + if any( + outputs[name].dtype != torch.bool for name in ("terminated", "truncated") + ): + raise TypeError("terminated and truncated must be boolean tensors") + return outputs + + def get_obs(self, **kwargs: Any) -> torch.Tensor: + """Read observations from graph state, or fresh native state during reset.""" + del kwargs + dynamics = self._ensure_dynamics() + state = ( + dynamics.create_state() + if self._resetting_functional + else self._functional_state + ) + return self._outputs(state)["obs"] + + def reset( + self, seed: int | None = None, options: dict | None = None + ) -> tuple[torch.Tensor, dict[str, Any]]: + """Run normal reset events then merge selected fresh rows into graph state.""" + dynamics = self._ensure_dynamics() + old_state = self._functional_state + options = {} if options is None else dict(options) + mask = torch.ones(self.num_envs, dtype=torch.bool, device=self.device) + if "reset_ids" in options: + ids = torch.as_tensor( + options["reset_ids"], dtype=torch.long, device=self.device + ).reshape(-1) + mask = torch.zeros_like(mask).scatter(0, ids, True) + # Base reset fills selected counters in place. Preserve any old counter + # tensor referenced by a terminal reward expression before entering it. + self._elapsed_steps = self._elapsed_steps.clone() + self._resetting_functional = True + try: + _, info = super().reset(seed=seed, options=options) + finally: + self._resetting_functional = False + self._functional_state = dynamics.reset(old_state, mask) + return self._outputs(self._functional_state)["obs"], info + + def step( + self, action: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, dict[str, Any]]: + """Advance configured physics macro steps without a second World.update.""" + dynamics = self._ensure_dynamics() + action = torch.as_tensor(action, dtype=torch.float32, device=self.device) + control = self._action_to_control(action, self._functional_state) + sensors = [ + sensor for sensor in self.sensors.values() if sensor.requires_substep_update + ] + for sensor in sensors: + sensor.begin_control_step() + for _ in range(int(self.cfg.sim_steps_per_control)): + self._functional_state = dynamics.step( + self._functional_state, control, physics_dt=self.physics_dt + ) + for sensor in sensors: + sensor.update_physics_step(self.physics_dt) + self._update_physical_objective() + # Never mutate a counter tensor retained by a task's reward expression. + self._elapsed_steps = self._elapsed_steps + 1 + with self.sim.render_frame(): + outputs = self._outputs(self._functional_state) + obs, reward = outputs["obs"], outputs["reward"] + terminated = outputs["terminated"] + if self.cfg.ignore_terminations: + terminated = torch.zeros_like(terminated) + truncated = outputs["truncated"] | ( + self._elapsed_steps >= self.max_episode_steps + ) + info = dict(outputs.get("info", {})) + done = terminated | truncated + self._hook_after_sim_step( + obs=obs, + action=action, + rewards=reward, + dones=done, + info=info, + terminateds=terminated, + truncateds=truncated, + ) + reset_ids = done.nonzero(as_tuple=False).squeeze(-1) + if reset_ids.numel(): + info["terminal_observation"] = obs + info["terminal_mask"] = done + fresh, _ = self.reset(options={"reset_ids": reset_ids}) + row_mask = done.reshape(self.num_envs, *([1] * (obs.ndim - 1))) + obs = torch.where(row_mask, fresh, obs) + return obs, reward, terminated, truncated, info + + def detach_state(self) -> torch.Tensor: + """Detach current physical/history state and observation without resetting.""" + self._ensure_dynamics() + self._functional_state = detach_dynamics_state(self._functional_state) + return self._outputs(self._functional_state)["obs"].detach() diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index b9f30254f..39dc5b32d 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -715,6 +715,7 @@ class ComponentCfg: viser_server = ViserServerCfg(**viser_server_config) env_cfg.sim_cfg = SimulationManagerCfg( + enable_ground_plane=config.get("enable_ground_plane", True), startup_summary=config.get("startup_summary", "compact"), dexsim_startup_info=config.get("dexsim_startup_info", False), scene_node_capacity=config.get("scene_node_capacity"), @@ -732,8 +733,11 @@ class ComponentCfg: ), ) - env_cfg.robot = _robot_cfg_from_dict(config["robot"]) - env_cfg.robot.enable_sensor = env_cfg.enable_sensor + env_cfg.robot = ( + None if config["robot"] is None else _robot_cfg_from_dict(config["robot"]) + ) + if env_cfg.robot is not None: + env_cfg.robot.enable_sensor = env_cfg.enable_sensor # Parse source-neutral expert trajectory settings owned by the environment. # Keep this under ``env`` so deployment configs do not duplicate runtime diff --git a/embodichain/lab/sim/cfg/physics.py b/embodichain/lab/sim/cfg/physics.py index 314eca508..88653124b 100644 --- a/embodichain/lab/sim/cfg/physics.py +++ b/embodichain/lab/sim/cfg/physics.py @@ -311,8 +311,9 @@ class NewtonPhysicsCfg(PhysicsBackendCfg): requires_grad: bool = False """Whether to finalize the Newton model with differentiable state enabled. - EmbodiChain currently requires the Semi-implicit solver for this mode and - disables CUDA graph capture when gradients are enabled. + Semi-implicit kinematics and registered differentiable solver providers + support this mode. Provider capabilities are validated after scene creation. + CUDA graph capture is disabled when gradients are enabled. """ use_cuda_graph: bool = True @@ -346,7 +347,8 @@ class NewtonPhysicsCfg(PhysicsBackendCfg): A mapping is converted to the matching DexSim Newton solver config. Include ``solver_type`` or ``class_type`` to select the solver, then add any - parameters accepted by that DexSim solver config. If omitted, EmbodiChain + parameters accepted by that DexSim solver config. Registered providers use + their entry-point name and an ``options`` mapping. If omitted, EmbodiChain preserves DexSim's scene-aware ``AutoSolverCfg`` default. A DexSim build exporting ``AutoSolverCfg`` is required; no concrete-solver fallback is used. """ @@ -382,6 +384,7 @@ def to_dexsim_cfg( SemiImplicitSolverCfg, VBDSolverCfg, XPBDSolverCfg, + RegisteredNewtonSolverCfg, ) torch_device = ( @@ -408,12 +411,16 @@ def to_dexsim_cfg( ) if self.requires_grad and ( - solver_cfg is None or solver_cfg.solver_type != "semi_implicit" + solver_cfg is None + or ( + solver_cfg.solver_type != "semi_implicit" + and not isinstance(solver_cfg, RegisteredNewtonSolverCfg) + ) ): logger.log_error( "Newton gradient mode requires an explicit " - "solver_type='semi_implicit'; AutoSolver does not select a " - "differentiable solver." + "solver_type='semi_implicit' or a registered differentiable " + "provider; AutoSolver does not select a differentiable solver." ) collision_pipeline_cfg = None @@ -458,48 +465,14 @@ def to_dexsim_cfg( def _normalize_newton_solver_type(solver_type: str) -> str: """Normalize public EmbodiChain and DexSim Newton solver aliases.""" - key = solver_type.replace("-", "_").lower() - aliases = { - "auto": "auto", - "autosolver": "auto", - "autosolvercfg": "auto", - "auto_solver": "auto", - "auto_solver_cfg": "auto", - "mjwarp": "mujoco_warp", - "mjwarpsolver": "mujoco_warp", - "mjwarpsolvercfg": "mujoco_warp", - "mjwarp_solver": "mujoco_warp", - "mjwarp_solver_cfg": "mujoco_warp", - "mujoco_warp": "mujoco_warp", - "mujocowarp": "mujoco_warp", - "mujocowarpsolver": "mujoco_warp", - "mujocowarpsolvercfg": "mujoco_warp", - "dexuni": "dexuni", - "dexunisolver": "dexuni", - "dexunisolvercfg": "dexuni", - "dexuni_solver": "dexuni", - "dexuni_solver_cfg": "dexuni", - "xpbdsolver": "xpbd", - "xpbdsolvercfg": "xpbd", - "xpbd": "xpbd", - "semiimplicit": "semi_implicit", - "semi_implicit": "semi_implicit", - "semiimplicitsolver": "semi_implicit", - "semiimplicitsolvercfg": "semi_implicit", - "featherstone": "featherstone", - "featherstonesolver": "featherstone", - "featherstonesolvercfg": "featherstone", - "vbd": "vbd", - "vbdsolver": "vbd", - "vbdsolvercfg": "vbd", - } - if key not in aliases: - logger.log_error( - f"Unsupported Newton solver type '{solver_type}'. " - "Expected one of 'auto', 'dexuni', 'mjwarp', 'xpbd', 'semi_implicit', " - "'featherstone', or 'vbd'." - ) - return aliases[key] + try: + from dexsim.engine.newton_physics import normalize_newton_solver_name + except ImportError as exc: + raise RuntimeError( + "Named Newton solvers require a DexSim build exposing " + "normalize_newton_solver_name and RegisteredNewtonSolverCfg." + ) from exc + return normalize_newton_solver_name(solver_type) def _newton_solver_cfg_to_dexsim( @@ -511,6 +484,21 @@ def _newton_solver_cfg_to_dexsim( return None if not isinstance(solver_cfg, Mapping): + from dexsim.engine.newton_physics import RegisteredNewtonSolverCfg + + if isinstance(solver_cfg, RegisteredNewtonSolverCfg): + name = _normalize_newton_solver_type(str(solver_cfg.solver_type)) + if name in solver_cfg_map: + if {"solver_type", "class_type"}.intersection(solver_cfg.options): + raise ValueError( + "Builtin solver options cannot override the registered name." + ) + return _newton_solver_cfg_to_dexsim( + {"solver_type": name, **dict(solver_cfg.options)}, solver_cfg_map + ) + return RegisteredNewtonSolverCfg( + solver_type=name, options=dict(solver_cfg.options) + ) if not hasattr(solver_cfg, "solver_type"): logger.log_error( "Newton solver_cfg must be a mapping or a DexSim Newton solver " @@ -525,6 +513,19 @@ def _newton_solver_cfg_to_dexsim( or "auto" ) normalized_solver_type = _normalize_newton_solver_type(str(configured_solver_type)) + if normalized_solver_type not in solver_cfg_map: + from dexsim.engine.newton_physics import RegisteredNewtonSolverCfg + + unexpected = set(solver_cfg_data) - {"options"} + if unexpected: + raise ValueError( + "Registered solver settings belong in options; unexpected fields: " + f"{sorted(unexpected)}." + ) + return RegisteredNewtonSolverCfg( + solver_type=normalized_solver_type, + options=solver_cfg_data.get("options", {}), + ) solver_cfg_type = solver_cfg_map[normalized_solver_type] if normalized_solver_type == "dexuni": try: diff --git a/embodichain/lab/sim/diff/__init__.py b/embodichain/lab/sim/diff/__init__.py index 1aafcd511..959788c40 100644 --- a/embodichain/lab/sim/diff/__init__.py +++ b/embodichain/lab/sim/diff/__init__.py @@ -13,11 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- -"""Differentiable Newton kinematics for EmbodiChain. +"""Differentiable Newton kinematics and registered solver dynamics. Bridges task-defined Warp kinematics into PyTorch autograd and exposes a :class:`tape_context` manager for advanced users composing their own kernels. -The package does not advance the Newton dynamics solver. +``FunctionalDynamics`` advances a registered solver through the public Scene +functional interface; the kinematics bridge retains its separate contract. """ from __future__ import annotations @@ -26,8 +27,11 @@ NewtonStepFunc, tape_context, ) +from .dynamics import FunctionalDynamics, detach_dynamics_state __all__ = [ "NewtonStepFunc", "tape_context", + "FunctionalDynamics", + "detach_dynamics_state", ] diff --git a/embodichain/lab/sim/diff/dynamics.py b/embodichain/lab/sim/diff/dynamics.py new file mode 100644 index 000000000..ebf5e577d --- /dev/null +++ b/embodichain/lab/sim/diff/dynamics.py @@ -0,0 +1,148 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Public functional dynamics facade over SimulationManager's Newton services.""" + +from __future__ import annotations +import math +from collections.abc import Mapping +from typing import Any, Protocol, TYPE_CHECKING +import torch + +if TYPE_CHECKING: + from dexsim.engine.newton_physics import FunctionalSolverState + +__all__ = ["FunctionalDynamics", "detach_dynamics_state"] + + +class _DynamicsManager(Protocol): + def create_functional_state(self) -> FunctionalSolverState: ... + def functional_parameters(self) -> Mapping[str, torch.Tensor]: ... + def functional_step( + self, + state: FunctionalSolverState, + control: Any, + parameters: Mapping[str, torch.Tensor] | None = None, + dt: float | None = None, + *, + publish: bool = True, + render_final_step: bool = True, + ) -> FunctionalSolverState: ... + def reset_functional_state( + self, state: FunctionalSolverState, world_mask: torch.Tensor | None = None + ) -> FunctionalSolverState: ... + + +def detach_dynamics_state(state: FunctionalSolverState) -> FunctionalSolverState: + """Create an explicit BPTT boundary without reset or state resampling. + + Args: + state: Owned functional state with tensor physical/history fields. + + Returns: + A new state containing detached views. Both states remain immutable; + native publication owns separate storage and cannot overwrite them. + """ + from dexsim.engine.newton_physics import FunctionalSolverState + + return FunctionalSolverState( + physical={k: v.detach() for k, v in state.physical.items()}, + history={k: v.detach() for k, v in state.history.items()}, + valid_history_lengths=( + None + if state.valid_history_lengths is None + else state.valid_history_lengths.detach() + ), + ) + + +class FunctionalDynamics: + """Advance differentiable state through a manager's public Newton interface. + + Args: + manager: SimulationManager or an object implementing the same public + functional dynamics methods. No backend or native buffer is accessed. + + A physics step means one configured macro interval. DexSim alone owns its + internal solver substeps. A control interval can contain several such steps. + Native state publication is detached; returned state retains its graph. + """ + + def __init__(self, manager: _DynamicsManager) -> None: + self._manager = manager + + def create_state(self) -> FunctionalSolverState: + """Snapshot native state and solver history into independent storage.""" + return self._manager.create_functional_state() + + def parameters(self) -> Mapping[str, torch.Tensor]: + """Return the live named solver parameters without copying or detaching.""" + return self._manager.functional_parameters() + + def step( + self, + state: FunctionalSolverState, + control: Any, + *, + physics_steps: int = 1, + physics_dt: float | None = None, + parameters: Mapping[str, torch.Tensor] | None = None, + publish: bool = True, + ) -> FunctionalSolverState: + """Hold control across physics_steps macro intervals and preserve BPTT. + + Args: + state: Initial functional state. + control: Provider-defined differentiable controls held for the interval. + physics_steps: Positive macro step count, independent of solver substeps. + physics_dt: Macro timestep; None uses the configured physics interval. + parameters: Optional explicit parameter mapping for functional evaluation. + publish: Whether to publish detached state and advance native time. + + Returns: + Final owned functional state. Rendering is left to the environment. + """ + if ( + isinstance(physics_steps, bool) + or not isinstance(physics_steps, int) + or physics_steps < 1 + ): + raise ValueError("physics_steps must be a positive integer") + if physics_dt is not None and ( + not math.isfinite(physics_dt) or physics_dt <= 0 + ): + raise ValueError("physics_dt must be positive and finite") + for _ in range(physics_steps): + state = self._manager.functional_step( + state, + control, + parameters=parameters, + dt=physics_dt, + publish=publish, + render_final_step=False, + ) + return state + + def reset( + self, state: FunctionalSolverState, world_mask: torch.Tensor | None = None + ) -> FunctionalSolverState: + """Merge freshly reset native rows into state while preserving other graphs. + + The caller first performs the manager's normal native reset/randomization. + This method snapshots those new rows and delegates ownership-aware merging + and detached publication to the manager. It does not advance time. + """ + return self._manager.reset_functional_state(state, world_mask=world_mask) diff --git a/embodichain/lab/sim/objects/articulation.py b/embodichain/lab/sim/objects/articulation.py index ccafd9dcf..c2e8049f0 100644 --- a/embodichain/lab/sim/objects/articulation.py +++ b/embodichain/lab/sim/objects/articulation.py @@ -2839,8 +2839,12 @@ def reset( # impulse when initialized away from their equality manifold. qpos = self._project_mimic_qpos(qpos) self.set_qpos(qpos, target=False, env_ids=local_env_ids) - # Set drive target to hold position. - self.set_qpos(qpos, target=True, env_ids=local_env_ids) + # Newton clear_dynamics synthesizes hold targets atomically with its + # state/history clear (either below or in the manager's coordinated + # clear). Do not send a user position-control command to effort-only + # providers while resetting them. + if not self.body_data.is_newton_backend: + self.set_qpos(qpos, target=True, env_ids=local_env_ids) if clear_dynamics: self.clear_dynamics(env_ids=local_env_ids) diff --git a/embodichain/lab/sim/physics/solvers/__init__.py b/embodichain/lab/sim/physics/solvers/__init__.py new file mode 100644 index 000000000..5e7d7036b --- /dev/null +++ b/embodichain/lab/sim/physics/solvers/__init__.py @@ -0,0 +1,25 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Physics solver providers discovered through DexSim entry points. + +Import providers explicitly; package discovery must not load neural weights or +initialize a physics model. +""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/embodichain/lab/sim/physics/solvers/_kinematics.py b/embodichain/lab/sim/physics/solvers/_kinematics.py new file mode 100644 index 000000000..f53488646 --- /dev/null +++ b/embodichain/lab/sim/physics/solvers/_kinematics.py @@ -0,0 +1,110 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""First-order Torch autograd boundary for native Newton FK.""" + +from __future__ import annotations + +from contextlib import nullcontext +from typing import Any + +import torch +from torch.autograd.function import once_differentiable +import warp as wp + +__all__: list[str] = [] + + +def _stream_scope(device: torch.device): + if device.type == "cuda": + return wp.ScopedStream(wp.stream_from_torch(torch.cuda.current_stream(device))) + return nullcontext() + + +class _NewtonFK(torch.autograd.Function): + """Record each call into independently owned Warp state and adjoints.""" + + @staticmethod + def forward(ctx: Any, q: torch.Tensor, qd: torch.Tensor, model: Any): + import newton + + if q.dtype != torch.float32 or qd.dtype != torch.float32: + raise TypeError("Newton FK requires float32 physical state tensors.") + if q.device != qd.device or q.device != torch.device( + wp.device_to_torch(model.device) + ): + raise ValueError("Newton model and FK inputs must share a device.") + if q.shape != (model.joint_coord_count,) or qd.shape != ( + model.joint_dof_count, + ): + raise ValueError( + "Newton FK requires flattened model-order joint coordinates." + ) + ctx.device = q.device + with _stream_scope(q.device): + # Never retain aliases to caller/native buffers: a later publication + # or reset may overwrite them before this transition's backward. + q_owned = q.detach().contiguous().clone() + qd_owned = qd.detach().contiguous().clone() + ctx.q_wp = wp.from_torch(q_owned, requires_grad=True) + ctx.qd_wp = wp.from_torch(qd_owned, requires_grad=True) + ctx.state = model.state(requires_grad=True) + ctx.tape = wp.Tape() + with ctx.tape: + newton.eval_fk(model, ctx.q_wp, ctx.qd_wp, ctx.state) + # Clones also keep the public outputs independent of Warp scratch. + body_q = wp.to_torch(ctx.state.body_q).clone() + body_qd = wp.to_torch(ctx.state.body_qd).clone() + return body_q, body_qd + + @staticmethod + @once_differentiable + def backward( + ctx: Any, body_q_grad: torch.Tensor | None, body_qd_grad: torch.Tensor | None + ): + with _stream_scope(ctx.device): + ctx.tape.zero() + grads = {} + if body_q_grad is not None: + grads[ctx.state.body_q] = wp.from_torch( + body_q_grad.contiguous(), dtype=wp.transform + ) + if body_qd_grad is not None: + grads[ctx.state.body_qd] = wp.from_torch( + body_qd_grad.contiguous(), dtype=wp.spatial_vector + ) + ctx.tape.backward(grads=grads) + q_grad = wp.to_torch(ctx.q_wp.grad).clone() + qd_grad = wp.to_torch(ctx.qd_wp.grad).clone() + return q_grad, qd_grad, None + + +def differentiable_fk( + model: Any, q: torch.Tensor, qd: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Return body poses and COM velocities with first-order input gradients. + + Args: + model: Finalized Newton model; its topology and physical parameters + must remain unchanged while a returned graph is alive. + q: Flattened float32 Newton joint coordinates. + qd: Flattened float32 Newton joint velocities. + + Returns: + Body transforms ``(N, 7)`` and COM linear/angular velocities ``(N, 6)``. + Storage belongs to this call, independently of live simulation state. + """ + return _NewtonFK.apply(q, qd, model) diff --git a/embodichain/lab/sim/physics/solvers/_nerd_binding.py b/embodichain/lab/sim/physics/solvers/_nerd_binding.py new file mode 100644 index 000000000..b0604606e --- /dev/null +++ b/embodichain/lab/sim/physics/solvers/_nerd_binding.py @@ -0,0 +1,476 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Validated Newton indexing and fixed-slot geometry for NeRD models.""" + +from __future__ import annotations + +import hashlib +import json +from typing import Any +from types import SimpleNamespace + +import numpy as np +import torch +import warp as wp + +from embodichain.compute.dynamics import plane_contact_features, quat_rotate + +__all__: list[str] = [] + +_PHYSICS_SIGNATURE_SCHEMA = "nerd-reference-physics-v2" + +_JOINTS = { + "cartpole": ("slider_to_cart", "cart_to_pole"), + "ant": tuple(name for i in range(1, 5) for name in (f"hip_{i}", f"ankle_{i}")), +} + + +def _numpy(value: Any) -> np.ndarray: + return ( + value.detach().cpu().numpy() + if isinstance(value, torch.Tensor) + else value.numpy() + ) + + +def _torch(value: Any) -> torch.Tensor: + return value if isinstance(value, torch.Tensor) else wp.to_torch(value) + + +def _local_name(label: str) -> str: + return label.rsplit("/", 1)[-1] + + +def _fingerprint(payload: dict[str, Any]) -> str: + # Infinities are legitimate unbounded joint limits. JSON strings retain + # their semantics without emitting nonstandard JSON numeric values. + def encode(value): + if isinstance(value, dict): + return {k: encode(v) for k, v in value.items()} + if isinstance(value, list): + return [encode(v) for v in value] + if isinstance(value, float) and not np.isfinite(value): + return str(value) + return value + + return hashlib.sha256( + json.dumps(encode(payload), sort_keys=True, allow_nan=False).encode() + ).hexdigest() + + +class NeRDBinding: + """Resolve homogeneous reference robots without assuming native row order. + + Construction reads model topology once. All per-step indexing is resident + on the model device and supports gradients through gather/scatter operations. + """ + + def __init__(self, model: Any, robot: str) -> None: + import newton + + self.model = model + self.robot = robot.lower() + if self.robot not in _JOINTS: + raise ValueError(f"Unsupported NeRD robot {robot!r}.") + if model.particle_count: + raise ValueError("NeRD does not integrate particle state.") + self.device = torch.device(wp.device_to_torch(model.device)) + self.num_envs = max(int(model.world_count), 1) + starts = _numpy(model.articulation_start) + worlds = _numpy(model.articulation_world) + if model.articulation_count != self.num_envs: + raise ValueError( + "NeRD requires exactly one reference articulation per world." + ) + joint_type = _numpy(model.joint_type) + parent = _numpy(model.joint_parent) + child = _numpy(model.joint_child) + q_starts = _numpy(model.joint_q_start) + qd_starts = _numpy(model.joint_qd_start) + shape_body = _numpy(model.shape_body) + shape_type = _numpy(model.shape_type) + flags = _numpy(model.shape_flags) + collision_flag = int(newton.ShapeFlags.COLLIDE_SHAPES) + shape_transform = _numpy(model.shape_transform) + shape_scale = _numpy(model.shape_scale) + shape_margin = _numpy(model.shape_margin) + body_com = _numpy(model.body_com) + planes = np.flatnonzero( + (shape_body == -1) & (shape_type == int(newton.GeoType.PLANE)) + ) + contact_pair_count = int(model.shape_contact_pair_count) + contact_pairs = ( + _numpy(model.shape_contact_pairs)[:contact_pair_count] + if contact_pair_count + else np.empty((0, 2), dtype=np.int32) + ) + if self.robot == "ant" and len(planes) != 1: + raise ValueError("NeRD Ant requires exactly one shared static plane.") + joint_xp, joint_xc = _numpy(model.joint_X_p), _numpy(model.joint_X_c) + identity = np.array([0, 0, 0, 0, 0, 0, 1], dtype=np.float32) + q_rows, qd_rows, root_bodies, ordered_bodies, world_ids = [], [], [], [], [] + slot_bodies, slot_points, slot_thickness = [], [], [] + signatures = [] + for world_id in range(self.num_envs): + match = np.flatnonzero(worlds == (world_id if model.world_count else -1)) + if len(match) != 1: + raise ValueError( + "NeRD worlds must contain one locally owned articulation." + ) + art = int(match[0]) + joints = list(range(int(starts[art]), int(starts[art + 1]))) + by_name = {_local_name(model.joint_label[j]): j for j in joints} + missing = set(_JOINTS[self.robot]) - by_name.keys() + if missing: + raise ValueError( + f"NeRD reference joint names are missing: {sorted(missing)}." + ) + scalar_joints = [by_name[n] for n in _JOINTS[self.robot]] + if any( + q_starts[j + 1] - q_starts[j] != 1 + or qd_starts[j + 1] - qd_starts[j] != 1 + for j in scalar_joints + ): + raise ValueError( + "NeRD reference joints must have one coordinate and velocity." + ) + roots = [j for j in joints if parent[j] == -1] + if len(roots) != 1: + raise ValueError("NeRD requires one world-root joint per articulation.") + root = roots[0] + if self.robot == "ant": + if int(joint_type[root]) != int(newton.JointType.FREE): + raise ValueError("Ant requires a free root joint.") + if not np.allclose(joint_xp[root], identity) or not np.allclose( + joint_xc[root], identity + ): + raise ValueError("NeRD Ant requires identity root joint frames.") + selected = [root, *scalar_joints] + else: + if int(joint_type[root]) != int(newton.JointType.FIXED): + raise ValueError("Cartpole requires a fixed root joint.") + selected = scalar_joints + q_ids = [ + i + for j in selected + for i in range(int(q_starts[j]), int(q_starts[j + 1])) + ] + qd_ids = [ + i + for j in selected + for i in range(int(qd_starts[j]), int(qd_starts[j + 1])) + ] + if sum(q_starts[j + 1] - q_starts[j] for j in joints) != len(q_ids): + raise ValueError("Scene contains unsupported articulation coordinates.") + bodies = [int(child[root]), *[int(child[j]) for j in scalar_joints]] + if len(set(bodies)) != len(bodies): + raise ValueError( + "Reference topology must use distinct root and joint-child bodies." + ) + shapes = [ + int(s) + for b in bodies + for s in np.flatnonzero( + (shape_body == b) & ((flags & collision_flag) != 0) + ) + ] + points, point_bodies, thickness = [], [], [] + if self.robot == "ant": + for s in shapes: + transform = torch.as_tensor(shape_transform[s]) + radius = float(shape_scale[s, 0]) + if int(shape_type[s]) == int(newton.GeoType.SPHERE): + centers = torch.zeros(1, 3) + elif int(shape_type[s]) == int(newton.GeoType.CAPSULE): + half_height = float(shape_scale[s, 1]) + centers = torch.tensor( + [[0.0, 0.0, half_height], [0.0, 0.0, -half_height]] + ) + else: + raise ValueError( + "Reference Ant contacts require sphere/capsule primitives." + ) + local = transform[:3] + quat_rotate(transform[3:], centers) + points.extend(local.tolist()) + point_bodies.extend([int(shape_body[s])] * len(centers)) + thickness.extend([radius + float(shape_margin[s])] * len(centers)) + if len(points) != 25: + raise ValueError( + f"Reference Ant requires 25 fixed contact slots, got {len(points)}." + ) + q_rows.append(q_ids) + qd_rows.append(qd_ids) + root_bodies.append(int(child[root])) + ordered_bodies.append(bodies) + world_ids.append(world_id if model.world_count else -1) + slot_bodies.append(point_bodies) + slot_points.append(points) + slot_thickness.append(thickness) + local_bodies = {body: index for index, body in enumerate(bodies)} + contract_joints = [root, *scalar_joints] + if any( + int(parent[joint]) not in local_bodies and int(parent[joint]) != -1 + for joint in contract_joints + ): + raise ValueError( + "Reference articulation contains unsupported joint-parent topology." + ) + local_shapes = {shape: index for index, shape in enumerate(shapes)} + effective_pairs = set() + for first, second in contact_pairs: + first, second = int(first), int(second) + if first not in local_shapes and second not in local_shapes: + continue + endpoints = [] + for shape in (first, second): + if shape in local_shapes: + endpoints.append(local_shapes[shape]) + elif self.robot == "ant" and shape == int(planes[0]): + endpoints.append(-1) + else: + raise ValueError( + "NeRD reference scenes cannot contain cross-world or unmodeled contacts." + ) + effective_pairs.add(tuple(sorted(endpoints))) + expected_pairs = ( + {(-1, index) for index in range(len(shapes))} + if self.robot == "ant" + else set() + ) + if effective_pairs != expected_pairs: + raise ValueError( + "NeRD reference collision policy must be plane-only for Ant and contact-free for Cartpole." + ) + signature = { + "schema": _PHYSICS_SIGNATURE_SCHEMA, + "robot": self.robot, + "joint_names": list(_JOINTS[self.robot]), + "body_names": [_local_name(model.body_label[b]) for b in bodies], + "joint_parent": [ + -1 if int(parent[j]) == -1 else local_bodies[int(parent[j])] + for j in contract_joints + ], + "joint_child": [local_bodies[int(child[j])] for j in contract_joints], + "root_frame_orientation": [ + joint_xp[root, 3:].tolist(), + joint_xc[root, 3:].tolist(), + ], + "allowed_contact_pairs": sorted(effective_pairs), + } + # Contact-free Cartpole may be translated with arena placement under + # uniform gravity. Fixed-root orientations still change its dynamics. + for name in ("body_mass", "body_inertia", "body_com", "body_flags"): + signature[name] = _numpy(getattr(model, name))[bodies].tolist() + for name in ( + "joint_type", + "joint_axis_dim", + "joint_dof_dim", + "joint_enabled", + ): + value = getattr(model, name, None) + if value is not None: + signature[name] = _numpy(value)[contract_joints].tolist() + for name in ("joint_X_p", "joint_X_c"): + signature[name] = _numpy(getattr(model, name))[scalar_joints].tolist() + for name in ( + "joint_axis", + "joint_limit_lower", + "joint_limit_upper", + "joint_limit_ke", + "joint_limit_kd", + "joint_friction", + "joint_armature", + "joint_damping", + "joint_target_ke", + "joint_target_kd", + "joint_target_mode", + "joint_effort_limit", + "joint_velocity_limit", + ): + value = getattr(model, name, None) + if value is not None: + signature[name] = _numpy(value)[qd_ids].tolist() + for name in ( + "shape_type", + "shape_scale", + "shape_transform", + "shape_margin", + "shape_gap", + "shape_is_solid", + "shape_material_mu", + "shape_material_ke", + "shape_material_kd", + "shape_material_kf", + "shape_material_ka", + "shape_material_kh", + "shape_material_restitution", + "shape_material_mu_torsional", + "shape_material_mu_rolling", + ): + value = getattr(model, name, None) + if value is not None: + signature[name] = _numpy(value)[shapes].tolist() + signatures.append(_fingerprint(signature)) + if len(set(signatures)) != 1: + raise ValueError("NeRD batch contains differing robot physics.") + self.physics_signature = signatures[0] + self.joint_names = list(_JOINTS[self.robot]) + self.q_indices = torch.tensor(q_rows, device=self.device, dtype=torch.long) + self.qd_indices = torch.tensor(qd_rows, device=self.device, dtype=torch.long) + self.root_bodies = torch.tensor( + root_bodies, device=self.device, dtype=torch.long + ) + self.world_ids = torch.tensor(world_ids, device=self.device, dtype=torch.long) + self.base_com = body_com[root_bodies[0]].tolist() + if not np.allclose(body_com[root_bodies], self.base_com): + raise ValueError("NeRD batch has differing base COM offsets.") + dynamic_bodies = set(np.flatnonzero(_numpy(model.body_inv_mass) > 0)) + if not dynamic_bodies <= {b for row in ordered_bodies for b in row}: + raise ValueError("NeRD cannot integrate additional dynamic bodies.") + self.contact_bodies = torch.tensor( + slot_bodies, device=self.device, dtype=torch.long + ) + self.contact_points = torch.tensor( + slot_points, device=self.device, dtype=torch.float32 + ).reshape(self.num_envs, -1, 3) + self.contact_thickness = torch.tensor( + slot_thickness, device=self.device, dtype=torch.float32 + ) + gravity = _torch(model.gravity)[self.world_ids].clone() + self.gravity = gravity + self.gravity_dir = gravity / gravity.norm(dim=-1, keepdim=True).clamp_min(1e-12) + self.plane_normal = torch.tensor([0.0, 0.0, 1.0], device=self.device) + self.plane_offset = torch.zeros((), device=self.device) + self.plane_margin = 0.0 + planes = np.flatnonzero( + (shape_body == -1) & (shape_type == int(newton.GeoType.PLANE)) + ) + if self.robot == "ant": + if len(planes) != 1: + raise ValueError("NeRD Ant requires exactly one shared static plane.") + plane = int(planes[0]) + transform = torch.tensor(shape_transform[plane], device=self.device) + self.plane_normal = quat_rotate(transform[3:], self.plane_normal) + self.plane_offset = (transform[:3] * self.plane_normal).sum() + self.plane_margin = float(shape_margin[plane]) + other_static = np.flatnonzero( + (shape_body == -1) + & (shape_type != int(newton.GeoType.PLANE)) + & ((flags & collision_flag) != 0) + ) + if len(other_static): + raise ValueError( + "NeRD reference scenes cannot contain other static colliders." + ) + owned_bodies = {b for row in ordered_bodies for b in row} + unowned_colliders = [ + s + for s, body in enumerate(shape_body) + if int(body) >= 0 + and int(body) not in owned_bodies + and int(flags[s]) & collision_flag + ] + if unowned_colliders: + raise ValueError( + "NeRD reference scenes cannot contain unmodeled fixed-body colliders." + ) + if self.robot == "cartpole" and any( + int(flags[s]) & collision_flag for s in planes + ): + raise ValueError( + "Contact-free Cartpole requires the ground plane to be disabled." + ) + environment = { + "schema": _PHYSICS_SIGNATURE_SCHEMA, + "robot_physics": self.physics_signature, + } + for name in ( + "shape_type", + "shape_transform", + "shape_margin", + "shape_gap", + "shape_material_mu", + "shape_material_ka", + "shape_material_kh", + "shape_material_restitution", + "shape_material_ke", + "shape_material_kd", + "shape_material_kf", + "shape_material_mu_torsional", + "shape_material_mu_rolling", + ): + values = getattr(model, name, None) + if values is not None: + environment[name] = _numpy(values)[planes].tolist() + self.physics_signature = _fingerprint(environment) + + @classmethod + def from_metadata(cls, metadata: dict[str, Any], robot: str) -> "NeRDBinding": + """Bind an owned public Scene metadata snapshot without native access.""" + values = dict(metadata) + values["device"] = wp.get_device(values["device"]) + return cls(SimpleNamespace(**values), robot) + + def read( + self, physical: dict[str, torch.Tensor] + ) -> tuple[torch.Tensor, torch.Tensor]: + return ( + physical["joint_q"][self.q_indices], + physical["joint_qd"][self.qd_indices], + ) + + def contacts(self, body_q: torch.Tensor) -> dict[str, torch.Tensor]: + if self.robot == "cartpole": + empty = body_q.new_empty((self.num_envs, 0)) + return { + name: empty + 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] + world_points = poses[..., :3] + quat_rotate(poses[..., 3:], self.contact_points) + result = plane_contact_features( + world_points, + self.contact_thickness, + self.plane_normal, + self.plane_offset, + self.plane_margin, + ) + result["contact_points_0"] = self.contact_points.flatten(1) + return result + + def metadata(self, dt: float) -> dict[str, Any]: + return { + "robot": self.robot, + "dt": dt, + "joint_names": self.joint_names, + "physics_signature": self.physics_signature, + "physics_signature_schema": _PHYSICS_SIGNATURE_SCHEMA, + "gravity": self.gravity[0].tolist(), + "codec": { + "robot": self.robot, + "base_com": self.base_com, + "continuous_q_indices": [1] if self.robot == "cartpole" else [], + }, + "contact_schema": "nerd-static-plane-v1", + } diff --git a/embodichain/lab/sim/physics/solvers/_reference_assets.py b/embodichain/lab/sim/physics/solvers/_reference_assets.py new file mode 100644 index 000000000..4bb93bc39 --- /dev/null +++ b/embodichain/lab/sim/physics/solvers/_reference_assets.py @@ -0,0 +1,350 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Translate pinned NeRD assets through public Newton import and Spawn data. + +Only a temporary CPU-side ModelBuilder is used for source interpretation. No +live World, backend, scene or model is mutated. Scene materialization remains +entirely owned by DexSim Spawn. +""" + +from __future__ import annotations + +from dataclasses import fields +import hashlib +from pathlib import Path +from typing import Any, TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + from dexsim.spawn import ArticulationDesc + +__all__ = [ + "make_reference_articulation", + "reference_asset_hash", + "reference_joint_names", + "reference_profile", +] + +_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 _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) -> Any: + """Interpret the pinned asset with upstream environment import defaults.""" + import newton + import warp as wp + + builder = newton.ModelBuilder() + if robot == "cartpole": + builder.default_joint_cfg = newton.ModelBuilder.JointDofConfig( + armature=0.01, limit_ke=1e4, limit_kd=10.0 + ) + builder.add_urdf( + str(_asset(robot)), + xform=wp.transform( + (0.0, 0.0, 0.0), + 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, + ) + 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 + ) + builder.add_mjcf( + str(_asset(robot)), + enable_self_collisions=False, + collapse_fixed_joints=True, + ignore_names=["floor", "ground"], + ) + return builder + + +def _pose(transform: Any) -> np.ndarray: + import warp as wp + + value = np.asarray(transform, dtype=np.float32) + pose = np.eye(4, dtype=np.float32) + pose[:3, :3] = np.asarray( + wp.quat_to_matrix(wp.quat(*value[3:])), dtype=np.float32 + ).reshape(3, 3) + pose[:3, 3] = value[:3] + return pose + + +def make_reference_articulation(robot: str) -> ArticulationDesc: + """Build an independent public descriptor with imported physical properties. + + Body/joint names retain their asset spelling, with only the importer's path + prefix removed. The free root is represented by ``fixed_base=False`` and is + created by Spawn. Default nonzero joint positions are in ``reference_profile`` + and must be applied through the host's public state interface after prepare. + """ + _asset(robot) + import newton + from dexsim.spawn import ( + ArticulationDesc, + CollisionDesc, + JointDesc, + LinkDesc, + NewtonJointDesc, + RenderDesc, + RigidBodyPhysicsDesc, + ) + from dexsim.engine import JointType + + builder = _import_reference(robot) + names = [str(name).rsplit("/", 1)[-1] for name in builder.body_label] + if len(names) != len(set(names)): + raise ValueError("Reference body names are not unique") + parents = { + int(child): int(parent) + for parent, child in zip(builder.joint_parent, builder.joint_child) + } + roots = [body for body, parent in parents.items() if parent < 0] + if len(roots) != 1: + raise ValueError("Reference asset must have one articulation root") + root = roots[0] + links = [ + LinkDesc( + name=name, + parent_name="" if parents[index] < 0 else names[parents[index]], + pose=_pose(builder.body_q[index]), + rigid_body=RigidBodyPhysicsDesc.dynamic( + mass=float(builder.body_mass[index]), + inertia=np.asarray(builder.body_inertia[index], dtype=np.float32) + .reshape(3, 3) + .copy(), + com_position=np.asarray( + builder.body_com[index], dtype=np.float32 + ).copy(), + ), + ) + for index, name in enumerate(names) + ] + if robot == "cartpole": + # Spawn exposes fixed/free articulation roots. Preserve the URDF's welded + # massless slideBar when Newton's collapse moved its prismatic child to + # the world; this adds no generalized coordinate or dynamic inertia. + for link in links: + if not link.parent_name: + link.parent_name = "slideBar" + links.append( + LinkDesc( + name="slideBar", + parent_name="", + pose=_pose(builder.body_q[root]), + rigid_body=RigidBodyPhysicsDesc.static(), + ) + ) + snapshot_names = ( + "shape_body", + "shape_flags", + "shape_type", + "shape_scale", + "shape_transform", + "shape_material_ke", + "shape_material_kd", + "shape_material_mu", + "shape_margin", + "shape_gap", + ) + snapshot = {key: np.asarray(getattr(builder, key)) for key in snapshot_names} + for index, body, collision in CollisionDesc.from_newton_snapshot(snapshot): + if body < 0: + raise ValueError( + "Reference articulation unexpectedly includes world geometry" + ) + collision.segment_name = str(builder.shape_label[index]).rsplit("/", 1)[-1] + # Snapshot conversion supplies geometry/basic contact fields. Preserve the + # rest explicitly to avoid Spawn's defaults changing source materials. + for field in fields(collision.newton): + source_name = "shape_material_" + field.name + if hasattr(builder, source_name): + setattr( + collision.newton, + field.name, + float(getattr(builder, source_name)[index]), + ) + collision.newton.is_solid = bool(builder.shape_is_solid[index]) + collision.newton.collision_group = int(builder.shape_collision_group[index]) + attribute = builder.custom_attributes.get("mujoco:condim") + if attribute is not None: + collision.newton.condim = int( + getattr( + attribute.values.get(index, attribute.default), + "value", + attribute.values.get(index, attribute.default), + ) + ) + links[body].collisions.append(collision) + links[body].visuals.append(RenderDesc.from_geometry(collision)) + joints = [] + for index, (parent, child) in enumerate( + zip(builder.joint_parent, builder.joint_child) + ): + if parent < 0 and robot != "cartpole": + continue + kind = builder.joint_type[index] + if kind not in (newton.JointType.PRISMATIC, newton.JointType.REVOLUTE): + raise ValueError(f"Unsupported reference joint type: {kind}") + offset = builder.joint_qd_start[index] + name = str(builder.joint_label[index]).rsplit("/", 1)[-1] + properties = NewtonJointDesc( + target_mode=int(builder.joint_target_mode[offset]), + **{ + key: float(getattr(builder, source)[offset]) + for key, source in ( + ("target_ke", "joint_target_ke"), + ("target_kd", "joint_target_kd"), + ("damping", "joint_damping"), + ("effort_limit", "joint_effort_limit"), + ("velocity_limit", "joint_velocity_limit"), + ("friction", "joint_friction"), + ("armature", "joint_armature"), + ("limit_ke", "joint_limit_ke"), + ("limit_kd", "joint_limit_kd"), + ) + }, + ) + joints.append( + JointDesc( + name=name, + parent_link_name=names[parent] if parent >= 0 else "slideBar", + child_link_name=names[child], + joint_type=( + JointType.PRISMATIC + if kind == newton.JointType.PRISMATIC + else ( + JointType.CONTINUOUS + if robot == "cartpole" and name == "cart_to_pole" + else JointType.REVOLUTE + ) + ), + origin_pose=( + _pose(builder.joint_X_p[index]) + if parent >= 0 + else np.linalg.inv(_pose(builder.body_q[root])) + @ _pose(builder.joint_X_p[index]) + ), + target_pose=_pose(builder.joint_X_c[index]), + axis=np.asarray(builder.joint_axis[offset], dtype=np.float32).copy(), + lower_limit=float(builder.joint_limit_lower[offset]), + upper_limit=float(builder.joint_limit_upper[offset]), + armature=properties.armature, + effort_limit=properties.effort_limit, + velocity_limit=properties.velocity_limit, + newton=properties, + ) + ) + if tuple(joint.name for joint in joints) != reference_joint_names(robot): + raise ValueError("Newton importer changed reference joint ordering") + pose = _pose(builder.body_q[root]) + if robot == "ant": + pose[:3, 3] = reference_profile(robot)["default_q"][:3] + return ArticulationDesc( + name=robot, + pose=pose, + links=([links[-1], *links[:-1]] if robot == "cartpole" else links), + joints=joints, + fixed_base=robot == "cartpole", + root_link_name=names[root] if robot == "ant" else "slideBar", + enable_self_collision=False, + ) diff --git a/embodichain/lab/sim/physics/solvers/collection.py b/embodichain/lab/sim/physics/solvers/collection.py new file mode 100644 index 000000000..7d246392b --- /dev/null +++ b/embodichain/lab/sim/physics/solvers/collection.py @@ -0,0 +1,1327 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Scene-owned analytical data collection and matched NeRD rollout evaluation.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from collections import OrderedDict +import hashlib +import math +from importlib import metadata +import json +import platform +from pathlib import Path +import time +from typing import Any + +import h5py +import numpy as np +import torch + +from embodichain.compute.dynamics import quat_rotate +from embodichain.compute.dynamics._evaluation import ( + _error_values, + _safe_number, + _teacher_row_sha256, +) +from embodichain.compute.dynamics.reference import reference_task_values as _task_values +from embodichain.data_pipeline.dynamics.shards import ( + DynamicsShardWriter, + load_manifest, + dataset_fingerprint, +) +from embodichain.utils import configclass +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.cfg import NewtonPhysicsCfg, RenderCfg + +from ._nerd_binding import NeRDBinding +from ._reference_assets import make_reference_articulation, reference_profile + +__all__ = [ + "DynamicsGenerationCfg", + "DynamicsSimulation", + "generate_dataset", + "evaluate_rollouts", +] + + +def _runtime_identity(device: torch.device) -> dict[str, Any]: + """Capture the actual neural executor separately from stored teacher metadata.""" + return { + "device": str(device), + "engine": "current_newton", + "device_name": ( + torch.cuda.get_device_name(device) + if device.type == "cuda" + else platform.machine() + ), + "python": platform.python_version(), + "platform": platform.platform(), + "state_dtype": "float32", + "cuda_runtime": torch.version.cuda, + "versions": { + name: metadata.version(name) + for name in ("embodichain", "dexsim-engine", "newton", "warp-lang", "torch") + }, + "source_sha256": { + name: _sha256(Path(__file__).with_name(name)) + for name in ( + "collection.py", + "nerd.py", + "_nerd_binding.py", + "_reference_assets.py", + ) + }, + } + + +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) + + +class DynamicsSimulation: + """Own one reference scene through SimulationManager and public Scene APIs. + + Args: + robot: Reference identity, ``cartpole`` or ``ant``. + device: Physics/tensor device. + num_envs: Number of independent replicated worlds. + solver: Named Newton solver configuration. + num_substeps: Teacher solver substeps per 1/60-second sample. + requires_grad: Allocate functional dynamics when supported by the provider. + """ + + def __init__( + self, + robot: str, + *, + device: str = "cpu", + num_envs: int = 1, + solver: Mapping[str, Any] | None = None, + num_substeps: int = 10, + requires_grad: bool = False, + ) -> None: + for name, value in (("num_envs", num_envs), ("num_substeps", num_substeps)): + _positive_integer(name, value) + if type(requires_grad) is not bool: + raise TypeError("requires_grad must be boolean") + self.profile = reference_profile(robot) + self.robot, self.device, self.num_envs = robot, torch.device(device), num_envs + self.dt = float(self.profile["dt"]) + solver_cfg = dict( + solver + or { + "solver_type": "mujoco_warp", + "use_mujoco_cpu": self.device.type == "cpu", + "separate_worlds": True, + } + ) + neural = solver_cfg.get("solver_type") == "nerd" + if neural and num_substeps != 1: + raise ValueError("The NeRD reference solver requires one macro substep.") + cfg = SimulationManagerCfg( + headless=True, + num_envs=num_envs, + arena_space=0.0, + enable_ground_plane=robot == "ant", + enable_entity_gizmo=False, + render_cfg=RenderCfg(renderer="no-render"), + startup_summary="off", + physics_cfg=NewtonPhysicsCfg( + device=device, + physics_dt=self.dt, + num_substeps=num_substeps, + requires_grad=requires_grad, + use_cuda_graph=False, + solver_cfg=solver_cfg, + collision_cfg=None if neural else {}, + ), + ) + self.sim = SimulationManager(cfg, requires_native_renderer=False) + self.articulation = None + try: + self.articulation = self.sim.add_articulation_descriptor( + make_reference_articulation(robot) + ) + self.sim.prepare() + self.binding = NeRDBinding.from_metadata( + dict(self.sim.spawn_result.dynamics_metadata()), robot + ) + self.joint_ids = [ + self.articulation.joint_names.index(name) + for name in self.binding.joint_names + ] + self.manifest = {**self.profile, **self.binding.metadata(self.dt)} + self.manifest["teacher"] = { + "solver": solver_cfg, + "num_substeps": num_substeps, + "versions": { + name: metadata.version(name) + for name in ("newton", "warp-lang", "torch", "dexsim-engine") + }, + } + self.manifest["teacher_origin"] = { + "kind": "current_newton", + "device": str(self.device), + "solver": solver_cfg, + "num_substeps": num_substeps, + "versions": self.manifest["teacher"]["versions"], + "dt": self.dt, + "physics_signature": self.binding.physics_signature, + "physics_signature_schema": self.binding.metadata(self.dt)[ + "physics_signature_schema" + ], + } + self.reset() + except Exception: + self.close() + raise + + def snapshot(self) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: + """Read independent physical state and aligned contact features.""" + physical = self.sim.spawn_result.dynamics_snapshot().physical + q, qd = self.binding.read(physical) + states = { + "q": q, + "qd": qd, + "root_body_q": physical["body_q"][self.binding.root_bodies], + } + contacts = self.binding.contacts(physical["body_q"]) + contacts.pop("contact_masks", None) # Recompute through the shared codec. + return states, contacts + + def reset( + self, q: torch.Tensor | None = None, qd: torch.Tensor | None = None + ) -> None: + """Set a complete reference batch through public articulation state writes.""" + self.sim.spawn_result.reset() + if q is None: + q = ( + torch.tensor(self.profile["default_q"], device=self.device) + .expand(self.num_envs, -1) + .clone() + ) + if qd is None: + qd = ( + torch.tensor(self.profile["default_qd"], device=self.device) + .expand(self.num_envs, -1) + .clone() + ) + 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.") + kwargs = {"joint_ids": self.joint_ids, "clear_dynamics": True} + if self.robot == "ant": + kwargs.update( + root_pose=q[:, :7], + root_velocity=qd[:, :6], + qpos=q[:, 7:], + qvel=qd[:, 6:], + ) + else: + kwargs.update(qpos=q, qvel=qd) + self.articulation.set_state(**kwargs) + + def step(self, effort: torch.Tensor) -> None: + """Apply recorded effort for one complete sampling interval.""" + if effort.shape != (self.num_envs, self.profile["qd_dim"]): + raise ValueError( + "Effort must use full reference generalized velocity layout." + ) + if self.robot == "ant": + if bool(torch.any(effort[:, :6] != 0)): + raise ValueError("Ant sampling cannot apply free-base forces.") + effort = effort[:, 6:] + self.articulation.set_qf(effort, joint_ids=self.joint_ids) + self.sim.update(step=1, render_final_step=False) + + def close(self) -> None: + """Release public child views before destroying the native World owner.""" + self.articulation = None + sim = getattr(self, "sim", None) + if sim is not None: + self.sim = None + sim.destroy(exit_process=False) + SimulationManager.flush_cleanup_queue() + + def __enter__(self) -> DynamicsSimulation: + """Return this scene owner for bounded context-managed execution.""" + return self + + def __exit__(self, *_args: Any) -> None: + """Release the owned simulation when leaving the context.""" + self.close() + + +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 + + 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: + raise TypeError("resume 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.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 DynamicsSimulation( + cfg.robot, + device=cfg.device, + num_envs=cfg.num_envs, + solver=cfg.teacher_solver, + num_substeps=cfg.num_substeps, + ) 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.binding.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 + + +class _FixedPolicy: + """One deterministic policy artifact, evaluated independently on each rollout.""" + + def __init__( + self, robot: str, config: Mapping[str, Any], device: torch.device + ) -> None: + self.robot, self.device = robot, device + self.config = dict(config) + self.policy = None + self.normalizer = None + kind = self.config.get("kind") + self.identity = { + "kind": kind, + "observation_schema": ( + "cartpole-state-4-v1" if robot == "cartpole" else "ant-com-state-27-v1" + ), + "deterministic": True, + "implementation": "embodichain-reference-feedback-v1", + "effort_limits": reference_profile(robot)["effort_limits"], + "scope": "fixed_baseline_not_paper_policy", + } + if kind == "pd": + allowed = {"kind", "gains"} if robot == "cartpole" else {"kind", "kp", "kd"} + if set(self.config) - allowed: + raise ValueError("Unknown PD policy setting") + if robot == "cartpole": + gains = self.config.get("gains", [5.0, 50.0, 5.0, 20.0]) + if ( + not isinstance(gains, Sequence) + or len(gains) != 4 + or any( + isinstance(v, bool) + or not isinstance(v, (int, float)) + or not math.isfinite(v) + for v in gains + ) + ): + raise ValueError( + "Cartpole PD gains must contain four finite numbers" + ) + self.gains = torch.tensor(gains, device=device) + self.identity["gains"] = list(gains) + else: + kp, kd = self.config.get("kp", 20.0), self.config.get("kd", 1.0) + if any( + isinstance(v, bool) + or not isinstance(v, (int, float)) + or not math.isfinite(v) + or v < 0 + for v in (kp, kd) + ): + raise ValueError("PD kp/kd must be finite nonnegative numbers") + self.kp, self.kd = float(kp), float(kd) + self.identity.update(kp=self.kp, kd=self.kd) + self.target = torch.tensor( + reference_profile(robot)["default_q"][7:], device=device + ) + elif kind == "actor_only": + if set(self.config) != {"kind", "config", "checkpoint"}: + raise ValueError( + "actor_only policy requires kind, config and checkpoint" + ) + from embodichain.learning.models import build_mlp_from_cfg + from embodichain.learning.rl.policies import build_policy + from embodichain.learning.rl.normalization import ( + RunningObservationNormalizer, + ) + from embodichain.utils.utility import load_config + + config_path, checkpoint_path = Path(self.config["config"]), Path( + self.config["checkpoint"] + ) + config_hash = _sha256(config_path) + checkpoint_hash = _sha256(checkpoint_path) + raw = load_config(config_path) + block = raw.get("policy", raw) + if block.get("name") != "actor_only" or block.get("obs_groups") is not None: + raise ValueError( + "Fixed-policy evaluation requires flat actor_only observations" + ) + obs_dim, action_dim = (4, 1) if robot == "cartpole" else (27, 8) + actor = build_mlp_from_cfg(block["actor"], obs_dim, action_dim) + self.policy = build_policy(block, obs_dim, action_dim, device, actor=actor) + payload = torch.load( + checkpoint_path, map_location=device, weights_only=True + ) + if payload.get("schema_version") != 1 or "policy" not in payload: + raise ValueError( + "Expected an EmbodiChain differentiable trainer policy checkpoint" + ) + if ( + _sha256(config_path) != config_hash + or _sha256(checkpoint_path) != checkpoint_hash + ): + raise ValueError("Policy artifact changed while loading") + if ( + raw.get("trainer", {}).get("normalize_observations", False) + and payload.get("observation_normalizer") is None + ): + raise ValueError( + "Policy config requires checkpoint observation normalization" + ) + self.policy.load_state_dict(payload["policy"], strict=True) + self.policy.eval().requires_grad_(False) + if payload.get("observation_normalizer") is not None: + self.normalizer = RunningObservationNormalizer(obs_dim, device) + self.normalizer.load_state_dict(payload["observation_normalizer"]) + self.identity.update( + config_sha256=config_hash, + checkpoint_sha256=checkpoint_hash, + scope="embodichain_fixed_actor_not_paper_policy", + ) + else: + raise ValueError("policy.kind must be pd or actor_only") + self.identity["sha256"] = hashlib.sha256( + json.dumps(self.identity, sort_keys=True).encode() + ).hexdigest() + + def __call__(self, q: torch.Tensor, qd: torch.Tensor) -> torch.Tensor: + obs, _, _ = _task_values(self.robot, q, qd) + finite = torch.isfinite(obs).all(-1) + obs = torch.nan_to_num(obs) + force = torch.zeros_like(qd) + if self.policy is not None: + from tensordict import TensorDict + + if self.normalizer is not None: + obs = self.normalizer.normalize(obs) + action = self.policy.get_action( + TensorDict({"obs": obs}, batch_size=[len(q)], device=q.device), + deterministic=True, + )["action"].clamp(-1, 1) + if self.robot == "cartpole": + force[:, 0] = action[:, 0] * 1500 + else: + force[:, 6:] = action * 400 + elif self.robot == "cartpole": + force[:, 0] = -(obs * self.gains).sum(-1).clamp(-1500, 1500) + else: + force[:, 6:] = ( + self.kp * (self.target - torch.nan_to_num(q[:, 7:])) + - self.kd * torch.nan_to_num(qd[:, 6:]) + ).clamp(-400, 400) + return torch.where(finite[:, None], force, 0) + + +def _evaluation_teacher( + config: Mapping[str, Any], + directory: Path, + horizon: int, + count: int, + batch: int, + policy: _FixedPolicy | None, +) -> dict[str, Any]: + """Persist every requested teacher trajectory, including invalid outcomes. + + These are evaluation traces, never a filtered training dataset. NaN/Inf arrays + are retained as evidence; no replacement candidate can hide teacher failures. + """ + if directory.exists(): + raise FileExistsError(f"Evaluation trace directory already exists: {directory}") + directory.mkdir(parents=True) + robot = config["robot"] + seed = config.get("seed", 500) + manifest = { + "schema_version": 1, + "trace_type": "evaluation_including_failures", + "split_seed": 0, + "shards": [], + } + with DynamicsSimulation( + robot, + device=config.get("device", "cpu"), + num_envs=batch, + solver=config.get("teacher_solver"), + num_substeps=config.get("num_substeps", 10), + ) as teacher: + manifest["metadata"] = teacher.manifest + manifest["evaluation"] = { + "seed": seed, + "trajectories": count, + "horizon": horizon, + "num_envs": batch, + "motion_mode": config.get("motion_mode", "random_effort"), + "policy": None if policy is None else policy.identity, + } + for start in range(0, count, batch): + n = min(batch, count - start) + samples = [ + _sample( + robot, + seed + i, + horizon, + "passive" if policy else config.get("motion_mode", "random_effort"), + config.get("initial_states", "sample"), + config.get("effort_scale", 1.0), + ) + for i in range(start, start + n) + ] + samples.extend([samples[-1]] * (batch - n)) + q, qd, efforts = [ + torch.stack([s[j] for s in samples]).to(teacher.device) + for j in range(3) + ] + teacher.reset(q, qd) + q_rows, qd_rows, force_rows = [], [], [] + with torch.no_grad(): + for step in range(horizon): + state, _ = teacher.snapshot() + q_rows.append(state["q"]) + qd_rows.append(state["qd"]) + force = ( + efforts[:, step] + if policy is None + else policy(state["q"], state["qd"]) + ) + force_rows.append(force) + teacher.step(force) + state, _ = teacher.snapshot() + q_rows.append(state["q"]) + qd_rows.append(state["qd"]) + q_array = torch.stack(q_rows, 1)[:n].cpu().numpy() + qd_array = torch.stack(qd_rows, 1)[:n].cpu().numpy() + force_array = torch.stack(force_rows, 1)[:n].cpu().numpy() + filename = f'teacher-{len(manifest["shards"]):06d}.h5' + rows = [] + with h5py.File(directory / filename, "w") as stream: + for row in range(n): + group = stream.create_group(str(row)) + group.create_dataset( + "states/q", data=q_array[row], compression="lzf" + ) + group.create_dataset( + "states/qd", data=qd_array[row], compression="lzf" + ) + group.create_dataset( + "transitions/joint_f", data=force_array[row], compression="lzf" + ) + rows.append( + { + "id": f"{seed}:{start+row}", + "group": str(row), + "length": horizon, + "split": "test", + "metadata": {"sampling_seed": seed + start + row}, + } + ) + manifest["shards"].append( + { + "file": filename, + "sha256": _sha256(directory / filename), + "trajectories": rows, + } + ) + _write_json(directory / "manifest.json", manifest) + return manifest + + +def _teacher_origin(metadata: Mapping[str, Any]) -> dict[str, Any]: + """Resolve a recorded domain, never infer official provenance from a robot name.""" + origin = metadata.get("teacher_origin") + if isinstance(origin, Mapping) and origin.get("kind") in { + "current_newton", + "official_warp", + }: + return dict(origin) + teacher = metadata.get("teacher") + if ( + metadata.get("teacher_profile") == "current_newton_import_generalized_effort" + and isinstance(teacher, Mapping) + and "newton" in teacher.get("versions", {}) + and isinstance(teacher.get("solver"), Mapping) + ): + return { + "kind": "current_newton", + **dict(teacher), + "dt": metadata.get("dt"), + "physics_signature": metadata.get("physics_signature"), + "physics_signature_schema": metadata.get("physics_signature_schema"), + } + raise ValueError("Teacher dataset has no verifiable recorded teacher_origin") + + +def _source_motion(manifest: Mapping[str, Any]) -> str | None: + """Read the persisted sampling contract, independently of the requested label.""" + evaluation = manifest.get("evaluation", {}) + if evaluation.get("motion_mode") is not None: + return str(evaluation["motion_mode"]) + metadata = manifest["metadata"] + generation = metadata.get("generation", {}) + if generation.get("motion_mode") is not None: + return str(generation["motion_mode"]) + sampling = metadata.get("sampling", {}) + if type(sampling.get("passive")) is bool: + return "passive" if sampling["passive"] else "random_effort" + return None + + +def _teacher_batches( + directory: Path, + manifest: Mapping[str, Any], + records: list[dict[str, Any]], + batch: int, + horizon: int, + device: torch.device, +): + """Read at most one batch of bounded prefixes and keep at most four files open.""" + handles: OrderedDict[str, Any] = 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"]] + row_values = ( + group["states/q"][: horizon + 1], + group["states/qd"][: horizon + 1], + group["transitions/joint_f"][:horizon], + ) + values.append(row_values) + row["reference_sha256"] = _teacher_row_sha256(row, row_values) + row["has_nonzero_effort"] = bool(np.any(row_values[2] != 0)) + values.extend([values[-1]] * (batch - len(values))) + arrays = [ + torch.from_numpy(np.stack([v[i] for v in values])).to(device) + for i in range(3) + ] + yield selected, *arrays + finally: + for stream in handles.values(): + stream.close() + + +def evaluate_rollouts(config: Mapping[str, Any]) -> dict[str, Any]: + """Evaluate batched independent rollouts with complete failure evidence. + + Fixed-policy actions are recomputed from each simulator's own state. Optional + immutable teacher_dataset replays canonical test trajectories for open-loop + comparisons; it cannot establish closed-loop policy behavior for that teacher. + Metrics synchronize once per batch, never once per metric or simulation step. + + Args: + config: Robot, model export, output path and evaluation settings. Optional + teacher_dataset selects immutable canonical traces; split defaults to + test. Fixed-policy mode instead needs a live current Newton teacher + and a policy mapping with kind pd or actor_only. Policy artifact paths + are already resolved by the CLI. num_envs is the actual replica count + for both teacher generation and neural replay. + + Returns: + Horizon/time error means, explicit valid/failure counts, policy returns, + artifact hashes and paths to per-trajectory JSONL and HDF5 evidence. + """ + allowed = { + "robot", + "model", + "device", + "num_envs", + "trajectories", + "horizons", + "seed", + "teacher_solver", + "num_substeps", + "motion_mode", + "initial_states", + "effort_scale", + "output", + "policy", + "teacher_dataset", + "split", + } + if set(config) - allowed: + raise ValueError(f"Unknown rollout settings: {sorted(set(config)-allowed)}") + robot = config["robot"] + reference_profile(robot) + count = _positive_integer("trajectories", config.get("trajectories", 2048)) + batch = _positive_integer("num_envs", config.get("num_envs", 64)) + seed = _positive_integer("seed", config.get("seed", 500), 0) + _positive_integer("num_substeps", config.get("num_substeps", 10)) + raw_horizons = config.get("horizons", [1, 10, 100, 1000]) + if ( + not isinstance(raw_horizons, Sequence) + or isinstance(raw_horizons, str) + or not raw_horizons + ): + raise ValueError("horizons must be a nonempty sequence") + horizons = sorted({_positive_integer("horizon", value) for value in raw_horizons}) + horizon = horizons[-1] + motion = config.get("motion_mode", "random_effort") + if motion not in {"passive", "random_effort", "fixed_policy"}: + raise ValueError("Unsupported evaluation motion_mode") + if config.get("initial_states", "sample") not in {"sample", "default"}: + raise ValueError("Unsupported initial_states") + scale = config.get("effort_scale", 1.0) + if ( + isinstance(scale, bool) + or not isinstance(scale, (int, float)) + or not math.isfinite(scale) + or scale < 0 + ): + raise ValueError("effort_scale must be finite and nonnegative") + if (motion == "fixed_policy") != ("policy" in config): + raise ValueError( + "fixed_policy mode requires an explicit policy configuration only in that mode" + ) + if motion == "fixed_policy" and config.get("teacher_dataset") is not None: + raise ValueError( + "Offline teacher_dataset cannot establish closed-loop fixed-policy behavior" + ) + if motion == "fixed_policy" and scale != 1.0: + raise ValueError( + "effort_scale applies to random_effort; configure policy gains for fixed_policy" + ) + device = torch.device(config.get("device", "cpu")) + policy = ( + _FixedPolicy(robot, config["policy"], device) + if motion == "fixed_policy" + else None + ) + output = Path(config["output"]) + model = Path(config["model"]) + model_sha = _sha256(model / "weights.pt") + model_manifest_sha = _sha256(model / "manifest.json") + evidence_path = output.with_name(output.stem + "-evidence.h5") + rows_path = output.with_name(output.stem + "-trajectories.jsonl") + if output.exists() or evidence_path.exists() or rows_path.exists(): + raise FileExistsError( + "Evaluation output/evidence already exists; choose a new output path" + ) + output.parent.mkdir(parents=True, exist_ok=True) + started = time.perf_counter() + if config.get("teacher_dataset") is None: + teacher_dir = output.with_name(output.stem + "-teacher-data") + dataset = _evaluation_teacher( + config, teacher_dir, horizon, count, batch, policy + ) + split = "test" + else: + teacher_dir = Path(config["teacher_dataset"]) + dataset = load_manifest(teacher_dir, verify=True) + split = config.get("split", "test") + if split not in {"train", "validation", "test", "all"}: + raise ValueError("split must be train, validation, test or all") + if dataset["metadata"].get("robot") != robot: + raise ValueError("Teacher dataset robot differs from evaluation robot") + records = [ + {**row, "file": shard["file"]} + for shard in dataset["shards"] + for row in shard["trajectories"] + if (split == "all" or row["split"] == split) and row["length"] >= horizon + ] + if len(records) < count: + raise ValueError( + f"Teacher dataset has only {len(records)} eligible trajectories, requested {count}" + ) + records = records[:count] + teacher_origin = _teacher_origin(dataset["metadata"]) + source_motion = _source_motion(dataset) + if config.get("teacher_dataset") is not None: + if source_motion == "fixed_policy": + raise ValueError( + "Fixed-policy traces cannot be labeled passive or random_effort replay" + ) + if motion == "random_effort" and source_motion != "random_effort": + raise ValueError( + "Teacher dataset does not establish a random_effort sampling contract" + ) + # A passive label is additionally checked against every selected effort + # value below. A claimed passive source cannot override nonzero evidence. + selected_teacher_digest = hashlib.sha256() + metric_names = ("q_mae", "qd_mae", "base_position_error", "base_orientation_error") + errors_sum = torch.zeros(horizon, 4, device=device, dtype=torch.float64) + valid_count = torch.zeros(horizon, device=device, dtype=torch.long) + teacher_nonfinite_count = torch.zeros_like(valid_count) + neural_nonfinite_count = torch.zeros_like(valid_count) + teacher_failure_count = torch.zeros_like(valid_count) + neural_failure_count = torch.zeros_like(valid_count) + horizon_sum = { + h: torch.zeros(4, device=device, dtype=torch.float64) for h in horizons + } + horizon_valid = { + h: torch.zeros((), device=device, dtype=torch.long) for h in horizons + } + teacher_returns = torch.zeros(horizon, device=device, dtype=torch.float64) + neural_returns = torch.zeros_like(teacher_returns) + with ( + DynamicsSimulation( + robot, + device=str(device), + num_envs=batch, + solver={"solver_type": "nerd", "options": {"checkpoint": str(model)}}, + num_substeps=1, + ) as neural, + h5py.File(evidence_path, "w") as evidence, + rows_path.open("w") as row_stream, + ): + runtime_identity = _runtime_identity(neural.binding.device) + if ( + _sha256(model / "weights.pt") != model_sha + or _sha256(model / "manifest.json") != model_manifest_sha + ): + raise ValueError("NeRD export changed while preparing evaluation") + for batch_index, (selected, target_q, target_qd, efforts) in enumerate( + _teacher_batches(teacher_dir, dataset, records, batch, horizon, device) + ): + n = len(selected) + if motion == "passive" and any( + record["has_nonzero_effort"] for record in selected + ): + raise ValueError( + "Cannot label nonzero teacher effort trajectories as passive" + ) + for record in selected: + selected_teacher_digest.update( + bytes.fromhex(record["reference_sha256"]) + ) + if ( + target_q.shape[-1] != neural.profile["q_dim"] + or target_qd.shape[-1] != neural.profile["qd_dim"] + or efforts.shape[-1] != neural.profile["qd_dim"] + ): + raise ValueError( + "Teacher data does not use canonical native q/qd/effort dimensions" + ) + neural.reset(target_q[:, 0], target_qd[:, 0]) + actual_initial, _ = neural.snapshot() + predicted_q = [actual_initial["q"]] + predicted_qd = [actual_initial["qd"]] + neural_efforts = [] + with torch.no_grad(): + for step in range(horizon): + force = ( + efforts[:, step] + if policy is None + else policy(predicted_q[-1], predicted_qd[-1]) + ) + neural_efforts.append(force) + neural.step(force) + current, _ = neural.snapshot() + predicted_q.append(current["q"]) + predicted_qd.append(current["qd"]) + q = torch.stack(predicted_q, 1)[:n] + qd = torch.stack(predicted_qd, 1)[:n] + tq = target_q[:n] + tqd = target_qd[:n] + teacher_nonfinite = ( + (~(torch.isfinite(tq).all(-1) & torch.isfinite(tqd).all(-1))) + .cummax(1) + .values[:, 1:] + ) + neural_nonfinite = ( + (~(torch.isfinite(q).all(-1) & torch.isfinite(qd).all(-1))) + .cummax(1) + .values[:, 1:] + ) + valid = ~(teacher_nonfinite | neural_nonfinite) + _, tr, tf = _task_values(robot, tq, tqd) + _, nr, nf = _task_values(robot, q, qd) + tf = ( + ( + tf + | ~torch.isfinite(tr) + | ~torch.isfinite(tq).all(-1) + | ~torch.isfinite(tqd).all(-1) + ) + .cummax(1) + .values + ) + nf = ( + ( + nf + | ~torch.isfinite(nr) + | ~torch.isfinite(q).all(-1) + | ~torch.isfinite(qd).all(-1) + ) + .cummax(1) + .values + ) + # Terminal transition reward counts once. Post-terminal continuation + # remains in error evidence but cannot inflate a policy's return. + tr = torch.where( + ~tf[:, :-1] & torch.isfinite(tr[:, 1:]), tr[:, 1:], 0 + ).cumsum(1) + nr = torch.where( + ~nf[:, :-1] & torch.isfinite(nr[:, 1:]), nr[:, 1:], 0 + ).cumsum(1) + error = _error_values(robot, q[:, 1:], qd[:, 1:], tq[:, 1:], tqd[:, 1:]) + error = torch.where(valid[..., None], error, 0) + errors_sum += error.double().sum(0) + valid_count += valid.sum(0) + teacher_nonfinite_count += teacher_nonfinite.sum(0) + neural_nonfinite_count += neural_nonfinite.sum(0) + teacher_failure_count += tf[:, 1:].sum(0) + neural_failure_count += nf[:, 1:].sum(0) + teacher_returns += tr.double().sum(0) + neural_returns += nr.double().sum(0) + cumulative = error.cumsum(1) + for h in horizons: + mask = valid[:, h - 1] + horizon_sum[h] += ( + torch.where(mask[:, None], cumulative[:, h - 1] / h, 0) + .double() + .sum(0) + ) + horizon_valid[h] += mask.sum() + # One batched transfer of each evidence array after the simulation loop. + arrays = { + k: v.detach().cpu().numpy() + for k, v in { + "q": q, + "qd": qd, + "joint_f": torch.stack(neural_efforts, 1)[:n], + "errors": error, + "valid": valid, + "teacher_nonfinite": teacher_nonfinite, + "neural_nonfinite": neural_nonfinite, + "teacher_failed": tf[:, 1:], + "neural_failed": nf[:, 1:], + "teacher_return": tr, + "neural_return": nr, + }.items() + } + group = evidence.create_group(str(batch_index)) + for key, value in arrays.items(): + group.create_dataset(key, data=value, compression="lzf") + for row, record in enumerate(selected): + item = { + "trajectory_id": record["id"], + "reference_sha256": record["reference_sha256"], + "sampling_seed": record.get("metadata", {}).get("sampling_seed"), + "evidence_group": str(batch_index), + "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 + ) + item["horizons"][str(h)] = { + **{ + key: _safe_number(value) + for key, value in zip(metric_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] + ), + } + row_stream.write(json.dumps(item, allow_nan=False) + "\n") + curve_mean = (errors_sum / valid_count.clamp_min(1)[:, None]).cpu().numpy() + counts = ( + torch.stack( + ( + valid_count, + teacher_nonfinite_count, + neural_nonfinite_count, + teacher_failure_count, + neural_failure_count, + ), + -1, + ) + .cpu() + .numpy() + ) + returns = ( + torch.stack((teacher_returns / count, neural_returns / count), -1).cpu().numpy() + ) + curve = [] + for step in range(horizon): + curve.append( + { + "step": step + 1, + **{ + key: (_safe_number(value) if counts[step, 0] else None) + for key, value in zip(metric_names, curve_mean[step]) + }, + "valid_trajectories": int(counts[step, 0]), + "teacher_nonfinite_count": int(counts[step, 1]), + "neural_nonfinite_count": int(counts[step, 2]), + "teacher_failed_count": int(counts[step, 3]), + "failed_trajectories": int(counts[step, 4]), + "teacher_nonfinite_rate": float(counts[step, 1] / count), + "nonfinite_rate": float(counts[step, 2] / count), + "teacher_failure_rate": float(counts[step, 3] / count), + "failure_rate": float(counts[step, 4] / count), + "teacher_policy_return": float(returns[step, 0]), + "policy_return": float(returns[step, 1]), + } + ) + metrics = {} + for h in horizons: + number = int(horizon_valid[h].cpu()) + mean = (horizon_sum[h] / max(number, 1)).cpu().tolist() + metrics[str(h)] = { + **curve[h - 1], + **{ + key: (_safe_number(value) if number else None) + for key, value in zip(metric_names, mean) + }, + "valid_trajectories": number, + "numerically_invalid_trajectories": count - number, + } + result = { + "status": "completed", + "scope": ( + "current_newton_closed_loop_not_paper_policy" + if policy + else "immutable_teacher_open_loop_not_official_training_qualification" + ), + "robot": robot, + "motion_mode": motion, + "seed": seed, + "model": str(model.resolve()), + "model_sha256": model_sha, + "model_manifest_sha256": model_manifest_sha, + "teacher_dataset": str(teacher_dir.resolve()), + "teacher_dataset_fingerprint": dataset_fingerprint(dataset), + "teacher_metadata": dataset["metadata"], + "teacher_origin": teacher_origin, + "teacher_source_motion_mode": source_motion, + "selected_teacher_sha256": selected_teacher_digest.hexdigest(), + "selected_teacher_hash_schema": "ordered-teacher-content-v1", + "teacher_split": split, + "trajectories": count, + "num_envs": batch, + "state_dtype": "float32", + "runtime_identity": runtime_identity, + "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": _sha256(rows_path), + "trace_evidence": str(evidence_path.resolve()), + "trace_evidence_sha256": _sha256(evidence_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 that horizon; physical failures remain included and separately counted. Policy returns stop after first task/numerical failure.", + "elapsed_seconds": time.perf_counter() - started, + } + _write_json(output, result) + return result diff --git a/embodichain/lab/sim/physics/solvers/nerd.py b/embodichain/lab/sim/physics/solvers/nerd.py new file mode 100644 index 000000000..df0c1ceca --- /dev/null +++ b/embodichain/lab/sim/physics/solvers/nerd.py @@ -0,0 +1,368 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Named NeRD solver using a shared native/functional state transition.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any +import math + +import newton +from newton.solvers import SolverBase +import numpy as np +import torch +import warp as wp + +from dexsim.engine.newton_physics import ( + FunctionalSolverState, + NewtonSolverCapabilities, + NewtonSolverContext, + NewtonSolverExtension, +) +from embodichain.compute.dynamics import NeRDStateCodec, append_history, reset_history +from embodichain.learning.models.nerd import load_model + +from ._kinematics import differentiable_fk, _stream_scope +from ._nerd_binding import NeRDBinding, _PHYSICS_SIGNATURE_SCHEMA + +__all__ = ["NeRDProvider", "NeuralRobotDynamicSolver"] + +_PHYSICAL_FIELDS = ( + "joint_q", + "joint_qd", + "body_q", + "body_qd", + "particle_q", + "particle_qd", +) + + +class NeuralRobotDynamicSolver(SolverBase): + """Advance homogeneous Cartpole/Ant worlds with a portable NeRD export. + + Args: + context: Model-owned construction resources supplied by DexSim. + checkpoint: Export directory containing weights and deployment manifest. + trainable: Whether the resident network parameters require gradients. + Functional callers may alternatively supply explicit parameter tensors. + """ + + def __init__( + self, context: NewtonSolverContext, checkpoint: str, *, trainable: bool = False + ) -> None: + super().__init__(context.model) + self.network, self.manifest = load_model( + checkpoint, device=wp.device_to_torch(context.device) + ) + if any( + value.is_floating_point() and value.dtype != torch.float32 + for value in self.network.state_dict().values() + ): + raise TypeError( + "NeRD Newton execution requires float32 model weights and buffers." + ) + self.network.requires_grad_(trainable) + codec_cfg = self.manifest.get("codec") + if not isinstance(codec_cfg, dict): + raise ValueError("NeRD export requires an explicit codec manifest.") + self.codec = NeRDStateCodec(**codec_cfg) + self.binding = NeRDBinding(context.model, self.codec.robot) + self._dt = float(self.manifest.get("dt", 0.0)) + self._check_dt(context.dt) + if self.manifest.get("joint_names") != self.binding.joint_names: + raise ValueError("NeRD exported joint order does not match the scene.") + if self.manifest.get("physics_signature_schema") != _PHYSICS_SIGNATURE_SCHEMA: + raise ValueError( + "Unsupported NeRD physics signature schema; regenerate data and exports " + "with the current physical contract instead of relabeling an old model." + ) + if self.manifest.get("physics_signature") != self.binding.physics_signature: + raise ValueError( + "NeRD exported physics signature does not match the scene." + ) + if not np.allclose(self.codec.base_com, self.binding.base_com, atol=1e-7): + raise ValueError("NeRD codec COM does not match the root body.") + if self.manifest.get("contact_schema") != "nerd-static-plane-v1": + raise ValueError("Unsupported NeRD contact feature schema.") + expected_gravity = self.manifest.get("gravity") + if expected_gravity is None or not torch.allclose( + self.binding.gravity, + self.binding.gravity.new_tensor(expected_gravity).expand_as( + self.binding.gravity + ), + ): + raise ValueError("NeRD scene gravity differs from the training manifest.") + for name in ("joint_target_ke", "joint_target_kd"): + values = getattr(self.model, name, None) + if values is not None and np.any(values.numpy() != 0): + raise ValueError( + "NeRD effort deployment requires zero native position/velocity drive gains." + ) + self.history_length = self.network.cfg.history_length + self._history = { + name: torch.zeros( + (self.binding.num_envs, self.history_length, width), + device=self.binding.device, + ) + for name, width in self.network.cfg.input_dims.items() + } + self._lengths = torch.zeros( + self.binding.num_envs, dtype=torch.long, device=self.binding.device + ) + self._closed = False + self._invalid = False + initial = self.initial_functional_state(self.model.state()) + q, qd = self.binding.read(initial.physical) + features = self.codec.encode( + q, + qd, + torch.zeros_like(qd), + self.binding.contacts(initial.physical["body_q"]), + self.binding.gravity_dir, + ) + if {k: v.shape[-1] for k, v in features.items()} != dict( + self.network.cfg.input_dims + ): + raise ValueError( + "NeRD network inputs do not match the scene feature schema." + ) + + def _check_dt(self, dt: float) -> None: + if self._dt <= 0 or not math.isclose( + float(dt), self._dt, rel_tol=0.0, abs_tol=1e-12 + ): + raise ValueError( + f"NeRD timestep must equal trained dt={self._dt}, got {dt}." + ) + + def initial_functional_state(self, native_state: Any) -> FunctionalSolverState: + """Clone current native state and resident history into an owned snapshot.""" + physical = { + name: wp.to_torch(getattr(native_state, name)).detach().clone() + for name in _PHYSICAL_FIELDS + if getattr(native_state, name, None) is not None + } + return FunctionalSolverState( + physical=physical, + history={k: v.detach().clone() for k, v in self._history.items()}, + valid_history_lengths=self._lengths.detach().clone(), + ) + + def functional_step( + self, + state: FunctionalSolverState, + control: Mapping[str, torch.Tensor], + parameters: Mapping[str, torch.Tensor] | None, + dt: float, + ) -> FunctionalSolverState: + """Return a new state while preserving all caller-owned graphs.""" + if self._closed: + raise RuntimeError("NeRD solver is closed.") + if self._invalid: + raise RuntimeError( + "NeRD physical model was invalidated; rebuild with a matching export." + ) + self._check_dt(dt) + if set(control) != {"joint_f"}: + raise ValueError("NeRD requires explicit joint_f effort controls only.") + force = control["joint_f"] + if ( + force.shape != (self.model.joint_dof_count,) + or force.dtype != torch.float32 + or force.device != self.binding.device + ): + raise ValueError( + "joint_f must use flattened float32 model-order DOFs on the model device." + ) + if self.codec.robot == "ant" and bool( + torch.any(force[self.binding.qd_indices[:, :6]] != 0) + ): + raise ValueError( + "NeRD Ant cannot model external generalized forces on its free base." + ) + if state.valid_history_lengths is None: + raise ValueError( + "NeRD functional state requires per-world history lengths." + ) + q, qd = self.binding.read(state.physical) + # q/qd are the independent physical coordinates. Recompute Ant contact + # geometry so the first step also differentiates through an initial q + # leaf supplied by the caller, not an unrelated detached body snapshot. + body_q = state.physical["body_q"] + if self.codec.robot == "ant": + body_q, _ = differentiable_fk( + self.model, state.physical["joint_q"], state.physical["joint_qd"] + ) + contacts = self.binding.contacts(body_q) + features = self.codec.encode( + q, qd, force[self.binding.qd_indices], contacts, self.binding.gravity_dir + ) + history, lengths = append_history( + state.history, state.valid_history_lengths, features + ) + if parameters is None: + prediction = self.network.evaluate(history, lengths).squeeze(1) + else: + output = torch.func.functional_call( + self.network, parameters, (history,), strict=False + ) + prediction = output.gather( + 1, (lengths - 1)[:, None, None].expand(-1, 1, output.shape[-1]) + ).squeeze(1) + next_q, next_qd = self.codec.decode(q, qd, prediction) + physical = dict(state.physical) + physical["joint_q"] = state.physical["joint_q"].scatter( + 0, self.binding.q_indices.flatten(), next_q.flatten() + ) + physical["joint_qd"] = state.physical["joint_qd"].scatter( + 0, self.binding.qd_indices.flatten(), next_qd.flatten() + ) + physical["body_q"], physical["body_qd"] = differentiable_fk( + self.model, physical["joint_q"], physical["joint_qd"] + ) + return FunctionalSolverState(physical, history, lengths) + + def commit_functional_state(self, state: FunctionalSolverState) -> None: + """Adopt detached history after an explicitly published functional step.""" + self._history = {k: v.detach().clone() for k, v in state.history.items()} + self._lengths = state.valid_history_lengths.detach().clone() + + def reset_functional_history( + self, + state: FunctionalSolverState, + reset_state: FunctionalSolverState, + world_mask: Any | None, + ): + """Clear selected histories without detaching unselected trajectories.""" + del reset_state + mask = self._selected_rows(world_mask) + return reset_history(state.history, state.valid_history_lengths, mask) + + def _selected_rows(self, world_mask: Any | None) -> torch.Tensor: + if world_mask is None: + return torch.ones( + self.binding.num_envs, dtype=torch.bool, device=self.binding.device + ) + mask = ( + wp.to_torch(world_mask) if isinstance(world_mask, wp.array) else world_mask + ) + return mask[self.binding.world_ids] + + def step( + self, state_in: Any, state_out: Any, control: Any, contacts: Any, dt: float + ) -> None: + """Advance the native double-buffer path with the same neural transition.""" + del contacts + with torch.no_grad(), _stream_scope(self.binding.device): + physical = { + name: wp.to_torch(getattr(state_in, name)) + for name in _PHYSICAL_FIELDS + if getattr(state_in, name, None) is not None + } + current = FunctionalSolverState(physical, self._history, self._lengths) + result = self.functional_step( + current, {"joint_f": wp.to_torch(control.joint_f)}, None, dt + ) + for name, value in result.physical.items(): + wp.copy( + getattr(state_out, name), + wp.from_torch( + value.contiguous(), dtype=getattr(state_out, name).dtype + ), + ) + self._history, self._lengths = ( + dict(result.history), + result.valid_history_lengths, + ) + + def reset( + self, state: Any, world_mask: Any | None = None, flags: Any | None = None + ) -> None: + """Clear selected resident history and synchronize FK after state reset.""" + del flags + with torch.no_grad(), _stream_scope(self.binding.device): + self._history, self._lengths = reset_history( + self._history, self._lengths, self._selected_rows(world_mask) + ) + newton.eval_fk(self.model, state.joint_q, state.joint_qd, state) + + def notify_model_changed(self, flags: Any) -> None: + """Reject changes invalidating a trained robot's physical contract.""" + del flags + try: + binding = NeRDBinding(self.model, self.codec.robot) + if ( + binding.physics_signature != self.binding.physics_signature + or not torch.equal(binding.gravity, self.binding.gravity) + ): + raise ValueError( + "Physical model changed; retrain or load a matching NeRD export." + ) + except Exception: + # Native mutation precedes Newton's notification. A caught error + # must not permit later stepping with incompatible learned physics. + self._invalid = True + raise + + def close(self) -> None: + """Release resident history after the owner stops stepping this solver.""" + self._closed = True + self._history = {} + + +class NeRDProvider: + """Entry-point provider for the registered ``nerd`` Newton solver.""" + + capabilities = NewtonSolverCapabilities( + differentiable=True, + reset=True, + contact_reports=False, + reduced_coordinates=True, + functional_contacts=False, + external_wrenches=False, + supported_control_fields=("joint_f",), + runtime_controls=False, + ) + + def register_custom_attributes( + self, builder: Any, options: Mapping[str, Any] + ) -> None: + """NeRD uses standard model state and requires no custom builder attributes.""" + del builder + if set(options) - {"checkpoint", "trainable"}: + raise ValueError("NeRD options support checkpoint and trainable only.") + if not isinstance(options.get("checkpoint"), str) or not options["checkpoint"]: + raise ValueError("NeRD requires an exported checkpoint directory.") + if not isinstance(options.get("trainable", False), bool): + raise TypeError("NeRD trainable must be boolean.") + + def create( + self, context: NewtonSolverContext, options: Mapping[str, Any] + ) -> NewtonSolverExtension: + """Build a native and functional adapter bound to the finalized scene.""" + solver = NeuralRobotDynamicSolver(context, **options) + return NewtonSolverExtension( + native_solver=solver, + capabilities=self.capabilities, + functional_step=solver.functional_step, + initial_functional_state=solver.initial_functional_state, + named_parameters=lambda: dict(solver.network.named_parameters()), + close=solver.close, + commit_functional_state=solver.commit_functional_state, + reset_functional_history=solver.reset_functional_history, + ) diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index 50e74d01f..0f2be661f 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -238,7 +238,11 @@ def __init__( startup_summary: Literal["compact", "full", "off"] = "compact", dexsim_startup_info: bool = False, scene_node_capacity: int | None = None, + enable_ground_plane: bool = True, ) -> None: + if not isinstance(enable_ground_plane, bool): + raise TypeError("enable_ground_plane must be boolean.") + self.enable_ground_plane = enable_ground_plane self.startup_summary = startup_summary self.dexsim_startup_info = dexsim_startup_info self.scene_node_capacity = scene_node_capacity @@ -290,6 +294,13 @@ def __init__( width: int = 1920 """The width of the simulation window.""" + enable_ground_plane: bool = True + """Declare the default physical ground plane during scene creation. + + Set false for contact-free dynamics, where hiding the ground would leave + its collision geometry active. This is a creation-time setting. + """ + startup_summary: Literal["compact", "full", "off"] = "compact" """Startup table detail. ``off`` disables both simulation and Gym summaries.""" @@ -969,7 +980,7 @@ def asset_uids(self) -> List[str]: Returns: List[str]: list of all assets uid. """ - uid_list = ["default_plane"] + uid_list = ["default_plane"] if self.sim_config.enable_ground_plane else [] uid_list.extend(list(self._lights.keys())) uid_list.extend(list(self._sensors.keys())) uid_list.extend(list(self._robots.keys())) @@ -1721,7 +1732,7 @@ def prepare(self) -> None: self._env = result.get_arena("default") self._arenas = [result.get_arena(name) for name in scene.arena_names] self.__dict__.pop("arena_offsets", None) - if self._default_plane is None: + if self._default_plane is None and self.sim_config.enable_ground_plane: self._bind_default_plane(scene.handles("default_plane")[0]) # Runtime readiness belongs to the SimulationManager. Keep this and @@ -1944,6 +1955,116 @@ def update( if step > 0: self._log_scene_summary() + def create_functional_state(self) -> Any: + """Snapshot independent state/history for a differentiable solver. + + Returns: + DexSim functional state owned independently of live buffers. + """ + self.prepare() + if not self.is_newton_backend: + raise ValueError("Functional dynamics requires the Newton backend.") + return self.spawn_result.create_functional_state() + + def functional_parameters(self) -> Mapping[str, torch.Tensor]: + """Return named trainable parameters of the current solver extension.""" + self.prepare() + if not self.is_newton_backend: + raise ValueError("Functional dynamics requires the Newton backend.") + return self.spawn_result.functional_parameters() + + def functional_step( + self, + state: Any, + control: Mapping[str, torch.Tensor], + parameters: Mapping[str, torch.Tensor] | None = None, + dt: float | None = None, + *, + publish: bool = True, + render_final_step: bool = True, + ) -> Any: + """Advance one functional macro step through the owning DexSim Scene. + + Args: + state: Independent physical and recurrent solver tensors. + control: Named differentiable Newton control tensors. + parameters: Optional named neural parameters for functional evaluation. + dt: Macro interval; defaults to the configured physics interval. + publish: Publish detached results and advance live time when true. + False evaluates an independent branch without touching live state. + render_final_step: Publish render consumers after a live step. + + Returns: + A new functional state retaining first-order input/parameter gradients. + """ + self.prepare() + if not self.is_newton_backend: + raise ValueError("Functional dynamics requires the Newton backend.") + resolved_dt = self.sim_config.physics_dt if dt is None else float(dt) + with self.profiler.section("functional_step", is_root=True): + result = self.spawn_result.functional_step( + state, control, parameters, dt=resolved_dt, publish=publish + ) + if publish: + self._visualization_sim_step += 1 + self._visualization_sim_time += resolved_dt + self._pending_record_dt += resolved_dt + if render_final_step: + with self.render_frame(): + pass + return result + + def reset_functional_state(self, state: Any, world_mask: Any | None = None) -> Any: + """Restore selected functional rows through the Scene reset contract. + + Args: + state: Functional state whose unselected rows keep their graphs. + world_mask: Optional boolean mask of the manager's environment rows. + A Newton mask with a false final global row is also accepted. + An all-false mask leaves the supplied state and graph unchanged. + + Returns: + Reset functional state with detached selected initial rows. + """ + self.prepare() + if not self.is_newton_backend: + raise ValueError("Functional dynamics requires the Newton backend.") + arenas = None + if world_mask is not None: + mask = ( + wp.to_torch(world_mask) + if isinstance(world_mask, wp.array) + else torch.as_tensor(world_mask, device=self.device) + ) + if mask.dtype != torch.bool or mask.ndim != 1: + raise ValueError( + "Functional reset requires a one-dimensional boolean mask." + ) + if mask.numel() == self.num_envs + 1: + if bool(mask[-1]): + if not bool(mask.all()): + raise ValueError( + "A partial manager reset cannot select global entities." + ) + mask = None + else: + mask = mask[:-1] + elif mask.numel() != self.num_envs: + raise ValueError("Functional reset mask must match environment count.") + if mask is not None: + arenas = [ + f"arena_{index}" + for index in mask.nonzero(as_tuple=False).flatten().tolist() + ] + if not arenas: + return state + # Environment reset/randomization has already written native state. + # Capture those values, not original model defaults, for selected rows. + fresh = self.spawn_result.create_functional_state() + return self.spawn_result.reset_functional_state( + state, arenas=arenas, reset_state=fresh, publish=True + ) + def get_env(self, arena_index: int = -1) -> dexsim.environment.Arena: """Get the arena or env by index. @@ -2148,6 +2269,9 @@ def set_emission_light( def _declare_spawn_default_plane(self) -> None: """Declare the global ground in the World's Spawn scene.""" + if not self.sim_config.enable_ground_plane: + return + from dexsim.spawn import ( CollisionApproximation, CollisionDesc, @@ -3190,6 +3314,42 @@ def add_articulation( self.prepare() return articulation + def add_articulation_descriptor(self, descriptor: ArticulationDesc) -> Articulation: + """Declare a procedural articulation through the existing Spawn owner. + + Args: + descriptor: Public DexSim descriptor with authored links and joints. + Geometry and physical properties retain descriptor ownership. + + Returns: + Articulation facade bound in place by :meth:`prepare`. + """ + from dexsim.spawn import ArticulationDesc + from embodichain.lab.sim.cfg import ArticulationRootPropertiesCfg + + if not isinstance(descriptor, ArticulationDesc): + raise TypeError("descriptor must be a DexSim ArticulationDesc.") + uid = descriptor.name + if uid in self._articulations or uid in self._robots: + raise ValueError(f"Articulation {uid!r} already exists.") + cfg = ArticulationCfg(uid=uid, fpath=descriptor.urdf_path) + cfg.init_local_pose = descriptor.pose.copy() + cfg.asset_physics_mode = "preserve" + cfg.root_props = ArticulationRootPropertiesCfg( + fixed_base=descriptor.fixed_base, + self_collision_enabled=descriptor.enable_self_collision, + ) + cfg.body_scale = tuple(float(v) for v in descriptor.body_scale) + cfg.build_pk_chain = False + facade = Articulation(cfg=cfg, device=self.device) + was_materialized = self.spawn_result is not None + self._spawn_scene.declare("articulation", uid, descriptor, facade=facade) + self._articulations[uid] = facade + self.notify_visualization_topology_changed() + if was_materialized: + self.prepare() + return facade + def get_articulation(self, uid: str) -> Articulation | None: """Get an articulation by its unique ID. diff --git a/embodichain/learning/__init__.py b/embodichain/learning/__init__.py index 373314c61..b83c68b8a 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 000000000..b1f8e18a5 --- /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/_runtime.py b/embodichain/learning/dynamics/_runtime.py new file mode 100644 index 000000000..ff0c3d164 --- /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/comparison.py b/embodichain/learning/dynamics/comparison.py new file mode 100644 index 000000000..5ffc2e78f --- /dev/null +++ b/embodichain/learning/dynamics/comparison.py @@ -0,0 +1,504 @@ +# ---------------------------------------------------------------------------- +# 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, + ) + + 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"]) + python = str(Path(config.get("python", sys.executable)).resolve()) + 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 = seed_dir / "order.npy" + order = np.lib.format.open_memmap( + order_path, + mode="w+", + dtype=np.int64, + shape=(cfg.total_steps, cfg.batch_size), + ) + generator = torch.Generator(device="cpu").manual_seed(seed) + for start in range(0, cfg.total_steps, 1024): + end = min(cfg.total_steps, start + 1024) + # rand_int produces the same flattened stream as consecutive + # minibatches; one generator is shared by both trainers. + order[start:end] = torch.randint( + len(dataset), (end - start, cfg.batch_size), generator=generator + ).numpy() + order.flush() + del order + 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", + "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]) -> 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))) + 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"]) + result = run_train( + { + "dataset": spec["dataset"], + "output_dir": str(directory), + "fine_tune": spec["initial_model"], + "sample_order": spec["sample_order"], + "train": spec["train"], + } + ) + 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("--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())) + 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 000000000..76ad6a6f1 --- /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 000000000..d49943a5d --- /dev/null +++ b/embodichain/learning/dynamics/reference.py @@ -0,0 +1,1072 @@ +# ---------------------------------------------------------------------------- +# 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 + from embodichain.compute.dynamics import quat_rotate + + if codec.robot == "cartpole": + return torch.cat((q, qd), -1) + com = q[..., :3] + quat_rotate(q[..., 3:7], q.new_tensor(codec.base_com)) + omega = qd[..., 3:6] + origin_velocity = qd[..., :3] - torch.linalg.cross(omega, com) + return torch.cat((q, omega, origin_velocity, qd[..., 6:]), -1) + + +def _official_batch(batch: dict[str, Any], codec: Any) -> dict[str, Any]: + """Transport native states to official world-frame arrays, before encoding.""" + import torch + + q, qd = batch["states"]["q"].float(), batch["states"]["qd"].float() + transitions = {key: value.float() for key, value in batch["transitions"].items()} + states = _legacy_states(q, qd, codec) + shape = q.shape[:2] + shape = (shape[0], shape[1] - 1) + root = batch["states"].get("root_body_q") + 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 +) -> 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") + 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 + from ._runtime import _training_runtime_identity + + cfg = DynamicsTrainCfg(**spec["train"]) + device = cfg.device + torch.set_num_threads(int(spec.get("cpu_threads", 1))) + common, manifest = load_model(spec["initial_model"], device=device) + 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") + order = np.load(order_path, mmap_mode="r", allow_pickle=False) + if order.shape != (cfg.total_steps, cfg.batch_size): + raise ValueError("Reference sample-order shape mismatch") + sample = training.get_batch([0]) + raw = _official_batch(sample, codec) + 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 + ) + 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) + destination = Path(spec["output_dir"]) + destination.mkdir(parents=True, exist_ok=True) + best = float("inf") + best_weights = None + + def transfer() -> None: + current = common.state_dict() + current.update(original.neural_model.state_dict()) + common.load_state_dict(current) + + try: + with (destination / "metrics.jsonl").open("w") as log: + for step in range(cfg.total_steps): + 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) + 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() + } + log.write(json.dumps(record, allow_nan=False) + "\n") + log.flush() + 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. + """ + 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": True, + "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, +) -> 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. + """ + 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" + ) + 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, + } + 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 and 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 "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 + ), + ) + 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( + "--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 + result = operation( + json.loads(args.spec.read_text()), trust_external_code=args.trust_external_code + ) + 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 000000000..bc27df9bc --- /dev/null +++ b/embodichain/learning/dynamics/reference_rollout.py @@ -0,0 +1,995 @@ +# ---------------------------------------------------------------------------- +# 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 = _positive( + "source sampling batch_size", manifest.get("sampling", {}).get("batch_size") + ) + 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. + 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", + } + 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 + 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 + ) + 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 in range(0, count, batch): + n = min(batch, count - start) + 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()) + student.reset(initial_states=initial_raw.clone()) + if start == 0: + portable.reset(initial_states=initial_raw.clone()) + teacher_raw = [teacher.states.clone()] + student_raw = [student.states.clone()] + teacher_controls = [] + student_controls = [] + for step in range(horizon): + tq, tqd = _canonical_states(teacher.states, codec, rotation) + nq, nqd = _canonical_states(student.states, codec, rotation) + tf = command[:, step] if policy is None else policy(tq, tqd) + nf = command[:, step] if policy is None else policy(nq, nqd) + # 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()) + 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] + nraw = torch.stack(student_raw, 1)[:n] + tq, tqd = _canonical_states(traw, codec, rotation) + 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", + "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 000000000..0f2d661a9 --- /dev/null +++ b/embodichain/learning/dynamics/report.py @@ -0,0 +1,1078 @@ +# ---------------------------------------------------------------------------- +# 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_REFERENCE = "476dc67fccc6a99be293bf611f897b1c06f13e30" +_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 result["reference_commit"] == _PARITY_REFERENCE + 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, "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 _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 metadata.get("source", {}).get("commit") == OFFICIAL_COMMIT + and bool(metadata.get("source_physical_arrays")) + and bool(metadata.get("source_raw_hdf5")) + ) + 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]]: + steps = [] + losses = [] + stride = max(1, (budget + max_points - 1) // max_points) + if path.exists(): + with path.open() as stream: + for line in stream: + record = json.loads(line) + if record["step"] % stride == 0 or record["step"] == 1: + steps.append(record["step"]) + losses.append(record["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 000000000..48f7c18dd --- /dev/null +++ b/embodichain/learning/dynamics/trainer.py @@ -0,0 +1,625 @@ +# ---------------------------------------------------------------------------- +# 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 + +__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. + """ + + 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 + + +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"] + inputs = codec.encode( + q[:, :-1], qd[:, :-1], effort, contacts or None, transitions.get("gravity_dir") + ) + target = codec.targets(q[:, :-1], qd[:, :-1], q[:, 1:], qd[:, 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: + error = prediction - target + else: + weights = torch.rsqrt(variance + 1e-5) + # Preserve official arithmetic ordering for numerical parity. + error = prediction * weights - target * weights + squared = error.square() + if mask is None: + return squared.mean() + 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)) + * 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() + if ( + min( + cfg.total_steps, + cfg.batch_size, + cfg.validation_interval, + cfg.checkpoint_interval, + cfg.lr_update_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 ( + 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) + 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 + 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 | 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 = 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") + 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 + prediction = self.model(inputs) + loss = normalized_mse(prediction, target, _model_variance(self.model), mask) + if not torch.isfinite(loss): + raise FloatingPointError("Non-finite dynamics training loss") + loss.backward() + norm = torch.nn.utils.clip_grad_norm_( + self.model.parameters(), self.cfg.max_grad_norm, error_if_nonfinite=True + ) + self.optimizer.step() + self.sampler.set_state(candidate_sampler.get_state()) + self.global_step += 1 + return { + "step": self.global_step, + "loss": float(loss.detach()), + "learning_rate": lr, + "grad_norm": float(norm), + } + + def save(self, path: str | Path) -> None: + """Atomically save weights and exact committed training/RNG progress.""" + numpy_state = np.random.get_state() + state = { + "format": "embodichain.nerd.training.v1", + "model_cfg": self.model.cfg.to_dict(), + "model": self.model.state_dict(), + "manifest": self.manifest, + "train_cfg": self.cfg.to_dict(), + "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": { + "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(self.cfg.device).type == "cuda" + else [] + ), + }, + } + _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") + 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, + }, + ), + ): + if state[key] != 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"]) + torch.set_rng_state(state["rng"]["torch"]) + random.setstate(state["rng"]["python"]) + name, keys, pos, gaussian, cached = state["rng"]["numpy"] + np.random.set_state( + (name, keys.numpy().astype(np.uint32), pos, gaussian, cached) + ) + if state["rng"]["cuda"]: + torch.cuda.set_rng_state_all(state["rng"]["cuda"]) + + 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") + 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 000000000..c5794b0c8 --- /dev/null +++ b/embodichain/learning/dynamics/workflow.py @@ -0,0 +1,245 @@ +# ---------------------------------------------------------------------------- +# 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", + } +) + + +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 == ".npy" + 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", + } + if config.get("output"): + Path(config["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 + + 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") + model = NeRDModel(NeRDModelCfg(**state["model_cfg"])) + model.load_state_dict(state["model"]) + manifest = { + **state["manifest"], + "training": { + "global_step": state["global_step"], + "dataset_fingerprint": state["dataset_fingerprint"], + "qualification": "unqualified", + }, + } + export_model(output, model, manifest) + return { + "output": str(output), + "global_step": state["global_step"], + "qualification": "unqualified", + } diff --git a/embodichain/learning/models/nerd/THIRD_PARTY_LICENSES b/embodichain/learning/models/nerd/THIRD_PARTY_LICENSES new file mode 100644 index 000000000..e30e79e71 --- /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 000000000..1d3ba033e --- /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 000000000..a31578627 --- /dev/null +++ b/embodichain/learning/models/nerd/checkpoint.py @@ -0,0 +1,188 @@ +# ---------------------------------------------------------------------------- +# 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 _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) + 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") + model = NeRDModel(NeRDModelCfg(**manifest["model_config"])) + 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 + } + 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, + ) + 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 000000000..5f4c3507b --- /dev/null +++ b/embodichain/learning/models/nerd/model.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. +# ---------------------------------------------------------------------------- + +"""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 +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. + """ + + 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 + + +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 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") + 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 000000000..eb025f85d --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/README.md @@ -0,0 +1,93 @@ +# Ant neural dynamics workflow + +These files configure `embodichain dynamics`; they are not Gym deployments and +have no runnable environment ID. The simulator host constructs the pinned NeRD +reference asset through EmbodiChain and DexSim Spawn. The generated manifest +provides the actual physics, asset and codec identity; training inherits it. + +Run from the repository root (replace `smoke` with `full` for the single-GPU profile): + +```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 +embodichain dynamics eval --config extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.rollout.smoke.yaml +``` + +Paths inside YAML resolve relative to the YAML file. CLI path overrides resolve +relative to the current directory. Output directories are under repository +`outputs/`, which is ignored by Git. Exports require a new destination directory. +Generation resumes with `--resume`; training resumes with +`--resume outputs/dynamics/ant/smoke/train/last.pt`. Exact resume preserves +budget, seed, device, configuration, immutable data fingerprint, optimizer, RNG +and the committed sampler cursor. Use `--fine-tune path/to/export` for a new run. + +The smoke profile uses CPU, 64 short trajectories, a small Transformer and +100 optimizer steps. It tests the complete plumbing and does not qualify model +accuracy. Ant smoke uses effort_scale=0.05 and default initial states; this is deliberately a different distribution from the official-range full profile. + +The full profile uses one CUDA device, the official six-layer, twelve-head +Transformer with embedding width 384, history 10, batch 512, and +4,995,000 Adam updates. This matches 999 trained epochs × 5000 updates in the +saved official configuration (epoch zero only evaluates). The learning rate is +held for 5000 updates, with offset one and a 1000-interval denominator, matching +the official epoch timing. Gradient clipping is a norm limit of 1.0; it does not +truncate temporal gradients. Validation uses the first 100 complete/partial +batches; final offline evaluation reads the entire test split. The expected +training budget is substantial and is not run automatically by tests. + +Default generation samples the full recorded force range: cart force ±1500 +with no pole actuation, or Ant scalar joints ±400. Teachers use MuJoCo-Warp with +10 substeps per 1/60-second macro step. CPU fixtures explicitly request separate +worlds. Source MJCF motor objects and solver XML options are not installed into +the teacher: the host applies actual generalized forces and records its solver +configuration. This profile is a current Newton teacher, not Warp 1.5.1 paper +reproduction. + +## Official comparison controls + +Default training uses deterministic sampling with replacement and streamed +population moments over overlapping training windows. These choices are explicit; +they do not reproduce the official dataloader's permutation or its running +normalizer initialization automatically. + +For a paired official comparison, pass the **same** immutable data, trajectory +split, initialized weights, batch order, normalizer artifact and update budget to +both trainers. `--statistics stats.pt` loads a tensor-only artifact containing +`schema_version: 1`, `inputs: {name: {mean, var}}`, and +`output: {mean, var}`. No validation/test samples may enter these statistics. +`--sample-order order.npy` reads a memory-mapped integer array of shape +`[total_steps, batch_size]`; each row indexes the full-window training dataset. +A `.pt` tensor is also accepted for short fixtures. The array hash and committed +cursor are checked on resume; large `.npy` orders are never loaded in full. + +Use `--seed 0` through `--seed 4` with distinct `--output-dir` values for the +five paired training seeds. Export each validation-selected checkpoint separately +and use evaluation's `--model` and `--output` overrides to retain every result. +Full rollout configurations request 2048 trajectories at horizons 1, 10, 100, +1000; repeat with `motion_mode: passive` and with the qualified fixed policy. +A common seed alone does not guarantee identical minibatches across trainers. + +No committed configuration is evidence that the full run passed. Required reports +include short numerical/Adam parity, both teachers × both trainers, five seed +pairs, confidence bounds, non-finite/failure rates, fixed-policy return and +software/asset/checkpoint hashes. `qualification.paired_noninferiority` implements +the seed-paired one-sided 95% bootstrap gate; fewer than five seeds is pending. +A 100-step smoke cannot replace the full budget or paper reproduction. + +## Registered differentiable task + +After exporting a matching model, `env.nerd.yaml` creates the registered +`AntDynamics-v1` task. Its checkpoint path resolves from the +working directory, unlike dynamics workflow file paths (config-relative). + +```bash +embodichain run-task --gym_config extensions/embodichain_tasks/configs/tasks/classic_control/ant/env.nerd.yaml --preview +embodichain train-rl --config extensions/embodichain_tasks/configs/tasks/classic_control/ant/agents/apg.nerd.yaml +``` + +Use preview for the Gym scene: this task does not implement an expert trajectory +collector. The APG configuration uses the existing learning-environment factory +and trainer, with graph-connected state/reward hooks. Its policy checkpoint is +separate from the NeRD dynamics-model export. diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/agents/apg.nerd.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/agents/apg.nerd.yaml new file mode 100644 index 000000000..4fbb692bb --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/agents/apg.nerd.yaml @@ -0,0 +1,44 @@ +# Requires a separately trained/qualified model at the checkpoint path below. +# Run from repository root: embodichain train-rl --config +trainer: + exp_name: ant_nerd_apg + learning_env: + name: AntDynamics-v1 + cfg: + checkpoint: outputs/dynamics/ant/full/export + max_episode_steps: 1000 + random_reset: true + seed: 42 + device: cuda:0 + num_envs: 64 + iterations: 500 + segment_length: 16 + update_horizon: 16 + enable_eval: true + eval_freq: 32000 + num_eval_envs: 16 + num_eval_episodes: 32 + eval_seed: 10042 + best_eval_metric: eval/avg_reward + best_eval_mode: max + save_frequency_updates: 50 + use_wandb: false +policy: + name: actor_only + initial_log_std: -2.0 + actor: + type: mlp + network_cfg: + hidden_sizes: [64, 64] + activation: tanh + last_activation: tanh + orthogonal_init: [1.414, 1.414, 0.01] +algorithm: + name: apg + cfg: + optimizer: + name: adam + learning_rate: 0.00025 + gamma: 0.99 + max_grad_norm: 0.5 + ent_coef: 0.0 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 000000000..2e4a7e59b --- /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.smoke.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.smoke.yaml new file mode 100644 index 000000000..0caca1bfb --- /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 000000000..da72a9f08 --- /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 000000000..c134f5c0d --- /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.rollout.full.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.rollout.full.yaml new file mode 100644 index 000000000..f9be1151f --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.rollout.full.yaml @@ -0,0 +1,21 @@ +# Dynamics workflow configuration; not a runnable Gym deployment. +mode: rollout +robot: ant +model: ../../../../../../outputs/dynamics/ant/full/export +device: cuda:0 +num_envs: 512 +trajectories: 2048 +horizons: +- 1 +- 10 +- 100 +- 1000 +seed: 123 +teacher_solver: + solver_type: mujoco_warp + use_mujoco_cpu: false +num_substeps: 10 +motion_mode: random_effort +initial_states: sample +effort_scale: 1.0 +output: ../../../../../../outputs/dynamics/ant/full/rollout.json diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.rollout.smoke.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.rollout.smoke.yaml new file mode 100644 index 000000000..2ac040d20 --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.rollout.smoke.yaml @@ -0,0 +1,20 @@ +# Dynamics workflow configuration; not a runnable Gym deployment. +mode: rollout +robot: ant +model: ../../../../../../outputs/dynamics/ant/smoke/export +device: cpu +num_envs: 4 +trajectories: 16 +horizons: +- 1 +- 10 +seed: 123 +teacher_solver: + solver_type: mujoco_warp + use_mujoco_cpu: true + separate_worlds: true +num_substeps: 10 +motion_mode: random_effort +initial_states: default +effort_scale: 0.05 +output: ../../../../../../outputs/dynamics/ant/smoke/rollout.json 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 000000000..73b720ba6 --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.train.full.yaml @@ -0,0 +1,36 @@ +# 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 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 000000000..853394306 --- /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/ant/env.nerd.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/env.nerd.yaml new file mode 100644 index 000000000..37b475213 --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/env.nerd.yaml @@ -0,0 +1,28 @@ +# Train and export the matching dynamics model first; run from the repository root. +id: AntDynamics-v1 +physics: newton +physics_config: + device: cpu + physics_dt: 0.016666666666666666 + num_substeps: 1 + requires_grad: true + use_cuda_graph: false + collision_cfg: null + solver_cfg: + solver_type: nerd + options: + checkpoint: outputs/dynamics/ant/full/export +num_envs: 16 +headless: true +renderer: auto +arena_space: 0.0 +enable_ground_plane: true +enable_entity_gizmo: false +robot_ik_gizmo: null +max_episode_steps: 1000 +robot: null +sensor: [] +env: + sim_steps_per_control: 1 + extensions: + random_reset: true 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 000000000..391bf6d5e --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/README.md @@ -0,0 +1,93 @@ +# Cartpole neural dynamics workflow + +These files configure `embodichain dynamics`; they are not Gym deployments and +have no runnable environment ID. The simulator host constructs the pinned NeRD +reference asset through EmbodiChain and DexSim Spawn. The generated manifest +provides the actual physics, asset and codec identity; training inherits it. + +Run from the repository root (replace `smoke` with `full` for the single-GPU profile): + +```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 dynamics eval --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.rollout.smoke.yaml +``` + +Paths inside YAML resolve relative to the YAML file. CLI path overrides resolve +relative to the current directory. Output directories are under repository +`outputs/`, which is ignored by Git. Exports require a new destination directory. +Generation resumes with `--resume`; training resumes with +`--resume outputs/dynamics/cartpole/smoke/train/last.pt`. Exact resume preserves +budget, seed, device, configuration, immutable data fingerprint, optimizer, RNG +and the committed sampler cursor. Use `--fine-tune path/to/export` for a new run. + +The smoke profile uses CPU, 64 short trajectories, a small Transformer and +100 optimizer steps. It tests the complete plumbing and does not qualify model +accuracy. Cartpole smoke uses default initial states; the full profile samples initial states. + +The full profile uses one CUDA device, the official six-layer, twelve-head +Transformer with embedding width 192, history 10, batch 512, and +4,995,000 Adam updates. This matches 999 trained epochs × 5000 updates in the +saved official configuration (epoch zero only evaluates). The learning rate is +held for 5000 updates, with offset one and a 1000-interval denominator, matching +the official epoch timing. Gradient clipping is a norm limit of 1.0; it does not +truncate temporal gradients. Validation uses the first 100 complete/partial +batches; final offline evaluation reads the entire test split. The expected +training budget is substantial and is not run automatically by tests. + +Default generation samples the full recorded force range: cart force ±1500 +with no pole actuation, or Ant scalar joints ±400. Teachers use MuJoCo-Warp with +10 substeps per 1/60-second macro step. CPU fixtures explicitly request separate +worlds. Source MJCF motor objects and solver XML options are not installed into +the teacher: the host applies actual generalized forces and records its solver +configuration. This profile is a current Newton teacher, not Warp 1.5.1 paper +reproduction. + +## Official comparison controls + +Default training uses deterministic sampling with replacement and streamed +population moments over overlapping training windows. These choices are explicit; +they do not reproduce the official dataloader's permutation or its running +normalizer initialization automatically. + +For a paired official comparison, pass the **same** immutable data, trajectory +split, initialized weights, batch order, normalizer artifact and update budget to +both trainers. `--statistics stats.pt` loads a tensor-only artifact containing +`schema_version: 1`, `inputs: {name: {mean, var}}`, and +`output: {mean, var}`. No validation/test samples may enter these statistics. +`--sample-order order.npy` reads a memory-mapped integer array of shape +`[total_steps, batch_size]`; each row indexes the full-window training dataset. +A `.pt` tensor is also accepted for short fixtures. The array hash and committed +cursor are checked on resume; large `.npy` orders are never loaded in full. + +Use `--seed 0` through `--seed 4` with distinct `--output-dir` values for the +five paired training seeds. Export each validation-selected checkpoint separately +and use evaluation's `--model` and `--output` overrides to retain every result. +Full rollout configurations request 2048 trajectories at horizons 1, 10, 100, +1000; repeat with `motion_mode: passive` and with the qualified fixed policy. +A common seed alone does not guarantee identical minibatches across trainers. + +No committed configuration is evidence that the full run passed. Required reports +include short numerical/Adam parity, both teachers × both trainers, five seed +pairs, confidence bounds, non-finite/failure rates, fixed-policy return and +software/asset/checkpoint hashes. `qualification.paired_noninferiority` implements +the seed-paired one-sided 95% bootstrap gate; fewer than five seeds is pending. +A 100-step smoke cannot replace the full budget or paper reproduction. + +## Registered differentiable task + +After exporting a matching model, `env.nerd.yaml` creates the registered +`CartpoleDynamics-v1` task. Its checkpoint path resolves from the +working directory, unlike dynamics workflow file paths (config-relative). + +```bash +embodichain run-task --gym_config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/env.nerd.yaml --preview +embodichain train-rl --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/agents/apg.nerd.yaml +``` + +Use preview for the Gym scene: this task does not implement an expert trajectory +collector. The APG configuration uses the existing learning-environment factory +and trainer, with graph-connected state/reward hooks. Its policy checkpoint is +separate from the NeRD dynamics-model export. diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/agents/apg.nerd.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/agents/apg.nerd.yaml new file mode 100644 index 000000000..c5e3e866c --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/agents/apg.nerd.yaml @@ -0,0 +1,44 @@ +# Requires a separately trained/qualified model at the checkpoint path below. +# Run from repository root: embodichain train-rl --config +trainer: + exp_name: cartpole_nerd_apg + learning_env: + name: CartpoleDynamics-v1 + cfg: + checkpoint: outputs/dynamics/cartpole/full/export + max_episode_steps: 500 + random_reset: true + seed: 42 + device: cuda:0 + num_envs: 64 + iterations: 500 + segment_length: 16 + update_horizon: 16 + enable_eval: true + eval_freq: 32000 + num_eval_envs: 16 + num_eval_episodes: 32 + eval_seed: 10042 + best_eval_metric: eval/avg_reward + best_eval_mode: max + save_frequency_updates: 50 + use_wandb: false +policy: + name: actor_only + initial_log_std: -2.0 + actor: + type: mlp + network_cfg: + hidden_sizes: [64, 64] + activation: tanh + last_activation: tanh + orthogonal_init: [1.414, 1.414, 0.01] +algorithm: + name: apg + cfg: + optimizer: + name: adam + learning_rate: 0.00025 + gamma: 0.99 + max_grad_norm: 0.5 + ent_coef: 0.0 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 000000000..7cd4eef56 --- /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.smoke.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.smoke.yaml new file mode 100644 index 000000000..6df395f4c --- /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 000000000..8ee89e2b3 --- /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 000000000..c5e4d1431 --- /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.rollout.full.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.rollout.full.yaml new file mode 100644 index 000000000..da5f40daa --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.rollout.full.yaml @@ -0,0 +1,21 @@ +# Dynamics workflow configuration; not a runnable Gym deployment. +mode: rollout +robot: cartpole +model: ../../../../../../outputs/dynamics/cartpole/full/export +device: cuda:0 +num_envs: 512 +trajectories: 2048 +horizons: +- 1 +- 10 +- 100 +- 1000 +seed: 123 +teacher_solver: + solver_type: mujoco_warp + use_mujoco_cpu: false +num_substeps: 10 +motion_mode: random_effort +initial_states: sample +effort_scale: 1.0 +output: ../../../../../../outputs/dynamics/cartpole/full/rollout.json diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.rollout.smoke.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.rollout.smoke.yaml new file mode 100644 index 000000000..bd12b84fc --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.rollout.smoke.yaml @@ -0,0 +1,20 @@ +# Dynamics workflow configuration; not a runnable Gym deployment. +mode: rollout +robot: cartpole +model: ../../../../../../outputs/dynamics/cartpole/smoke/export +device: cpu +num_envs: 4 +trajectories: 16 +horizons: +- 1 +- 10 +seed: 123 +teacher_solver: + solver_type: mujoco_warp + use_mujoco_cpu: true + separate_worlds: true +num_substeps: 10 +motion_mode: random_effort +initial_states: default +effort_scale: 1.0 +output: ../../../../../../outputs/dynamics/cartpole/smoke/rollout.json 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 000000000..76df09dee --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.train.full.yaml @@ -0,0 +1,32 @@ +# 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 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 000000000..cf1815fde --- /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/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/env.nerd.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/env.nerd.yaml new file mode 100644 index 000000000..bbb3bca96 --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/env.nerd.yaml @@ -0,0 +1,28 @@ +# Train and export the matching dynamics model first; run from the repository root. +id: CartpoleDynamics-v1 +physics: newton +physics_config: + device: cpu + physics_dt: 0.016666666666666666 + num_substeps: 1 + requires_grad: true + use_cuda_graph: false + collision_cfg: null + solver_cfg: + solver_type: nerd + options: + checkpoint: outputs/dynamics/cartpole/full/export +num_envs: 16 +headless: true +renderer: auto +arena_space: 0.0 +enable_ground_plane: false +enable_entity_gizmo: false +robot_ik_gizmo: null +max_episode_steps: 500 +robot: null +sensor: [] +env: + sim_steps_per_control: 1 + extensions: + random_reset: true diff --git a/extensions/embodichain_tasks/embodichain_tasks/classic_control/ant.py b/extensions/embodichain_tasks/embodichain_tasks/classic_control/ant.py new file mode 100644 index 000000000..8a5e6569a --- /dev/null +++ b/extensions/embodichain_tasks/embodichain_tasks/classic_control/ant.py @@ -0,0 +1,244 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Ant control through the registered NeRD dynamics solver.""" + +from __future__ import annotations +from collections.abc import Sequence +from typing import Any, TYPE_CHECKING +import gymnasium as gym +import numpy as np +import torch +from embodichain.lab.gym.envs.differentiable_dynamics_env import ( + DifferentiableDynamicsEnv, +) +from embodichain.lab.gym.envs.embodied_env import EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env +from embodichain.lab.sim.cfg import NewtonPhysicsCfg, RenderCfg +from embodichain.lab.sim.sim_manager import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.objects.articulation import Articulation +from embodichain.lab.sim.physics.solvers._reference_assets import ( + make_reference_articulation, + reference_profile, +) +from embodichain.lab.sim.physics.solvers._nerd_binding import NeRDBinding +from embodichain.compute.dynamics import quat_rotate, reference_task_values +from embodichain.learning.rl.env import register_learning_env + +if TYPE_CHECKING: + from dexsim.engine.newton_physics import FunctionalSolverState + +__all__ = ["AntDynamicsEnv", "make_ant_dynamics_env"] + + +@register_env("AntDynamics-v1", max_episode_steps=1000, supports_rl=True) +class AntDynamicsEnv(DifferentiableDynamicsEnv): + """Differentiable ant task with normalized effort actions. + + The pinned reference articulation is task-owned. A compatible exported model + must be selected by physics_config.solver_cfg; model accuracy is determined + by its independent validation, not by this task's registration. + """ + + def __init__(self, cfg: EmbodiedEnvCfg, **kwargs: Any) -> None: + self.random_reset = (cfg.extensions or {}).get("random_reset", True) + if not isinstance(self.random_reset, bool): + raise TypeError("random_reset must be boolean") + self._profile = reference_profile("ant") + self._binding: NeRDBinding | None = None + if cfg.robot is not None: + raise ValueError( + "AntDynamics owns its reference robot; configure robot: null" + ) + if cfg.sim_cfg is not None and bool(cfg.sim_cfg.enable_ground_plane) != True: + raise ValueError("AntDynamics requires enable_ground_plane: true") + super().__init__(cfg, **kwargs) + + def _declare_robot(self, **kwargs: Any) -> Articulation: + del kwargs + return self.sim.add_articulation_descriptor(make_reference_articulation("ant")) + + def _setup_robot(self, **kwargs: Any) -> Articulation: + del kwargs + self._binding = NeRDBinding.from_metadata( + dict(self.sim.spawn_result.dynamics_metadata()), "ant" + ) + self._joint_ids = [ + self.robot.joint_names.index(name) for name in self._binding.joint_names + ] + self.active_joint_ids = list(self._joint_ids) + self.single_action_space = gym.spaces.Box( + -1.0, 1.0, shape=(8,), dtype=np.float32 + ) + self.single_observation_space = gym.spaces.Box( + -np.inf, np.inf, shape=(27,), dtype=np.float32 + ) + return self.robot + + def _action_to_control( + self, action: torch.Tensor, state: FunctionalSolverState + ) -> dict[str, torch.Tensor]: + if action.shape != (self.num_envs, 8): + raise ValueError( + f"Expected action shape {(self.num_envs,8)}, got {tuple(action.shape)}" + ) + indices = self._binding.qd_indices[:, 6:] + effort = action.clamp(-1.0, 1.0) * 400.0 + return { + "joint_f": torch.zeros_like(state.physical["joint_qd"]).scatter( + 0, indices.flatten(), effort.flatten() + ) + } + + def _initialize_episode( + self, env_ids: Sequence[int] | torch.Tensor | None = None, **kwargs: Any + ) -> None: + super()._initialize_episode(env_ids, **kwargs) + ids = ( + torch.arange(self.num_envs, device=self.device) + if env_ids is None + else torch.as_tensor(env_ids, device=self.device, dtype=torch.long) + ) + count = ids.numel() + generator = self.get_generator("ant_dynamics_reset") + q = ( + torch.tensor(self._profile["default_q"], device=self.device) + .expand(count, -1) + .clone() + ) + qd = torch.zeros(count, self._profile["qd_dim"], device=self.device) + if self.random_reset: + q[:, :3] += ( + 2 * torch.rand(count, 3, device=self.device, generator=generator) - 1 + ) * q.new_tensor((0.2, 0.2, 0.1)) + axis = 2 * torch.rand(count, 3, device=self.device, generator=generator) - 1 + axis = axis / axis.norm(dim=-1, keepdim=True).clamp_min(1e-9) + angle = ( + 2 * torch.rand(count, 1, device=self.device, generator=generator) - 1 + ) * (torch.pi / 6) + q[:, 3:7] = torch.cat( + (axis * torch.sin(angle / 2), torch.cos(angle / 2)), -1 + ) + q[:, 7:] += 0.2 * ( + 2 * torch.rand(count, 8, device=self.device, generator=generator) - 1 + ) + qd[:, :3] = quat_rotate( + q[:, 3:7], + 0.1 + * ( + 2 * torch.rand(count, 3, device=self.device, generator=generator) + - 1 + ), + ) + qd[:, 3:6] = quat_rotate( + q[:, 3:7], + 0.25 + * ( + 2 * torch.rand(count, 3, device=self.device, generator=generator) + - 1 + ), + ) + qd[:, 6:] = 0.25 * ( + 2 * torch.rand(count, 8, device=self.device, generator=generator) - 1 + ) + self.robot.set_state( + root_pose=q[:, :7], + root_velocity=qd[:, :6], + qpos=q[:, 7:], + qvel=qd[:, 6:], + env_ids=ids, + joint_ids=self._joint_ids, + clear_dynamics=True, + ) + + def _read_outputs(self, state: FunctionalSolverState) -> dict[str, Any]: + q, qd = self._binding.read(state.physical) + observation, reward, failed = reference_task_values("ant", q, qd) + return { + "obs": observation, + "reward": reward, + "terminated": failed, + "truncated": torch.zeros_like(failed), + "info": { + "fail": failed, + "metrics": { + "height": q[:, 2].detach(), + "forward_velocity": qd[:, 1].detach(), + }, + }, + } + + def close(self, *, exit_process: bool = False) -> None: + """Close the task without terminating the trainer process by default.""" + super().close(exit_process=exit_process) + if not exit_process: + SimulationManager.flush_cleanup_queue() + + +@register_learning_env("AntDynamics-v1") +def make_ant_dynamics_env( + num_envs: int, + device: torch.device | str, + *, + checkpoint: str, + max_episode_steps: int = 1000, + random_reset: bool = True, + seed: int | None = None, +) -> AntDynamicsEnv: + """Build the same task for the existing differentiable RL trainer. + + Args: + num_envs: Number of parallel homogeneous robots. + device: Simulation and policy device. + checkpoint: Export directory, absolute or relative to the working directory. + max_episode_steps: Control-step episode limit. + random_reset: Sample task-specific initial states when true. + seed: Optional initial RNG seed; explicit reset seeds take precedence. + + Returns: + A registered task implementing the DifferentiableVecEnv protocol. + """ + physics = NewtonPhysicsCfg( + device=device, + physics_dt=1 / 60, + num_substeps=1, + requires_grad=True, + use_cuda_graph=False, + collision_cfg=None, + solver_cfg={"solver_type": "nerd", "options": {"checkpoint": checkpoint}}, + ) + sim_cfg = SimulationManagerCfg( + headless=True, + num_envs=num_envs, + device=device, + arena_space=0.0, + enable_ground_plane=True, + enable_entity_gizmo=False, + robot_ik_gizmo=None, + render_cfg=RenderCfg(renderer="no-render"), + startup_summary="off", + physics_cfg=physics, + ) + cfg = EmbodiedEnvCfg( + num_envs=num_envs, + sim_cfg=sim_cfg, + sim_steps_per_control=1, + max_episode_steps=max_episode_steps, + seed=seed, + robot=None, + extensions={"random_reset": random_reset}, + ) + return AntDynamicsEnv(cfg) diff --git a/extensions/embodichain_tasks/embodichain_tasks/classic_control/cartpole.py b/extensions/embodichain_tasks/embodichain_tasks/classic_control/cartpole.py new file mode 100644 index 000000000..cb3cc108f --- /dev/null +++ b/extensions/embodichain_tasks/embodichain_tasks/classic_control/cartpole.py @@ -0,0 +1,208 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Cartpole control through the registered NeRD dynamics solver.""" + +from __future__ import annotations +from collections.abc import Sequence +from typing import Any, TYPE_CHECKING +import gymnasium as gym +import numpy as np +import torch +from embodichain.lab.gym.envs.differentiable_dynamics_env import ( + DifferentiableDynamicsEnv, +) +from embodichain.lab.gym.envs.embodied_env import EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env +from embodichain.lab.sim.cfg import NewtonPhysicsCfg, RenderCfg +from embodichain.lab.sim.sim_manager import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.objects.articulation import Articulation +from embodichain.lab.sim.physics.solvers._reference_assets import ( + make_reference_articulation, + reference_profile, +) +from embodichain.lab.sim.physics.solvers._nerd_binding import NeRDBinding +from embodichain.compute.dynamics import reference_task_values +from embodichain.learning.rl.env import register_learning_env + +if TYPE_CHECKING: + from dexsim.engine.newton_physics import FunctionalSolverState + +__all__ = ["CartpoleDynamicsEnv", "make_cartpole_dynamics_env"] + + +@register_env("CartpoleDynamics-v1", max_episode_steps=500, supports_rl=True) +class CartpoleDynamicsEnv(DifferentiableDynamicsEnv): + """Differentiable cartpole task with normalized effort actions. + + The pinned reference articulation is task-owned. A compatible exported model + must be selected by physics_config.solver_cfg; model accuracy is determined + by its independent validation, not by this task's registration. + """ + + def __init__(self, cfg: EmbodiedEnvCfg, **kwargs: Any) -> None: + self.random_reset = (cfg.extensions or {}).get("random_reset", True) + if not isinstance(self.random_reset, bool): + raise TypeError("random_reset must be boolean") + self._profile = reference_profile("cartpole") + self._binding: NeRDBinding | None = None + if cfg.robot is not None: + raise ValueError( + "CartpoleDynamics owns its reference robot; configure robot: null" + ) + if cfg.sim_cfg is not None and bool(cfg.sim_cfg.enable_ground_plane) != False: + raise ValueError("CartpoleDynamics requires enable_ground_plane: false") + super().__init__(cfg, **kwargs) + + def _declare_robot(self, **kwargs: Any) -> Articulation: + del kwargs + return self.sim.add_articulation_descriptor( + make_reference_articulation("cartpole") + ) + + def _setup_robot(self, **kwargs: Any) -> Articulation: + del kwargs + self._binding = NeRDBinding.from_metadata( + dict(self.sim.spawn_result.dynamics_metadata()), "cartpole" + ) + self._joint_ids = [ + self.robot.joint_names.index(name) for name in self._binding.joint_names + ] + self.active_joint_ids = list(self._joint_ids) + self.single_action_space = gym.spaces.Box( + -1.0, 1.0, shape=(1,), dtype=np.float32 + ) + self.single_observation_space = gym.spaces.Box( + -np.inf, np.inf, shape=(4,), dtype=np.float32 + ) + return self.robot + + def _action_to_control( + self, action: torch.Tensor, state: FunctionalSolverState + ) -> dict[str, torch.Tensor]: + if action.shape != (self.num_envs, 1): + raise ValueError( + f"Expected action shape {(self.num_envs,1)}, got {tuple(action.shape)}" + ) + indices = self._binding.qd_indices[:, :1] + effort = action.clamp(-1.0, 1.0) * 1500.0 + return { + "joint_f": torch.zeros_like(state.physical["joint_qd"]).scatter( + 0, indices.flatten(), effort.flatten() + ) + } + + def _initialize_episode( + self, env_ids: Sequence[int] | torch.Tensor | None = None, **kwargs: Any + ) -> None: + super()._initialize_episode(env_ids, **kwargs) + ids = ( + torch.arange(self.num_envs, device=self.device) + if env_ids is None + else torch.as_tensor(env_ids, device=self.device, dtype=torch.long) + ) + count = ids.numel() + generator = self.get_generator("cartpole_dynamics_reset") + q = ( + torch.tensor(self._profile["default_q"], device=self.device) + .expand(count, -1) + .clone() + ) + qd = torch.zeros(count, self._profile["qd_dim"], device=self.device) + if self.random_reset: + q = ( + 2 * torch.rand(count, 2, device=self.device, generator=generator) - 1 + ) * q.new_tensor((1.0, torch.pi)) + qd = 2 * torch.rand(count, 2, device=self.device, generator=generator) - 1 + self.robot.set_state( + qpos=q, qvel=qd, env_ids=ids, joint_ids=self._joint_ids, clear_dynamics=True + ) + + def _read_outputs(self, state: FunctionalSolverState) -> dict[str, Any]: + q, qd = self._binding.read(state.physical) + observation, reward, failed = reference_task_values("cartpole", q, qd) + return { + "obs": observation, + "reward": reward, + "terminated": failed, + "truncated": torch.zeros_like(failed), + "info": { + "fail": failed, + "metrics": {"pole_angle": observation[:, 1].abs().detach()}, + }, + } + + def close(self, *, exit_process: bool = False) -> None: + """Close the task without terminating the trainer process by default.""" + super().close(exit_process=exit_process) + if not exit_process: + SimulationManager.flush_cleanup_queue() + + +@register_learning_env("CartpoleDynamics-v1") +def make_cartpole_dynamics_env( + num_envs: int, + device: torch.device | str, + *, + checkpoint: str, + max_episode_steps: int = 500, + random_reset: bool = True, + seed: int | None = None, +) -> CartpoleDynamicsEnv: + """Build the same task for the existing differentiable RL trainer. + + Args: + num_envs: Number of parallel homogeneous robots. + device: Simulation and policy device. + checkpoint: Export directory, absolute or relative to the working directory. + max_episode_steps: Control-step episode limit. + random_reset: Sample task-specific initial states when true. + seed: Optional initial RNG seed; explicit reset seeds take precedence. + + Returns: + A registered task implementing the DifferentiableVecEnv protocol. + """ + physics = NewtonPhysicsCfg( + device=device, + physics_dt=1 / 60, + num_substeps=1, + requires_grad=True, + use_cuda_graph=False, + collision_cfg=None, + solver_cfg={"solver_type": "nerd", "options": {"checkpoint": checkpoint}}, + ) + sim_cfg = SimulationManagerCfg( + headless=True, + num_envs=num_envs, + device=device, + arena_space=0.0, + enable_ground_plane=False, + enable_entity_gizmo=False, + robot_ik_gizmo=None, + render_cfg=RenderCfg(renderer="no-render"), + startup_summary="off", + physics_cfg=physics, + ) + cfg = EmbodiedEnvCfg( + num_envs=num_envs, + sim_cfg=sim_cfg, + sim_steps_per_control=1, + max_episode_steps=max_episode_steps, + seed=seed, + robot=None, + extensions={"random_reset": random_reset}, + ) + return CartpoleDynamicsEnv(cfg) diff --git a/pyproject.toml b/pyproject.toml index 20930f944..fba16a989 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -85,6 +85,9 @@ embodichain_tasks = "embodichain_tasks" [project.entry-points."embodichain.robots"] official = "embodichain_robots._catalog:ROBOT_CONFIGS" +[project.entry-points."dexsim.newton_solvers"] +nerd = "embodichain.lab.sim.physics.solvers.nerd:NeRDProvider" + [tool.uv.sources] bpy = { index = "blender" } diff --git a/scripts/benchmark/learning/dynamics_input.py b/scripts/benchmark/learning/dynamics_input.py new file mode 100644 index 000000000..89772bd6f --- /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_training.py b/scripts/benchmark/learning/dynamics_training.py new file mode 100644 index 000000000..69dd5e9c6 --- /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/setup.py b/setup.py index c1b01fc26..d86f48047 100644 --- a/setup.py +++ b/setup.py @@ -146,6 +146,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.gen_sim.simready_pipeline.configs": ["*.json"], "embodichain.lab.scripts._atomic_action": ["README.md"], "embodichain_tasks.configs": [ diff --git a/tests/cli/test_dynamics.py b/tests/cli/test_dynamics.py new file mode 100644 index 000000000..f047841ca --- /dev/null +++ b/tests/cli/test_dynamics.py @@ -0,0 +1,332 @@ +# ---------------------------------------------------------------------------- +# 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_rollout_routing_and_path_resolution(tmp_path, monkeypatch): + host = ModuleType("embodichain.lab.sim.physics.solvers.collection") + host.generate_dataset = Mock(return_value={"status": "ok"}) + host.evaluate_rollouts = 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 + (tmp_path / "export").mkdir() + config.write_text( + "mode: rollout\nrobot: ant\nmodel: export\nmotion_mode: passive\noutput: report.json\n" + ) + main(["eval", "--config", str(config)]) + received = host.evaluate_rollouts.call_args.args[0] + assert "mode" not in received + assert received["motion_mode"] == "passive" + assert received["model"] == str(tmp_path / "export") + + +def test_rollout_teacher_and_policy_paths_are_config_relative(tmp_path, monkeypatch): + host = ModuleType("embodichain.lab.sim.physics.solvers.collection") + host.evaluate_rollouts = Mock(return_value={"status": "ok"}) + monkeypatch.setitem(sys.modules, host.__name__, host) + for name in ("export", "teacher"): + (tmp_path / name).mkdir() + for name in ("actor.yaml", "actor.pt"): + (tmp_path / name).write_text("mock artifact") + config = tmp_path / "eval.yaml" + config.write_text( + yaml.safe_dump( + { + "mode": "rollout", + "model": "export", + "teacher_dataset": "teacher", + "policy": { + "kind": "actor_only", + "config": "actor.yaml", + "checkpoint": "actor.pt", + }, + } + ) + ) + main(["eval", "--config", str(config)]) + received = host.evaluate_rollouts.call_args.args[0] + assert received["teacher_dataset"] == str(tmp_path / "teacher") + assert received["policy"]["config"] == str(tmp_path / "actor.yaml") + assert received["policy"]["checkpoint"] == str(tmp_path / "actor.pt") + (tmp_path / "actor.pt").unlink() + with pytest.raises(SystemExit) as error: + main(["eval", "--config", str(config)]) + assert error.value.code == 2 + + +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": "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 / "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.lab.sim.physics.solvers.collection") + host.generate_dataset = Mock(return_value={"status": "not_run"}) + host.evaluate_rollouts = 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 = [] + for robot in ("cartpole", "ant"): + files.extend( + ( + repo + / "extensions" + / "embodichain_tasks" + / "configs" + / "tasks" + / "classic_control" + / robot + ).glob("dynamics.*.yaml") + ) + assert len(files) == 16 + 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(["eval" if operation == "rollout" else 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/test_codec.py b/tests/compute/dynamics/test_codec.py new file mode 100644 index 000000000..44d1a65fa --- /dev/null +++ b/tests/compute/dynamics/test_codec.py @@ -0,0 +1,161 @@ +# ---------------------------------------------------------------------------- +# 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 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_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_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_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 000000000..ba93266d3 --- /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 000000000..6f65154d5 --- /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_shards.py b/tests/data_pipeline/dynamics/test_shards.py new file mode 100644 index 000000000..6ab6caa9b --- /dev/null +++ b/tests/data_pipeline/dynamics/test_shards.py @@ -0,0 +1,274 @@ +# ---------------------------------------------------------------------------- +# 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 +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"} + + +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/gym/envs/tasks/test_neural_dynamics_tasks.py b/tests/gym/envs/tasks/test_neural_dynamics_tasks.py new file mode 100644 index 000000000..fd0134fb6 --- /dev/null +++ b/tests/gym/envs/tasks/test_neural_dynamics_tasks.py @@ -0,0 +1,171 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Registered neural tasks, file-owned configs and actual APG routing.""" + +from __future__ import annotations +import importlib +import importlib.util +import json +import os +from pathlib import Path +import subprocess +import sys +import pytest +import torch +from embodichain.lab.gym.utils.gym_utils import config_to_cfg +from embodichain.lab.gym.utils.registration import REGISTERED_ENVS, build_env +from embodichain.learning.rl.env import ( + get_registered_learning_env_names, + build_learning_env, +) +from embodichain.utils.utility import load_config + +_ROOT = Path(__file__).resolve().parents[4] +_CONFIGS = _ROOT / "extensions/embodichain_tasks/configs/tasks/classic_control" + + +@pytest.mark.parametrize("robot,title", [("cartpole", "Cartpole"), ("ant", "Ant")]) +def test_registered_task_and_yaml_ownership(robot: str, title: str) -> None: + importlib.import_module(f"embodichain_tasks.classic_control.{robot}") + assert f"{title}Dynamics-v1" in REGISTERED_ENVS + assert f"{title.lower()}dynamics-v1" in get_registered_learning_env_names() + path = _CONFIGS / robot / "env.nerd.yaml" + config = load_config(path) + parsed = config_to_cfg(config, source_path=path) + assert parsed.sim_steps_per_control == 1 + assert parsed.robot is None + assert parsed.sim_cfg.enable_ground_plane is (robot == "ant") + assert parsed.sim_cfg.physics_cfg.requires_grad + assert parsed.sim_cfg.physics_cfg.solver_cfg["solver_type"] == "nerd" + agent = load_config(_CONFIGS / robot / "agents/apg.nerd.yaml") + assert agent["trainer"]["learning_env"]["name"] == f"{title}Dynamics-v1" + assert "gym_config" not in agent["trainer"] + assert agent["algorithm"]["name"] == "apg" + + +def _check_task_and_trainer(tmp_path: Path, robot: str, title: str) -> None: + importlib.import_module(f"embodichain_tasks.classic_control.{robot}") + from embodichain.lab.sim.physics.solvers.collection import DynamicsSimulation + from embodichain.learning.rl.train import train_from_config + + helper_path = _ROOT / "tests/sim/physics/solvers/test_nerd_solver.py" + spec = importlib.util.spec_from_file_location("_nerd_export_helper", helper_path) + helper = importlib.util.module_from_spec(spec) + spec.loader.exec_module(helper) + with DynamicsSimulation( + robot, + num_envs=2, + device="cpu", + num_substeps=1, + ) as teacher: + checkpoint = helper._small_export( + teacher.sim.spawn_result, tmp_path / "model", robot + ) + deployment_path = _CONFIGS / robot / "env.nerd.yaml" + deployment = load_config(deployment_path) + deployment.update(num_envs=2, max_episode_steps=4, renderer="no-render") + deployment["physics_config"]["solver_cfg"]["options"]["checkpoint"] = str( + checkpoint + ) + env = build_env( + deployment["id"], config_to_cfg(deployment, source_path=deployment_path) + ) + try: + obs1, _ = env.reset(seed=9) + saved = obs1.detach().clone() + obs2, _ = env.reset(seed=9) + torch.testing.assert_close(saved, obs2) + width = 1 if robot == "cartpole" else 8 + observation_width = 4 if robot == "cartpole" else 27 + assert obs2.shape == (2, observation_width) + action = torch.full((2, width), 0.1, requires_grad=True) + obs, reward, terminated, truncated, info = env.step(action) + assert reward.shape == (2,) and obs.shape == (2, observation_width) + (reward.sum() + 0.001 * obs.square().sum()).backward() + assert ( + action.grad is not None + and torch.isfinite(action.grad).all() + and action.grad.abs().sum() > 0 + ) + assert not env.detach_state().requires_grad + # Selected episode reset clears only selected history, uses a continuing + # named RNG stream, and keeps the other row's physical state unchanged. + before = env.functional_state + unchanged = ( + before.physical["joint_q"][env._binding.q_indices[1]].detach().clone() + ) + env.reset(options={"reset_ids": torch.tensor([0])}) + assert env.functional_state.valid_history_lengths.tolist() == [0, 1] + torch.testing.assert_close( + env.functional_state.physical["joint_q"][env._binding.q_indices[1]], + unchanged, + ) + finally: + env.close() + cfg = load_config(_CONFIGS / robot / "agents/apg.nerd.yaml") + cfg["trainer"].update( + device="cpu", + num_envs=2, + iterations=1, + segment_length=2, + update_horizon=2, + enable_eval=False, + use_wandb=False, + save_frequency_updates=0, + ) + cfg["trainer"]["learning_env"]["cfg"].update( + checkpoint=str(checkpoint), max_episode_steps=4, random_reset=False + ) + cfg["policy"]["actor"]["network_cfg"] = { + "hidden_sizes": [8], + "activation": "tanh", + "last_activation": "tanh", + } + path = tmp_path / "apg.json" + path.write_text(json.dumps(cfg)) + os.chdir(tmp_path) + result = train_from_config(path) + assert result is not None + assert result["global_step"] == 4 + assert list((tmp_path / "outputs").rglob("*.pt")) + + +@pytest.mark.subprocess_sim +@pytest.mark.parametrize("robot,title", [("cartpole", "Cartpole"), ("ant", "Ant")]) +def test_actual_registered_task_and_one_apg_update( + tmp_path: Path, robot: str, title: str +) -> None: + code = "import importlib.util,sys;from pathlib import Path;s=importlib.util.spec_from_file_location('neural_task_test',sys.argv[1]);m=importlib.util.module_from_spec(s);s.loader.exec_module(m);m._check_task_and_trainer(Path(sys.argv[2]),sys.argv[3],sys.argv[4])" + subprocess.run( + [ + sys.executable, + "-c", + code, + str(Path(__file__).resolve()), + str(tmp_path), + robot, + title, + ], + check=True, + env={ + **os.environ, + "OMP_NUM_THREADS": "2", + "MKL_NUM_THREADS": "2", + "PYTHONPATH": str(_ROOT), + "EMBODICHAIN_SIM_EXIT_PROCESS": "0", + }, + ) diff --git a/tests/gym/envs/test_differentiable_dynamics_env.py b/tests/gym/envs/test_differentiable_dynamics_env.py new file mode 100644 index 000000000..f68aeda4f --- /dev/null +++ b/tests/gym/envs/test_differentiable_dynamics_env.py @@ -0,0 +1,176 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Environment graph/reset semantics over actual CPU functional dynamics.""" + +from __future__ import annotations +from types import SimpleNamespace +import torch +import warp as wp +import pytest +from embodichain.lab.gym.envs.differentiable_dynamics_env import ( + DifferentiableDynamicsEnv, +) +from embodichain.lab.gym.envs.embodied_env import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.sim.cfg import NewtonPhysicsCfg, DefaultPhysicsCfg +from embodichain.lab.sim.sim_manager import SimulationManagerCfg +from embodichain.lab.sim.diff.dynamics import FunctionalDynamics + +# Reuse a real registered-provider fixture without importing tutorial code. +import importlib.util +from pathlib import Path + +_fixture_spec = importlib.util.spec_from_file_location( + "_dynamics_fixture", + Path(__file__).resolve().parents[2] / "sim" / "diff" / "test_dynamics.py", +) +_fixture_module = importlib.util.module_from_spec(_fixture_spec) +_fixture_spec.loader.exec_module(_fixture_module) +manager = _fixture_module.manager + + +class _Task(DifferentiableDynamicsEnv): + def _action_to_control(self, action, state): + return action.flatten() + + def _read_outputs(self, state): + q = state.physical["joint_q"] + return { + "obs": q[:, None], + "reward": q.square(), + "terminated": q > self.threshold, + "truncated": torch.zeros_like(q, dtype=torch.bool), + } + + +def make_env( + manager, monkeypatch: pytest.MonkeyPatch, *, threshold: float = 10.0, steps: int = 2 +): + # Bypass only the scene/render construction already covered by EmbodiedEnv; + # production step/reset/facade and actual solver execution remain exercised. + env = object.__new__(_Task) + env.sim = manager + env._num_envs = 2 + env.cfg = SimpleNamespace(sim_steps_per_control=steps, ignore_terminations=False) + env.sim_cfg = SimpleNamespace(physics_dt=0.1, device="cpu") + env.threshold = threshold + env._dynamics = FunctionalDynamics(manager) + env._functional_state = env._dynamics.create_state() + env._resetting_functional = False + env._elapsed_steps = torch.zeros(2, dtype=torch.int32) + env.max_episode_steps = 100 + env.sensors = {} + env._update_physical_objective = lambda: None + env._hook_after_sim_step = lambda **kwargs: None + + def native_reset(self, seed=None, options=None): + ids = options.get("reset_ids", torch.arange(2)) + selected = torch.zeros(3, dtype=torch.bool) + selected[ids] = True + self.sim.backend.reset_state(wp.from_torch(selected, dtype=wp.bool)) + self._elapsed_steps[ids] = 0 + return self.get_obs(), {} + + monkeypatch.setattr(EmbodiedEnv, "reset", native_reset) + return env + + +def test_control_interval_and_multistep_action_gradient(manager, monkeypatch) -> None: + env = make_env(manager, monkeypatch) + first = torch.tensor([[1.0], [2.0]], requires_grad=True) + second = torch.tensor([[0.5], [1.0]], requires_grad=True) + obs, reward, *_ = env.step(first) + obs2, reward2, *_ = env.step(second) + torch.testing.assert_close(obs2, torch.tensor([[0.6], [1.2]])) + assert manager.backend.step_index == 4 + (reward.sum() + reward2.sum()).backward() + torch.testing.assert_close(first.grad, torch.tensor([[0.8], [1.6]])) + torch.testing.assert_close(second.grad, torch.tensor([[0.48], [0.96]])) + + +def test_terminal_reset_keeps_reward_graph_and_other_rows(manager, monkeypatch) -> None: + env = make_env(manager, monkeypatch, threshold=0.3, steps=1) + action = torch.tensor([[2.0], [1.0]], requires_grad=True) + obs, reward, terminated, truncated, info = env.step(action) + assert terminated.tolist() == [True, False] + torch.testing.assert_close(obs, torch.tensor([[0.0], [0.2]])) + assert env.functional_state.valid_history_lengths.tolist() == [0, 2] + assert manager.backend.step_index == 1 + assert info["terminal_mask"].tolist() == [True, False] + loss = reward.sum() + obs.sum() + env.functional_state.history["memory"].sum() + loss.backward() + torch.testing.assert_close(action.grad, torch.tensor([[0.16], [0.38]])) + + +def test_detach_state_stops_prior_control_gradient_without_reset( + manager, monkeypatch +) -> None: + env = make_env(manager, monkeypatch, steps=1) + first = torch.ones(2, 1, requires_grad=True) + env.step(first) + before = env.functional_state.physical["joint_q"].detach().clone() + obs = env.detach_state() + assert not obs.requires_grad + torch.testing.assert_close(env.functional_state.physical["joint_q"], before) + second = torch.ones(2, 1, requires_grad=True) + _, reward, *_ = env.step(second) + reward.sum().backward() + assert first.grad is None + assert second.grad is not None + assert manager.backend.step_index == 2 + + +def test_max_episode_steps_truncate_and_reset(manager, monkeypatch) -> None: + env = make_env(manager, monkeypatch, steps=1) + env.max_episode_steps = 1 + obs, reward, terminated, truncated, _ = env.step( + torch.ones(2, 1, requires_grad=True) + ) + assert truncated.tolist() == [True, True] and not terminated.any() + assert not obs.any() and env._elapsed_steps.tolist() == [0, 0] + reward.sum().backward() + + +def test_configuration_rejects_nonfunctional_paths_before_scene_creation() -> None: + cfg = EmbodiedEnvCfg(sim_cfg=SimulationManagerCfg(physics_cfg=DefaultPhysicsCfg())) + with pytest.raises(ValueError, match="NewtonPhysicsCfg"): + DifferentiableDynamicsEnv._validate_dynamics_cfg(cfg) + cfg.sim_cfg.physics_cfg = NewtonPhysicsCfg(requires_grad=False) + with pytest.raises(ValueError, match="requires_grad"): + DifferentiableDynamicsEnv._validate_dynamics_cfg(cfg) + cfg.sim_cfg.physics_cfg.requires_grad = True + cfg.actions = {"unsupported": object()} + with pytest.raises(ValueError, match="functional hooks"): + DifferentiableDynamicsEnv._validate_dynamics_cfg(cfg) + + +def test_sensor_sampling_follows_macro_steps_and_ignore_terminations( + manager, monkeypatch +) -> None: + env = make_env(manager, monkeypatch, threshold=0.0, steps=3) + env.cfg.ignore_terminations = True + calls = [] + env.sensors = { + "sensor": SimpleNamespace( + requires_substep_update=True, + begin_control_step=lambda: calls.append("begin"), + update_physics_step=lambda dt: calls.append(dt), + ) + } + obs, _, terminated, _, _ = env.step(torch.ones(2, 1)) + assert calls == ["begin", 0.1, 0.1, 0.1] + assert not terminated.any() + torch.testing.assert_close(obs, torch.tensor([[0.6], [0.6]])) diff --git a/tests/learning/dynamics/test_comparison.py b/tests/learning/dynamics/test_comparison.py new file mode 100644 index 000000000..820890683 --- /dev/null +++ b/tests/learning/dynamics/test_comparison.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. +# ---------------------------------------------------------------------------- + +"""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 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 000000000..15042ae47 --- /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 000000000..a8f843497 --- /dev/null +++ b/tests/learning/dynamics/test_reference.py @@ -0,0 +1,454 @@ +# ---------------------------------------------------------------------------- +# 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 os +from pathlib import Path +import subprocess +import sys + +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({}) + + +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 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 diff --git a/tests/learning/dynamics/test_reference_rollout.py b/tests/learning/dynamics/test_reference_rollout.py new file mode 100644 index 000000000..c9802d24e --- /dev/null +++ b/tests/learning/dynamics/test_reference_rollout.py @@ -0,0 +1,306 @@ +# ---------------------------------------------------------------------------- +# 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"] + # 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": "", + }, + ) diff --git a/tests/learning/dynamics/test_report.py b/tests/learning/dynamics/test_report.py new file mode 100644 index 000000000..2fe72422b --- /dev/null +++ b/tests/learning/dynamics/test_report.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. +# ---------------------------------------------------------------------------- + +"""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 ( + _matched_rollouts, + _numerical_evidence, + _rollout_evidence, + annotate_rollout_evidence, +) +from embodichain.learning.dynamics.reference import _hash + + +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( + "field,value", + [ + ("max_forward_error", float("nan")), + ("max_codec_decode_error", -1.0), + ("atol", 1.0), + ("reference_commit", "unrelated"), + ], +) +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_trainer.py b/tests/learning/dynamics/test_trainer.py new file mode 100644 index 000000000..8d8fccb5c --- /dev/null +++ b/tests/learning/dynamics/test_trainer.py @@ -0,0 +1,475 @@ +# ---------------------------------------------------------------------------- +# 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 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): + 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, + ) + 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 + ) + + +def test_exact_resume_with_dropout_optimizer_sampler_rng(tmp_path): + make_data(tmp_path / "data") + full = build(tmp_path / "data", tmp_path / "full", steps=20, dropout=0.1) + 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) + assert interrupted.train(stop_after=7)["status"] == "interrupted" + restored = build(tmp_path / "data", tmp_path / "resume", steps=20, dropout=0.1) + 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_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 = torch.rsqrt(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 + + +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_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() + + +def test_training_annotations_cannot_invent_missing_physics_contract(tmp_path): + metadata = make_data(tmp_path / "data") + with pytest.raises(ValueError, match="physics_signature mismatch"): + run_train( + { + "dataset": str(tmp_path / "data"), + "output_dir": str(tmp_path / "run"), + "manifest": {**metadata, "physics_signature": "invented"}, + "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" diff --git a/tests/learning/models/_nerd_reference_check.py b/tests/learning/models/_nerd_reference_check.py new file mode 100644 index 000000000..2ab1b72d2 --- /dev/null +++ b/tests/learning/models/_nerd_reference_check.py @@ -0,0 +1,303 @@ +# ---------------------------------------------------------------------------- +# 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) -> 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. + """ + 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 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"], + ) + + 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()}) + 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:] + 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)) + 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 000000000..bf3a7faad --- /dev/null +++ b/tests/learning/models/test_nerd.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. +# ---------------------------------------------------------------------------- + +from __future__ import annotations +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, {}) + + +@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": str(Path(__file__).resolve().parents[3]), + }, + ) diff --git a/tests/sim/diff/test_dynamics.py b/tests/sim/diff/test_dynamics.py new file mode 100644 index 000000000..42316c2b6 --- /dev/null +++ b/tests/sim/diff/test_dynamics.py @@ -0,0 +1,218 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Functional facade tests using a real CPU Newton model and installed provider.""" + +from __future__ import annotations +from contextlib import nullcontext +from pathlib import Path +from types import SimpleNamespace +import sys +import pytest +import torch +import warp as wp +from dexsim.engine.newton_physics import NewtonCfg, RegisteredNewtonSolverCfg +from dexsim.engine.newton_physics.newton_backend import NewtonBackend +from dexsim.scene import Scene +from embodichain.lab.sim.diff.dynamics import FunctionalDynamics, detach_dynamics_state + +_PROVIDER = """ +import torch +import warp as wp +from newton.solvers import SolverBase +from dexsim.engine.newton_physics import FunctionalSolverState, NewtonSolverCapabilities, NewtonSolverExtension +class Native(SolverBase): + def __init__(self, model): + super().__init__(model) + self.memory = torch.zeros(model.world_count) + self.lengths = torch.zeros(model.world_count,dtype=torch.long) + def step(self,state_in,state_out,control,contacts,dt): + raise AssertionError('Native solver must not be stepped by functional env') + def reset(self,state,world_mask=None,flags=None): + mask=torch.ones_like(self.lengths,dtype=torch.bool) if world_mask is None else wp.to_torch(world_mask)[:-1] + self.memory=torch.where(mask,0,self.memory) + self.lengths=torch.where(mask,0,self.lengths) +class Provider: + def register_custom_attributes(self,builder,options): + pass + def create(self,context,options): + native=Native(context.model) + gain=torch.tensor(2.,requires_grad=True) + def initial(state): + physical={k:wp.to_torch(getattr(state,k)).clone() for k in ('joint_q','joint_qd','body_q','body_qd')} + return FunctionalSolverState(physical,{'memory':native.memory.clone()},native.lengths.clone()) + def transition(state,control,parameters,dt): + q=state.physical['joint_q']+control*parameters['gain']*dt + velocity=control*parameters['gain'] + body=torch.cat((q[:,None],state.physical['body_q'][:,1:]),-1) + body_velocity=torch.cat((velocity[:,None],torch.zeros_like(state.physical['body_qd'][:,1:])), -1) + return FunctionalSolverState({'joint_q':q,'joint_qd':velocity,'body_q':body,'body_qd':body_velocity},{'memory':state.history['memory']+control*dt},state.valid_history_lengths+1) + def commit(state): + native.memory=state.history['memory'].clone() + native.lengths=state.valid_history_lengths.clone() + def reset_history(state,fresh,world_mask): + mask=torch.ones_like(state.valid_history_lengths,dtype=torch.bool) if world_mask is None else wp.to_torch(world_mask)[:-1] + return {'memory':torch.where(mask,torch.zeros_like(state.history['memory']),state.history['memory'])},torch.where(mask,0,state.valid_history_lengths) + return NewtonSolverExtension(native_solver=native,capabilities=NewtonSolverCapabilities(differentiable=True,reset=True),functional_step=transition,initial_functional_state=initial,named_parameters=lambda:{'gain':gain},commit_functional_state=commit,reset_functional_history=reset_history) +""" + + +class _SceneManager: + """Renderer-free manager shell; physics, ownership merging and time are real.""" + + def __init__(self, backend: NewtonBackend) -> None: + self.device = torch.device("cpu") + self.backend = backend + self.scene = Scene("newton", SimpleNamespace(_newton_backend=backend)) + self.published_steps = 0 + self.elapsed = 0.0 + + def create_functional_state(self): + return self.scene.create_functional_state() + + def functional_parameters(self): + return self.scene.functional_parameters() + + def functional_step( + self, + state, + control, + parameters=None, + dt=None, + *, + publish=True, + render_final_step=True, + ): + assert render_final_step is False + result = self.scene.functional_step( + state, control, parameters, dt, publish=publish + ) + if publish: + self.published_steps += 1 + self.elapsed += self.backend.cfg.dt if dt is None else dt + return result + + def reset_functional_state(self, state, world_mask=None): + # This shell has no native arena registry: exercise the public backend + # counterpart with its actual Newton ownership map for selective rows. + selected = ( + None + if world_mask is None + else wp.from_torch( + torch.cat((world_mask, world_mask.new_zeros(1))), dtype=wp.bool + ) + ) + fresh = self.scene.create_functional_state() + return self.backend.reset_functional_state(state, selected, reset_state=fresh) + + def render_frame(self): + return nullcontext() + + +@pytest.fixture +def manager(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + info = tmp_path / "functional_dynamics_fixture-1.0.dist-info" + info.mkdir() + (info / "METADATA").write_text("Name: functional-dynamics-fixture\nVersion: 1.0\n") + (info / "entry_points.txt").write_text( + "[dexsim.newton_solvers]\nfunctional_fixture = functional_provider_fixture:Provider\n" + ) + (tmp_path / "functional_provider_fixture.py").write_text(_PROVIDER) + monkeypatch.syspath_prepend(str(tmp_path)) + monkeypatch.delitem(sys.modules, "functional_provider_fixture", raising=False) + cfg = NewtonCfg( + device="cpu", + dt=0.1, + num_substeps=2, + requires_grad=True, + sync_to_dexsim=False, + collision_pipeline_cfg=None, + solver_cfg=RegisteredNewtonSolverCfg("functional_fixture"), + ) + cfg.use_cuda_graph = False + backend = NewtonBackend(cfg) + from dexsim.engine.newton_physics._runtime import create_model_builder + + source = create_model_builder() + body = source.add_link(mass=1.0) + joint = source.add_joint_prismatic(-1, body, axis=wp.vec3(1.0, 0.0, 0.0)) + source.add_articulation([joint]) + source.add_shape_sphere(body, radius=0.1) + for _ in range(2): + backend.builder.add_world(source) + backend.finalize() + try: + yield _SceneManager(backend) + finally: + backend.close() + + +def test_macro_steps_do_not_double_solver_substeps_or_time(manager) -> None: + dynamics = FunctionalDynamics(manager) + state = dynamics.create_state() + action = torch.tensor([1.0, 2.0], requires_grad=True) + out = dynamics.step(state, action, physics_steps=3, physics_dt=0.1) + torch.testing.assert_close(out.physical["joint_q"], torch.tensor([0.6, 1.2])) + assert manager.backend.sim_time == pytest.approx(0.3) + assert manager.backend.step_index == 3 and manager.published_steps == 3 + assert out.valid_history_lengths.tolist() == [6, 6] + out.physical["joint_q"].sum().backward() + torch.testing.assert_close(action.grad, torch.tensor([0.6, 0.6])) + torch.testing.assert_close(dynamics.parameters()["gain"].grad, torch.tensor(0.9)) + # A later native overwrite cannot alter retained observations or old graphs. + manager.backend.state_0.joint_q.zero_() + torch.testing.assert_close(out.physical["joint_q"], torch.tensor([0.6, 1.2])) + + +def test_isolated_rollout_and_detach_do_not_mutate_live_state(manager) -> None: + dynamics = FunctionalDynamics(manager) + action = torch.ones(2, requires_grad=True) + result = dynamics.step( + dynamics.create_state(), action, physics_steps=2, publish=False + ) + assert manager.backend.step_index == 0 and manager.elapsed == 0.0 + assert manager.backend.solver.lengths.tolist() == [0, 0] + saved = result.physical["joint_q"] + detached = detach_dynamics_state(result) + assert saved.requires_grad and not detached.physical["joint_q"].requires_grad + assert not detached.history["memory"].requires_grad + torch.testing.assert_close(saved, detached.physical["joint_q"]) + + +def test_partial_reset_preserves_other_physical_and_history_gradients(manager) -> None: + dynamics = FunctionalDynamics(manager) + action = torch.tensor([1.0, 2.0], requires_grad=True) + state = dynamics.step(dynamics.create_state(), action) + # Normal manager reset occurred first, including native provider reset. + manager.backend.reset_state( + wp.array([True, False, False], dtype=wp.bool, device="cpu") + ) + merged = dynamics.reset(state, torch.tensor([True, False])) + assert manager.backend.step_index == 1 + torch.testing.assert_close(merged.physical["joint_q"], torch.tensor([0.0, 0.4])) + assert merged.valid_history_lengths.tolist() == [0, 2] + loss = merged.physical["joint_q"].sum() + merged.history["memory"].sum() + loss.backward() + torch.testing.assert_close(action.grad, torch.tensor([0.0, 0.3])) + + +def test_invalid_macro_steps_fail_before_stepping(manager) -> None: + dynamics = FunctionalDynamics(manager) + state = dynamics.create_state() + for count in [0, -1, True, 1.5]: + with pytest.raises(ValueError, match="positive integer"): + dynamics.step(state, torch.ones(2), physics_steps=count) + assert manager.backend.step_index == 0 diff --git a/tests/sim/physics/solvers/test_collection.py b/tests/sim/physics/solvers/test_collection.py new file mode 100644 index 000000000..4510f588d --- /dev/null +++ b/tests/sim/physics/solvers/test_collection.py @@ -0,0 +1,438 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Actual bounded generation and batched evaluation, isolated from renderer fixtures.""" + +from __future__ import annotations +import importlib.util +import json +import os +from pathlib import Path +import subprocess +import sys +import h5py +import numpy as np +import pytest +import torch +from embodichain.lab.sim.physics.solvers.collection import ( + DynamicsGenerationCfg, + generate_dataset, + evaluate_rollouts, + _sample, + _task_values, + _FixedPolicy, +) +from embodichain.data_pipeline.dynamics.shards import load_manifest, trajectory_split + +_ROOT = Path(__file__).resolve().parents[4] + + +@pytest.mark.parametrize( + "key,value", + [ + ("num_envs", True), + ("horizon", 2.5), + ("trajectories", 0), + ("resume", "false"), + ("seed", False), + ("num_substeps", -1), + ("max_attempts", 1), + ("effort_scale", float("nan")), + ], +) +def test_generation_config_rejects_invalid_values_before_runtime(key, value) -> None: + with pytest.raises((ValueError, TypeError)): + DynamicsGenerationCfg(**{key: value}) + + +def test_sampling_is_candidate_seed_local_and_passive_zero() -> None: + first = _sample("ant", 9, 5, "random_effort", "sample", 0.1) + _sample("cartpole", 73, 2, "random_effort", "sample", 1.0) + second = _sample("ant", 9, 5, "random_effort", "sample", 0.1) + for a, b in zip(first, second): + torch.testing.assert_close(a, b, atol=0, rtol=0) + assert not first[2][:, :6].any() + assert not _sample("cartpole", 3, 5, "passive", "default", 1.0)[2].any() + + +def test_fixed_policy_adapter_uses_own_state_and_records_hash() -> None: + policy = _FixedPolicy("cartpole", {"kind": "pd"}, torch.device("cpu")) + q = torch.tensor([[0.0, 0.0], [0.1, 0.2]]) + qd = torch.zeros_like(q) + force = policy(q, qd) + assert force[0, 0] == 0 and force[1, 0] != 0 + assert len(policy.identity["sha256"]) == 64 + assert not force[:, 1].any() + + +def _helper(): + path = Path(__file__).with_name("test_nerd_solver.py") + spec = importlib.util.spec_from_file_location("_collection_export_fixture", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _check_pipeline(tmp_path: Path, robot: str) -> None: + from embodichain.lab.sim.physics.solvers.collection import DynamicsSimulation + + cfg = { + "robot": robot, + "output_dir": str(tmp_path / "data"), + "num_envs": 2, + "trajectories": 3, + "horizon": 3, + "seed": 7, + "split_seed": 11, + "num_substeps": 2, + "trajectories_per_shard": 2, + "device": "cpu", + "initial_states": "default", + "motion_mode": "random_effort", + "effort_scale": 0.01, + } + result = generate_dataset(cfg) + assert result["trajectories"] == 3 and result["transitions"] == 9 + manifest = load_manifest(tmp_path / "data", verify=True) + original = [row for shard in manifest["shards"] for row in shard["trajectories"]] + assert [row["id"] for row in original] == ["7:0", "7:1", "7:2"] + assert [row["split"] for row in original] == [ + trajectory_split(row["id"], 11) for row in original + ] + resumed = generate_dataset({**cfg, "trajectories": 5, "resume": True}) + assert resumed["trajectories"] == 5 + manifest = load_manifest(tmp_path / "data", verify=True) + all_rows = [row for shard in manifest["shards"] for row in shard["trajectories"]] + assert [row["id"] for row in all_rows] == [f"7:{i}" for i in range(5)] + assert all_rows[:3] == original + first_shard = manifest["shards"][0] + with h5py.File(tmp_path / "data" / first_shard["file"], "r") as stream: + row = stream[first_shard["trajectories"][0]["group"]] + q = torch.tensor(row["states/q"][()]) + qd = torch.tensor(row["states/qd"][()]) + force = torch.tensor(row["transitions/joint_f"][()]) + saved_contacts = { + key: torch.tensor(value[0:1]) + for key, value in row["transitions"].items() + if key.startswith("contact_") + } + assert q.shape[0] == 4 and force.shape[0] == 3 + if robot == "cartpole": + assert row["transitions/contact_depths"].shape == (3, 0) + with DynamicsSimulation(robot, device="cpu", num_envs=1, num_substeps=2) as host: + host.reset(q[:1], qd[:1]) + state, contacts = host.snapshot() + for key, value in saved_contacts.items(): + torch.testing.assert_close(contacts[key], value) + host.step(force[:1]) + next_state, _ = host.snapshot() + torch.testing.assert_close(next_state["q"], q[1:2], atol=1e-5, rtol=1e-4) + torch.testing.assert_close(next_state["qd"], qd[1:2], atol=1e-5, rtol=1e-4) + checkpoint = _helper()._small_export( + host.sim.spawn_result, tmp_path / "model", robot + ) + common = { + "robot": robot, + "model": str(checkpoint), + "device": "cpu", + "num_envs": 2, + "trajectories": 3, + "horizons": [1, 3], + "num_substeps": 2, + "initial_states": "default", + } + report = evaluate_rollouts( + { + **common, + "teacher_dataset": str(tmp_path / "data"), + "split": "all", + "motion_mode": "random_effort", + "output": str(tmp_path / "replay.json"), + } + ) + assert report["num_envs"] == 2 and report["trajectories"] == 3 + assert len(report["error_vs_time"]) == 3 + assert report["teacher_origin"]["kind"] == "current_newton" + assert report["runtime_identity"]["device"] == "cpu" + assert report["runtime_identity"]["versions"]["torch"] == torch.__version__ + assert len(report["runtime_identity"]["source_sha256"]["nerd.py"]) == 64 + assert len(report["selected_teacher_sha256"]) == 64 + # Same immutable teacher trajectories have the same selected-content digest + # even when replay batching and output destination change. + rebatched = evaluate_rollouts( + { + **common, + "num_envs": 1, + "teacher_dataset": str(tmp_path / "data"), + "split": "all", + "motion_mode": "random_effort", + "output": str(tmp_path / "rebatched.json"), + } + ) + assert rebatched["selected_teacher_sha256"] == report["selected_teacher_sha256"] + with pytest.raises(ValueError, match="nonzero.*passive"): + evaluate_rollouts( + { + **common, + "teacher_dataset": str(tmp_path / "data"), + "split": "all", + "motion_mode": "passive", + "output": str(tmp_path / "mislabeled.json"), + } + ) + rows = [ + json.loads(line) + for line in Path(report["trajectory_evidence"]).read_text().splitlines() + ] + assert [row["trajectory_id"] for row in rows] == ["7:0", "7:1", "7:2"] + with h5py.File(report["trace_evidence"], "r") as stream: + assert stream["0/q"].shape[0] == 2 and stream["1/q"].shape[0] == 1 + fixed = evaluate_rollouts( + { + **common, + "motion_mode": "fixed_policy", + "policy": {"kind": "pd"}, + "output": str(tmp_path / "closed.json"), + } + ) + assert ( + fixed["policy_return"] is not None + and fixed["teacher_policy_return"] is not None + ) + assert len(fixed["policy_sha256"]) == 64 + teacher_manifest = load_manifest(fixed["teacher_dataset"], verify=True) + teacher_file = ( + Path(fixed["teacher_dataset"]) / teacher_manifest["shards"][0]["file"] + ) + with ( + h5py.File(teacher_file, "r") as teacher, + h5py.File(fixed["trace_evidence"], "r") as neural, + ): + # The same feedback policy recomputes actions from the deviating states. + assert not np.allclose( + teacher["0/transitions/joint_f"][1:], neural["0/joint_f"][0, 1:], atol=1e-8 + ) + assert ( + fixed["horizons"]["3"]["valid_trajectories"] + + fixed["horizons"]["3"]["numerically_invalid_trajectories"] + == 3 + ) + + +def _check_actor_and_failures(tmp_path: Path) -> None: + from embodichain.lab.sim.physics.solvers.collection import DynamicsSimulation + from embodichain.learning.models import build_mlp_from_cfg + from embodichain.learning.rl.policies import build_policy + from embodichain.learning.models.nerd import load_model, export_model + + with DynamicsSimulation( + "cartpole", device="cpu", num_envs=2, num_substeps=2 + ) as host: + checkpoint = _helper()._small_export( + host.sim.spawn_result, tmp_path / "model", "cartpole" + ) + block = { + "name": "actor_only", + "actor": { + "type": "mlp", + "network_cfg": { + "hidden_sizes": [8], + "activation": "tanh", + "last_activation": "tanh", + }, + }, + } + torch.manual_seed(11) + actor = build_mlp_from_cfg(block["actor"], 4, 1) + policy = build_policy(block, 4, 1, torch.device("cpu"), actor=actor) + config_path = tmp_path / "policy.json" + config_path.write_text(json.dumps({"policy": block})) + weights = tmp_path / "policy.pt" + torch.save({"schema_version": 1, "policy": policy.state_dict()}, weights) + common = { + "robot": "cartpole", + "model": str(checkpoint), + "device": "cpu", + "num_envs": 2, + "trajectories": 3, + "horizons": [1, 3], + "num_substeps": 2, + "initial_states": "default", + } + result = evaluate_rollouts( + { + **common, + "motion_mode": "fixed_policy", + "policy": { + "kind": "actor_only", + "config": str(config_path), + "checkpoint": str(weights), + }, + "output": str(tmp_path / "actor.json"), + } + ) + assert result["policy"]["checkpoint_sha256"] and result["policy_return"] is not None + model, manifest = load_model(checkpoint) + with torch.no_grad(): + # Finite FP32 weights whose velocity predictions overflow during rollout. + model.model.output_net.bias[2:] = torch.finfo(torch.float32).max + divergent = tmp_path / "divergent" + export_model(divergent, model, manifest) + report = evaluate_rollouts( + { + **common, + "model": str(divergent), + "motion_mode": "passive", + "output": str(tmp_path / "nonfinite.json"), + } + ) + assert report["nonfinite_rate"] == 1.0 + assert report["failure_rate"] == 1.0 + assert report["horizons"]["3"]["valid_trajectories"] == 0 + assert report["horizons"]["3"]["q_mae"] is None + assert len(Path(report["trajectory_evidence"]).read_text().splitlines()) == 3 + + +def _check_rejection_limit(tmp_path: Path) -> None: + import embodichain.lab.sim.physics.solvers.collection as collection + + original = collection._sample + + def fast_velocity(*args): + q, qd, effort = original(*args) + qd[:] = 2e5 + return q, qd, effort + + collection._sample = fast_velocity + try: + with pytest.raises(RuntimeError, match="attempt limit"): + generate_dataset( + { + "robot": "cartpole", + "output_dir": str(tmp_path / "rejected"), + "num_envs": 2, + "trajectories": 1, + "horizon": 1, + "max_attempts": 2, + "device": "cpu", + "initial_states": "default", + "motion_mode": "passive", + } + ) + progress = json.loads( + (tmp_path / "rejected/generation-progress.json").read_text() + ) + assert progress["rejected"] == 2 and len(progress["rejections"]) == 2 + assert not load_manifest(tmp_path / "rejected")["shards"] + finally: + collection._sample = original + + +def _subprocess(tmp_path: Path, operation: str, robot: str = "cartpole") -> None: + code = "import importlib.util,sys;from pathlib import Path;s=importlib.util.spec_from_file_location('collection_test',sys.argv[1]);m=importlib.util.module_from_spec(s);s.loader.exec_module(m);getattr(m,sys.argv[3])(Path(sys.argv[2]),sys.argv[4]) if sys.argv[3]=='_check_pipeline' else getattr(m,sys.argv[3])(Path(sys.argv[2]))" + subprocess.run( + [ + sys.executable, + "-c", + code, + str(Path(__file__).resolve()), + str(tmp_path), + operation, + robot, + ], + check=True, + env={ + **os.environ, + "OMP_NUM_THREADS": "2", + "MKL_NUM_THREADS": "2", + "PYTHONPATH": str(_ROOT), + "EMBODICHAIN_SIM_EXIT_PROCESS": "0", + }, + ) + + +@pytest.mark.subprocess_sim +@pytest.mark.parametrize("robot", ["cartpole", "ant"]) +def test_actual_generation_resume_alignment_and_batched_closed_loop( + tmp_path: Path, robot: str +) -> None: + _subprocess(tmp_path, "_check_pipeline", robot) + + +@pytest.mark.subprocess_sim +def test_actual_actor_artifact_and_nonfinite_evidence(tmp_path: Path) -> None: + _subprocess(tmp_path, "_check_actor_and_failures") + + +@pytest.mark.subprocess_sim +def test_actual_rejection_attempt_limit(tmp_path: Path) -> None: + _subprocess(tmp_path, "_check_rejection_limit") + + +def test_error_metric_wraps_only_continuous_coordinates() -> None: + from embodichain.lab.sim.physics.solvers.collection import _error_values + + q = torch.zeros(1, 15) + q[:, 6] = 1.0 + target = q.clone() + q[:, 7] = 2 * torch.pi + error = _error_values("ant", q, torch.zeros(1, 14), target, torch.zeros(1, 14)) + torch.testing.assert_close( + error[:, 0], torch.tensor([2 * torch.pi / 8], dtype=torch.double) + ) + cart = torch.tensor([[0.0, 2 * torch.pi]]) + error = _error_values( + "cartpole", cart, torch.zeros(1, 2), torch.zeros(1, 2), torch.zeros(1, 2) + ) + assert error[0, 0] < 1e-6 + + +def test_teacher_content_hash_uses_values_and_logical_identity_not_paths() -> None: + from embodichain.lab.sim.physics.solvers.collection import ( + _teacher_row_sha256, + _teacher_origin, + ) + + record = { + "id": "sample:7", + "file": "shard-a.h5", + "group": "0", + "metadata": {"sampling_seed": 7}, + } + values = [ + np.zeros((4, 2), dtype=np.float32), + np.zeros((4, 2), dtype=np.float32), + np.ones((3, 2), dtype=np.float32), + ] + expected = _teacher_row_sha256(record, values) + assert ( + _teacher_row_sha256( + {**record, "file": "/different/location.h5", "group": "17"}, values + ) + == expected + ) + changed = [array.copy() for array in values] + changed[0][2, 1] = 1.0 + assert _teacher_row_sha256(record, changed) != expected + assert _teacher_row_sha256({**record, "id": "sample:8"}, values) != expected + official = { + "kind": "official_warp", + "warp_version": "1.5.1", + "source_commit": "pinned", + } + assert _teacher_origin({"teacher_origin": official}) == official + with pytest.raises(ValueError, match="teacher_origin"): + _teacher_origin({"robot": "ant"}) diff --git a/tests/sim/physics/solvers/test_kinematics.py b/tests/sim/physics/solvers/test_kinematics.py new file mode 100644 index 000000000..5394e7fe2 --- /dev/null +++ b/tests/sim/physics/solvers/test_kinematics.py @@ -0,0 +1,92 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Numerical checks of the native FK autograd boundary.""" + +from __future__ import annotations + +import pytest +import torch +import warp as wp + +newton = pytest.importorskip("newton") + +from embodichain.lab.sim.physics.solvers._kinematics import differentiable_fk + + +def _model(): + builder = newton.ModelBuilder() + body = builder.add_link( + mass=1.0, + com=wp.vec3(0.0, 0.0, -0.5), + inertia=wp.mat33(1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0), + ) + joint = builder.add_joint_revolute( + parent=-1, child=body, axis=wp.vec3(0.0, 1.0, 0.0) + ) + builder.add_articulation([joint]) + return builder.finalize(device="cpu") + + +def test_fk_pose_and_velocity_gradients_match_finite_difference(): + model = _model() + q = torch.tensor([0.4], requires_grad=True) + qd = torch.tensor([0.7], requires_grad=True) + pose, velocity = differentiable_fk(model, q, qd) + pose_weights = torch.arange(1.0, 8.0).view(1, 7) + velocity_weights = torch.arange(1.0, 7.0).view(1, 6) + loss = (pose * pose_weights).sum() + (velocity * velocity_weights).sum() + gradients = torch.autograd.grad(loss, (q, qd), retain_graph=True) + repeated = torch.autograd.grad(loss, (q, qd)) + assert all(torch.allclose(a, b) for a, b in zip(gradients, repeated)) + + def evaluate(position, speed): + a, b = differentiable_fk(model, position, speed) + return (a * pose_weights).sum() + (b * velocity_weights).sum() + + epsilon = 1.0e-3 + expected_q = ( + evaluate(q.detach() + epsilon, qd.detach()) + - evaluate(q.detach() - epsilon, qd.detach()) + ) / (2 * epsilon) + expected_qd = ( + evaluate(q.detach(), qd.detach() + epsilon) + - evaluate(q.detach(), qd.detach() - epsilon) + ) / (2 * epsilon) + assert torch.allclose(gradients[0], expected_q.reshape(1), atol=2.0e-3, rtol=2.0e-3) + assert torch.allclose( + gradients[1], expected_qd.reshape(1), atol=2.0e-3, rtol=2.0e-3 + ) + + +def test_later_fk_does_not_overwrite_retained_outputs_or_graph(): + model = _model() + q = torch.tensor([0.2], requires_grad=True) + qd = torch.tensor([0.3], requires_grad=True) + first, _ = differentiable_fk(model, q, qd) + saved = first.detach().clone() + differentiable_fk(model, q + 1.0, qd + 2.0) + assert torch.equal(first, saved) + first.sum().backward() + assert q.grad is not None and torch.isfinite(q.grad).all() + + +def test_fk_rejects_unmatched_coordinate_layout(): + model = _model() + with pytest.raises(ValueError, match="flattened"): + differentiable_fk(model, torch.zeros(1, 1), torch.zeros(1)) + with pytest.raises(TypeError, match="float32"): + differentiable_fk(model, torch.zeros(1, dtype=torch.float64), torch.zeros(1)) diff --git a/tests/sim/physics/solvers/test_nerd_solver.py b/tests/sim/physics/solvers/test_nerd_solver.py new file mode 100644 index 000000000..6e04ab46e --- /dev/null +++ b/tests/sim/physics/solvers/test_nerd_solver.py @@ -0,0 +1,900 @@ +# ---------------------------------------------------------------------------- +# 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 NeRD integration contracts with exact assets and installed entry points. + +Small seeded networks prove execution/gradient contracts, not trained accuracy. +""" + +from __future__ import annotations +from contextlib import contextmanager +from pathlib import Path +from types import SimpleNamespace +from typing import Any +import json +import os +import subprocess +import sys +import numpy as np +import pytest +import torch +import warp as wp +from dexsim.engine.newton_physics import ( + NewtonCfg, + RegisteredNewtonSolverCfg, + FunctionalSolverState, +) +from dexsim.engine.newton_physics.newton_backend import NewtonBackend +from dexsim.scene import Scene +from embodichain.lab.sim.physics.solvers._reference_assets import ( + make_reference_articulation, + reference_profile, +) +from embodichain.lab.sim.physics.solvers._nerd_binding import NeRDBinding +from embodichain.learning.models.nerd import NeRDModel, NeRDModelCfg, export_model + +DT = 1 / 60 + + +def _source(robot: str): + """Use production descriptor lowering, with no render allocations.""" + from dexsim.spawn.adapters.newton_articulation_adapter import ( + _translate_explicit_articulation, + ) + from dexsim.engine.newton_physics.articulation.articulation import ( + NewtonArticulation, + ) + from dexsim.engine.newton_physics._runtime import create_model_builder + + descriptor = make_reference_articulation(robot) + builder = create_model_builder() + specs, joints = _translate_explicit_articulation( + SimpleNamespace( + get_render_body=lambda name: None, + get_link_names=lambda *_: [link.name for link in descriptor.links], + ), + None, + descriptor, + {}, + "mujoco_warp", + reuse_render=True, + ) + handle = NewtonArticulation( + builder, + None, + descriptor.name, + descriptor.root_link_name, + specs, + joints, + build_device="cpu", + ) + handle.finalize( + descriptor.pose, fixed_base=descriptor.fixed_base, disable_self_collision=True + ) + builder.joint_q[:] = reference_profile(robot)["default_q"] + builder.joint_qd[:] = reference_profile(robot)["default_qd"] + return builder + + +@contextmanager +def _scene( + robot: str, + checkpoint: Path | None = None, + *, + num_envs: int = 2, + dt: float = DT, + mutate=None, + device: str = "cpu", +): + cfg = NewtonCfg( + device=device, + dt=dt, + num_substeps=1, + requires_grad=checkpoint is not None, + sync_to_dexsim=False, + collision_pipeline_cfg=None, + solver_cfg=( + RegisteredNewtonSolverCfg("semi_implicit") + if checkpoint is None + else RegisteredNewtonSolverCfg( + "nerd", {"checkpoint": str(checkpoint), "trainable": True} + ) + ), + ) + cfg.use_cuda_graph = False + backend = NewtonBackend(cfg) + try: + source = _source(robot) + for _ in range(num_envs): + backend.builder.add_world(source) + if robot == "ant": + backend.builder.add_ground_plane() + if mutate is not None: + mutate(backend.builder) + backend.finalize() + yield Scene("newton", SimpleNamespace(_newton_backend=backend)), backend + finally: + backend.close() + + +def _small_export(scene, destination: Path, robot: str) -> Path: + binding = NeRDBinding.from_metadata(dict(scene.dynamics_metadata()), robot) + physical = scene.dynamics_snapshot().physical + q, qd = binding.read(physical) + from embodichain.compute.dynamics import NeRDStateCodec + + codec = NeRDStateCodec(**binding.metadata(DT)["codec"]) + features = codec.encode( + q, + qd, + torch.zeros_like(qd), + binding.contacts(physical["body_q"]), + binding.gravity_dir, + ) + torch.manual_seed(42) + model = NeRDModel( + NeRDModelCfg( + input_dims={k: v.shape[-1] for k, v in features.items()}, + output_dim=q.shape[-1] + qd.shape[-1], + n_embd=24, + n_head=3, + n_layer=2, + 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_model(destination, model, binding.metadata(DT)) + return destination + + +def _export(tmp_path: Path, robot: str) -> Path: + with _scene(robot) as (scene, _): + return _small_export(scene, tmp_path / robot, robot) + + +@pytest.mark.parametrize("robot", ["cartpole", "ant"]) +def test_native_matches_functional_and_commits_history( + tmp_path: Path, robot: str +) -> None: + checkpoint = _export(tmp_path, robot) + with _scene(robot, checkpoint) as (scene, backend): + initial = scene.create_functional_state() + binding = backend.solver.binding + force = torch.zeros(backend.model.joint_dof_count) + force[binding.qd_indices[:, -1]] = 0.3 + expected = scene.functional_step(initial, {"joint_f": force}, publish=False) + assert backend.step_index == 0 + backend.control.joint_f.assign(force.numpy()) + backend.step() + actual = scene.create_functional_state() + for name in expected.physical: + torch.testing.assert_close( + actual.physical[name], expected.physical[name], atol=1e-6, rtol=1e-5 + ) + assert actual.valid_history_lengths.tolist() == [1, 1] + # Publish a functional transition, then native stepping must continue its + # history rather than reusing the history before publication. + published = scene.functional_step(actual, {"joint_f": force}) + expected_next = scene.functional_step( + published, {"joint_f": force}, publish=False + ) + backend.step() + after = scene.create_functional_state() + torch.testing.assert_close( + after.physical["joint_q"], + expected_next.physical["joint_q"], + atol=1e-6, + rtol=1e-5, + ) + assert after.valid_history_lengths.tolist() == [3, 3] + assert backend.step_index == 3 + + +@pytest.mark.parametrize("robot", ["cartpole", "ant"]) +@pytest.mark.parametrize( + "device", ["cpu", pytest.param("cuda:0", marks=pytest.mark.gpu)] +) +def test_multistep_gradients_and_native_storage_independence( + tmp_path: Path, robot: str, device: str +) -> None: + if device.startswith("cuda") and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + checkpoint = _export(tmp_path, robot) + with _scene(robot, checkpoint, device=device) as (scene, backend): + state = scene.create_functional_state() + initial_q = state.physical["joint_q"].requires_grad_() + initial_qd = state.physical["joint_qd"].requires_grad_() + controls = [] + for _ in range(5): + force = torch.zeros( + backend.model.joint_dof_count, device=device, requires_grad=True + ) + controls.append(force) + state = scene.functional_step(state, {"joint_f": force}, publish=True) + saved_q = state.physical["joint_q"].detach().clone() + saved_body = state.physical["body_q"].detach().clone() + backend.state_0.joint_q.zero_() + backend.state_0.body_q.zero_() + torch.testing.assert_close(state.physical["joint_q"], saved_q) + torch.testing.assert_close(state.physical["body_q"], saved_body) + loss = ( + state.physical["joint_q"].square().sum() + + state.physical["body_q"].square().sum() + + state.physical["body_qd"].square().sum() + ) + loss.backward() + for value in [initial_q, initial_qd, *controls]: + assert value.grad is not None and torch.isfinite(value.grad).all() + assert value.grad.abs().sum() > 0 + parameters = scene.functional_parameters() + assert parameters and all( + p.grad is not None and torch.isfinite(p.grad).all() + for p in parameters.values() + ) + assert sum(p.grad.abs().sum() for p in parameters.values()) > 0 + + +@pytest.mark.parametrize("robot", ["cartpole", "ant"]) +@pytest.mark.parametrize( + "device", ["cpu", pytest.param("cuda:0", marks=pytest.mark.gpu)] +) +def test_partial_reset_preserves_unselected_history_graph( + tmp_path: Path, robot: str, device: str +) -> None: + if device.startswith("cuda") and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + checkpoint = _export(tmp_path, robot) + with _scene(robot, checkpoint, device=device) as (scene, backend): + before = scene.create_functional_state() + force = torch.zeros( + backend.model.joint_dof_count, device=device, requires_grad=True + ) + state = scene.functional_step(before, {"joint_f": force}) + mask = wp.array([True, False, False], dtype=wp.bool, device=device) + # Exercise real public backend ownership mapping in the renderer-free + # Scene shell; actual manager arena selection has its own integration test. + reset = backend.reset_functional_state(state, mask) + assert reset.valid_history_lengths.tolist() == [0, 1] + for value in reset.history.values(): + assert not value[0].any() + state2 = scene.functional_step( + reset, {"joint_f": torch.zeros_like(force)}, publish=False + ) + assert state2.valid_history_lengths.tolist() == [1, 2] + q = state2.physical["joint_q"][backend.solver.binding.q_indices] + q.sum().backward() + assert force.grad is not None + selected = backend.solver.binding.qd_indices[0] + unselected = backend.solver.binding.qd_indices[1] + assert not force.grad[selected].any() + assert force.grad[unselected].abs().sum() > 0 + + +@pytest.mark.parametrize( + "field,value,message", + [ + ("dt", 0.02, "timestep"), + ("contact_schema", "wrong", "contact"), + ("physics_signature", "wrong", "physics signature"), + ("physics_signature_schema", None, "signature schema"), + ("physics_signature_schema", "unsupported", "signature schema"), + ("joint_names", ["wrong"], "joint order"), + ], +) +def test_incompatible_manifest_rejected_at_creation( + tmp_path: Path, field: str, value: Any, message: str +) -> None: + checkpoint = _export(tmp_path, "cartpole") + path = checkpoint / "manifest.json" + manifest = json.loads(path.read_text()) + manifest[field] = value + path.write_text(json.dumps(manifest)) + with pytest.raises(ValueError, match=message): + with _scene("cartpole", checkpoint): + pass + + +def test_wrong_timestep_and_control_schema_fail_without_advance(tmp_path: Path) -> None: + checkpoint = _export(tmp_path, "cartpole") + with _scene("cartpole", checkpoint) as (scene, backend): + state = scene.create_functional_state() + with pytest.raises(ValueError, match="timestep"): + scene.functional_step(state, {"joint_f": torch.zeros(4)}, dt=0.02) + with pytest.raises(ValueError, match="joint_f"): + scene.functional_step(state, {"joint_target_pos": torch.zeros(4)}) + assert backend.step_index == 0 + + +def test_non_float32_export_rejected_before_newton_execution(tmp_path: Path) -> None: + from embodichain.learning.models.nerd import load_model + + original = _export(tmp_path, "cartpole") + model, manifest = load_model(original) + checkpoint = tmp_path / "float64-export" + export_model(checkpoint, model.double(), manifest) + with pytest.raises(TypeError, match="float32 model"): + with _scene("cartpole", checkpoint): + pass + + +def test_ant_rejects_unmodeled_free_base_force(tmp_path: Path) -> None: + checkpoint = _export(tmp_path, "ant") + with _scene("ant", checkpoint) as (scene, backend): + force = torch.zeros(backend.model.joint_dof_count) + force[backend.solver.binding.qd_indices[0, 0]] = 1.0 + with pytest.raises( + (ValueError, NotImplementedError), match="base|root|force|effort" + ): + scene.functional_step(scene.create_functional_state(), {"joint_f": force}) + assert backend.step_index == 0 + + +def test_joint_axis_change_invalidates_physics_signature(tmp_path: Path) -> None: + checkpoint = _export(tmp_path, "cartpole") + + def mutate(builder): + builder.joint_axis[0] = wp.vec3(0.0, 0.0, 1.0) + builder.joint_axis[2] = wp.vec3(0.0, 0.0, 1.0) + + with pytest.raises(ValueError, match="physics signature"): + with _scene("cartpole", checkpoint, mutate=mutate): + pass + + +def test_ant_plane_friction_change_invalidates_physics_signature( + tmp_path: Path, +) -> None: + checkpoint = _export(tmp_path, "ant") + + def mutate(builder): + builder.shape_material_mu[-1] = 0.12 + + with pytest.raises(ValueError, match="physics signature"): + with _scene("ant", checkpoint, mutate=mutate): + pass + + +def _check_real_manager_public_scene_and_partial_arena_reset( + tmp_path: Path, robot: str +) -> None: + from embodichain.lab.sim.sim_manager import SimulationManager, SimulationManagerCfg + from embodichain.lab.sim.cfg import NewtonPhysicsCfg, RenderCfg + + def make(checkpoint=None): + physics = NewtonPhysicsCfg( + device="cpu", + physics_dt=DT, + num_substeps=1, + requires_grad=checkpoint is not None, + use_cuda_graph=False, + collision_cfg=None, + solver_cfg=( + {"solver_type": "semi_implicit"} + if checkpoint is None + else { + "solver_type": "nerd", + "options": {"checkpoint": str(checkpoint), "trainable": True}, + } + ), + ) + cfg = SimulationManagerCfg( + headless=True, + num_envs=2, + device="cpu", + physics_cfg=physics, + render_cfg=RenderCfg(renderer="no-render"), + enable_ground_plane=robot == "ant", + enable_entity_gizmo=False, + robot_ik_gizmo=None, + startup_summary="off", + ) + sim = SimulationManager(cfg) + try: + sim.add_articulation_descriptor(make_reference_articulation(robot)) + sim.prepare() + except BaseException: + sim.destroy(exit_process=False) + SimulationManager.flush_cleanup_queue() + raise + return sim + + teacher = make() + try: + checkpoint = _small_export( + teacher.spawn_result, tmp_path / "manager-model", robot + ) + finally: + teacher.destroy(exit_process=False) + SimulationManager.flush_cleanup_queue() + sim = make(checkpoint) + try: + initial = sim.create_functional_state() + control = torch.zeros_like(initial.physical["joint_qd"], requires_grad=True) + state = sim.functional_step( + initial, {"joint_f": control}, render_final_step=False + ) + sim.reset_objects_state(env_ids=torch.tensor([0])) + reset = sim.reset_functional_state(state, torch.tensor([True, False])) + assert reset.valid_history_lengths.tolist() == [0, 1] + assert sim._visualization_sim_step == 1 + reset.physical["joint_q"].sum().backward() + assert control.grad is not None and torch.isfinite(control.grad).all() + binding = NeRDBinding.from_metadata( + dict(sim.spawn_result.dynamics_metadata()), robot + ) + assert not control.grad[binding.qd_indices[0]].any() + assert control.grad[binding.qd_indices[1]].abs().sum() > 0 + finally: + sim.destroy(exit_process=False) + SimulationManager.flush_cleanup_queue() + + +def test_fk_vjp_matches_directional_difference(tmp_path: Path) -> None: + from embodichain.lab.sim.physics.solvers._kinematics import differentiable_fk + + with _scene("ant") as (scene, backend): + physical = scene.dynamics_snapshot().physical + q = physical["joint_q"].clone().requires_grad_() + qd = physical["joint_qd"].clone().requires_grad_() + # Perturb scalar joints only so the root quaternion remains unit. + direction = torch.zeros_like(q) + for row in NeRDBinding(backend.model, "ant").q_indices: + direction[row[7:]] = torch.linspace(0.2, 0.8, 8) + velocity_direction = torch.linspace(-0.3, 0.3, qd.numel()) + weights = torch.linspace(0.1, 1.0, backend.model.body_count * 7).reshape(-1, 7) + + def objective(q, qd): + body, velocity = differentiable_fk(backend.model, q, qd) + return (body * weights).sum() + velocity.square().sum() + + loss = objective(q, qd) + gq, gv = torch.autograd.grad(loss, (q, qd)) + epsilon = 1e-3 + finite = ( + objective( + q.detach() + epsilon * direction, + qd.detach() + epsilon * velocity_direction, + ) + - objective( + q.detach() - epsilon * direction, + qd.detach() - epsilon * velocity_direction, + ) + ) / (2 * epsilon) + analytic = (gq * direction).sum() + (gv * velocity_direction).sum() + torch.testing.assert_close(analytic, finite, rtol=2e-3, atol=2e-4) + + +def test_ant_initial_joint_position_contact_gradient_matches_fd(tmp_path: Path) -> None: + checkpoint = _export(tmp_path, "ant") + with _scene("ant", checkpoint) as (scene, backend): + initial = scene.create_functional_state() + q = initial.physical["joint_q"].clone().requires_grad_() + state = FunctionalSolverState( + {**initial.physical, "joint_q": q}, + initial.history, + initial.valid_history_lengths, + ) + force = torch.zeros_like(initial.physical["joint_qd"]) + + def feature_loss(value): + # Deliberately retain the initial detached body snapshot: q is the + # authoritative reduced state and must rebuild contact geometry. + current = FunctionalSolverState( + {**state.physical, "joint_q": value}, + state.history, + state.valid_history_lengths, + ) + result = scene.functional_step(current, {"joint_f": force}, publish=False) + return result.history["contact_depths"][:, 0].sum() + + grad = torch.autograd.grad(feature_loss(q), q)[0] + direction = torch.zeros_like(q) + direction[backend.solver.binding.q_indices[:, 2]] = 1.0 + epsilon = 1e-3 + finite = ( + feature_loss(q.detach() + epsilon * direction) + - feature_loss(q.detach() - epsilon * direction) + ) / (2 * epsilon) + analytic = (grad * direction).sum() + assert analytic.abs() > 0.5 + torch.testing.assert_close(analytic, finite, rtol=1e-3, atol=1e-4) + + +def test_model_change_invalidates_future_steps_even_after_failed_notification( + tmp_path: Path, +) -> None: + checkpoint = _export(tmp_path, "cartpole") + with _scene("cartpole", checkpoint) as (scene, backend): + state = scene.create_functional_state() + original = backend.model.joint_axis.numpy().copy() + changed = original.copy() + changed[:, :] = [0.0, 0.0, 1.0] + backend.model.joint_axis.assign(changed) + with pytest.raises(ValueError, match="changed|signature"): + backend.solver.notify_model_changed(0) + backend.model.joint_axis.assign(original) + with pytest.raises(RuntimeError, match="invalid|rebuild|changed"): + scene.functional_step( + state, {"joint_f": torch.zeros_like(state.physical["joint_qd"])} + ) + assert backend.step_index == 0 + + +@pytest.mark.subprocess_sim +@pytest.mark.parametrize("robot", ["cartpole", "ant"]) +def test_real_manager_public_scene_and_partial_arena_reset( + tmp_path: Path, robot: str +) -> None: + # DexSim renderer is process-global; isolate no-render tests from the main + # suite's hybrid renderer fixture without resetting native global state. + code = "import importlib.util,sys; from pathlib import Path; spec=importlib.util.spec_from_file_location('nerd_provider_test',sys.argv[1]); module=importlib.util.module_from_spec(spec); spec.loader.exec_module(module); module._check_real_manager_public_scene_and_partial_arena_reset(Path(sys.argv[2]),sys.argv[3])" + subprocess.run( + [ + sys.executable, + "-c", + code, + str(Path(__file__).resolve()), + str(tmp_path), + robot, + ], + check=True, + env={ + **os.environ, + "OMP_NUM_THREADS": "2", + "MKL_NUM_THREADS": "2", + "PYTHONPATH": str(Path(__file__).resolve().parents[4]), + }, + ) + + +def _check_real_differentiable_environment(tmp_path: Path) -> None: + import gymnasium as gym + from embodichain.lab.gym.envs.differentiable_dynamics_env import ( + DifferentiableDynamicsEnv, + ) + from embodichain.lab.gym.envs.embodied_env import EmbodiedEnvCfg + from embodichain.lab.sim.sim_manager import SimulationManager, SimulationManagerCfg + from embodichain.lab.sim.cfg import NewtonPhysicsCfg, RenderCfg + + def sim_cfg(checkpoint=None): + physics = NewtonPhysicsCfg( + device="cpu", + physics_dt=DT, + num_substeps=1, + requires_grad=checkpoint is not None, + use_cuda_graph=False, + collision_cfg=None, + solver_cfg=( + {"solver_type": "semi_implicit"} + if checkpoint is None + else { + "solver_type": "nerd", + "options": {"checkpoint": str(checkpoint), "trainable": True}, + } + ), + ) + return SimulationManagerCfg( + headless=True, + num_envs=2, + device="cpu", + physics_cfg=physics, + render_cfg=RenderCfg(renderer="no-render"), + enable_ground_plane=False, + enable_entity_gizmo=False, + robot_ik_gizmo=None, + startup_summary="off", + ) + + teacher = SimulationManager(sim_cfg()) + try: + teacher.add_articulation_descriptor(make_reference_articulation("cartpole")) + teacher.prepare() + checkpoint = _small_export( + teacher.spawn_result, tmp_path / "env-model", "cartpole" + ) + finally: + teacher.destroy(exit_process=False) + SimulationManager.flush_cleanup_queue() + + class Task(DifferentiableDynamicsEnv): + def _setup_scene(self, **kwargs): + self.sim = SimulationManager(self.sim_cfg) + self.sim.add_articulation_descriptor( + make_reference_articulation("cartpole") + ) + + def _setup_robot(self, **kwargs): + self.single_action_space = gym.spaces.Box( + -1.0, 1.0, shape=(2,), dtype=np.float32 + ) + self.single_observation_space = gym.spaces.Box( + -np.inf, np.inf, shape=(2,), dtype=np.float32 + ) + return self.sim.get_articulation("cartpole") + + def _action_to_control(self, action, state): + return {"joint_f": action.flatten()} + + def _read_outputs(self, state): + q = state.physical["joint_q"].reshape(2, 2) + done = (torch.arange(2, device=q.device) == 0) & (self._elapsed_steps >= 1) + return { + "obs": q, + "reward": q.square().sum(-1), + "terminated": done, + "truncated": torch.zeros(2, dtype=torch.bool, device=q.device), + } + + env = Task( + EmbodiedEnvCfg( + num_envs=2, + sim_cfg=sim_cfg(checkpoint), + sim_steps_per_control=1, + max_episode_steps=5, + ) + ) + try: + obs, _ = env.reset(seed=7) + assert obs.shape == (2, 2) + action = torch.ones(2, 2, requires_grad=True) + obs, reward, terminated, truncated, info = env.step(action) + assert terminated.tolist() == [True, False] + assert env.functional_state.valid_history_lengths.tolist() == [0, 1] + assert env.sim._visualization_sim_step == 1 + loss = reward.sum() + obs.square().sum() + loss.backward() + assert action.grad is not None and torch.isfinite(action.grad).all() + assert action.grad.abs().sum() > 0 + assert all( + value.grad is not None for value in env.dynamics_parameters().values() + ) + assert not env.detach_state().requires_grad + finally: + env.sim.destroy(exit_process=False) + SimulationManager.flush_cleanup_queue() + + +@pytest.mark.subprocess_sim +def test_actual_differentiable_environment_constructor_reset_and_backward( + tmp_path: Path, +) -> None: + code = "import importlib.util,sys; from pathlib import Path; spec=importlib.util.spec_from_file_location('nerd_provider_test',sys.argv[1]); module=importlib.util.module_from_spec(spec); spec.loader.exec_module(module); module._check_real_differentiable_environment(Path(sys.argv[2]))" + subprocess.run( + [sys.executable, "-c", code, str(Path(__file__).resolve()), str(tmp_path)], + check=True, + env={ + **os.environ, + "OMP_NUM_THREADS": "2", + "MKL_NUM_THREADS": "2", + "PYTHONPATH": str(Path(__file__).resolve().parents[4]), + }, + ) + + +def test_frozen_network_accepts_explicit_trainable_parameter_mapping( + tmp_path: Path, +) -> None: + checkpoint = _export(tmp_path, "cartpole") + with _scene("cartpole", checkpoint) as (scene, backend): + backend.solver.network.requires_grad_(False) + state = scene.create_functional_state() + force = torch.ones_like(state.physical["joint_qd"], requires_grad=True) + baseline = scene.functional_step(state, {"joint_f": force}, publish=False) + parameters = { + name: value.detach().clone().requires_grad_() + for name, value in scene.functional_parameters().items() + } + explicit = scene.functional_step( + state, {"joint_f": force}, parameters=parameters, publish=False + ) + torch.testing.assert_close( + explicit.physical["joint_q"], baseline.physical["joint_q"] + ) + explicit.physical["joint_q"].square().sum().backward() + assert all( + value.grad is None for value in scene.functional_parameters().values() + ) + assert all( + value.grad is not None and torch.isfinite(value.grad).all() + for value in parameters.values() + ) + assert force.grad is not None and force.grad.abs().sum() > 0 + altered = dict(parameters) + altered["model.output_net.bias"] = parameters["model.output_net.bias"] + 0.1 + changed = scene.functional_step( + state, {"joint_f": force}, parameters=altered, publish=False + ) + assert not torch.allclose( + changed.physical["joint_q"], baseline.physical["joint_q"] + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize("robot", ["cartpole", "ant"]) +def test_cuda_native_functional_and_contact_backward_on_torch_stream( + tmp_path: Path, robot: str +) -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + checkpoint = _export(tmp_path, robot) + with _scene(robot, checkpoint, device="cuda:0") as (scene, backend): + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream), wp.ScopedStream(wp.stream_from_torch(stream)): + state = scene.create_functional_state() + force = torch.zeros_like(state.physical["joint_qd"], requires_grad=True) + force_data = force.detach() + expected = scene.functional_step(state, {"joint_f": force}, publish=False) + wp.copy(backend.control.joint_f, wp.from_torch(force_data)) + backend.step() + actual = scene.create_functional_state() + torch.testing.assert_close( + actual.physical["joint_q"], + expected.physical["joint_q"], + atol=1e-6, + rtol=1e-5, + ) + expected.physical["body_q"].square().sum().backward() + assert force.grad is not None and torch.isfinite(force.grad).all() + if robot == "ant": + q = state.physical["joint_q"].detach().clone() + q[backend.solver.binding.q_indices[:, 2]] = 0.3 + q.requires_grad_() + lowered = FunctionalSolverState( + {**state.physical, "joint_q": q}, + state.history, + state.valid_history_lengths, + ) + result = scene.functional_step( + lowered, {"joint_f": torch.zeros_like(force)}, publish=False + ) + # Foot penetration yields active continuous plane features; masks + # are discrete and held fixed in the local derivative contract. + features = result.history["contact_depths"][:, 0] + assert (features < 0).any() + features.sum().backward() + assert q.grad is not None and torch.isfinite(q.grad).all() + assert q.grad[backend.solver.binding.q_indices[:, 7:]].abs().sum() > 0 + stream.synchronize() + + +@pytest.mark.parametrize("robot", ["cartpole", "ant"]) +@pytest.mark.parametrize( + "device", ["cpu", pytest.param("cuda:0", marks=pytest.mark.gpu)] +) +def test_explicit_incoming_history_gradient( + tmp_path: Path, robot: str, device: str +) -> None: + if device.startswith("cuda") and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + checkpoint = _export(tmp_path, robot) + with _scene(robot, checkpoint, device=device) as (scene, backend): + first = scene.create_functional_state() + force = torch.zeros_like(first.physical["joint_qd"]) + warm = scene.functional_step(first, {"joint_f": force}, publish=False) + history = { + name: value.detach().clone().requires_grad_() + for name, value in warm.history.items() + } + state = FunctionalSolverState( + {name: value.detach() for name, value in warm.physical.items()}, + history, + warm.valid_history_lengths, + ) + result = scene.functional_step(state, {"joint_f": force}, publish=False) + ( + result.physical["joint_q"].square().sum() + + result.physical["body_q"].square().sum() + ).backward() + for value in history.values(): + assert value.grad is not None and torch.isfinite(value.grad).all() + assert value.grad[:, 0].abs().sum() > 0 + assert not value.grad[:, 1:].any() + + +def test_cartpole_fixed_root_orientation_is_bound(tmp_path: Path) -> None: + checkpoint = _export(tmp_path, "cartpole") + + def invert_root(builder): + for joint, parent in enumerate(builder.joint_parent): + if parent == -1: + builder.joint_X_p[joint] = wp.transform( + (0.0, 0.0, 0.0), wp.quat(1.0, 0.0, 0.0, 0.0) + ) + + with pytest.raises(ValueError, match="physics signature"): + with _scene("cartpole", checkpoint, mutate=invert_root): + pass + + +@pytest.mark.parametrize( + "field", + [ + "shape_material_ke", + "shape_material_kd", + "shape_material_restitution", + "shape_material_ka", + "shape_material_mu_torsional", + "shape_gap", + ], +) +def test_ant_robot_contact_properties_are_bound(tmp_path: Path, field: str) -> None: + checkpoint = _export(tmp_path, "ant") + + def change_material(builder): + # Modify every robot equally: distinguish incompatible export from mere + # heterogeneous-batch rejection. Leave the shared plane unchanged. + values = getattr(builder, field) + for shape, body in enumerate(builder.shape_body): + if body >= 0: + values[shape] = float(values[shape]) + 0.123 + + with pytest.raises(ValueError, match="physics signature"): + with _scene("ant", checkpoint, mutate=change_material): + pass + + +@pytest.mark.parametrize("mutation", ["self_collision", "disabled_plane", "topology"]) +def test_ant_contact_policy_and_topology_are_bound( + tmp_path: Path, mutation: str +) -> None: + checkpoint = _export(tmp_path, "ant") + + def change(builder): + if mutation == "self_collision": + builder.shape_collision_filter_pairs.clear() + elif mutation == "disabled_plane": + builder.shape_collision_group[-1] = 0 + else: + roots = [ + joint + for joint, parent in enumerate(builder.joint_parent) + if parent == -1 + ] + for root in roots: + builder.joint_parent[root + 2] = builder.joint_child[root] + + with pytest.raises(ValueError, match="plane-only|physics signature"): + with _scene("ant", checkpoint, mutate=change): + pass + + +@pytest.mark.parametrize("robot", ["cartpole", "ant"]) +def test_complete_signature_matches_public_snapshot_and_replica_count( + robot: str, +) -> None: + with _scene(robot, num_envs=1) as (one, backend): + direct = NeRDBinding(backend.model, robot) + snapshot = NeRDBinding.from_metadata(dict(one.dynamics_metadata()), robot) + assert direct.physics_signature == snapshot.physics_signature + assert ( + direct.metadata(DT)["physics_signature_schema"] + == "nerd-reference-physics-v2" + ) + signature = direct.physics_signature + with _scene(robot, num_envs=2) as (two, backend): + assert NeRDBinding(backend.model, robot).physics_signature == signature + assert ( + NeRDBinding.from_metadata( + dict(two.dynamics_metadata()), robot + ).physics_signature + == signature + ) diff --git a/tests/sim/physics/solvers/test_reference_assets.py b/tests/sim/physics/solvers/test_reference_assets.py new file mode 100644 index 000000000..0b5e12b16 --- /dev/null +++ b/tests/sim/physics/solvers/test_reference_assets.py @@ -0,0 +1,182 @@ +# ---------------------------------------------------------------------------- +# 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 checks against the official Newton asset importer and Spawn lowering.""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import newton +import numpy as np +import pytest +import warp as wp + +from embodichain.lab.sim.physics.solvers._reference_assets import ( + _import_reference, + make_reference_articulation, + reference_asset_hash, + reference_joint_names, + reference_profile, +) + + +def _lower_descriptor(descriptor): + # Exercise production Spawn translation without creating a renderer/World. + from dexsim.spawn.adapters.newton_articulation_adapter import ( + _translate_explicit_articulation, + ) + from dexsim.engine.newton_physics.articulation.articulation import ( + NewtonArticulation, + ) + from dexsim.engine.newton_physics._runtime import create_model_builder + + builder = create_model_builder() + specs, joints = _translate_explicit_articulation( + MagicMock(), None, descriptor, {}, "mujoco_warp", reuse_render=True + ) + handle = NewtonArticulation( + builder, + None, + descriptor.name, + descriptor.root_link_name, + specs, + joints, + build_device="cpu", + ) + handle.finalize( + descriptor.pose, fixed_base=descriptor.fixed_base, disable_self_collision=True + ) + return builder + + +@pytest.mark.parametrize("robot", ["cartpole", "ant"]) +def test_spawn_lowering_preserves_reference_physical_properties(robot): + reference = _import_reference(robot) + descriptor = make_reference_articulation(robot) + lowered = _lower_descriptor(descriptor) + lower_bodies = {name: index for index, name in enumerate(lowered.body_label)} + for source_index, label in enumerate(reference.body_label): + target_index = lower_bodies[label.rsplit("/", 1)[-1]] + for field in ("body_mass", "body_com", "body_inertia"): + np.testing.assert_allclose( + np.asarray(getattr(lowered, field)[target_index]), + np.asarray(getattr(reference, field)[source_index]), + rtol=1e-6, + atol=1e-8, + err_msg=f"{robot} {label} {field}", + ) + # Compare physical shapes; the source also imports non-colliding visuals. + active_shapes = [ + i + for i, flag in enumerate(reference.shape_flags) + if int(flag) + & int(newton.ShapeFlags.COLLIDE_SHAPES | newton.ShapeFlags.COLLIDE_PARTICLES) + ] + # Shape order and exact primitive scales/materials survive public descriptors. + for field in ( + "shape_type", + "shape_scale", + "shape_transform", + "shape_margin", + "shape_gap", + "shape_material_ke", + "shape_material_kd", + "shape_material_kf", + "shape_material_mu", + "shape_material_mu_torsional", + "shape_material_mu_rolling", + ): + np.testing.assert_allclose( + np.asarray(getattr(lowered, field)), + np.asarray(getattr(reference, field))[active_shapes], + rtol=1e-6, + atol=1e-7, + err_msg=f"{robot} {field}", + ) + lower_joints = {name: index for index, name in enumerate(lowered.joint_label)} + for source_joint, label in enumerate(reference.joint_label): + name = label.rsplit("/", 1)[-1] + if name not in reference_joint_names(robot): + continue + source_dof = reference.joint_qd_start[source_joint] + target_dof = lowered.joint_qd_start[lower_joints[name]] + for field in ( + "joint_armature", + "joint_damping", + "joint_friction", + "joint_limit_ke", + "joint_limit_kd", + "joint_limit_lower", + "joint_limit_upper", + "joint_effort_limit", + "joint_velocity_limit", + "joint_target_ke", + "joint_target_kd", + ): + np.testing.assert_allclose( + getattr(lowered, field)[target_dof], + getattr(reference, field)[source_dof], + rtol=1e-6, + atol=1e-7, + err_msg=f"{robot} {name} {field}", + ) + np.testing.assert_allclose( + np.asarray(lowered.joint_axis[target_dof]), + np.asarray(reference.joint_axis[source_dof]), + ) + assert len(lowered.joint_q) == reference_profile(robot)["q_dim"] + assert len(lowered.joint_qd) == reference_profile(robot)["qd_dim"] + + +@pytest.mark.parametrize("robot", ["cartpole", "ant"]) +def test_cpu_fk_matches_reference_importer(robot): + reference = _import_reference(robot) + lowered = _lower_descriptor(make_reference_articulation(robot)) + q = reference_profile(robot)["default_q"] + q[-1] += 0.17 + reference.joint_q[:] = q + lowered.joint_q[:] = q + models = [builder.finalize(device="cpu") for builder in (reference, lowered)] + states = [model.state() for model in models] + for model, state in zip(models, states): + newton.eval_fk(model, model.joint_q, model.joint_qd, state) + lower_indices = {name: index for index, name in enumerate(lowered.body_label)} + for index, name in enumerate(reference.body_label): + actual = states[1].body_q.numpy()[lower_indices[name.rsplit("/", 1)[-1]]] + expected = states[0].body_q.numpy()[index] + np.testing.assert_allclose(actual[:3], expected[:3], rtol=1e-5, atol=1e-6) + assert abs(float(np.dot(actual[3:], expected[3:]))) == pytest.approx( + 1.0, abs=1e-6 + ) + + +def test_reference_hash_and_profiles(): + assert ( + reference_asset_hash("ant") + == "23e9a720465a71fde3fd9fd2d5a18f8c11ec642f3dca8921c66cb4e169421fa7" + ) + assert ( + reference_asset_hash("cartpole") + == "f31276531c03050a3ac3a3e24443899f38a4424586db17140978383433dec237" + ) + first = make_reference_articulation("ant") + second = make_reference_articulation("ant") + first.links[0].rigid_body.mass = 99.0 + assert second.links[0].rigid_body.mass == pytest.approx(0.48388024694610887) + assert reference_profile("ant")["effort_limits"] == [400.0] * 8 + with pytest.raises(ValueError, match="Unsupported"): + make_reference_articulation("unknown") diff --git a/tests/sim/test_sim_manager.py b/tests/sim/test_sim_manager.py index a369772a3..d1219bec4 100644 --- a/tests/sim/test_sim_manager.py +++ b/tests/sim/test_sim_manager.py @@ -66,6 +66,28 @@ (0.0, 0.0, 1.0), ) + +@pytest.mark.parametrize("mask_size", [2, 3]) +def test_empty_functional_reset_preserves_graph_without_native_publication(mask_size): + manager = SimpleNamespace( + prepare=MagicMock(), + is_newton_backend=True, + device=torch.device("cpu"), + num_envs=2, + spawn_result=MagicMock(), + ) + initial = torch.tensor([1.0, 2.0], requires_grad=True) + state = SimpleNamespace(physical={"joint_q": initial * 2}) + result = SimulationManager.reset_functional_state( + manager, state, torch.zeros(mask_size, dtype=torch.bool) + ) + assert result is state + result.physical["joint_q"].sum().backward() + torch.testing.assert_close(initial.grad, torch.full_like(initial, 2)) + manager.spawn_result.create_functional_state.assert_not_called() + manager.spawn_result.reset_functional_state.assert_not_called() + + pytestmark = pytest.mark.no_sim @@ -2036,6 +2058,7 @@ def test_robot_mounted_camera_backend_lifecycle(tmp_path: Path, backend: str) -> def test_default_plane_authors_repeated_uv_before_spawn() -> None: sim = object.__new__(SimulationManager) + sim.sim_config = SimulationManagerCfg() sim._spawn_scene = MagicMock() sim._spawn_scene.handles.return_value = [] sim._spawn_default_plane_material = object() diff --git a/tests/sim/test_sim_manager_cfg.py b/tests/sim/test_sim_manager_cfg.py index 6adf8c294..d581acf70 100644 --- a/tests/sim/test_sim_manager_cfg.py +++ b/tests/sim/test_sim_manager_cfg.py @@ -41,6 +41,58 @@ def test_simulation_manager_cfg_uses_default_physics_cfg() -> None: assert type(cfg.physics_cfg) is DefaultPhysicsCfg +def test_contact_free_manager_explicitly_disables_physical_ground() -> None: + config = SimulationManagerCfg(enable_ground_plane=False) + assert config.to_dict()["enable_ground_plane"] is False + manager = object.__new__(SimulationManager) + manager.sim_config = config + manager._spawn_scene = MagicMock() + manager._declare_spawn_default_plane() + manager._spawn_scene.declare.assert_not_called() + with pytest.raises(TypeError, match="boolean"): + SimulationManagerCfg(enable_ground_plane="false") + + +def test_registered_newton_solver_preserves_name_and_options() -> None: + from dexsim.engine.newton_physics import RegisteredNewtonSolverCfg + + config = NewtonPhysicsCfg( + solver_cfg={ + "solver_type": "nerd", + "options": {"checkpoint": "/models/example"}, + }, + requires_grad=True, + ) + native = config.to_dexsim_cfg(gpu_id=0) + assert isinstance(native.solver_cfg, RegisteredNewtonSolverCfg) + assert native.solver_cfg.solver_type == "nerd" + assert native.solver_cfg.options == {"checkpoint": "/models/example"} + assert not native.use_cuda_graph + with pytest.raises(ValueError, match="belong in options"): + NewtonPhysicsCfg( + solver_cfg={"solver_type": "nerd", "checkpoint": "model"} + ).to_dexsim_cfg(gpu_id=0) + + +def test_registered_config_cannot_bypass_builtin_gradient_restriction() -> None: + from dexsim.engine.newton_physics import RegisteredNewtonSolverCfg + + config = NewtonPhysicsCfg( + requires_grad=True, + solver_cfg=RegisteredNewtonSolverCfg( + solver_type="mjwarp", options={"use_mujoco_cpu": True} + ), + ) + with pytest.raises(RuntimeError, match="gradient mode"): + config.to_dexsim_cfg(gpu_id=0) + with pytest.raises(ValueError, match="cannot override"): + NewtonPhysicsCfg( + solver_cfg=RegisteredNewtonSolverCfg( + solver_type="mujoco_warp", options={"solver_type": "xpbd"} + ) + ).to_dexsim_cfg(gpu_id=0) + + @pytest.mark.no_sim def test_newton_physics_cfg_owns_backend_device_default() -> None: """The concrete Newton default shadows the generic backend default.""" diff --git a/tests/test_agent_context_map.py b/tests/test_agent_context_map.py index e915a49c8..94934ce95 100644 --- a/tests/test_agent_context_map.py +++ b/tests/test_agent_context_map.py @@ -74,6 +74,7 @@ def test_map_registers_the_supported_context_domains() -> None: "scene-expansion", "scene-services", "differentiable-env", + "neural-dynamics", } assert set(topics) == expected_topic_ids From 0e0915bc3f8c69b8cdd469098263fced2eccda9e Mon Sep 17 00:00:00 2001 From: yuecideng Date: Wed, 21 Oct 2026 11:41:51 +0000 Subject: [PATCH 2/7] fix(dynamics): align official sampling and numerical execution --- .../topics/neural-dynamics/neural-dynamics.md | 76 ++- .../simulation-system/simulation-system.md | 3 + design/neural-dynamics-solvers.md | 66 ++- embodichain/compute/dynamics/codec.py | 249 +++++++-- .../data_pipeline/dynamics/_resident.py | 131 +++++ embodichain/data_pipeline/dynamics/shards.py | 27 +- .../sim/physics/solvers/_reference_assets.py | 73 ++- .../lab/sim/physics/solvers/collection.py | 149 ++++- embodichain/lab/sim/physics/solvers/nerd.py | 4 +- embodichain/lab/sim/sim_manager.py | 39 ++ embodichain/lab/sim/spawn/scene.py | 10 +- embodichain/learning/dynamics/_numerics.py | 121 +++++ .../learning/dynamics/_reference_traces.py | 198 +++++++ .../learning/dynamics/_sample_order.py | 376 +++++++++++++ embodichain/learning/dynamics/comparison.py | 61 ++- embodichain/learning/dynamics/reference.py | 253 ++++++++- .../learning/dynamics/reference_rollout.py | 183 +++++-- embodichain/learning/dynamics/report.py | 84 ++- embodichain/learning/dynamics/trainer.py | 216 ++++++-- embodichain/learning/dynamics/workflow.py | 13 +- .../learning/models/nerd/checkpoint.py | 39 +- embodichain/learning/models/nerd/model.py | 24 +- .../tasks/classic_control/ant/README.md | 33 ++ .../ant/dynamics.generate.warp-reference.yaml | 18 + .../ant/dynamics.train.full.yaml | 2 + .../tasks/classic_control/cartpole/README.md | 33 ++ .../dynamics.generate.warp-reference.yaml | 18 + .../cartpole/dynamics.train.full.yaml | 2 + pyproject.toml | 2 +- .../benchmark/learning/dynamics_sampling.py | 193 +++++++ scripts/experiments/neural_dynamics/README.md | 262 +++++++++ .../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 ++++++++++ .../neural_dynamics/physics_alignment.py | 283 ++++++++++ tests/cli/test_dynamics.py | 16 +- .../dynamics/data/ant-official-frame.json | 64 +++ .../dynamics/data/ant-official-frame.npz | Bin 0 -> 34898 bytes tests/compute/dynamics/test_codec.py | 170 ++++++ tests/data_pipeline/dynamics/test_shards.py | 69 +++ tests/learning/dynamics/test_comparison.py | 20 + tests/learning/dynamics/test_reference.py | 481 ++++++++++++++++ .../dynamics/test_reference_rollout.py | 227 ++++++++ .../dynamics/test_reference_traces.py | 247 +++++++++ tests/learning/dynamics/test_report.py | 141 ++++- tests/learning/dynamics/test_sample_order.py | 328 +++++++++++ .../dynamics/test_small_budget_observation.py | 169 ++++++ tests/learning/dynamics/test_trainer.py | 444 ++++++++++++++- .../learning/models/_nerd_reference_check.py | 63 ++- tests/learning/models/test_nerd.py | 114 +++- .../solvers/data/warp_reference_arrays.json | 467 ++++++++++++++++ tests/sim/physics/solvers/test_collection.py | 78 +++ .../physics/solvers/test_reference_assets.py | 89 ++- tests/sim/spawn/test_scene.py | 15 + tests/sim/test_sim_manager.py | 47 ++ 62 files changed, 9398 insertions(+), 239 deletions(-) create mode 100644 embodichain/data_pipeline/dynamics/_resident.py create mode 100644 embodichain/learning/dynamics/_numerics.py create mode 100644 embodichain/learning/dynamics/_reference_traces.py create mode 100644 embodichain/learning/dynamics/_sample_order.py 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/cartpole/dynamics.generate.warp-reference.yaml create mode 100644 scripts/benchmark/learning/dynamics_sampling.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 scripts/experiments/neural_dynamics/physics_alignment.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/learning/dynamics/test_reference_traces.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/sim/physics/solvers/data/warp_reference_arrays.json diff --git a/agent_context/topics/neural-dynamics/neural-dynamics.md b/agent_context/topics/neural-dynamics/neural-dynamics.md index 5c0d89b60..3ab993bf8 100644 --- a/agent_context/topics/neural-dynamics/neural-dynamics.md +++ b/agent_context/topics/neural-dynamics/neural-dynamics.md @@ -31,6 +31,18 @@ APIs, not native backend mutation. ## Data and state contract +`DynamicsSimulation`, generation and rollout configs accept an explicit +`physics_profile: native|warp_reference`. The default retains modern importer +semantics. `warp_reference` reconstructs measured legacy asset mass/inertia, +body armature, force-mode damping and ordered capsule contact slots before +Spawn materialization; its shared ground has explicit material parameters. +It still runs current Newton and records that teacher identity. Featherstone +teachers use complete public articulation writes at macro boundaries; their +mass-matrix update interval must divide the macro substep count. Old and new +contact algorithms can differ even after parameter matching. The isolated +`scripts/experiments/neural_dynamics/physics_alignment.py` runner preserves +identical canonical inputs and hash-verified traces for that comparison. + NeRD validates supported homogeneous articulations, named joint mapping, trained timestep, gravity, physics signature and contact schema. Reference Cartpole is contact-free: disable the physical default plane at manager creation, not only @@ -40,17 +52,57 @@ are generalized efforts; unsupported base/external forces are rejected. Newton COM velocities and NeRD spatial twists differ in both component ordering and reference point. The codec owns both conversions. Relative targets use the frame at t, not t+1. Training and runtime share the codec and frozen normalizers. +Quaternion arithmetic and scripted body-transform execution are part of FP32 +reference alignment. Training targets and native decoding reuse the current +unnormalized embedding rather than transforming that state twice. CUDA graph +specialization can otherwise change low-order bits within one transition. +Periodic positions are wrapped before differencing, matching the original +trainer; loss uses its literal reciprocal-sqrt and native MSE reduction. +Reference adapters reuse the codec's velocity transport on the model device, +so raw CPU/CUDA residency cannot change preprocessing precision. +For original-data training, `import_legacy_datasets` can explicitly preserve +`states/legacy_velocity` before COM conversion. This unrotated FP32 channel +requires bit-exact transition continuity and supplies both current and next +velocity directly to the codec; native `q`/`qd` and the default import retain +their contracts. Source fidelity does not qualify a model for modern physics. +Compare both implementations in the same floating-point world frame. Comparing +an original Y-up result with its Z-up input rotation also measures the reference's +own finite-precision frame non-invariance. Committed HDF5 shards hold T+1 states and T controls/contact inputs with provenance. Splits are trajectory-level; windows cannot cross episodes or splits. Workers own -independent file handles and bounded caches. Offline training does not initialize -simulation. Teacher/neural deployments separately own physics configuration. +independent file handles and bounded caches. Long single-process learners can +explicitly budget raw CPU/CUDA residency through `resident_cache_bytes`; cache +gathers preserve precision, padding and returned ownership and fail before +exceeding the field budget. Serialized worker datasets drop that cache. Offline +training does not initialize simulation. Teacher/neural deployments separately +own physics configuration. Training checkpoints retain optimizer, scheduler, RNG and committed sample progress. Portable exports contain tensor weights and JSON metadata without executable model objects. Legacy conversion is explicitly trusted and offline. Official comparisons share exact statistics and sample order; overlapping-window statistical weighting must not be confused with unique-transition weighting. +`_sample_order.py` supports checkpointed random-with-replacement streams and +explicit original-DataLoader shuffle descriptors. The latter bind source pass +seeds and full-permutation hashes, preserve epoch-tail dropping, and cache at +most one permutation; they must not be replaced by independent seeded shuffles. +Old descriptors and `.npy` matrices remain readable. Both trainers checkpoint +optimizer/RNG/best-model progress and +bind the preprocessing execution contract. Three real-layout CUDA preprocessing +warmups precede updates and RNG restoration, preserving exact resumed execution. +Sparse `log_interval` controls observation volume without changing updates or +selection. Keep virtualenv executable paths intact: resolving their symlinks +can silently select the base interpreter instead. + +The optional training `numerical_profile="reference"` selects math attention +and ordered FP64 accumulation of FP32 per-parameter gradient norms. Defaults +retain their existing Torch behavior. Checkpoints bind this policy; portable +exports persist the model's attention backend so inference uses the same choice. +Changing policy requires a new fine-tuning run. The original-method reference +adapter accepts only its default policy and rejects unsupported math exports. +An explicit common policy enables controlled comparisons; it does not promise +bitwise equivalence across CPU and GPU execution. ## Differentiable lifecycle @@ -85,6 +137,17 @@ thresholds are frozen in comparison manifests; missing runs remain pending. The acceptance contract is `design/neural-dynamics-solvers.md`; user commands are documented in `docs/source/guides/neural_dynamics.md`. +Equal seeds do not prove identical GPU teacher trajectories. Freeze shared +evaluation corpora and verify their ordered content hashes for paired runs. +Verify held-out seeds against the actual stored trajectory seeds: collection +uses the generation seed plus candidate index, so comparing top-level seeds +alone does not exclude training/evaluation overlap. Check actual canonical +initial states when establishing independent evaluation corpora. +Recorded live PD feedback traces can be reused only with matching policy +identity and action-by-action checks against teacher states. Neural feedback +still uses each neural trajectory's own state; training or random-effort files +cannot stand in for a verified closed-loop teacher. + Historical rollout runs in a fresh process with the original pinned Warp and source checkout. It restores either trainer's portable weights into the same original runtime and checks recorded physical arrays including gravity. This is @@ -93,3 +156,12 @@ Newton. Shared metrics remain simulation-independent so old Warp never imports DexSim/Newton. The controlled matrix's shared sampling/statistics/validation protocol is distinct from literal paper-recipe reproduction; its report labels must retain that distinction. + +Fixed historical teacher data can be generated with the standalone +`scripts/experiments/neural_dynamics/official_evaluation_data.py` in Python3.8 / +Torch2.2 without importing the EmbodiChain package or SDK; its pure coordinate +transport modules are explicitly hashed. Historical `teacher_dataset` replay +validates generation provenance, source/model physics, raw/canonical state +agreement and feedback actions. It retains all requested outcomes while the +neural state evolves independently. Data-generation and model-runtime versions +are recorded separately. diff --git a/agent_context/topics/simulation-system/simulation-system.md b/agent_context/topics/simulation-system/simulation-system.md index 24f807375..b02071e9a 100644 --- a/agent_context/topics/simulation-system/simulation-system.md +++ b/agent_context/topics/simulation-system/simulation-system.md @@ -25,6 +25,9 @@ and readiness for each committed topology revision. DexSim `SceneBuilder` and `Scene` own descriptor revisions, native materialization, replicated arenas, and backend handles. `add_*()` returns a declared facade; `prepare()` binds that same object in place. UIDs identify resources in the manager registry. +Procedural articulation/object descriptors enter through the manager's public +descriptor methods. Spawn owns their replication count: `per_env=False` binds +one global facade instance even in a replicated scene. `BaseEnv` owns Gym reset/step timing; `EmbodiedEnv` translates task composition into manager declarations. Scene replication and collision isolation belong diff --git a/design/neural-dynamics-solvers.md b/design/neural-dynamics-solvers.md index 1553b6120..543c127f9 100644 --- a/design/neural-dynamics-solvers.md +++ b/design/neural-dynamics-solvers.md @@ -11,6 +11,27 @@ contract fixes, refreshed artifacts and remaining requirements. The latest [bounded five-seed comparison](reports/neural-dynamics-bounded-five-seed-20261020.md) and [historical rollout evidence](reports/neural-dynamics-historical-rollout-20261020.md) cover the subsequent implementation work without claiming full qualification. +The [trained-model and clean-environment report](reports/neural-dynamics-trained-comparison-20261020.md) +records the Ant conversion correction, scratch-training comparison and the active +full-budget matrix at that report's snapshot time using official-version data. +The subsequent [CUDA training and full-v2 report](reports/neural-dynamics-cuda-alignment-20261020.md) +records actual optimizer/resume alignment and the frozen common-teacher, +disjoint-seed long-rollout protocol; full accuracy acceptance remains pending. +The [supported reference-evaluation report](reports/neural-dynamics-supported-reference-evaluation-20261020.md) +records the data provenance: both original training and held-out +teacher data are generated in the supported old environment and replayed through +separately recorded model runtimes. +The latest [original-entry small-budget comparison](reports/neural-dynamics-small-budget-20261020.md) +supersedes the full-budget execution plan for the current task. The user requested +four 1,000-update main runs and a workflow check; the forty-cell queues and their +evaluation waiters are stopped with artifacts preserved. That round does not +claim project-wide accuracy qualification or public-model convergence. +The [numerical attribution follow-up](reports/neural-dynamics-numerical-attribution-20261020.md) +adds 1,000-update same-environment parity for Ant, identifies CPU attention and +gradient-norm reduction differences, validates a separate common numerical policy +through the actual migrated trainer, and isolates Ant transport effects during +training. The official default baseline and production numerical defaults remain +unchanged. ## Objective and decisions @@ -109,7 +130,8 @@ boundary. Optimizer updates occur after backward and before a detached next roll Pin reference source: main `370e89d2b8ce68aa3578aa7812859c681b7502bc` for generation/training; Newton preview `476dc67fccc6a99be293bf611f897b1c06f13e30` -for solver integration. Target integration stack is Newton 1.6 / Warp 1.17. +for solver integration. The requested SDK MR !1483 artifact at `b61e920f` runs +Newton 1.7.0 / Warp 1.17.0; the initial integration used Newton 1.6. Record the actual DexSim build, Torch, CUDA and driver identities in every report. Use upstream-matching Cartpole/Ant assets and explicit joint/shape mapping. Bind @@ -173,7 +195,12 @@ actual generalized forces, contact data, trajectory IDs and valid lengths. Record schema/codec versions, units, asset/config hashes, teacher versions and sampling seeds. Use bounded staging/writer buffers and commit complete shards; resume only committed work. Dataset workers open independent HDF5 handles and -read windows lazily. Never retain the complete dataset on GPU or in host memory. +read windows lazily by default. Do not materialize the expanded overlapping-window +dataset or a dense full-budget sampling index matrix. A single-process learner +may explicitly opt into raw CPU/CUDA field residency with a strict byte limit +and bounded host staging; it must fail before exceeding that budget. Repeated +window gathers retain independent ownership, and serialized worker datasets drop +resident caches. Cache placement cannot change preprocessing precision or samples. Default dataset splits are trajectory-level 80/10/10. Windows cannot cross resets, trajectories or splits. Normalization uses training data only and is frozen for @@ -280,6 +307,14 @@ recomputes fixed-policy actions from its own state. Record which reference runtime executed the network; only current-domain tests establish the native Newton/functional adapter behavior. +Historical comparison teacher data must also be generated in the declared +supported original environment. Preserve raw and canonical states/controls +without filtering failures, and reuse immutable traces across models. Verify +generation versions, source/asset/physical arrays, coordinate transport and +teacher feedback. Regenerating reference data under a different framework +version requires a separate explicit scope. The model evaluation runtime is +identified independently. + ### Evaluation and acceptance For Cartpole and Ant evaluate passive motion, random effort and fixed-policy @@ -324,12 +359,21 @@ assert equivalence to paper results from a different teacher or training budget. All items remain pending until backed by authoritative code and executed results. The implementation log must distinguish smoke, numerical-parity, full-budget -quality, performance and regression evidence. The goal is not complete while any -required capability or qualification artifact is unimplemented or unverified. - -Current evidence status (2026-10-20): packages 1–4 and 6 have implementation, -focused runtime/tests and local-artifact evidence in the linked reports. Package 5 -has executed end-to-end flows, numerical parity, 40 bounded training cells and -both-domain short rollouts. Its full-size data, declared full training budget, -2048-trajectory long-rollout and final acceptance gates remain pending. The -controlled comparison is not substituted for literal paper-recipe reproduction. +quality, performance and regression evidence. Project-wide qualification requires +all its artifacts. A user-selected bounded round has its own completion conditions +and does not confer project-wide qualification. + +Current evidence status (2026-10-21): packages 1–4 and 6 have implementation, +focused runtime/tests and scoped artifact evidence. Package 5 includes native +end-to-end smoke, numerical parity, original-entry small-budget controls and +supported-environment historical rollouts. The user stopped the 40-cell +full-budget matrix and dependent waiters; its checkpoints, logs and configurations +remain preserved. The subsequent one-seed-per-robot round uses original supported +Warp data, captured original shuffle order and explicit reference/math numerics. +Cartpole completed 4,995,000 actual updates and passed fixed offline acceptance; +its validation-selected best is update 4,940,000. Ant final-budget acceptance is +still pending in this source snapshot. Historical-domain success does not qualify +current-Newton deployment. Larger multi-seed/long-rollout qualification remains +separate from this round, and neither comparison establishes same-data public +model convergence reproduction. The [current evidence index](reports/neural-dynamics-pr-status-20261020.md) +links the reports and records their scopes. diff --git a/embodichain/compute/dynamics/codec.py b/embodichain/compute/dynamics/codec.py index 6eb353a02..57a7236f0 100644 --- a/embodichain/compute/dynamics/codec.py +++ b/embodichain/compute/dynamics/codec.py @@ -43,23 +43,93 @@ def quat_inverse(q: Tensor) -> Tensor: 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.""" - av, bv = torch.broadcast_tensors(a[..., :3], b[..., :3]) - return torch.cat( - ( - a[..., 3:] * bv + b[..., 3:] * av + torch.linalg.cross(av, bv), - a[..., 3:] * b[..., 3:] - (av * bv).sum(-1, keepdim=True), - ), - -1, + """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 broadcastable unit xyzw quaternions.""" + """Rotate vectors by unit xyzw quaternions in official NeRD FP32 order.""" xyz, v = torch.broadcast_tensors(q[..., :3], v) - uv = torch.linalg.cross(xyz, v) - return v + 2 * (q[..., 3:] * uv + torch.linalg.cross(xyz, uv)) + 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: @@ -71,7 +141,7 @@ def _normalize(q: Tensor) -> Tensor: def _wrap(x: Tensor) -> Tensor: - return torch.remainder(x + torch.pi, 2 * torch.pi) - torch.pi + return x - torch.floor((x + torch.pi) / (2 * torch.pi)) * (2 * torch.pi) class NeRDStateCodec: @@ -161,19 +231,41 @@ def legacy_to_native(self, q: Tensor, velocity: Tensor) -> Tensor: -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]: - inv = quat_inverse(anchor[..., 3:7]) - pos = quat_rotate(inv, q[..., :3] - anchor[..., :3]) - rot = quat_multiply(inv, q[..., 3:7]) - omega = quat_rotate(inv, velocity[..., :3]) - nu = quat_rotate( - inv, - velocity[..., 3:6] - torch.linalg.cross(anchor[..., :3], velocity[..., :3]), + 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, rot, q[..., 7:]), -1), torch.cat( - (omega, nu, velocity[..., 6:]), -1 + 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( @@ -183,13 +275,27 @@ def encode( 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.native_to_legacy(q, qd) + velocity = self._resolve_legacy_velocity(q, qd, legacy_velocity) if self.robot == "cartpole": if effort.shape != qd.shape: raise ValueError( @@ -233,14 +339,57 @@ def encode( "gravity_dir": quat_rotate(inv.squeeze(-2), gravity_dir), } - def targets(self, q: Tensor, qd: Tensor, next_q: Tensor, next_qd: Tensor) -> Tensor: - """Encode relative next state in the frame at t, including quaternion delta.""" - velocity = self.native_to_legacy(q, qd) - next_velocity = self.native_to_legacy(next_q, next_qd) + 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 - q, velocity = self._body(anchor, q, velocity) + 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])) ) @@ -248,18 +397,54 @@ def targets(self, q: Tensor, qd: Tensor, next_q: Tensor, next_qd: Tensor) -> Ten return torch.cat( (diff[..., :3], delta, diff[..., 7:], next_velocity - velocity), -1 ) - return torch.cat((self._wrap_q(next_q - q), 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 + 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.""" + """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(q + prediction[..., :2]), qd + prediction[..., 2:] - bq, bv = self._body(q, q, velocity) + 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]) ) diff --git a/embodichain/data_pipeline/dynamics/_resident.py b/embodichain/data_pipeline/dynamics/_resident.py new file mode 100644 index 000000000..9663efcab --- /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/shards.py b/embodichain/data_pipeline/dynamics/shards.py index aa3d1748a..061778784 100644 --- a/embodichain/data_pipeline/dynamics/shards.py +++ b/embodichain/data_pipeline/dynamics/shards.py @@ -328,6 +328,27 @@ def __init__( 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 @@ -434,6 +455,8 @@ def get_batch( 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] = [] @@ -514,7 +537,8 @@ def get_batch( } def close(self) -> None: - """Close this worker's bounded file-handle cache.""" + """Release resident fields and file handles; returned batches stay owned.""" + self._resident = None for handle in self._handles.values(): handle.close() self._handles.clear() @@ -522,6 +546,7 @@ def close(self) -> None: def __getstate__(self) -> dict[str, Any]: state = self.__dict__.copy() state["_handles"] = OrderedDict() + state["_resident"] = None return state def __del__(self) -> None: diff --git a/embodichain/lab/sim/physics/solvers/_reference_assets.py b/embodichain/lab/sim/physics/solvers/_reference_assets.py index 4bb93bc39..fb1730426 100644 --- a/embodichain/lab/sim/physics/solvers/_reference_assets.py +++ b/embodichain/lab/sim/physics/solvers/_reference_assets.py @@ -27,6 +27,7 @@ import hashlib from pathlib import Path from typing import Any, TYPE_CHECKING +import xml.etree.ElementTree as ET import numpy as np @@ -121,26 +122,48 @@ def reference_profile(robot: str) -> dict[str, Any]: } -def _import_reference(robot: str) -> Any: +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_from_axis_angle(wp.vec3(0.0, 0.0, 1.0), wp.pi * 0.5), + ( + 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 @@ -148,12 +171,48 @@ def _import_reference(robot: str) -> Any: 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( - str(_asset(robot)), + 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 @@ -169,9 +228,13 @@ def _pose(transform: Any) -> np.ndarray: return pose -def make_reference_articulation(robot: str) -> ArticulationDesc: +def make_reference_articulation( + robot: str, physics_profile: str = "native" +) -> ArticulationDesc: """Build an independent public descriptor with imported physical properties. + ``warp_reference`` explicitly reproduces measured old-Warp asset parameters; + it does not select the old simulator or assert identical contact algorithms. Body/joint names retain their asset spelling, with only the importer's path prefix removed. The free root is represented by ``fixed_base=False`` and is created by Spawn. Default nonzero joint positions are in ``reference_profile`` @@ -190,7 +253,7 @@ def make_reference_articulation(robot: str) -> ArticulationDesc: ) from dexsim.engine import JointType - builder = _import_reference(robot) + builder = _import_reference(robot, physics_profile) names = [str(name).rsplit("/", 1)[-1] for name in builder.body_label] if len(names) != len(set(names)): raise ValueError("Reference body names are not unique") diff --git a/embodichain/lab/sim/physics/solvers/collection.py b/embodichain/lab/sim/physics/solvers/collection.py index 7d246392b..727d375c9 100644 --- a/embodichain/lab/sim/physics/solvers/collection.py +++ b/embodichain/lab/sim/physics/solvers/collection.py @@ -118,30 +118,51 @@ def __init__( solver: Mapping[str, Any] | None = None, num_substeps: int = 10, requires_grad: bool = False, + physics_profile: str = "native", ) -> None: for name, value in (("num_envs", num_envs), ("num_substeps", num_substeps)): _positive_integer(name, value) if type(requires_grad) is not bool: raise TypeError("requires_grad must be boolean") + if physics_profile not in {"native", "warp_reference"}: + raise ValueError( + f"Unsupported reference physics profile: {physics_profile}" + ) self.profile = reference_profile(robot) + self.profile["physics_profile"] = physics_profile self.robot, self.device, self.num_envs = robot, torch.device(device), num_envs self.dt = float(self.profile["dt"]) solver_cfg = dict( solver - or { - "solver_type": "mujoco_warp", - "use_mujoco_cpu": self.device.type == "cpu", - "separate_worlds": True, - } + 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, + } + ) ) neural = solver_cfg.get("solver_type") == "nerd" + self._featherstone_teacher = solver_cfg.get("solver_type") == "featherstone" + if self._featherstone_teacher: + interval = solver_cfg.get("update_mass_matrix_interval", 1) + if type(interval) is not int or interval < 1 or num_substeps % interval: + raise ValueError( + "Featherstone teacher mass-matrix interval must divide num_substeps " + "so every reset starts with a fresh mass matrix." + ) if neural and num_substeps != 1: raise ValueError("The NeRD reference solver requires one macro substep.") cfg = SimulationManagerCfg( headless=True, num_envs=num_envs, arena_space=0.0, - enable_ground_plane=robot == "ant", + enable_ground_plane=robot == "ant" and physics_profile == "native", enable_entity_gizmo=False, render_cfg=RenderCfg(renderer="no-render"), startup_summary="off", @@ -152,14 +173,52 @@ def __init__( requires_grad=requires_grad, use_cuda_graph=False, solver_cfg=solver_cfg, - collision_cfg=None if neural else {}, + collision_cfg=( + None + if neural + else ( + {"update_interval": None} + if physics_profile == "warp_reference" + else {} + ) + ), ), ) self.sim = SimulationManager(cfg, requires_native_renderer=False) self.articulation = None try: + if robot == "ant" and physics_profile == "warp_reference": + # The old ModelBuilder's shared plane has different material + # defaults from SimulationManager's ordinary default plane. + from dexsim.spawn import ( + CollisionDesc, + GeometryDesc, + NewtonCollisionDesc, + ObjectDesc, + RigidBodyPhysicsDesc, + ) + + collision = CollisionDesc.from_geometry(GeometryDesc.plane(1000.0)) + collision.newton = NewtonCollisionDesc( + ke=1e5, + kd=1e3, + kf=1e3, + mu=0.5, + margin=1e-5, + gap=0.05, + mu_torsional=0.0, + mu_rolling=0.0, + ) + self.sim.add_rigid_object_descriptor( + ObjectDesc( + name="reference_ground", + collisions=[collision], + physics=RigidBodyPhysicsDesc.static(), + per_env=False, + ) + ) self.articulation = self.sim.add_articulation_descriptor( - make_reference_articulation(robot) + make_reference_articulation(robot, physics_profile) ) self.sim.prepare() self.binding = NeRDBinding.from_metadata( @@ -173,6 +232,7 @@ def __init__( self.manifest["teacher"] = { "solver": solver_cfg, "num_substeps": num_substeps, + "physics_profile": physics_profile, "versions": { name: metadata.version(name) for name in ("newton", "warp-lang", "torch", "dexsim-engine") @@ -212,7 +272,12 @@ def reset( self, q: torch.Tensor | None = None, qd: torch.Tensor | None = None ) -> None: """Set a complete reference batch through public articulation state writes.""" - self.sim.spawn_result.reset() + # Featherstone has no Scene.reset capability in the named SDK. Complete + # articulation state writes are supported and clear both native buffers + # and inputs below. Calls occur at macro boundaries; the mass-matrix + # interval divides each macro step for the reference teacher. + if not self._featherstone_teacher: + self.sim.spawn_result.reset() if q is None: q = ( torch.tensor(self.profile["default_q"], device=self.device) @@ -359,6 +424,7 @@ class DynamicsGenerationCfg: initial_states: str = "sample" max_attempts: int | None = None effort_scale: float = 1.0 + physics_profile: str = "native" def __post_init__(self) -> None: for name in ( @@ -388,6 +454,8 @@ def __post_init__(self) -> None: 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 ): @@ -427,6 +495,7 @@ def generate_dataset(config: Mapping[str, Any]) -> dict[str, Any]: num_envs=cfg.num_envs, solver=cfg.teacher_solver, num_substeps=cfg.num_substeps, + physics_profile=cfg.physics_profile, ) as host: manifest = dict(host.manifest) manifest["generation"] = { @@ -617,6 +686,7 @@ def __init__( self.target = torch.tensor( reference_profile(robot)["default_q"][7:], device=device ) + self.identity["target"] = self.target.detach().cpu().tolist() elif kind == "actor_only": if set(self.config) != {"kind", "config", "checkpoint"}: raise ValueError( @@ -736,6 +806,7 @@ def _evaluation_teacher( num_envs=batch, solver=config.get("teacher_solver"), num_substeps=config.get("num_substeps", 10), + physics_profile=config.get("physics_profile", "native"), ) as teacher: manifest["metadata"] = teacher.manifest manifest["evaluation"] = { @@ -858,6 +929,33 @@ def _source_motion(manifest: Mapping[str, Any]) -> str | None: return None +def _verify_feedback_trace(manifest: Mapping[str, Any], policy: _FixedPolicy) -> None: + """Require an actual evaluation trace bound to the declared PD feedback.""" + if policy.identity["kind"] != "pd": + raise ValueError("Cached fixed-policy evaluation currently supports PD only") + evaluation = manifest.get("evaluation", {}) + if ( + manifest.get("trace_type") != "evaluation_including_failures" + or evaluation.get("motion_mode") != "fixed_policy" + or evaluation.get("policy") != policy.identity + or _teacher_origin(manifest["metadata"])["kind"] != "current_newton" + ): + raise ValueError("Teacher trace does not verify the declared feedback policy") + + +@torch.no_grad() +def _verify_feedback_actions( + policy: _FixedPolicy, q: torch.Tensor, qd: torch.Tensor, efforts: torch.Tensor +) -> None: + """Check all recorded actions on device and synchronize once per batch.""" + matched = torch.ones((), dtype=torch.bool, device=q.device) + for step in range(efforts.shape[1]): + expected = policy(q[:, step].contiguous(), qd[:, step].contiguous()) + matched = matched & (expected == efforts[:, step]).all() + if not bool(matched): + raise ValueError("Recorded teacher actions differ from its own-state feedback") + + def _teacher_batches( directory: Path, manifest: Mapping[str, Any], @@ -905,14 +1003,16 @@ def evaluate_rollouts(config: Mapping[str, Any]) -> dict[str, Any]: Fixed-policy actions are recomputed from each simulator's own state. Optional immutable teacher_dataset replays canonical test trajectories for open-loop - comparisons; it cannot establish closed-loop policy behavior for that teacher. + comparisons. Recorded live PD feedback traces support closed-loop evaluation + after their policy identity and all teacher actions have been verified. Metrics synchronize once per batch, never once per metric or simulation step. Args: config: Robot, model export, output path and evaluation settings. Optional teacher_dataset selects immutable canonical traces; split defaults to - test. Fixed-policy mode instead needs a live current Newton teacher - and a policy mapping with kind pd or actor_only. Policy artifact paths + test. Fixed-policy mode needs live current Newton feedback, or a + recorded live PD trace with verified policy/actions. A live teacher + supports kind pd or actor_only. Policy artifact paths are already resolved by the CLI. num_envs is the actual replica count for both teacher generation and neural replay. @@ -937,6 +1037,7 @@ def evaluate_rollouts(config: Mapping[str, Any]) -> dict[str, Any]: "policy", "teacher_dataset", "split", + "physics_profile", } if set(config) - allowed: raise ValueError(f"Unknown rollout settings: {sorted(set(config)-allowed)}") @@ -972,10 +1073,6 @@ def evaluate_rollouts(config: Mapping[str, Any]) -> dict[str, Any]: raise ValueError( "fixed_policy mode requires an explicit policy configuration only in that mode" ) - if motion == "fixed_policy" and config.get("teacher_dataset") is not None: - raise ValueError( - "Offline teacher_dataset cannot establish closed-loop fixed-policy behavior" - ) if motion == "fixed_policy" and scale != 1.0: raise ValueError( "effort_scale applies to random_effort; configure policy gains for fixed_policy" @@ -1026,7 +1123,17 @@ def evaluate_rollouts(config: Mapping[str, Any]) -> dict[str, Any]: teacher_origin = _teacher_origin(dataset["metadata"]) source_motion = _source_motion(dataset) if config.get("teacher_dataset") is not None: - if source_motion == "fixed_policy": + if motion == "fixed_policy": + assert policy is not None + _verify_feedback_trace(dataset, policy) + exported = json.loads((model / "manifest.json").read_text()) + if dataset["metadata"].get("physics_signature") != exported.get( + "physics_signature" + ): + raise ValueError( + "Feedback teacher trace/model physics signatures differ" + ) + elif source_motion == "fixed_policy": raise ValueError( "Fixed-policy traces cannot be labeled passive or random_effort replay" ) @@ -1059,6 +1166,7 @@ def evaluate_rollouts(config: Mapping[str, Any]) -> dict[str, Any]: num_envs=batch, solver={"solver_type": "nerd", "options": {"checkpoint": str(model)}}, num_substeps=1, + physics_profile=config.get("physics_profile", "native"), ) as neural, h5py.File(evidence_path, "w") as evidence, rows_path.open("w") as row_stream, @@ -1073,6 +1181,8 @@ def evaluate_rollouts(config: Mapping[str, Any]) -> dict[str, Any]: _teacher_batches(teacher_dir, dataset, records, batch, horizon, device) ): n = len(selected) + if policy is not None and config.get("teacher_dataset") is not None: + _verify_feedback_actions(policy, target_q, target_qd, efforts) if motion == "passive" and any( record["has_nonzero_effort"] for record in selected ): @@ -1299,6 +1409,11 @@ def evaluate_rollouts(config: Mapping[str, Any]) -> dict[str, Any]: "teacher_metadata": dataset["metadata"], "teacher_origin": teacher_origin, "teacher_source_motion_mode": source_motion, + "teacher_feedback_verification": ( + "all_selected_actions_verified_against_teacher_states" + if policy is not None and config.get("teacher_dataset") is not None + else "live_feedback_execution" if policy is not None else None + ), "selected_teacher_sha256": selected_teacher_digest.hexdigest(), "selected_teacher_hash_schema": "ordered-teacher-content-v1", "teacher_split": split, diff --git a/embodichain/lab/sim/physics/solvers/nerd.py b/embodichain/lab/sim/physics/solvers/nerd.py index df0c1ceca..e262b845b 100644 --- a/embodichain/lab/sim/physics/solvers/nerd.py +++ b/embodichain/lab/sim/physics/solvers/nerd.py @@ -224,7 +224,9 @@ def functional_step( prediction = output.gather( 1, (lengths - 1)[:, None, None].expand(-1, 1, output.shape[-1]) ).squeeze(1) - next_q, next_qd = self.codec.decode(q, qd, prediction) + next_q, next_qd = self.codec.decode( + q, qd, prediction, current_embedding=features["states_embedding"] + ) physical = dict(state.physical) physical["joint_q"] = state.physical["joint_q"].scatter( 0, self.binding.q_indices.flatten(), next_q.flatten() diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index 0f2be661f..1e1023e82 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -137,6 +137,7 @@ def _is_usd_path(path: object | None) -> bool: ) if TYPE_CHECKING: + from dexsim.spawn import ArticulationDesc, ObjectDesc from dexsim.engine import PhysicsScene from dexsim.scene import Scene from dexsim.interaction import EntityGizmoConfig, EntityGizmoManipulator @@ -2694,6 +2695,44 @@ def add_rigid_object( """ return self._declare_rigid_object(cfg, replace=False) + def add_rigid_object_descriptor(self, descriptor: ObjectDesc) -> RigidObject: + """Declare procedural rigid geometry through the existing Spawn owner. + + Args: + descriptor: Public DexSim object descriptor owning geometry, physical + parameters and per-environment or global replication intent. + + Returns: + Rigid object facade bound by :meth:`prepare`. + """ + from dexsim.spawn import ObjectDesc + + if not isinstance(descriptor, ObjectDesc) or descriptor.physics is None: + raise TypeError("descriptor must be a physical DexSim ObjectDesc.") + uid = descriptor.name + if uid in self.asset_uids: + raise ValueError(f"Asset {uid!r} already exists.") + body_type = { + ActorType.DYNAMIC: "dynamic", + ActorType.KINEMATIC: "kinematic", + ActorType.STATIC: "static", + }[descriptor.physics.actor_type] + cfg = RigidObjectCfg( + uid=uid, + init_local_pose=descriptor.pose.copy(), + body_type=body_type, + body_scale=tuple(float(value) for value in descriptor.body_scale), + asset_physics_mode="preserve", + ) + facade = RigidObject(cfg=cfg, device=self.device) + was_materialized = self.spawn_result is not None + self._spawn_scene.declare("rigid_object", uid, descriptor, facade=facade) + self._rigid_objects[uid] = facade + self.notify_visualization_topology_changed() + if was_materialized: + self.prepare() + return facade + def replace_rigid_object(self, cfg: RigidObjectCfg) -> RigidObject: """Replace a rigid object with the same UID in one preparation. diff --git a/embodichain/lab/sim/spawn/scene.py b/embodichain/lab/sim/spawn/scene.py index 73df30664..52de7e863 100644 --- a/embodichain/lab/sim/spawn/scene.py +++ b/embodichain/lab/sim/spawn/scene.py @@ -87,7 +87,7 @@ def declare( """Add a descriptor and associate it with an EmbodiChain facade.""" if uid in self._assets: raise ValueError(f"Spawn asset uid is already declared: {uid!r}.") - self._initialize_facade_declaration(facade) + self._initialize_facade_declaration(facade, descriptor) declaration = _AssetDeclaration( kind=kind, descriptor=descriptor, @@ -165,7 +165,7 @@ def track( """Track a descriptor that was already added to ``SceneBuilder``.""" if uid in self._assets: raise ValueError(f"Spawn asset uid is already declared: {uid!r}.") - self._initialize_facade_declaration(facade) + self._initialize_facade_declaration(facade, descriptor) declaration = _AssetDeclaration(kind, descriptor, facade) self._assets[uid] = declaration handles = self.handles(uid) @@ -279,7 +279,9 @@ def handles(self, uid: str) -> tuple[Any, ...]: return () return tuple(result.handles[path] for path in paths) - def _initialize_facade_declaration(self, facade: Any | None) -> None: + def _initialize_facade_declaration( + self, facade: Any | None, descriptor: Any + ) -> None: """Give a Spawn facade the instance count owned by this scene. The duck-typed fallback keeps ``SpawnScene`` usable by lightweight @@ -290,7 +292,7 @@ def _initialize_facade_declaration(self, facade: Any | None) -> None: return initialize = getattr(facade, "_initialize_spawn_declaration", None) if initialize is not None: - initialize(self._num_envs) + initialize(self._num_envs if getattr(descriptor, "per_env", True) else 1) def _resolve_articulation_source(self, descriptor: Any) -> None: """Resolve one descriptor through the available DexSim boundary.""" diff --git a/embodichain/learning/dynamics/_numerics.py b/embodichain/learning/dynamics/_numerics.py new file mode 100644 index 000000000..85d127da0 --- /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 000000000..def714dcf --- /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/_sample_order.py b/embodichain/learning/dynamics/_sample_order.py new file mode 100644 index 000000000..593ae1b98 --- /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 index 5ffc2e78f..74b5f5414 100644 --- a/embodichain/learning/dynamics/comparison.py +++ b/embodichain/learning/dynamics/comparison.py @@ -64,6 +64,7 @@ def prepare_comparison(config: dict[str, Any]) -> Path: save_statistics, load_statistics, ) + from ._sample_order import _write_checkpointed_order allowed = { "output_dir", @@ -100,7 +101,9 @@ def prepare_comparison(config: dict[str, Any]) -> Path: if any(root.iterdir()): raise FileExistsError("Comparison output must be a new empty directory") source = reference_identity(config["source"]) - python = str(Path(config.get("python", sys.executable)).resolve()) + # 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") @@ -163,23 +166,14 @@ def prepare_comparison(config: dict[str, Any]) -> Path: for seed in seeds: seed_dir = teacher_dir / f"seed-{seed}" seed_dir.mkdir() - order_path = seed_dir / "order.npy" - order = np.lib.format.open_memmap( - order_path, - mode="w+", - dtype=np.int64, - shape=(cfg.total_steps, cfg.batch_size), + 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, ) - generator = torch.Generator(device="cpu").manual_seed(seed) - for start in range(0, cfg.total_steps, 1024): - end = min(cfg.total_steps, start + 1024) - # rand_int produces the same flattened stream as consecutive - # minibatches; one generator is shared by both trainers. - order[start:end] = torch.randint( - len(dataset), (end - start, cfg.batch_size), generator=generator - ).numpy() - order.flush() - del order + 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) @@ -255,6 +249,7 @@ def prepare_comparison(config: dict[str, Any]) -> Path: "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, @@ -300,7 +295,12 @@ def _read_plan(path: str | Path) -> tuple[Path, dict[str, Any]]: return root, plan -def run_candidate(spec: dict[str, Any]) -> dict[str, Any]: +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 ( @@ -310,6 +310,7 @@ def run_candidate(spec: dict[str, Any]) -> dict[str, Any]: 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"] @@ -323,15 +324,29 @@ def run_candidate(spec: dict[str, Any]) -> dict[str, Any]: ): 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), - "fine_tune": spec["initial_model"], + "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( @@ -476,13 +491,19 @@ 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())) + 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()))) diff --git a/embodichain/learning/dynamics/reference.py b/embodichain/learning/dynamics/reference.py index d49943a5d..6cfdff3db 100644 --- a/embodichain/learning/dynamics/reference.py +++ b/embodichain/learning/dynamics/reference.py @@ -87,26 +87,39 @@ def reference_identity(source: str | Path) -> dict[str, Any]: def _legacy_states(q: Any, qd: Any, codec: Any) -> Any: import torch - from embodichain.compute.dynamics import quat_rotate - if codec.robot == "cartpole": - return torch.cat((q, qd), -1) - com = q[..., :3] + quat_rotate(q[..., 3:7], q.new_tensor(codec.base_com)) - omega = qd[..., 3:6] - origin_velocity = qd[..., :3] - torch.linalg.cross(omega, com) - return torch.cat((q, omega, origin_velocity, qd[..., 6:]), -1) + # 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) -> dict[str, Any]: - """Transport native states to official world-frame arrays, before encoding.""" +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, qd = batch["states"]["q"].float(), batch["states"]["qd"].float() - transitions = {key: value.float() for key, value in batch["transitions"].items()} - states = _legacy_states(q, qd, codec) + 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] @@ -272,7 +285,11 @@ def _restore_official_model( def run_reference( - spec: dict[str, Any], *, trust_external_code: bool = False + 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. @@ -282,6 +299,17 @@ def run_reference( """ 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 @@ -291,13 +319,25 @@ def run_reference( 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 + 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") @@ -317,11 +357,40 @@ def run_reference( order_path = Path(spec["sample_order"]) if _hash(order_path) != spec["sample_order_sha256"]: raise ValueError("Reference sample order checksum mismatch") - order = np.load(order_path, mmap_mode="r", allow_pickle=False) + 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) + raw = _official_batch(sample, codec, device=device) integrator = _reference_integrator( codec.robot, raw["contact_depths"].shape[-1], device ) @@ -343,6 +412,8 @@ def run_reference( 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 @@ -350,23 +421,69 @@ def run_reference( 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: - with (destination / "metrics.jsonl").open("w") as log: - for step in range(cfg.total_steps): + 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) + 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, @@ -400,8 +517,25 @@ def transfer() -> None: key: value.detach().cpu().clone() for key, value in common.state_dict().items() } - log.write(json.dumps(record, allow_nan=False) + "\n") - log.flush() + 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) @@ -448,6 +582,7 @@ def export_legacy_dataset( 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 @@ -514,7 +649,7 @@ def export_legacy_dataset( "path": str(output.resolve()), "sha256": _hash(output), "layout": "T,B,features", - "native_to_legacy_qd": True, + "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)) @@ -529,6 +664,7 @@ def import_legacy_datasets( 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. @@ -537,6 +673,22 @@ def import_legacy_datasets( 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 @@ -552,6 +704,10 @@ def import_legacy_datasets( 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 @@ -577,6 +733,25 @@ def import_legacy_datasets( "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 @@ -603,12 +778,22 @@ def import_legacy_datasets( 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 and not np.allclose( - current[1:], following[:-1], atol=1e-6, rtol=1e-5 - ): - raise ValueError( - "Legacy trajectory crosses an undeclared reset/discontinuity" - ) + 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) ) @@ -632,6 +817,8 @@ def import_legacy_datasets( -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: @@ -1032,6 +1219,7 @@ def generate_official_dataset( if spec.get("repair_uninitialized_gravity", False) else None ), + preserve_legacy_velocity=spec.get("preserve_legacy_velocity", False), ) result.update( status="completed", @@ -1055,6 +1243,8 @@ 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", @@ -1062,8 +1252,13 @@ def main() -> None: ) 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 + json.loads(args.spec.read_text()), + trust_external_code=args.trust_external_code, + **extra, ) print(json.dumps(result, indent=2)) diff --git a/embodichain/learning/dynamics/reference_rollout.py b/embodichain/learning/dynamics/reference_rollout.py index bc27df9bc..3226dd71f 100644 --- a/embodichain/learning/dynamics/reference_rollout.py +++ b/embodichain/learning/dynamics/reference_rollout.py @@ -206,8 +206,14 @@ def _verify_historical_physics( ) 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", manifest.get("sampling", {}).get("batch_size") + "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: @@ -375,6 +381,8 @@ def evaluate_reference_rollouts( 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. @@ -401,6 +409,7 @@ def evaluate_reference_rollouts( "warp_version", "cpu_threads", "verify_portable_steps", + "teacher_dataset", } if set(spec) - allowed: raise ValueError( @@ -576,6 +585,10 @@ def init_rnn(self, batch_size): 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: @@ -632,55 +645,132 @@ def init_rnn(self, batch_size): 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 in range(0, count, batch): + for start, recorded in zip(range(0, count, batch), batches): n = min(batch, count - start) - 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()) + 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()] + 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) + 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 else policy(tq, tqd) + 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()) - teacher_raw.append( - teacher.step_with_joint_act( - tf, env_mode="ground-truth" - ).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() ) @@ -707,9 +797,17 @@ def init_rnn(self, batch_size): ).max() ), ) - traw = torch.stack(teacher_raw, 1)[:n] + 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) + 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] @@ -919,6 +1017,29 @@ def init_rnn(self, batch_size): "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", diff --git a/embodichain/learning/dynamics/report.py b/embodichain/learning/dynamics/report.py index 0f2d661a9..406fe3378 100644 --- a/embodichain/learning/dynamics/report.py +++ b/embodichain/learning/dynamics/report.py @@ -30,7 +30,10 @@ __all__ = ["write_comparison_report", "annotate_rollout_evidence"] -_PARITY_REFERENCE = "476dc67fccc6a99be293bf611f897b1c06f13e30" +_PARITY_REFERENCES = { + OFFICIAL_COMMIT: "official_warp_main", + "476dc67fccc6a99be293bf611f897b1c06f13e30": "newton_preview", +} _HORIZONS = (1, 10, 100, 1000) @@ -93,7 +96,8 @@ def _numerical_evidence(path: str | Path, robot: str) -> dict[str, Any]: } identity_matches = ( str(result["robot"]).lower() == robot - and result["reference_commit"] == _PARITY_REFERENCE + 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" @@ -133,7 +137,13 @@ def _numerical_evidence(path: str | Path, robot: str) -> dict[str, Any]: "pending", "Residuals need per-element relative-tolerance evidence", ) - return {**result, "checks": checks, "status": status, "reason": reason} + return { + **result, + "reference_kind": _PARITY_REFERENCES[result["reference_commit"]], + "checks": checks, + "status": status, + "reason": reason, + } def _rollout_evidence( @@ -314,6 +324,47 @@ def _matched_rollouts(left: dict[str, Any], right: dict[str, Any]) -> bool: ) +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": @@ -335,9 +386,14 @@ def _dataset_evidence(teacher: str, declaration: dict[str, Any]) -> dict[str, An verified = ( origin.get("kind") == teacher and origin.get("source_commit") == OFFICIAL_COMMIT - and metadata.get("source", {}).get("commit") == OFFICIAL_COMMIT and bool(metadata.get("source_physical_arrays")) - and bool(metadata.get("source_raw_hdf5")) + 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", {}) @@ -454,16 +510,30 @@ def _full_profile_issues(plan: dict[str, Any]) -> list[str]: 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 - 1) // max_points) + 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) - if record["step"] % stride == 0 or record["step"] == 1: + 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 diff --git a/embodichain/learning/dynamics/trainer.py b/embodichain/learning/dynamics/trainer.py index 48f7c18dd..e31c13ec5 100644 --- a/embodichain/learning/dynamics/trainer.py +++ b/embodichain/learning/dynamics/trainer.py @@ -32,6 +32,15 @@ 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", @@ -53,6 +62,10 @@ class DynamicsTrainCfg: 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 @@ -67,6 +80,9 @@ class DynamicsTrainCfg: 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( @@ -88,10 +104,27 @@ def encode_batch( 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") + 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:] + ), ) - target = codec.targets(q[:, :-1], qd[:, :-1], q[:, 1:], qd[:, 1:]) return inputs, target, batch["valid_mask"].to(device) @@ -103,14 +136,16 @@ def normalized_mse( ) -> Tensor: """Official MSE over normalized physical targets, excluding padded steps.""" if variance is None: - error = prediction - target + scaled_prediction, scaled_target = prediction, target else: - weights = torch.rsqrt(variance + 1e-5) + weights = 1.0 / torch.sqrt(variance + 1e-5) # Preserve official arithmetic ordering for numerical parity. - error = prediction * weights - target * weights - squared = error.square() + scaled_prediction, scaled_target = prediction * weights, target * weights if mask is None: - return squared.mean() + 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) @@ -207,6 +242,35 @@ def _atomic_save(path: Path, value: Any) -> None: 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 @@ -250,7 +314,14 @@ def evaluate_model( raise FloatingPointError("Non-finite evaluation prediction") n = int(mask.sum()) accumulated["normalized_mse"] += ( - float(normalized_mse(prediction, target, _model_variance(model), mask)) + float( + normalized_mse( + prediction, + target, + _model_variance(model), + mask if dataset.include_prefix else None, + ) + ) * n ) accumulated["target_mse"] += ( @@ -324,6 +395,7 @@ def __init__( sample_order: Tensor | str | Path | None = None, ) -> None: cfg.validate() + _validate_numerical_profile(cfg.numerical_profile) if ( min( cfg.total_steps, @@ -331,6 +403,7 @@ def __init__( cfg.validation_interval, cfg.checkpoint_interval, cfg.lr_update_interval, + cfg.log_interval, ) < 1 ): @@ -339,6 +412,8 @@ def __init__( 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 @@ -352,6 +427,8 @@ def __init__( 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) @@ -377,6 +454,26 @@ def __init__( ) 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) @@ -386,7 +483,7 @@ def __init__( def _prepare_order( self, source: Tensor | str | Path | None - ) -> Tensor | np.ndarray | None: + ) -> Tensor | np.ndarray | _CheckpointedOrder | _OfficialShuffleOrder | None: if source is None: return None if isinstance(source, (str, Path)): @@ -395,9 +492,11 @@ def _prepare_order( with path.open("rb") as stream: for block in iter(lambda: stream.read(1024 * 1024), b""): digest.update(block) - 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") + 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(), @@ -448,33 +547,53 @@ def step(self) -> dict[str, float | int]: lr = self._lr() for group in self.optimizer.param_groups: group["lr"] = lr - prediction = self.model(inputs) - loss = normalized_mse(prediction, target, _model_variance(self.model), mask) - if not torch.isfinite(loss): - raise FloatingPointError("Non-finite dynamics training loss") - loss.backward() - norm = torch.nn.utils.clip_grad_norm_( - self.model.parameters(), self.cfg.max_grad_norm, error_if_nonfinite=True - ) + 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()), + "loss": float(loss.detach()) if observed is None else observed[0], "learning_rate": lr, - "grad_norm": float(norm), + "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.""" - numpy_state = np.random.get_state() 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}, @@ -487,22 +606,7 @@ def save(self, path: str | Path) -> None: "include_prefix": self.train_dataset.include_prefix, }, "statistics": self.statistics, - "rng": { - "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(self.cfg.device).type == "cuda" - else [] - ), - }, + "rng": _capture_rng_state(self.cfg.device), } _atomic_save(Path(path), state) @@ -511,6 +615,15 @@ def resume(self, path: str | Path) -> None: 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()), @@ -524,7 +637,12 @@ def resume(self, path: str | Path) -> None: }, ), ): - if state[key] != expected: + 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" ) @@ -553,14 +671,7 @@ def resume(self, path: str | Path) -> None: ) self.statistics = state["statistics"] self.sampler.set_state(state["sampler_state"]) - torch.set_rng_state(state["rng"]["torch"]) - random.setstate(state["rng"]["python"]) - name, keys, pos, gaussian, cached = state["rng"]["numpy"] - np.random.set_state( - (name, keys.numpy().astype(np.uint32), pos, gaussian, cached) - ) - if state["rng"]["cuda"]: - torch.cuda.set_rng_state_all(state["rng"]["cuda"]) + _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.""" @@ -605,8 +716,13 @@ def train(self, *, stop_after: int | None = None) -> dict[str, Any]: if score < self.best_validation: self.best_validation = score self.save(self.output_dir / "best.pt") - log.write(json.dumps(last, allow_nan=False) + "\n") - log.flush() + 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") diff --git a/embodichain/learning/dynamics/workflow.py b/embodichain/learning/dynamics/workflow.py index c5794b0c8..d06fc13ce 100644 --- a/embodichain/learning/dynamics/workflow.py +++ b/embodichain/learning/dynamics/workflow.py @@ -44,6 +44,7 @@ "world_rotation_xyzw", "gravity_transport_override", "source_physical_arrays", + "legacy_velocity", } ) @@ -138,7 +139,7 @@ def run_train(config: dict[str, Any]) -> dict[str, Any]: order_path = Path(config["sample_order"]) order = ( order_path - if order_path.suffix == ".npy" + if order_path.suffix in {".npy", ".json"} else torch.load(order_path, map_location="cpu", weights_only=True) ) try: @@ -209,6 +210,8 @@ def run_eval(config: dict[str, Any]) -> dict[str, Any]: ), "dataset_fingerprint": dataset.fingerprint, "qualification": "offline_only", + "numerical_profile": manifest.get("numerical_profile", "default"), + "attention_backend": model.cfg.attention_backend, } if config.get("output"): Path(config["output"]).write_text( @@ -223,18 +226,24 @@ 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) @@ -242,4 +251,6 @@ def run_export(checkpoint: str | Path, output: str | Path) -> dict[str, Any]: "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/checkpoint.py b/embodichain/learning/models/nerd/checkpoint.py index a31578627..f161a55d1 100644 --- a/embodichain/learning/models/nerd/checkpoint.py +++ b/embodichain/learning/models/nerd/checkpoint.py @@ -30,6 +30,29 @@ _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: @@ -50,6 +73,10 @@ def export_model(path: str | Path, model: NeRDModel, manifest: dict[str, Any]) - 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) @@ -73,7 +100,9 @@ def load_model( raise ValueError("Unsupported NeRD export format_version") if _digest(path / "weights.pt") != manifest.get("weights_sha256"): raise ValueError("NeRD weights checksum mismatch") - model = NeRDModel(NeRDModelCfg(**manifest["model_config"])) + 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()) @@ -153,6 +182,11 @@ def rename(name: str) -> str: 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()), @@ -169,6 +203,9 @@ def rename(name: str) -> str: 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() diff --git a/embodichain/learning/models/nerd/model.py b/embodichain/learning/models/nerd/model.py index 5f4c3507b..3abd86c07 100644 --- a/embodichain/learning/models/nerd/model.py +++ b/embodichain/learning/models/nerd/model.py @@ -24,6 +24,7 @@ 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 @@ -38,6 +39,8 @@ class NeRDModelCfg: 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] = {} @@ -53,6 +56,19 @@ class NeRDModelCfg: 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): @@ -196,6 +212,11 @@ class NeRDModel(nn.Module): 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): @@ -266,7 +287,8 @@ def forward(self, inputs: Mapping[str, Tensor]) -> Tensor: ) if not 1 <= shape[1] <= self.cfg.block_size: raise ValueError("Sequence length is outside the model block_size") - value = self.model(self.transformer_model(torch.cat(values, -1))) + 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 ( diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/README.md b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/README.md index eb025f85d..309550b30 100644 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/README.md +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/README.md @@ -91,3 +91,36 @@ Use preview for the Gym scene: this task does not implement an expert trajectory collector. The APG configuration uses the existing learning-environment factory and trainer, with graph-connected state/reward hooks. Its policy checkpoint is separate from the NeRD dynamics-model export. + +## Legacy Warp parameter comparison + +`dynamics.generate.warp-reference.yaml` selects `physics_profile: warp_reference` +and Featherstone with 16 substeps per 1/60-second frame. It preserves measured +legacy asset masses/inertia, body armature, force-mode damping and contact-slot +order. Ant also uses an explicitly authored shared ground material. The simulator +is still current Newton; contact algorithms can differ from old Warp. + +Generate this bounded 512-trajectory dataset with: + +```bash +embodichain dynamics generate --config extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.warp-reference.yaml +``` + +Keep the same `physics_profile` in rollout evaluation. Existing exports remain +bound to their recorded physical signature; generate new data and train or +fine-tune a new export for this profile. The default `native` profile retains +modern importer semantics. Both profiles need an installed DexSim with named +solver extensions; the measured SDK branch is MR !1483 at `b61e920f`. + +The isolated historical/current replay tool is +`scripts/experiments/neural_dynamics/physics_alignment.py`. Its shared input and +trace hashes, plus finite/failure counts, are required when comparing trajectories. + + +## Supported-version reference data + +The [original data and paired evaluation workflow](../../../../../../scripts/experiments/neural_dynamics/README.md) +uses isolated Python3.8/Torch2.2 generation, immutable raw/canonical teacher +traces and independently evolving model feedback. Training and evaluation +runtimes are recorded separately; a current-Newton corpus cannot be relabeled +as original-Warp data. 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 000000000..c0f0f5453 --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.warp-reference.yaml @@ -0,0 +1,18 @@ +# 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 +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 index 73b720ba6..212b3e951 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 @@ -34,3 +34,5 @@ train: 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/README.md b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/README.md index 391bf6d5e..0c2a73b08 100644 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/README.md +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/README.md @@ -91,3 +91,36 @@ Use preview for the Gym scene: this task does not implement an expert trajectory collector. The APG configuration uses the existing learning-environment factory and trainer, with graph-connected state/reward hooks. Its policy checkpoint is separate from the NeRD dynamics-model export. + +## Legacy Warp parameter comparison + +`dynamics.generate.warp-reference.yaml` selects `physics_profile: warp_reference` +and Featherstone with 5 substeps per 1/60-second frame. It preserves measured +legacy asset masses/inertia, body armature, force-mode damping and contact-slot +order. Ant also uses an explicitly authored shared ground material. The simulator +is still current Newton; contact algorithms can differ from old Warp. + +Generate this bounded 512-trajectory dataset with: + +```bash +embodichain dynamics generate --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.warp-reference.yaml +``` + +Keep the same `physics_profile` in rollout evaluation. Existing exports remain +bound to their recorded physical signature; generate new data and train or +fine-tune a new export for this profile. The default `native` profile retains +modern importer semantics. Both profiles need an installed DexSim with named +solver extensions; the measured SDK branch is MR !1483 at `b61e920f`. + +The isolated historical/current replay tool is +`scripts/experiments/neural_dynamics/physics_alignment.py`. Its shared input and +trace hashes, plus finite/failure counts, are required when comparing trajectories. + + +## Supported-version reference data + +The [original data and paired evaluation workflow](../../../../../../scripts/experiments/neural_dynamics/README.md) +uses isolated Python3.8/Torch2.2 generation, immutable raw/canonical teacher +traces and independently evolving model feedback. Training and evaluation +runtimes are recorded separately; a current-Newton corpus cannot be relabeled +as original-Warp data. 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 000000000..a7fd7e8b3 --- /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 index 76df09dee..cb287d7f1 100644 --- 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 @@ -30,3 +30,5 @@ train: checkpoint_interval: 5000 lr_update_interval: 5000 lr_offset: 1 + log_interval: 500 + resident_cache_bytes: 4294967296 diff --git a/pyproject.toml b/pyproject.toml index fba16a989..fb95c299d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ dynamic = ["version"] # Core install dependencies (kept from requirements.txt). dependencies = [ - "dexsim_engine==0.5.1rc1", + "dexsim_engine==0.5.1", "shapely>=2.0", "setuptools>=78.1.1", "gymnasium>=0.29.1", diff --git a/scripts/benchmark/learning/dynamics_sampling.py b/scripts/benchmark/learning/dynamics_sampling.py new file mode 100644 index 000000000..62360ee3a --- /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/experiments/neural_dynamics/README.md b/scripts/experiments/neural_dynamics/README.md new file mode 100644 index 000000000..0e863c6fd --- /dev/null +++ b/scripts/experiments/neural_dynamics/README.md @@ -0,0 +1,262 @@ +# Original NeRD data and matched training + +Use two isolated environments. The original data worker uses Python 3.8.20, +Torch 2.2.2, NumPy 1.24.4 and Warp 1.8.0 for Cartpole or Warp 1.5.1 for Ant. The +EmbodiChain environment owns current-Newton collection, controlled training, +export and native inference. Use the matching SDK build for the named solver +and functional Scene contracts. Keep the unchanged original checkout at commit +`370e89d2b8ce68aa3578aa7812859c681b7502bc`. + +## Original-entry comparison with a small budget + +`official_small_budget.py` runs the actual upstream `train/train.py` through +`runpy`, including its constructor, DataLoaders, FP32 RunningMeanStd, loss, +epoch-zero validation, rollout evaluation and independent checkpoint selectors. +It observes the emitted batches rather than generating an approximate shuffle. +Use the supported original environment above. CPU execution maps the upstream +CUDA-stream request to `None`; the original numerical methods remain unchanged. + +Provide a JSON plan with `source`, `robot`, `output_dir`, `training_config`, +`seed: 0`, `steps: 1000` and `cpu_threads: 8`. Start from the original robot YAML, +change the existing dataset paths and set `algorithm.num_iters_per_epoch: 1000`. +Keep `num_epochs: 1000` and the original optimizer/model fields. The observer +stops before the second trained epoch, after the original validation, simulation +evaluation and checkpoint selection complete. Thus all 1,000 updates retain the +original first-trained-epoch rate, 0.0009991. Existing output directories are +rejected. The original source checkout stays unmodified. + +```bash +CUDA_VISIBLE_DEVICES="" PYTHONPATH="$NERD_WARP_ROOT" "$NERD_DATA_PYTHON" scripts/experiments/neural_dynamics/official_small_budget.py --config original-small-plan.json --trust-external-code +``` + +Import the same raw split files for the migrated trainer. For Ant, preserve the +original Y-up world coordinates in this experiment (`world_rotation_xyzw=None`) +so extra world rotation does not enter the comparison. This does not change the +native Newton deployment coordinate contract. The official run supplies the +actual initial weights, normalization statistics and consumed sample indices. + +```bash +CUDA_VISIBLE_DEVICES="" "$EMBODICHAIN_PYTHON" scripts/experiments/neural_dynamics/matched_small_budget.py --official-dir /workspace/exp_data/ant/official --dataset /workspace/exp_data/ant/original-frame-data --output-dir /workspace/exp_data/ant/migrated --device cpu +``` + +The migrated experiment calls production `DynamicsTrainer.step` and export APIs; +its outer loop reproduces the captured original sampling and validation selector. +This differs from the default CLI sampling/statistics protocol. Compare both CPU +runs first, and record any separate CUDA run as an environment diagnostic. + +New first-batch observations are saved **before** the original model forward, +with `first-preprocessed-batch.json` identifying physical features and binding +the tensor file's SHA256. The original model normalizes its input dictionary in +place. Earlier observations without that sidecar contain already-normalized +features; the matched reader preserves that legacy interpretation and records +the observed stage. Rebuild physical inputs from raw data for gradient probes, +rather than passing a legacy normalized snapshot through the model again. + +`official_small_budget_evaluate.py` evaluates the selected original checkpoint on +its untouched test split, records fixed validation curves and produces forward +golden tensors under Python3.8/Torch2.2. `official_small_budget_rollout.py` reuses +fixed prefixes of the supported-environment evaluation corpora below, evaluates +both trained weight sets through the original integrator in the same supported +runtime, and records endpoint/prefix errors at 1/10/100 steps. A migrated export +must first be restored into the original model format, preserving all weights +and normalizers. All failed trajectories remain in the measurements, and each +feedback trajectory uses its own predicted state. + +Historical-data model exports and native SDK models have different deployment +provenance. Verify the current SDK generate/train/export/infer workflow with +native-data models; never replace a historical model's missing native physics +signature to bypass deployment checks. + +## Reference/math exported-model inference + +`official_math_rollout.py` evaluates an explicitly `reference`/math portable +export through the unchanged original model and Transformer integrator. This is +the reusable entry used for the full-budget model's supported-environment +accuracy checks. The original-default baseline remains separately evaluated by +`official_small_budget_rollout.py` with its native official checkpoint. + +Use the original Python 3.8.20/Torch 2.2.2 environment and the robot's pinned Warp. +Provide fixed official teacher trajectories, the portable export and the exact +pure transport modules whose hashes were recorded when the teachers were +created. Source, asset, raw-state/action, physics and transport identities are +checked; changing a teacher or transport requires its own experiment identity. + +```bash +CUDA_VISIBLE_DEVICES="" PYTHONDONTWRITEBYTECODE=1 PYTHONPATH="$NERD_WARP_ROOT" "$NERD_DATA_PYTHON" scripts/experiments/neural_dynamics/official_math_rollout.py --source "$NERD_REFERENCE_SOURCE" --portable "$NERD_PORTABLE_MODEL" --scope checkpoint_accuracy --teacher "$NERD_TEACHER_DATA" --transport-root "$NERD_FROZEN_TRANSPORT" --output "$NERD_NEW_RESULT_JSON" --count 16 --trust-external-code +``` + +The output contains H1/10/100 endpoint and prefix metrics, complete NPZ traces, +source/runtime/physics provenance and the actual attention operator. Passive and +random-effort runs replay the same controls; PD recomputes each neural branch's +controls from its own predicted state. Failed trajectories remain included. +`--scope smoke` is reserved for zero-update pipeline checks. This tool neither +starts training nor qualifies historical weights for current-Newton deployment. + +`_portable_reference.py` and `_historical_physics.py` are standalone helpers for +this old Python environment. They do not import the installed EmbodiChain package +or SDK, and do not modify the original source or default baseline. + +## Cross-version numerical attribution + +Treat training implementation, runtime kernels and data transport as separate +controls. On shared canonical data, the current-environment original numerical +methods and migrated trainer match through 1,000 updates for both robots. This +does not remove the Ant raw-spatial/native-COM target-rounding distinction. + +CPU SDPA dispatch and the final FP32 gradient-norm reduction can differ between +Torch 2.2 and 2.7. A jointly specified math-attention and fixed-order total-norm +policy is a separate controlled variant. It changes the original default +numerical path and must not replace the unchanged official default baseline or +silently change production settings. Record optimizer tensor equality separately +from version-specific metadata such as `decoupled_weight_decay=False`. + +The [numerical attribution report](../../../design/reports/neural-dynamics-numerical-attribution-20261020.md) +records bounded operator probes, the 1,000-update controls and their limits. + +## Full-budget continuation with an explicit numerical policy + +Production `DynamicsTrainCfg.numerical_profile="reference"` selects math SDPA +and ordered FP64 accumulation of FP32 per-parameter gradient norms. Portable +exports persist `attention_backend="math"`; changing the policy requires a new +fine-tuning run. The original-method `run_reference` adapter supports only its +default policy and rejects math initial exports rather than silently ignoring +their execution metadata. Keep the original default baseline separate. + +For source-fidelity imports, `import_legacy_datasets` accepts +`preserve_legacy_velocity=True` with `world_rotation_xyzw=None`. It preserves +original FP32 velocity before COM conversion, verifies bit-exact transition +continuity, and supplies current/next raw velocity directly to production +encoding. This option does not change native deployment physics qualification. + +Keep `train.total_steps=4995000` for both initial and resumed phases. Use only +top-level `stop_after=5000` for the first checkpoint; then remove `fine_tune` and +`statistics`, provide `resume`, and retain the entire training configuration. +The exact source-shuffle descriptor includes the original DataLoader's epoch +tail dropping and source permutation hashes. Independent seeded permutations +or random-with-replacement orders do not reproduce that stream. + +The [full-budget progress report](../../../design/reports/neural-dynamics-full-budget-20261021.md) +binds the two formal runs, official zero-step initialization, GPU recovery gates, +fixed-test metrics and original-supported-runtime rollout comparisons. Both +actual committed cursors must reach the full budget before completion is claimed. + +## Training data and paired training + +`official_model_reference.py --generate-only` executes the unmodified original +sampler without a model. `NERD_DATA_PYTHON` is the original environment's +interpreter; `NERD_WARP_ROOT` is the robot's pinned Warp dependency directory: + +```bash +CUDA_VISIBLE_DEVICES="" PYTHONPATH="$NERD_WARP_ROOT" "$NERD_DATA_PYTHON" scripts/experiments/neural_dynamics/official_model_reference.py --source /path/to/neural-robot-dynamics --robot ant --output /workspace/exp_data/ant/reference-train --count 10000 --num-envs 512 --horizon 100 --seed 0 --generate-only --trust-external-code +``` + +Generate validation/test separately with different seeds and output paths. The +original sampler rounds counts to replica batches; inspect the measured count. +`import_legacy_datasets` in `embodichain.learning.dynamics.reference` imports raw +split files with canonical transport and explicit gravity-transport repair. +Preserve unchanged raw files and actual split identities, versions and hashes. + +Prepare the controlled matrix with +`embodichain.learning.dynamics.comparison.prepare_comparison`, supplying both +immutable teacher datasets, original source/dependency paths, five seeds and +complete model/train configuration. `run_comparison` executes frozen commands. +Both trainers share initialized tensors, normalizers, sample indices, updates +and the validation selector. Use distinct output directories and explicitly +resume incomplete work from its own optimizer/RNG checkpoint. This controlled +protocol does not establish literal paper-recipe reproduction. + +## Freeze supported-environment evaluation data + +`official_evaluation_data.py` retains all requested outcomes, including physical +failures and NaN/Inf, with no rejection/resampling. It runs original physics and +sampling unchanged. Pure Torch coordinate/observation files load directly under +an isolated namespace; their paths/hashes are recorded separately. The worker +does not import the EmbodiChain package, DexSim or Newton. + +Example `original-eval-data.json`: + +```json +{ + "source": "/path/to/neural-robot-dynamics", + "robot": "ant", + "output_dir": "/workspace/exp_data/ant/original-eval-passive", + "num_envs": 64, + "trajectories": 2048, + "horizon": 1000, + "seed": 1000000, + "initial_states": "sample", + "motion_mode": "passive" +} +``` + +```bash +CUDA_VISIBLE_DEVICES="" PYTHONPATH="$NERD_WARP_ROOT" "$NERD_DATA_PYTHON" scripts/experiments/neural_dynamics/official_evaluation_data.py --config original-eval-data.json --trust-external-code +``` + +Use new directories for the other scenarios. Set `motion_mode: random_effort`, +or set `motion_mode: fixed_policy` with Ant policy +`{"kind": "pd", "kp": 20.0, "kd": 1.0}`. Cartpole PD uses +`{"kind": "pd", "gains": [5.0, 50.0, 5.0, 20.0]}`. These are declared baselines, +not the paper's learned policy. Original raw states/efforts are preserved beside +canonical q/qd/efforts and measured physics. + +Choose seeds outside every stored training/validation/test candidate range and +audit actual initial states. Current collection uses generation seed plus +candidate index, so comparing top-level seeds alone is insufficient. Reuse one +immutable corpus per robot/scenario for both trainers and all model seeds. + +## Evaluate both exports against one corpus + +Run the historical evaluator in a fresh EmbodiChain Python process, with the +old Warp dependency directory prepended to `PYTHONPATH`. Do not replace installed +current Warp or import DexSim/Newton into that process. Example spec: + +```json +{ + "source": "/path/to/neural-robot-dynamics", + "robot": "ant", + "model": "/workspace/exp_data/ant/official-seed0/best-export", + "teacher_dataset": "/workspace/exp_data/ant/original-eval-passive", + "output": "/workspace/exp_data/ant/official-seed0/passive.json", + "device": "cpu", + "num_envs": 64, + "trajectories": 2048, + "horizons": [1, 10, 100, 1000], + "seed": 1000000, + "initial_states": "sample", + "motion_mode": "passive", + "cpu_threads": 4 +} +``` + +```bash +CUDA_VISIBLE_DEVICES="" PYTHONPATH="$NERD_WARP_ROOT" "$EMBODICHAIN_PYTHON" -m embodichain.learning.dynamics.reference_rollout --spec historical-eval.json --trust-external-code +``` + +The evaluator verifies generation versions/source/physics, raw/canonical state +transport, ordered teacher contents and feedback actions. The neural simulator +starts from the recorded initial state, evolves independently, and computes +feedback from its own state. Later teacher states are never copied into it. +Model-runtime versions and teacher-generation versions are recorded separately. +Omitting `teacher_dataset` retains live teacher generation for diagnostics. + +For the migrated export change only model/output. Repeat for both robots, all +three scenarios and five seeds, retaining complete budgets. Use +`embodichain dynamics` for separate current-Newton deployment evaluation; it also +accepts a verified `teacher_dataset`, including fixed-PD feedback traces. + +Attach results with `annotate_rollout_evidence`, then call +`write_comparison_report`. Training completion, numerical parity, hashes and all +error/failure/nonfinite/return gates are separate checks. Completed commands and +short pilots cannot substitute for full trained-model accuracy qualification. + +## Optional reference tests + +For the opt-in checkpoint tests, set `NERD_REFERENCE_ROOT` to a pinned checkout +with real LFS weights. When that checkout is original Warp main, also set +`NERD_REFERENCE_DEPS` to the matching isolated Warp dependency directory. A +checkout containing LFS pointer files cannot supply the pretrained checkpoint. +For the fixed-teacher tests, `NERD_SUPPORTED_EVAL_DATA` contains directories named +`cartpole-passive`, `ant-passive`, and the corresponding `random_effort` and +`fixed_policy` directories. These tests inspect original generation provenance +separately from the model-inference runtime. diff --git a/scripts/experiments/neural_dynamics/_historical_physics.py b/scripts/experiments/neural_dynamics/_historical_physics.py new file mode 100644 index 000000000..f75718a72 --- /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 000000000..c6cffc09c --- /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 000000000..81404b197 --- /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 000000000..ec4d2d8d3 --- /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 000000000..b40933776 --- /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 000000000..0f28e8e23 --- /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 000000000..7d04d4080 --- /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 000000000..25b7af7b3 --- /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 000000000..8ab2ac998 --- /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 000000000..f44fe077e --- /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/scripts/experiments/neural_dynamics/physics_alignment.py b/scripts/experiments/neural_dynamics/physics_alignment.py new file mode 100644 index 000000000..1c95bf845 --- /dev/null +++ b/scripts/experiments/neural_dynamics/physics_alignment.py @@ -0,0 +1,283 @@ +# ---------------------------------------------------------------------------- +# 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 identical canonical states/efforts in isolated Warp and Newton. + +Run one worker per engine/version; compare the saved NPZ arrays, never live +simulators in one process. The old source remains pinned and unmodified. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path +import sys + +import numpy as np + +__all__: list[str] = [] + + +def _compare(args: argparse.Namespace) -> dict: + import torch + from embodichain.compute.dynamics._evaluation import _error_values + + if args.compare_with is None: + raise ValueError("compare requires --compare-with reference.npz") + paths = (args.compare_with, args.inputs) + receipts = [json.loads(path.with_suffix(".json").read_text()) for path in paths] + if any(row["robot"] != args.robot for row in receipts): + raise ValueError("Trace robot identities differ") + if receipts[0]["inputs_sha256"] != receipts[1]["inputs_sha256"]: + raise ValueError("Traces do not share identical input content") + for path, receipt in zip(paths, receipts): + if hashlib.sha256(path.read_bytes()).hexdigest() != receipt["trace_sha256"]: + raise ValueError(f"Trace checksum mismatch: {path}") + with np.load(paths[0]) as ref, np.load(paths[1]) as actual: + q, qd, rq, rqd = [ + torch.from_numpy(value.copy()) + for value in (actual["q"], actual["qd"], ref["q"], ref["qd"]) + ] + if q.shape != rq.shape or qd.shape != rqd.shape: + raise ValueError("Trace shapes differ") + torch.testing.assert_close(q[0], rq[0], rtol=1e-5, atol=1e-6) + torch.testing.assert_close(qd[0], rqd[0], rtol=1e-5, atol=1e-6) + errors = _error_values(args.robot, q, qd, rq, rqd) + finite = torch.isfinite(torch.cat((q, qd, rq, rqd), -1)).all(-1) + results = {} + for h in (1, 10, 100, 1000): + if h >= q.shape[0]: + continue + valid = finite[: h + 1].all(0) + results[str(h)] = { + "finite_trajectories": int(valid.sum()), + "nonfinite_trajectories": int((~valid).sum()), + "errors": ( + None + if not bool(valid.all()) + else dict( + zip( + ( + "q_mae", + "qd_mae", + "base_position_error", + "base_orientation_error_rad", + ), + errors[1 : h + 1].mean((0, 1)).tolist(), + ) + ) + ), + } + return { + "robot": args.robot, + "reference": receipts[0], + "actual": receipts[1], + "horizons": results, + "population": "All requested trajectories; errors remain undefined when any trajectory becomes nonfinite.", + } + + +def _inputs(robot: str, count: int, horizon: int) -> dict[str, np.ndarray]: + from embodichain.lab.sim.physics.solvers.collection import _sample + + rows = [ + _sample(robot, 42 + i, horizon, "random_effort", "default", 0.02) + for i in range(count) + ] + q, qd, efforts = [np.stack([row[j].numpy() for row in rows]) for j in range(3)] + rng = np.random.default_rng(42) + if robot == "cartpole": + q[:, 0] = rng.uniform(-0.25, 0.25, count) + q[:, 1] = rng.uniform(-0.4, 0.4, count) + qd[:] = rng.uniform(-0.3, 0.3, qd.shape) + else: + q[:, 2] = np.resize([0.55, 1.0, 3.0, 1.0], count) + q[:, 7:] += rng.uniform(-0.03, 0.03, (count, 8)) + qd[:, 6:] = rng.uniform(-0.3, 0.3, (count, 8)) + return {"q": q, "qd": qd, "efforts": efforts} + + +def _legacy( + args: argparse.Namespace, data: dict[str, np.ndarray] +) -> tuple[np.ndarray, np.ndarray, dict]: + from embodichain.learning.dynamics.reference import reference_identity + + source = reference_identity(args.source) + sys.path.insert(0, source["root"]) + import warp as wp + from scipy.spatial.transform import Rotation + from utils.env_utils import create_abstract_contact_env + + expected = "1.8.0" if args.robot == "cartpole" else "1.5.1" + if wp.__version__ != expected: + raise ValueError(f"Expected Warp {expected}, found {wp.__version__}") + env = create_abstract_contact_env( + "Cartpole" if args.robot == "cartpole" else "Ant", + args.count, + device="cpu", + random_reset=False, + ) + q, qd = data["q"].copy(), data["qd"].copy() + rotation = Rotation.from_euler("x", np.pi / 2) + if args.robot == "ant": + q[:, :3] = rotation.inv().apply(q[:, :3]) + q[:, 3:7] = (rotation.inv() * Rotation.from_quat(q[:, 3:7])).as_quat() + omega = rotation.inv().apply(qd[:, 3:6]) + linear = rotation.inv().apply(qd[:, :3]) - np.cross(omega, q[:, :3]) + qd[:, :6] = np.concatenate((omega, linear), axis=-1) + env.state_0.joint_q.assign(q.astype(np.float32).ravel()) + env.state_0.joint_qd.assign(qd.astype(np.float32).ravel()) + wp.sim.eval_fk( + env.model, env.state_0.joint_q, env.state_0.joint_qd, None, env.state_0 + ) + qs, vs = [], [] + for step in range(args.horizon + 1): + a = env.state_0.joint_q.numpy().reshape(args.count, -1).copy() + v = env.state_0.joint_qd.numpy().reshape(args.count, -1).copy() + if args.robot == "ant": + v[:, :6] = np.concatenate( + ( + rotation.apply(v[:, 3:6] + np.cross(v[:, :3], a[:, :3])), + rotation.apply(v[:, :3]), + ), + axis=-1, + ) + a[:, :3] = rotation.apply(a[:, :3]) + a[:, 3:7] = (rotation * Rotation.from_quat(a[:, 3:7])).as_quat() + qs.append(a) + vs.append(v) + if step < args.horizon: + force = data["efforts"][:, step] + env.control.joint_act.assign( + (force[:, 6:] if args.robot == "ant" else force).ravel() + ) + env.update() + return ( + np.stack(qs), + np.stack(vs), + { + "source": source, + "warp": wp.__version__, + "substeps": env.sim_substeps, + "mass_update_interval": env.integrator.update_mass_matrix_every, + "gravity": list(env.model.gravity), + }, + ) + + +def _current( + args: argparse.Namespace, data: dict[str, np.ndarray] +) -> tuple[np.ndarray, np.ndarray, dict]: + import torch + import dexsim + from embodichain.lab.sim.physics.solvers.collection import DynamicsSimulation + + substeps = args.substeps or (5 if args.robot == "cartpole" else 16) + solver = ( + None + if args.solver == "featherstone" + else {"solver_type": "mujoco_warp", "use_mujoco_cpu": True} + ) + if args.profile == "native" and solver is None: + solver = { + "solver_type": "featherstone", + "update_mass_matrix_interval": substeps, + } + with DynamicsSimulation( + args.robot, + num_envs=args.count, + physics_profile=args.profile, + num_substeps=substeps, + solver=solver, + ) as env: + env.reset(torch.from_numpy(data["q"]), torch.from_numpy(data["qd"])) + qs, vs = [], [] + for step in range(args.horizon + 1): + state, _ = env.snapshot() + qs.append(state["q"].numpy()) + vs.append(state["qd"].numpy()) + if step < args.horizon: + env.step(torch.from_numpy(data["efforts"][:, step])) + metadata = {**env.manifest, "dexsim_commit": dexsim.__commit_id__} + return np.stack(qs), np.stack(vs), metadata + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--engine", choices=["legacy", "newton", "compare"], required=True + ) + parser.add_argument("--robot", choices=["cartpole", "ant"], required=True) + parser.add_argument("--source", type=Path) + parser.add_argument("--trust-external-code", action="store_true") + parser.add_argument( + "--profile", choices=["native", "warp_reference"], default="warp_reference" + ) + parser.add_argument( + "--solver", choices=["featherstone", "mujoco_warp"], default="featherstone" + ) + parser.add_argument("--substeps", type=int) + parser.add_argument("--count", type=int, default=8) + parser.add_argument("--horizon", type=int, default=1000) + parser.add_argument("--inputs", type=Path, required=True) + parser.add_argument("--compare-with", type=Path) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if args.engine == "compare": + if args.output.exists(): + raise FileExistsError(args.output) + args.output.write_text(json.dumps(_compare(args), indent=2, allow_nan=False)) + print(args.output) + return + if min(args.count, args.horizon) < 1: + parser.error("count and horizon must be positive") + if args.engine == "legacy" and ( + not args.trust_external_code or args.source is None + ): + parser.error("legacy execution requires --source and --trust-external-code") + if args.output.exists(): + raise FileExistsError(args.output) + if args.inputs.exists(): + with np.load(args.inputs) as rows: + data = {k: rows[k] for k in rows.files} + else: + if args.engine != "newton": + parser.error("generate shared inputs with the current worker first") + data = _inputs(args.robot, args.count, args.horizon) + args.inputs.parent.mkdir(parents=True, exist_ok=True) + np.savez(args.inputs, **data) + if data["efforts"].shape[:2] != (args.count, args.horizon): + raise ValueError("input shape differs from count/horizon") + q, qd, meta = (_legacy if args.engine == "legacy" else _current)(args, data) + args.output.parent.mkdir(parents=True, exist_ok=True) + np.savez(args.output, q=q, qd=qd) + meta.update( + engine=args.engine, + profile=args.profile, + robot=args.robot, + inputs_sha256=hashlib.sha256(args.inputs.read_bytes()).hexdigest(), + trace_sha256=hashlib.sha256(args.output.read_bytes()).hexdigest(), + horizon=args.horizon, + count=args.count, + ) + args.output.with_suffix(".json").write_text(json.dumps(meta, indent=2, default=str)) + print(args.output) + + +if __name__ == "__main__": + main() diff --git a/tests/cli/test_dynamics.py b/tests/cli/test_dynamics.py index f047841ca..2170f6016 100644 --- a/tests/cli/test_dynamics.py +++ b/tests/cli/test_dynamics.py @@ -244,8 +244,19 @@ def train(config): (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.rollout.smoke.yaml", + "dynamics.rollout.full.yaml", + } for robot in ("cartpole", "ant"): - files.extend( + robot_configs = list( ( repo / "extensions" @@ -256,7 +267,8 @@ def train(config): / robot ).glob("dynamics.*.yaml") ) - assert len(files) == 16 + 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. 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 000000000..8737a5f7c --- /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 index 44d1a65fa..f2dbb6cdb 100644 --- a/tests/compute/dynamics/test_codec.py +++ b/tests/compute/dynamics/test_codec.py @@ -15,6 +15,9 @@ # ---------------------------------------------------------------------------- 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 ( @@ -34,6 +37,28 @@ def ant_state(batch: int = 2) -> tuple[torch.Tensor, torch.Tensor]: 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() @@ -46,6 +71,111 @@ def test_com_reference_shift_and_inverse() -> None: 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() @@ -65,6 +195,46 @@ def test_ant_targets_decode_in_frame_at_t() -> None: 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]]) diff --git a/tests/data_pipeline/dynamics/test_shards.py b/tests/data_pipeline/dynamics/test_shards.py index 6ab6caa9b..380a3401b 100644 --- a/tests/data_pipeline/dynamics/test_shards.py +++ b/tests/data_pipeline/dynamics/test_shards.py @@ -21,6 +21,7 @@ import pickle import numpy as np import pytest +import torch from torch.utils.data import DataLoader from embodichain.data_pipeline.dynamics.shards import ( @@ -114,6 +115,74 @@ def test_stable_split_not_ingestion_order(): 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 diff --git a/tests/learning/dynamics/test_comparison.py b/tests/learning/dynamics/test_comparison.py index 820890683..c7da56531 100644 --- a/tests/learning/dynamics/test_comparison.py +++ b/tests/learning/dynamics/test_comparison.py @@ -168,3 +168,23 @@ def test_rollout_attachment_preserves_measurements_and_requires_identity(tmp_pat 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_reference.py b/tests/learning/dynamics/test_reference.py index a8f843497..470e14e1b 100644 --- a/tests/learning/dynamics/test_reference.py +++ b/tests/learning/dynamics/test_reference.py @@ -19,10 +19,12 @@ 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 @@ -46,6 +48,107 @@ def test_external_execution_requires_explicit_trust(): 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: @@ -236,6 +339,254 @@ def test_legacy_roundtrip_preserves_native_qd_and_declared_split(tmp_path): 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 @@ -452,3 +803,133 @@ def test_portable_to_official_restoration_preserves_state_dtype_and_ownership(): 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 index c9802d24e..98a62860c 100644 --- a/tests/learning/dynamics/test_reference_rollout.py +++ b/tests/learning/dynamics/test_reference_rollout.py @@ -254,6 +254,43 @@ def _check_original_runtime(tmp_path: Path, robot: str, source: str) -> None: ) 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] @@ -304,3 +341,193 @@ def test_real_historical_teacher_neural_runtime_and_source_contract( "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 000000000..7257b0c7b --- /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 index 2fe72422b..788bca3ab 100644 --- a/tests/learning/dynamics/test_report.py +++ b/tests/learning/dynamics/test_report.py @@ -25,12 +25,129 @@ 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 _hash +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: @@ -83,6 +200,26 @@ def test_numerical_receipt_requires_measurements_and_matching_identity(tmp_path) 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", [ @@ -90,6 +227,8 @@ def test_numerical_receipt_requires_measurements_and_matching_identity(tmp_path) ("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): diff --git a/tests/learning/dynamics/test_sample_order.py b/tests/learning/dynamics/test_sample_order.py new file mode 100644 index 000000000..7886d24b0 --- /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 000000000..9f711f329 --- /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 index 8d8fccb5c..195fb75f1 100644 --- a/tests/learning/dynamics/test_trainer.py +++ b/tests/learning/dynamics/test_trainer.py @@ -19,6 +19,7 @@ from __future__ import annotations import json +import math import subprocess import sys @@ -95,7 +96,17 @@ def architecture(dropout=0.0): ) -def build(root, out, *, steps=30, dropout=0.0): +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") @@ -109,6 +120,10 @@ def build(root, out, *, steps=30, dropout=0.0): 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"] @@ -133,14 +148,33 @@ def test_statistics_only_training_windows(tmp_path): ) -def test_exact_resume_with_dropout_optimizer_sampler_rng(tmp_path): +@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) + 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) + 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) + 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(): @@ -212,11 +246,25 @@ def test_workflow_fine_tune_and_resume_contract(tmp_path): 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 = torch.rsqrt(variance + 1e-5) + weights = 1.0 / torch.sqrt(variance + 1e-5) expected = torch.nn.functional.mse_loss( prediction[:, :1] * weights, target[:, :1] * weights ) @@ -226,6 +274,25 @@ def test_official_loss_order_and_mask(): 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" @@ -339,6 +406,64 @@ def test_memory_mapped_order_resume_and_integrity(tmp_path): 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") @@ -434,14 +559,23 @@ def test_training_manifest_cannot_relabel_dataset_physics(tmp_path, mutation): assert not (tmp_path / "run").exists() -def test_training_annotations_cannot_invent_missing_physics_contract(tmp_path): +@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="physics_signature mismatch"): + with pytest.raises(ValueError, match=f"{field} mismatch"): run_train( { "dataset": str(tmp_path / "data"), "output_dir": str(tmp_path / "run"), - "manifest": {**metadata, "physics_signature": "invented"}, + "manifest": {**metadata, field: value}, "model": architecture().to_dict(), "train": {"total_steps": 1, "batch_size": 2}, } @@ -473,3 +607,295 @@ def test_unrelated_annotations_and_fine_tune_to_new_measured_physics(tmp_path): 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 index 2ab1b72d2..5c31b58f2 100644 --- a/tests/learning/models/_nerd_reference_check.py +++ b/tests/learning/models/_nerd_reference_check.py @@ -27,15 +27,29 @@ from embodichain.learning.models.nerd import convert_legacy_checkpoint, load_model -def compare_codec(robot: str, destination: str) -> dict[str, float]: +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. """ - from newton import JointType - from solvers.solver_neural import NeuralSolver + 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" @@ -48,7 +62,7 @@ def compare_codec(robot: str, destination: str) -> dict[str, float]: 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 0 + 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 @@ -99,10 +113,18 @@ def compare_codec(robot: str, destination: str) -> dict[str, float]: "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"], + 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: @@ -129,6 +151,17 @@ def legacy(q: torch.Tensor, qd: torch.Tensor) -> torch.Tensor: } 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(): @@ -181,6 +214,14 @@ def legacy(q: torch.Tensor, qd: torch.Tensor) -> torch.Tensor: 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 = ( @@ -277,7 +318,11 @@ def reference_inputs(values: dict[str, torch.Tensor]) -> dict[str, torch.Tensor] result["max_adam_parameter_error"] = max( result["max_adam_parameter_error"], (p - rp).abs().max().item() ) - result.update(compare_codec(robot, destination)) + 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() diff --git a/tests/learning/models/test_nerd.py b/tests/learning/models/test_nerd.py index bf3a7faad..77ce6d093 100644 --- a/tests/learning/models/test_nerd.py +++ b/tests/learning/models/test_nerd.py @@ -15,6 +15,7 @@ # ---------------------------------------------------------------------------- from __future__ import annotations +import json import os from pathlib import Path import subprocess @@ -127,6 +128,109 @@ def test_legacy_pickle_requires_explicit_trust(tmp_path: Path) -> None: 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 @@ -144,6 +248,14 @@ def test_explicit_official_checkpoint_forward_gradient_and_adam( **os.environ, "OMP_NUM_THREADS": "2", "MKL_NUM_THREADS": "2", - "PYTHONPATH": str(Path(__file__).resolve().parents[3]), + "PYTHONPATH": os.pathsep.join( + filter( + None, + ( + os.environ.get("NERD_REFERENCE_DEPS"), + str(Path(__file__).resolve().parents[3]), + ), + ) + ), }, ) diff --git a/tests/sim/physics/solvers/data/warp_reference_arrays.json b/tests/sim/physics/solvers/data/warp_reference_arrays.json new file mode 100644 index 000000000..5d1691d75 --- /dev/null +++ b/tests/sim/physics/solvers/data/warp_reference_arrays.json @@ -0,0 +1,467 @@ +{ + "reference_commit": "370e89d2b8ce68aa3578aa7812859c681b7502bc", + "scope": "Measured unmodified official Environment model arrays; legacy Y-up body-local geometry; not a learned-model quality fixture.", + "cartpole": { + "body_name": [ + "cart", + "pole" + ], + "body_mass": [ + 50.0, + 2.5 + ], + "body_com": [ + [ + 0.0, + 0.0, + 0.0 + ], + [ + 0.0, + 0.0, + 0.5 + ] + ], + "body_inertia": [ + [ + [ + 1.21833336353302, + 0.0, + 0.0 + ], + [ + 0.0, + 1.21833336353302, + 0.0 + ], + [ + 0.0, + 0.0, + 2.0933332443237305 + ] + ], + [ + [ + 0.2188541740179062, + 0.0, + 0.0 + ], + [ + 0.0, + 0.2188541740179062, + 0.0 + ], + [ + 0.0, + 0.0, + 0.011041666381061077 + ] + ] + ], + "joint_armature": [ + 0.009999999776482582, + 0.009999999776482582 + ], + "joint_limit_ke": [ + 10000.0, + 10000.0 + ], + "joint_limit_kd": [ + 10.0, + 10.0 + ], + "warp_version": "1.8.0", + "source_importer_sha256": "f0d9efedb89487eb0bb75cf8fe10dfca5a8e950ad1731f2bf46f5f7ad3af8c10" + }, + "ant": { + "body_name": [ + "torso", + "front_left_leg", + "front_left_foot", + "front_right_leg", + "front_right_foot", + "left_back_leg", + "left_back_foot", + "right_back_leg", + "right_back_foot" + ], + "body_mass": [ + 31.653451919555664, + 7.831550598144531, + 13.518441200256348, + 7.831550598144531, + 13.518441200256348, + 7.831550598144531, + 13.518441200256348, + 7.831550598144531, + 13.518441200256348 + ], + "body_com": [ + [ + 0.0, + 0.0, + 0.0 + ], + [ + 0.10000000149011612, + 0.10000000149011612, + 0.0 + ], + [ + 0.20000001788139343, + 0.20000001788139343, + 0.0 + ], + [ + -0.10000000149011612, + 0.10000000149011612, + 0.0 + ], + [ + -0.20000001788139343, + 0.20000001788139343, + 0.0 + ], + [ + -0.10000000149011612, + -0.10000000149011612, + 0.0 + ], + [ + -0.20000001788139343, + -0.20000001788139343, + 0.0 + ], + [ + 0.10000000149011612, + -0.10000000149011612, + 0.0 + ], + [ + 0.20000001788139343, + -0.20000001788139343, + 0.0 + ] + ], + "body_inertia": [ + [ + [ + 0.5960065126419067, + 0.0, + 0.0 + ], + [ + 0.0, + 0.5960065126419067, + 0.0 + ], + [ + 0.0, + 0.0, + 1.0890781879425049 + ] + ], + [ + [ + 0.07864078879356384, + -0.04495241865515709, + 0.0 + ], + [ + -0.04495241865515709, + 0.07864080369472504, + 0.0 + ], + [ + 0.0, + 0.0, + 0.12359321862459183 + ] + ], + [ + [ + 0.298422247171402, + -0.24653588235378265, + 0.0 + ], + [ + -0.24653588235378265, + 0.29842230677604675, + 0.0 + ], + [ + 0.0, + 0.0, + 0.5449581742286682 + ] + ], + [ + [ + 0.07864078879356384, + 0.04495241865515709, + 0.0 + ], + [ + 0.04495241865515709, + 0.07864080369472504, + 0.0 + ], + [ + 0.0, + 0.0, + 0.12359321862459183 + ] + ], + [ + [ + 0.298422247171402, + 0.24653588235378265, + 0.0 + ], + [ + 0.24653588235378265, + 0.29842230677604675, + 0.0 + ], + [ + 0.0, + 0.0, + 0.5449581742286682 + ] + ], + [ + [ + 0.07864080369472504, + -0.044952426105737686, + 0.0 + ], + [ + -0.04495241865515709, + 0.07864079624414444, + 0.0 + ], + [ + 0.0, + 0.0, + 0.12359318882226944 + ] + ], + [ + [ + 0.29842233657836914, + -0.24653591215610504, + 0.0 + ], + [ + -0.24653591215610504, + 0.29842230677604675, + 0.0 + ], + [ + 0.0, + 0.0, + 0.5449580550193787 + ] + ], + [ + [ + 0.07864080369472504, + 0.044952426105737686, + 0.0 + ], + [ + 0.04495241865515709, + 0.07864079624414444, + 0.0 + ], + [ + 0.0, + 0.0, + 0.12359318882226944 + ] + ], + [ + [ + 0.29842233657836914, + 0.24653591215610504, + 0.0 + ], + [ + 0.24653591215610504, + 0.29842230677604675, + 0.0 + ], + [ + 0.0, + 0.0, + 0.5449580550193787 + ] + ] + ], + "joint_armature": [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.009999999776482582, + 0.009999999776482582, + 0.009999999776482582, + 0.009999999776482582, + 0.009999999776482582, + 0.009999999776482582, + 0.009999999776482582, + 0.009999999776482582 + ], + "joint_limit_ke": [ + 1000.0, + 1000.0, + 1000.0, + 1000.0, + 1000.0, + 1000.0, + 1000.0, + 1000.0 + ], + "joint_limit_kd": [ + 100.0, + 100.0, + 100.0, + 100.0, + 100.0, + 100.0, + 100.0, + 100.0 + ], + "warp_version": "1.5.1", + "ordered_contact_points": [ + [ + 0.0, + 0.0, + 0.0 + ], + [ + 0.20000001262410383, + 0.1999999995773283, + 0.0 + ], + [ + -9.64387160407032e-09, + 3.4029039308602904e-09, + 0.0 + ], + [ + -0.20000001262410383, + 0.1999999995773283, + 0.0 + ], + [ + 9.64387160407032e-09, + 3.4029039308602904e-09, + 0.0 + ], + [ + -0.2000000071173526, + -0.20000000508408, + 0.0 + ], + [ + 4.137120354630319e-09, + 2.10384776266892e-09, + 0.0 + ], + [ + 0.2000000071173526, + -0.20000000508408, + 0.0 + ], + [ + -4.137120354630319e-09, + 2.10384776266892e-09, + 0.0 + ], + [ + 0.20000001262410383, + 0.1999999995773283, + 0.0 + ], + [ + -9.64387160407032e-09, + 3.4029039308602904e-09, + 0.0 + ], + [ + 0.40000002524820766, + 0.3999999991546566, + 0.0 + ], + [ + -1.928774320814064e-08, + 6.805807861720581e-09, + 0.0 + ], + [ + -0.20000001262410383, + 0.1999999995773283, + 0.0 + ], + [ + 9.64387160407032e-09, + 3.4029039308602904e-09, + 0.0 + ], + [ + -0.40000002524820766, + 0.3999999991546566, + 0.0 + ], + [ + 1.928774320814064e-08, + 6.805807861720581e-09, + 0.0 + ], + [ + -0.2000000071173526, + -0.20000000508408, + 0.0 + ], + [ + 4.137120354630319e-09, + 2.10384776266892e-09, + 0.0 + ], + [ + -0.4000000142347052, + -0.40000001016816, + 0.0 + ], + [ + 8.274240709260638e-09, + 4.20769552533784e-09, + 0.0 + ], + [ + 0.2000000071173526, + -0.20000000508408, + 0.0 + ], + [ + -4.137120354630319e-09, + 2.10384776266892e-09, + 0.0 + ], + [ + 0.4000000142347052, + -0.40000001016816, + 0.0 + ], + [ + -8.274240709260638e-09, + 4.20769552533784e-09, + 0.0 + ] + ], + "source_importer_sha256": "55d12045b75eed0cce05f96e10794987681ce7c2d060c95539ae6edf15348221" + } +} diff --git a/tests/sim/physics/solvers/test_collection.py b/tests/sim/physics/solvers/test_collection.py index 4510f588d..05efac3b1 100644 --- a/tests/sim/physics/solvers/test_collection.py +++ b/tests/sim/physics/solvers/test_collection.py @@ -51,6 +51,7 @@ ("num_substeps", -1), ("max_attempts", 1), ("effort_scale", float("nan")), + ("physics_profile", "unregistered"), ], ) def test_generation_config_rejects_invalid_values_before_runtime(key, value) -> None: @@ -68,6 +69,31 @@ def test_sampling_is_candidate_seed_local_and_passive_zero() -> None: assert not _sample("cartpole", 3, 5, "passive", "default", 1.0)[2].any() +@pytest.mark.parametrize("robot,substeps", [("cartpole", 5), ("ant", 16)]) +def test_warp_parameter_profile_preserves_reset_replay(robot, substeps) -> None: + from embodichain.lab.sim.physics.solvers.collection import DynamicsSimulation + + with DynamicsSimulation( + robot, num_envs=2, num_substeps=substeps, physics_profile="warp_reference" + ) as host: + q, qd, effort = _sample(robot, 5, 2, "random_effort", "default", 0.02) + q, qd = q.repeat(2, 1), qd.repeat(2, 1) + host.reset(q, qd) + host.step(effort[0].repeat(2, 1)) + first, _ = host.snapshot() + host.step(effort[1].repeat(2, 1)) + host.reset(q, qd) + host.step(effort[0].repeat(2, 1)) + repeated, _ = host.snapshot() + for field in ("q", "qd"): + torch.testing.assert_close( + first[field], repeated[field], rtol=1e-5, atol=1e-6 + ) + assert host.manifest["physics_profile"] == "warp_reference" + assert host.manifest["teacher_origin"]["kind"] == "current_newton" + assert host.manifest["teacher"]["solver"]["solver_type"] == "featherstone" + + def test_fixed_policy_adapter_uses_own_state_and_records_hash() -> None: policy = _FixedPolicy("cartpole", {"kind": "pd"}, torch.device("cpu")) q = torch.tensor([[0.0, 0.0], [0.1, 0.2]]) @@ -230,6 +256,58 @@ def _check_pipeline(tmp_path: Path, robot: str) -> None: + fixed["horizons"]["3"]["numerically_invalid_trajectories"] == 3 ) + replay = evaluate_rollouts( + { + **common, + "motion_mode": "fixed_policy", + "policy": {"kind": "pd"}, + "teacher_dataset": fixed["teacher_dataset"], + "output": str(tmp_path / "closed-replay.json"), + } + ) + assert replay["selected_teacher_sha256"] == fixed["selected_teacher_sha256"] + assert replay["horizons"] == fixed["horizons"] + assert replay["teacher_feedback_verification"] == ( + "all_selected_actions_verified_against_teacher_states" + ) + with pytest.raises(ValueError, match="declared feedback policy"): + evaluate_rollouts( + { + **common, + "motion_mode": "fixed_policy", + "policy": ( + {"kind": "pd", "gains": [1.0, 2.0, 3.0, 4.0]} + if robot == "cartpole" + else {"kind": "pd", "kp": 10.0, "kd": 1.0} + ), + "teacher_dataset": fixed["teacher_dataset"], + "output": str(tmp_path / "wrong-policy.json"), + } + ) + # A self-consistent file hash cannot excuse actions unrelated to the stored + # teacher's own observations and declared feedback configuration. + import shutil + + changed = tmp_path / "changed-feedback" + shutil.copytree(fixed["teacher_dataset"], changed) + source_manifest = load_manifest(changed, verify=True) + first = source_manifest["shards"][0] + with h5py.File(changed / first["file"], "r+") as stream: + stream["0/transitions/joint_f"][0, 0] += 1.0 + import hashlib + + first["sha256"] = hashlib.sha256((changed / first["file"]).read_bytes()).hexdigest() + (changed / "manifest.json").write_text(json.dumps(source_manifest)) + with pytest.raises(ValueError, match="own-state feedback"): + evaluate_rollouts( + { + **common, + "motion_mode": "fixed_policy", + "policy": {"kind": "pd"}, + "teacher_dataset": str(changed), + "output": str(tmp_path / "changed-feedback.json"), + } + ) def _check_actor_and_failures(tmp_path: Path) -> None: diff --git a/tests/sim/physics/solvers/test_reference_assets.py b/tests/sim/physics/solvers/test_reference_assets.py index 0b5e12b16..b9a061079 100644 --- a/tests/sim/physics/solvers/test_reference_assets.py +++ b/tests/sim/physics/solvers/test_reference_assets.py @@ -19,6 +19,8 @@ from __future__ import annotations from unittest.mock import MagicMock +import json +from pathlib import Path import newton import numpy as np @@ -64,9 +66,10 @@ def _lower_descriptor(descriptor): @pytest.mark.parametrize("robot", ["cartpole", "ant"]) -def test_spawn_lowering_preserves_reference_physical_properties(robot): - reference = _import_reference(robot) - descriptor = make_reference_articulation(robot) +@pytest.mark.parametrize("physics_profile", ["native", "warp_reference"]) +def test_spawn_lowering_preserves_reference_physical_properties(robot, physics_profile): + reference = _import_reference(robot, physics_profile) + descriptor = make_reference_articulation(robot, physics_profile) lowered = _lower_descriptor(descriptor) lower_bodies = {name: index for index, name in enumerate(lowered.body_label)} for source_index, label in enumerate(reference.body_label): @@ -180,3 +183,83 @@ def test_reference_hash_and_profiles(): assert reference_profile("ant")["effort_limits"] == [400.0] * 8 with pytest.raises(ValueError, match="Unsupported"): make_reference_articulation("unknown") + + +def test_warp_cartpole_import_preserves_measured_mass_and_body_armature(): + """Check measured Warp 1.8 arrays, rather than URDF inertial declarations.""" + builder = _import_reference("cartpole", "warp_reference") + np.testing.assert_allclose(builder.body_mass, [50.0, 2.5], rtol=1e-6) + np.testing.assert_allclose( + np.asarray(builder.body_inertia).reshape(-1, 3, 3), + [ + np.diag([1.21833336353302, 1.21833336353302, 2.0933332443237305]), + np.diag([0.2188541740179062, 0.2188541740179062, 0.011041666381061077]), + ], + rtol=1e-6, + atol=1e-8, + ) + np.testing.assert_allclose(builder.joint_armature, [0.01, 0.01]) + # The modern interpretation stays available as its own physical domain. + np.testing.assert_allclose(_import_reference("cartpole").body_mass, [1.0, 10.0]) + + +def test_warp_ant_import_preserves_geometry_density_and_force_mode(): + """Check measured Warp 1.5.1 masses, inertia regularization and friction.""" + builder = _import_reference("ant", "warp_reference") + np.testing.assert_allclose( + builder.body_mass, + [31.653451919555664, *[7.831550598144531, 13.518441200256348] * 4], + rtol=1e-6, + ) + np.testing.assert_allclose( + np.asarray(builder.body_inertia[0]).reshape(3, 3), + np.diag([0.5960065126419067, 0.5960065126419067, 1.0890781879425049]), + rtol=1e-6, + atol=1e-7, + ) + np.testing.assert_allclose(builder.joint_damping, 0.0) + np.testing.assert_allclose(builder.joint_target_kd, 0.0) + np.testing.assert_allclose(builder.shape_material_mu, 0.75) + np.testing.assert_allclose(builder.shape_margin, 0.0) + transform = builder.shape_transform[1] + first_slot = wp.transform_point( + transform, wp.vec3(0.0, 0.0, builder.shape_scale[1][1]) + ) + np.testing.assert_allclose(first_slot, [0.2, 0.2, 0.0], atol=5e-8) + assert reference_asset_hash("ant") == ( + "23e9a720465a71fde3fd9fd2d5a18f8c11ec642f3dca8921c66cb4e169421fa7" + ) + + +def test_unknown_reference_physics_profile_rejected(): + with pytest.raises(ValueError, match="physics profile"): + _import_reference("ant", "unregistered") + + +@pytest.mark.parametrize("robot", ["cartpole", "ant"]) +def test_warp_profile_matches_recorded_body_arrays_and_ordered_contacts(robot): + fixture = json.loads( + (Path(__file__).parent / "data/warp_reference_arrays.json").read_text() + )[robot] + builder = _import_reference(robot, "warp_reference") + for name in ("body_mass", "body_com", "body_inertia", "joint_armature"): + actual = np.asarray(getattr(builder, name)) + expected = np.asarray(fixture[name]).reshape(actual.shape) + np.testing.assert_allclose(actual, expected, rtol=2e-6, atol=2e-7, err_msg=name) + inertia = np.asarray(builder.body_inertia).reshape(-1, 3, 3) + inverse = np.asarray(builder.body_inv_inertia).reshape(-1, 3, 3) + np.testing.assert_allclose( + inertia @ inverse, np.broadcast_to(np.eye(3), inertia.shape), atol=1e-6 + ) + if robot == "ant": + points = [] + for scale, transform in zip(builder.shape_scale, builder.shape_transform): + offsets = ( + [wp.vec3(0.0)] + if scale[1] == 0 + else [wp.vec3(0.0, 0.0, scale[1]), wp.vec3(0.0, 0.0, -scale[1])] + ) + points.extend( + list(wp.transform_point(transform, value)) for value in offsets + ) + np.testing.assert_allclose(points, fixture["ordered_contact_points"], atol=2e-7) diff --git a/tests/sim/spawn/test_scene.py b/tests/sim/spawn/test_scene.py index aeea7eb87..a653bca68 100644 --- a/tests/sim/spawn/test_scene.py +++ b/tests/sim/spawn/test_scene.py @@ -80,6 +80,21 @@ def _prepare_spawn_runtime_config(self, _result: object) -> None: self.events.append("runtime_config") +def test_global_descriptor_declares_one_facade_instance_in_replicated_scene() -> None: + handle = object() + scene = _make_scene({"reference_ground": handle}) + scene._num_envs = 8 + facade = _RetryableFacade() + scene.track( + "rigid_object", + "reference_ground", + SimpleNamespace(name="reference_ground", per_env=False), + facade=facade, + ) + assert facade.declared_num_instances == 1 + assert facade._entities == [handle] + + def test_bind_retries_only_incomplete_declarations() -> None: first_handle = object() second_handle = object() diff --git a/tests/sim/test_sim_manager.py b/tests/sim/test_sim_manager.py index d1219bec4..edbcc66fa 100644 --- a/tests/sim/test_sim_manager.py +++ b/tests/sim/test_sim_manager.py @@ -67,6 +67,53 @@ ) +def test_procedural_object_descriptor_retains_global_material_and_scene_owner( + monkeypatch, +): + from dexsim.spawn import ( + CollisionDesc, + GeometryDesc, + NewtonCollisionDesc, + ObjectDesc, + RigidBodyPhysicsDesc, + ) + + collision = CollisionDesc.from_geometry(GeometryDesc.plane(1000.0)) + collision.newton = NewtonCollisionDesc(ke=1e5, kd=1e3, mu=0.5, margin=0.0) + descriptor = ObjectDesc( + name="reference_ground", + collisions=[collision], + physics=RigidBodyPhysicsDesc.static(), + per_env=False, + ) + facade = object() + monkeypatch.setattr(sim_manager_module, "RigidObject", lambda **kwargs: facade) + manager = SimpleNamespace( + asset_uids=[], + device=torch.device("cpu"), + spawn_result=None, + _spawn_scene=MagicMock(), + _rigid_objects={}, + notify_visualization_topology_changed=MagicMock(), + prepare=MagicMock(), + ) + result = SimulationManager.add_rigid_object_descriptor(manager, descriptor) + assert result is facade + manager._spawn_scene.declare.assert_called_once_with( + "rigid_object", "reference_ground", descriptor, facade=facade + ) + assert not descriptor.per_env and descriptor.collisions[0].newton.ke == 1e5 + assert manager._rigid_objects["reference_ground"] is facade + manager.prepare.assert_not_called() + manager.asset_uids.append("reference_ground") + with pytest.raises(ValueError, match="already exists"): + SimulationManager.add_rigid_object_descriptor(manager, descriptor) + with pytest.raises(TypeError, match="physical"): + SimulationManager.add_rigid_object_descriptor( + manager, ObjectDesc(name="visual") + ) + + @pytest.mark.parametrize("mask_size", [2, 3]) def test_empty_functional_reset_preserves_graph_without_native_publication(mask_size): manager = SimpleNamespace( From 0d5cffe87bf915af7e3e9386b28f79a45f10f9ae Mon Sep 17 00:00:00 2001 From: yuecideng Date: Wed, 21 Oct 2026 11:42:09 +0000 Subject: [PATCH 3/7] docs(dynamics): organize alignment evidence and full-budget status --- ...ural-dynamics-cuda-alignment-20261020.json | 1040 ++ ...neural-dynamics-cuda-alignment-20261020.md | 159 + .../neural-dynamics-full-budget-20261021.md | 150 + ...namics-numerical-attribution-20261020.json | 195 + ...dynamics-numerical-attribution-20261020.md | 86 + ...ics-official-model-alignment-20261020.json | 507 + ...amics-official-model-alignment-20261020.md | 113 + ...dynamics-parameter-alignment-20261020.json | 12367 ++++++++++++++++ .../neural-dynamics-pr-status-20261020.md | 365 +- .../neural-dynamics-publication-20261021.json | 113 + ...neural-dynamics-small-budget-20261020.json | 1445 ++ .../neural-dynamics-small-budget-20261020.md | 92 + ...pported-reference-evaluation-20261020.json | 1247 ++ ...supported-reference-evaluation-20261020.md | 132 + ...-dynamics-trained-comparison-20261020.json | 7928 ++++++++++ ...al-dynamics-trained-comparison-20261020.md | 164 + 16 files changed, 25835 insertions(+), 268 deletions(-) create mode 100644 design/reports/neural-dynamics-cuda-alignment-20261020.json create mode 100644 design/reports/neural-dynamics-cuda-alignment-20261020.md create mode 100644 design/reports/neural-dynamics-full-budget-20261021.md create mode 100644 design/reports/neural-dynamics-numerical-attribution-20261020.json create mode 100644 design/reports/neural-dynamics-numerical-attribution-20261020.md create mode 100644 design/reports/neural-dynamics-official-model-alignment-20261020.json create mode 100644 design/reports/neural-dynamics-official-model-alignment-20261020.md create mode 100644 design/reports/neural-dynamics-parameter-alignment-20261020.json create mode 100644 design/reports/neural-dynamics-publication-20261021.json create mode 100644 design/reports/neural-dynamics-small-budget-20261020.json create mode 100644 design/reports/neural-dynamics-small-budget-20261020.md create mode 100644 design/reports/neural-dynamics-supported-reference-evaluation-20261020.json create mode 100644 design/reports/neural-dynamics-supported-reference-evaluation-20261020.md create mode 100644 design/reports/neural-dynamics-trained-comparison-20261020.json create mode 100644 design/reports/neural-dynamics-trained-comparison-20261020.md diff --git a/design/reports/neural-dynamics-cuda-alignment-20261020.json b/design/reports/neural-dynamics-cuda-alignment-20261020.json new file mode 100644 index 000000000..572796502 --- /dev/null +++ b/design/reports/neural-dynamics-cuda-alignment-20261020.json @@ -0,0 +1,1040 @@ +{ + "status": "active_full_accuracy_qualification_pending", + "observed_utc": "2026-10-20T09:33:46.805535+00:00", + "scope": "Actual CUDA entrypoint arithmetic and exactresume diagnostics plus unchanged full-budget training; no final model accuracy inferred from partial work", + "numerical_tolerance_absolute": 1e-05, + "cuda_arithmetic": { + "before_actual_first_update": { + "path": "/root/.cache/embodichain/pr822-accuracy/gpu-real-data-prefix-ant/step1-comparison.json", + "sha256": "fd0dd30b25f5799d56bd7a877ae99ae8fac5d0857a8695f850b05e5d2e2fa120" + }, + "after_actual_first_update": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/numerics/production-prefix-ant/step1-comparison.json", + "sha256": "240d47a07b1c5ccf64425c7b4a4e96df39b3784614d01d311d9530784f47e2b6" + }, + "actual_ten_updates_and_exact_resume": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/numerics/warmed-training-verified.json", + "sha256": "db0caed9cea8b9bf1397b63c34538b57ce56a13934a09dd9d69b9cc7ab069e52", + "measured": { + "ant_trainers": 4.4321350287646055e-06, + "official_resume": 0.0, + "migrated_resume": 0.0, + "cartpole_trainers": 0.0 + } + }, + "recomputed_checkpoint_tensors_and_optimizer": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/numerics/warmed-training-checkpoint-evidence.json", + "sha256": "f235d1c1d4b1cb8f5b4c5596bba66f965231b487c02a10da9eda47ce9e6648e8" + } + }, + "immutable_full_v2_implementation": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/implementation-v2.json", + "sha256": "fd4a22182c97cadbd36be64194109ef19c766f64d8afc91d8f825fd4e69124ea" + }, + "unchanged_training_inputs_verification": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/full-v2-readiness.json", + "sha256": "a2406dd89bf4c7f880dcaeb8754ab7277da117fcf42ad39c7f53a724ab7b941c" + }, + "full_v2_readiness": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/full-v2-readiness-verification.json", + "sha256": "66998c2f6ee15d552529cf094100d56b7f8f6c2bff2636fe60a1a6a1b604e5d4" + }, + "training_plans": { + "cartpole": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/experiment.json", + "sha256": "c64d88ee5973b20182c07858e31b065ccaa98b713fb9afcb0de4ac5edc120f4e" + }, + "ant": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/experiment.json", + "sha256": "d0fe75e124dca02be358241cb148470db8fe5200feb4569fb6a6ec3460f91c3c" + } + }, + "full_matrix_cells": 40, + "completed_full_cells": 0, + "active_training_processes": 8, + "allowed_physical_gpus": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6 + ], + "gpu7_reallocation_verification": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/gpu-reallocation-verification.json", + "sha256": "2d69e5fc09a58f86ee7a30b310a498d2d3f018bde296c415901ead1b87775658" + }, + "training_progress": [ + { + "robot": "cartpole", + "id": "official_warp-official-0", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/official_warp/seed-0/official.json", + "sha256": "d089d5ccd1d8b9c7ca121bd904b1c05611d5fde8d7d9d2915cb4f5998efeded2" + }, + "declared_updates": 4995000, + "status": "running", + "actual_process": { + "pid": 1129489, + "physical_gpu": 4, + "proc_start_ticks": "164966376", + "command": [ + "/root/.cache/embodichain/pr822-accuracy/official-ec-env/bin/python", + "-m", + "embodichain.learning.dynamics.reference", + "--spec", + "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/official_warp/seed-0/official.json", + "--trust-external-code" + ] + }, + "last_observed_metric": { + "step": 230500, + "loss": 4.7971367166610435e-05, + "learning_rate": 0.0009576999999999999, + "grad_norm": 0.021629858762025833 + } + }, + { + "robot": "cartpole", + "id": "official_warp-embodichain-0", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/official_warp/seed-0/embodichain.json", + "sha256": "5e97e27325b18efe9358ca8d880eef7eb7b23e24a47ee5e5d98795842820b4da" + }, + "declared_updates": 4995000, + "status": "running", + "actual_process": { + "pid": 1129515, + "physical_gpu": 5, + "proc_start_ticks": "164966511", + "command": [ + "/root/.cache/embodichain/pr822-accuracy/official-ec-env/bin/python", + "-m", + "embodichain.learning.dynamics.comparison", + "--spec", + "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/official_warp/seed-0/embodichain.json", + "--candidate" + ] + }, + "last_observed_metric": { + "step": 270500, + "loss": 1.8255217582918704e-05, + "learning_rate": 0.0009505, + "grad_norm": 0.008207005448639393 + } + }, + { + "robot": "cartpole", + "id": "official_warp-official-1", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/official_warp/seed-1/official.json", + "sha256": "232896c4d40ff0450b505b7644c886a0e06217e46e99ea95e1986c4b397f531b" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "cartpole", + "id": "official_warp-embodichain-1", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/official_warp/seed-1/embodichain.json", + "sha256": "5695ccc78ff1d2e5f551c342720a2fdff95c7deb0183159d209fc6ea63caf4df" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "cartpole", + "id": "official_warp-official-2", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/official_warp/seed-2/official.json", + "sha256": "2244d1bfd9ad22b80818ae91c3d46bb3bc19ab8fe3c4ab2b491b4206c223d3da" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "cartpole", + "id": "official_warp-embodichain-2", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/official_warp/seed-2/embodichain.json", + "sha256": "a7864c50db4ba2dac809bc725e50ddff752031cad2203046362a23e128352843" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "cartpole", + "id": "official_warp-official-3", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/official_warp/seed-3/official.json", + "sha256": "8117193c9ce928ac8f785b5d7ef6497c36e48625afc066ab5d46546a2cc97ee5" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "cartpole", + "id": "official_warp-embodichain-3", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/official_warp/seed-3/embodichain.json", + "sha256": "1103cef1063c1f3cc30a4a69d44e685bbdb2bfbaf4a0c4ef2723f1a6646684f8" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "cartpole", + "id": "official_warp-official-4", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/official_warp/seed-4/official.json", + "sha256": "2817017c5698e5e59c7f825aaab3a5f6b125bf589041a32bcfc1e77108d17281" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "cartpole", + "id": "official_warp-embodichain-4", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/official_warp/seed-4/embodichain.json", + "sha256": "c6ffe002537c98b6fc76670a28b0f999a10108d737a084d0e5caaedea0507474" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "cartpole", + "id": "current_newton-official-0", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/current_newton/seed-0/official.json", + "sha256": "10d3c322cca4ed5a4794cd8ea8ae3aabfa700d038f51dd12f11510e9062153f9" + }, + "declared_updates": 4995000, + "status": "running", + "actual_process": { + "pid": 1129548, + "physical_gpu": 6, + "proc_start_ticks": "164966611", + "command": [ + "/root/.cache/embodichain/pr822-accuracy/official-ec-env/bin/python", + "-m", + "embodichain.learning.dynamics.reference", + "--spec", + "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/current_newton/seed-0/official.json", + "--trust-external-code" + ] + }, + "last_observed_metric": { + "step": 151000, + "loss": 8.002131653483957e-05, + "learning_rate": 0.0009721, + "grad_norm": 0.029288411140441895 + } + }, + { + "robot": "cartpole", + "id": "current_newton-embodichain-0", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/current_newton/seed-0/embodichain.json", + "sha256": "fe3280e1a6d29671c52d7b00e811854b6391257b1883a076cb48d4cf3f204882" + }, + "declared_updates": 4995000, + "status": "running", + "actual_process": { + "pid": 1157557, + "physical_gpu": 6, + "proc_start_ticks": "165081259", + "command": [ + "/root/.cache/embodichain/pr822-accuracy/official-ec-env/bin/python", + "-m", + "embodichain.learning.dynamics.comparison", + "--spec", + "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/current_newton/seed-0/embodichain.json", + "--candidate", + "--resume", + "/workspace/exp_data/embodichain-pr822-accuracy-20261020/gpu7-recovery-evidence/last.pt" + ] + }, + "last_observed_metric": { + "step": 184000, + "loss": 0.00019227356824558228, + "learning_rate": 0.0009667, + "grad_norm": 0.058301541954278946 + } + }, + { + "robot": "cartpole", + "id": "current_newton-official-1", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/current_newton/seed-1/official.json", + "sha256": "fcf89f97e73f5d0598e44b4a2ac6702bfd791380c96b7b4cb349d3b45182e98e" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "cartpole", + "id": "current_newton-embodichain-1", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/current_newton/seed-1/embodichain.json", + "sha256": "00fa98dcc3e9824ce8ade1e7dc369b2821c4ed53255f85de4ac7dcbd599b88d0" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "cartpole", + "id": "current_newton-official-2", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/current_newton/seed-2/official.json", + "sha256": "5e4f4c9a0d6e1f1276ae52adff6d939728f255935fe4684486ca78ccbb77ac31" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "cartpole", + "id": "current_newton-embodichain-2", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/current_newton/seed-2/embodichain.json", + "sha256": "a90ace548874dbb7c0da8686cc52dc3e28873f58b0f69394418177ce9e27b219" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "cartpole", + "id": "current_newton-official-3", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/current_newton/seed-3/official.json", + "sha256": "b74e04d0c402e95996c53af6ffea024e11d7230f334784ff892d2afa83e22465" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "cartpole", + "id": "current_newton-embodichain-3", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/current_newton/seed-3/embodichain.json", + "sha256": "4ef07df9dfe4e9b2184b1eb7521c64620034bb77d12c883bc62ad715b756fd95" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "cartpole", + "id": "current_newton-official-4", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/current_newton/seed-4/official.json", + "sha256": "c1fdf4c996860175ea999e422233760262aec49e811ae288b5c7fc6dd7b0c89a" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "cartpole", + "id": "current_newton-embodichain-4", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-cartpole-full-v2/current_newton/seed-4/embodichain.json", + "sha256": "a2664de788eee886353af172e704d6ca8b02cf94c115e9340e6946b7c2f47771" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "ant", + "id": "official_warp-official-0", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/official_warp/seed-0/official.json", + "sha256": "e2a1820292e4da78e97967918370491ae9aa0ae16c4a15da85403c075038f843" + }, + "declared_updates": 4995000, + "status": "running", + "actual_process": { + "pid": 1129370, + "physical_gpu": 0, + "proc_start_ticks": "164965920", + "command": [ + "/root/.cache/embodichain/pr822-accuracy/official-ec-env/bin/python", + "-m", + "embodichain.learning.dynamics.reference", + "--spec", + "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/official_warp/seed-0/official.json", + "--trust-external-code" + ] + }, + "last_observed_metric": { + "step": 106000, + "loss": 0.0007118280045688152, + "learning_rate": 0.0009801999999999999, + "grad_norm": 0.0392836332321167 + } + }, + { + "robot": "ant", + "id": "official_warp-embodichain-0", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/official_warp/seed-0/embodichain.json", + "sha256": "cf378ef67c4cb00d152bfff34a4d9a909da94e93a9cc80df77a5ac9fc0812ad2" + }, + "declared_updates": 4995000, + "status": "running", + "actual_process": { + "pid": 1129398, + "physical_gpu": 1, + "proc_start_ticks": "164966033", + "command": [ + "/root/.cache/embodichain/pr822-accuracy/official-ec-env/bin/python", + "-m", + "embodichain.learning.dynamics.comparison", + "--spec", + "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/official_warp/seed-0/embodichain.json", + "--candidate" + ] + }, + "last_observed_metric": { + "step": 142000, + "loss": 0.0004192662891000509, + "learning_rate": 0.0009739, + "grad_norm": 0.01704532466828823 + } + }, + { + "robot": "ant", + "id": "official_warp-official-1", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/official_warp/seed-1/official.json", + "sha256": "9d3b13803328dcb91f1e8d8341789f741acf4c1d043ed8fddc65e02f60ccdf56" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "ant", + "id": "official_warp-embodichain-1", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/official_warp/seed-1/embodichain.json", + "sha256": "3a7acf310a43be3dc1035d6ed13c22b240c75c303fad97e2b5a973bdb3553f55" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "ant", + "id": "official_warp-official-2", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/official_warp/seed-2/official.json", + "sha256": "ea6cb6b5583c92eb93bd15b3c764a7e68323ff6a65332508f9ffcc0f84fff741" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "ant", + "id": "official_warp-embodichain-2", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/official_warp/seed-2/embodichain.json", + "sha256": "8de97933a2c1f9f8770dcd1066d762c7a7bcbb62a3ede583484892f7fd8baa5d" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "ant", + "id": "official_warp-official-3", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/official_warp/seed-3/official.json", + "sha256": "61742cb13a665723b2fc490d48f2145257bb92e3f25ed09b14c37c1aad5ca1ac" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "ant", + "id": "official_warp-embodichain-3", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/official_warp/seed-3/embodichain.json", + "sha256": "df26acceb3cacc4b521817c1faa88cd186fd94f789045e9934c04dbd8daa2be0" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "ant", + "id": "official_warp-official-4", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/official_warp/seed-4/official.json", + "sha256": "2a8cc8eb695700c25417f86d6dae306b42c538dfdea2056d7824aab545a2c15c" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "ant", + "id": "official_warp-embodichain-4", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/official_warp/seed-4/embodichain.json", + "sha256": "c04c495eee4561477e18f4cde6a7f2957d1f38c70dbc4403f55e6b08f0132903" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "ant", + "id": "current_newton-official-0", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/current_newton/seed-0/official.json", + "sha256": "e50d58278fe9f5b4e37de32648f6782377d5d027e1a553f47943a87e27f0158f" + }, + "declared_updates": 4995000, + "status": "running", + "actual_process": { + "pid": 1129429, + "physical_gpu": 2, + "proc_start_ticks": "164966147", + "command": [ + "/root/.cache/embodichain/pr822-accuracy/official-ec-env/bin/python", + "-m", + "embodichain.learning.dynamics.reference", + "--spec", + "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/current_newton/seed-0/official.json", + "--trust-external-code" + ] + }, + "last_observed_metric": { + "step": 130000, + "loss": 0.0004277752304915339, + "learning_rate": 0.0009766, + "grad_norm": 0.022283406928181648, + "validation": { + "normalized_mse": 0.001994298987556249, + "target_mse": 0.015075126811861992, + "q_error_norm": 0.004691861733794212, + "qd_error_norm": 0.47554288983345033, + "base_position_error": 0.00027830247185193, + "base_orientation_error_rad": 0.0007081530699506402, + "valid_positions": 512000, + "windows": 51200, + "split": "validation" + } + } + }, + { + "robot": "ant", + "id": "current_newton-embodichain-0", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/current_newton/seed-0/embodichain.json", + "sha256": "356748328e09e18ec1445fd2544b106e4c640eb7cfefe1b7f1a30df2dba90291" + }, + "declared_updates": 4995000, + "status": "running", + "actual_process": { + "pid": 1129457, + "physical_gpu": 3, + "proc_start_ticks": "164966294", + "command": [ + "/root/.cache/embodichain/pr822-accuracy/official-ec-env/bin/python", + "-m", + "embodichain.learning.dynamics.comparison", + "--spec", + "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/current_newton/seed-0/embodichain.json", + "--candidate" + ] + }, + "last_observed_metric": { + "step": 143500, + "loss": 0.0004127389402128756, + "learning_rate": 0.0009739, + "grad_norm": 0.01879895105957985 + } + }, + { + "robot": "ant", + "id": "current_newton-official-1", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/current_newton/seed-1/official.json", + "sha256": "91af60a6c687fb4e3d18957d8d6010d2e149c1f42d38c898a9c10935af989b25" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "ant", + "id": "current_newton-embodichain-1", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/current_newton/seed-1/embodichain.json", + "sha256": "bf9f025d00dd0115d58c6ec7ae02a7d36e4ba8af6959efa79e1c01582fb6785c" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "ant", + "id": "current_newton-official-2", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/current_newton/seed-2/official.json", + "sha256": "cf6c9a4274b3c290e328884f1e881055b343d65fbaded3ed9488baa58e8bbcb3" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "ant", + "id": "current_newton-embodichain-2", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/current_newton/seed-2/embodichain.json", + "sha256": "4038f09c118a58b1ff40b27d2e374dcbe4d07f3a94b28b6fe8ccfbac58055186" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "ant", + "id": "current_newton-official-3", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/current_newton/seed-3/official.json", + "sha256": "8a5fa8fd4697383d0d49208fc1668539977ed5819a1260177736639ded69863f" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "ant", + "id": "current_newton-embodichain-3", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/current_newton/seed-3/embodichain.json", + "sha256": "f0a3b4de6a50c23f76cfc8e3390e2b69cf22c4d5c73a606a38d47e04b052e456" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "ant", + "id": "current_newton-official-4", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/current_newton/seed-4/official.json", + "sha256": "54eafa0de09ec5b978251f0c02cfe0cd618053623f61487dfa972b2dc689df49" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + }, + { + "robot": "ant", + "id": "current_newton-embodichain-4", + "spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/paired-training-ant-full-v2/current_newton/seed-4/embodichain.json", + "sha256": "6f69921a24597227b6374e59a40f947cc2407d5132e91db4e9f1c0f6cd2c15fd" + }, + "declared_updates": 4995000, + "status": "queued", + "actual_process": null, + "last_observed_metric": null + } + ], + "teacher_repeatability": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/numerics/teacher-repeatability/results.json", + "sha256": "2e42c703982bf3936c7f1c605da59bf4ca7a158741d7a3d8ccc9b5189aa12e13" + }, + "evaluation_protocol": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/full-rollout-protocol.json", + "sha256": "f17d90dd91202998d96b208373adaf551ab75c03e990cd9442851ed32548a033" + }, + "evaluation_worker": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/qualification_rollout_worker.py", + "sha256": "f06b6075bb3256f513ca355ad3b03eda047437782fab9d1ab161ff3184cd5100", + "pid": 1164825, + "session": 97030, + "proc_exists": true, + "latest_actual_status": "waiting_for_verified_training_or_device" + }, + "evaluation_guards": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/rollout-worker-guard-checks.json", + "sha256": "c5cde671eaaf28eeca2e9d2a8809a3c2419a413f8724043a14a5b38733cdc674" + }, + "heldout_seed_audit": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/evaluation-holdout-audit.json", + "sha256": "3bb1eb9cbc3279268b4a59a52d0d01270521ba2d6081a8fc0c67c750fa3e3108" + }, + "actual_initial_state_overlap_audit": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/initial-state-holdout-audit.json", + "sha256": "fb2b885f2e3a0b1689939e161fd3abde24145f6b4270522411407120f93350bf" + }, + "common_teacher_corpora": [ + { + "id": "cartpole-passive", + "generation_spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/teacher-generation-cartpole.json", + "sha256": "35738f9e39e60bdd43f8daee85ffb4a4e687b190540a99676712a72529887b3a" + }, + "status": "completed", + "receipt": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/common-teachers/cartpole/passive.receipt.json", + "sha256": "d594ef22cff8945124dfeccd4ecd571342c4bfd1b5fb9c0ec5f375757b5900de" + }, + "measured": { + "format": "task-specific-frozen-common-teacher.v1", + "status": "completed", + "id": "cartpole-passive", + "robot": "cartpole", + "teacher": "current_newton", + "motion_mode": "passive", + "spec": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/teacher-generation-cartpole.json", + "spec_sha256": "35738f9e39e60bdd43f8daee85ffb4a4e687b190540a99676712a72529887b3a", + "holdout_audit_sha256": "3bb1eb9cbc3279268b4a59a52d0d01270521ba2d6081a8fc0c67c750fa3e3108", + "implementation_manifest_sha256": "fd4a22182c97cadbd36be64194109ef19c766f64d8afc91d8f825fd4e69124ea", + "trajectories": 2048, + "horizon": 1000, + "sampling_seed": 1000000, + "manifest": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/common-teachers/cartpole/passive/manifest.json", + "manifest_sha256": "2a812ffa1db4609076bfbfb7efebff801a4093942f6356a32259abefee77517d", + "selected_teacher_sha256": "bfd38a041d250e97a84132f5f47bcce82996a4f5b045b33ce53c22cbdd01d116", + "selected_teacher_hash_schema": "ordered-teacher-content-v1", + "feedback_actions_verified": false, + "policy_identity": null, + "physical_gpu": 4, + "elapsed_seconds": 103.45478200912476, + "completed_utc": "2026-10-20T09:06:34.451250+00:00", + "meaning": "All requested common teacher traces generated; not a trained model or accuracy acceptance" + } + }, + { + "id": "cartpole-random_effort", + "generation_spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/teacher-generation-cartpole.json", + "sha256": "35738f9e39e60bdd43f8daee85ffb4a4e687b190540a99676712a72529887b3a" + }, + "status": "completed", + "receipt": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/common-teachers/cartpole/random_effort.receipt.json", + "sha256": "ac10e0cda1d44e6cc4e4d9d3b251a2865be27e04eefea331ba7a0138457a517c" + }, + "measured": { + "format": "task-specific-frozen-common-teacher.v1", + "status": "completed", + "id": "cartpole-random_effort", + "robot": "cartpole", + "teacher": "current_newton", + "motion_mode": "random_effort", + "spec": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/teacher-generation-cartpole.json", + "spec_sha256": "35738f9e39e60bdd43f8daee85ffb4a4e687b190540a99676712a72529887b3a", + "holdout_audit_sha256": "3bb1eb9cbc3279268b4a59a52d0d01270521ba2d6081a8fc0c67c750fa3e3108", + "implementation_manifest_sha256": "fd4a22182c97cadbd36be64194109ef19c766f64d8afc91d8f825fd4e69124ea", + "trajectories": 2048, + "horizon": 1000, + "sampling_seed": 1000000, + "manifest": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/common-teachers/cartpole/random_effort/manifest.json", + "manifest_sha256": "6cc42cadf8d47a2390740e0a7c4f80f9f2803fabde728dd257de20226d851fe2", + "selected_teacher_sha256": "9b4ab66f5e38373d628a44f6539c9ff6dc39acf1f6eee5da775bc7765bb82cb6", + "selected_teacher_hash_schema": "ordered-teacher-content-v1", + "feedback_actions_verified": false, + "policy_identity": null, + "physical_gpu": 4, + "elapsed_seconds": 103.04208588600159, + "completed_utc": "2026-10-20T09:08:17.493681+00:00", + "meaning": "All requested common teacher traces generated; not a trained model or accuracy acceptance" + } + }, + { + "id": "cartpole-fixed_policy", + "generation_spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/teacher-generation-cartpole.json", + "sha256": "35738f9e39e60bdd43f8daee85ffb4a4e687b190540a99676712a72529887b3a" + }, + "status": "completed", + "receipt": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/common-teachers/cartpole/fixed_policy.receipt.json", + "sha256": "50566d069807755e668e853452024bc1a0e2450e4406c80901f948436704aabe" + }, + "measured": { + "format": "task-specific-frozen-common-teacher.v1", + "status": "completed", + "id": "cartpole-fixed_policy", + "robot": "cartpole", + "teacher": "current_newton", + "motion_mode": "fixed_policy", + "spec": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/teacher-generation-cartpole.json", + "spec_sha256": "35738f9e39e60bdd43f8daee85ffb4a4e687b190540a99676712a72529887b3a", + "holdout_audit_sha256": "3bb1eb9cbc3279268b4a59a52d0d01270521ba2d6081a8fc0c67c750fa3e3108", + "implementation_manifest_sha256": "fd4a22182c97cadbd36be64194109ef19c766f64d8afc91d8f825fd4e69124ea", + "trajectories": 2048, + "horizon": 1000, + "sampling_seed": 1000000, + "manifest": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/common-teachers/cartpole/fixed_policy/manifest.json", + "manifest_sha256": "329af9438debfc0051cec2f5df59290d5a6fb7a99efbca92047dbfcf9faa66ab", + "selected_teacher_sha256": "7645c8c81e3ef60e492d0b687906eeaedff6e675b2fa3c005472aa1266e54ad5", + "selected_teacher_hash_schema": "ordered-teacher-content-v1", + "feedback_actions_verified": true, + "policy_identity": { + "kind": "pd", + "observation_schema": "cartpole-state-4-v1", + "deterministic": true, + "implementation": "embodichain-reference-feedback-v1", + "effort_limits": [ + 1500.0, + 0.0 + ], + "scope": "fixed_baseline_not_paper_policy", + "gains": [ + 5.0, + 50.0, + 5.0, + 20.0 + ], + "sha256": "5f373a5c1aab23985abc978b15ed04891beed8a47e5350e6064f2d9d5ab5930a" + }, + "physical_gpu": 4, + "elapsed_seconds": 149.14963150024414, + "completed_utc": "2026-10-20T09:10:46.645395+00:00", + "meaning": "All requested common teacher traces generated; not a trained model or accuracy acceptance" + } + }, + { + "id": "ant-passive", + "generation_spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/teacher-generation-ant.json", + "sha256": "4c8def70e041e89c326a92a377ec1378e17e08cf700267a4d1a0771a142c4a62" + }, + "status": "completed", + "receipt": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/common-teachers/ant/passive.receipt.json", + "sha256": "a8dbbdc8f83829ce374699d0b83638f53ab39a71fe6b52f5b2c6ef783f5f0702" + }, + "measured": { + "format": "task-specific-frozen-common-teacher.v1", + "status": "completed", + "id": "ant-passive", + "robot": "ant", + "teacher": "current_newton", + "motion_mode": "passive", + "spec": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/teacher-generation-ant.json", + "spec_sha256": "4c8def70e041e89c326a92a377ec1378e17e08cf700267a4d1a0771a142c4a62", + "holdout_audit_sha256": "3bb1eb9cbc3279268b4a59a52d0d01270521ba2d6081a8fc0c67c750fa3e3108", + "implementation_manifest_sha256": "fd4a22182c97cadbd36be64194109ef19c766f64d8afc91d8f825fd4e69124ea", + "trajectories": 2048, + "horizon": 1000, + "sampling_seed": 1000000, + "manifest": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/common-teachers/ant/passive/manifest.json", + "manifest_sha256": "0715fa59472d48a9223ec1e46d3ba6c0d8fb539842f8b3a9fc5bf1974e63b6e7", + "selected_teacher_sha256": "b7a432233e101c81df8e73df7af9cf1766f75610127a1b818f1ed86874914dd2", + "selected_teacher_hash_schema": "ordered-teacher-content-v1", + "feedback_actions_verified": false, + "policy_identity": null, + "physical_gpu": 0, + "elapsed_seconds": 371.0439872741699, + "completed_utc": "2026-10-20T09:11:00.696762+00:00", + "meaning": "All requested common teacher traces generated; not a trained model or accuracy acceptance" + } + }, + { + "id": "ant-random_effort", + "generation_spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/teacher-generation-ant.json", + "sha256": "4c8def70e041e89c326a92a377ec1378e17e08cf700267a4d1a0771a142c4a62" + }, + "status": "completed", + "receipt": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/common-teachers/ant/random_effort.receipt.json", + "sha256": "e64c2e05a23b22069dc9580e79412a076cc9b3bbdf52739b8165ffa8d67bb1bb" + }, + "measured": { + "format": "task-specific-frozen-common-teacher.v1", + "status": "completed", + "id": "ant-random_effort", + "robot": "ant", + "teacher": "current_newton", + "motion_mode": "random_effort", + "spec": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/teacher-generation-ant.json", + "spec_sha256": "4c8def70e041e89c326a92a377ec1378e17e08cf700267a4d1a0771a142c4a62", + "holdout_audit_sha256": "3bb1eb9cbc3279268b4a59a52d0d01270521ba2d6081a8fc0c67c750fa3e3108", + "implementation_manifest_sha256": "fd4a22182c97cadbd36be64194109ef19c766f64d8afc91d8f825fd4e69124ea", + "trajectories": 2048, + "horizon": 1000, + "sampling_seed": 1000000, + "manifest": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/common-teachers/ant/random_effort/manifest.json", + "manifest_sha256": "9e0f89f3e8d543ee26e7a3007f61f3ef76f431e9c0206435b8299f737dbf2a3e", + "selected_teacher_sha256": "49563912ac214982a78fcbd7463023f2d9e9caeaf62c4bdb53f7b2dc9aec735a", + "selected_teacher_hash_schema": "ordered-teacher-content-v1", + "feedback_actions_verified": false, + "policy_identity": null, + "physical_gpu": 0, + "elapsed_seconds": 365.61407947540283, + "completed_utc": "2026-10-20T09:17:06.311246+00:00", + "meaning": "All requested common teacher traces generated; not a trained model or accuracy acceptance" + } + }, + { + "id": "ant-fixed_policy", + "generation_spec": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/teacher-generation-ant.json", + "sha256": "4c8def70e041e89c326a92a377ec1378e17e08cf700267a4d1a0771a142c4a62" + }, + "status": "completed", + "receipt": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/common-teachers/ant/fixed_policy.receipt.json", + "sha256": "5f323b0b7c723740e269f57ad358e7ee7362b6bec89b6c3f8dc544df293fbb9a" + }, + "measured": { + "format": "task-specific-frozen-common-teacher.v1", + "status": "completed", + "id": "ant-fixed_policy", + "robot": "ant", + "teacher": "current_newton", + "motion_mode": "fixed_policy", + "spec": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/teacher-generation-ant.json", + "spec_sha256": "4c8def70e041e89c326a92a377ec1378e17e08cf700267a4d1a0771a142c4a62", + "holdout_audit_sha256": "3bb1eb9cbc3279268b4a59a52d0d01270521ba2d6081a8fc0c67c750fa3e3108", + "implementation_manifest_sha256": "fd4a22182c97cadbd36be64194109ef19c766f64d8afc91d8f825fd4e69124ea", + "trajectories": 2048, + "horizon": 1000, + "sampling_seed": 1000000, + "manifest": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/common-teachers/ant/fixed_policy/manifest.json", + "manifest_sha256": "4451176a4627c729fc56c0cb4ff2b5f02d9f2d4972e395414adc8eb8917fd81e", + "selected_teacher_sha256": "99f2d39df0c805364e6f01265b222a9fea01039e4e9a3f180747a88802aedb69", + "selected_teacher_hash_schema": "ordered-teacher-content-v1", + "feedback_actions_verified": true, + "policy_identity": { + "kind": "pd", + "observation_schema": "ant-com-state-27-v1", + "deterministic": true, + "implementation": "embodichain-reference-feedback-v1", + "effort_limits": [ + 400.0, + 400.0, + 400.0, + 400.0, + 400.0, + 400.0, + 400.0, + 400.0 + ], + "scope": "fixed_baseline_not_paper_policy", + "kp": 20.0, + "kd": 1.0, + "target": [ + 0.0, + 1.0, + 0.0, + -1.0, + 0.0, + -1.0, + 0.0, + 1.0 + ], + "sha256": "2dfde4898331fc07566d72fdfd90e78cbf60eeb6f006e9d151c5d636992f937a" + }, + "physical_gpu": 0, + "elapsed_seconds": 510.4218785762787, + "completed_utc": "2026-10-20T09:25:36.734695+00:00", + "meaning": "All requested common teacher traces generated; not a trained model or accuracy acceptance" + } + } + ], + "all_common_teachers_verified": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/common-teachers-verification.json", + "sha256": "9ebb437344a645ba5518b88603edffdacb4ec4bf0b0455126223f1a75e26d3fb" + }, + "native_long_replay_pilot_diagnostics": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/diagnostics/native-long-replay-verification.json", + "sha256": "8596c61ffec20f1065d6183433a23eddd6eceb1e456ae8fcb733d304ed746d36" + }, + "actual_partial_validation_curves": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/training-validation-progress.json", + "sha256": "7f15c1593f02c476858848f181bfc59f2f160b96b0739e0660c5920d2186440a" + }, + "rollout_storage": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/rollout-storage.json", + "sha256": "13972aab7cb36707ea4a954f69d401bde9a87cab425f8064e04e234b27e5da51" + }, + "focused_regression": { + "passed": 276, + "skipped": 8, + "failed": 0, + "elapsed_seconds": 193.14, + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/numerics/affected-focused-regression-v2-final.log", + "sha256": "5914db3446d06b27ba6e8024614c33ab8cb9a503aeec27b78a0247be86930b14" + }, + "remaining": [ + "complete all40 unchanged4995000update cells", + "all120 full-population evaluations and per-pair teacher-content verification", + "paired5seed confidence bounds and failure/nonfinite/policy-return gates", + "full learning/error curves and final qualification report" + ], + "active_evaluation": { + "protocol": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v3/full-rollout-protocol.json", + "sha256": "a8abbadf56d9d2400d371b5e88633b295c5dcc2c6d6655419f4c4086c776f949" + }, + "worker_pid": 1182660, + "worker_session": 76376, + "latest_report": { + "path": "/root/.codex/worktrees/pr822-analysis/EmbodiChain/design/reports/neural-dynamics-supported-reference-evaluation-20261020.json", + "sha256": "fb98b989840e60a613ce3125201b8e1da0a0abb55d108f5f191b518d75d4a9e4" + }, + "scope": "V3supported-environment immutable teacher replay; earlier v2evaluation waiter superseded, training unchanged" + } +} diff --git a/design/reports/neural-dynamics-cuda-alignment-20261020.md b/design/reports/neural-dynamics-cuda-alignment-20261020.md new file mode 100644 index 000000000..54e6f116b --- /dev/null +++ b/design/reports/neural-dynamics-cuda-alignment-20261020.md @@ -0,0 +1,159 @@ +# NeRD CUDA alignment and full-budget evaluation — 2026-10-20 + +**Full model-accuracy qualification is still pending.** The real Ant CUDA +training-entrypoint residual is corrected, exact resumed execution is verified, +and the unchanged 40-cell full-budget training matrix is running. The official +supported environment generated the reference data; both trainers execute in +the clean EmbodiChain environment to control framework differences. + +The [supported reference-evaluation follow-up](neural-dynamics-supported-reference-evaluation-20261020.md) +now records the active v3 evaluation protocol: historical teacher trajectories +are generated once in the supported environment and replayed, replacing the +v2 evaluator's live Torch2.7 teacher regeneration. Running training is unchanged. + +## Actual training arithmetic + +The earlier CPU comparison used byte-identical canonical inputs and obtained +zero Ant feature/target/prediction differences. A subsequent real CUDA probe +used the actual full-data training entry points, B512/H10 and original-width +network. Although the first loss was identical, the first Adam update differed +by **2.563931e-5** in one parameter. Target arithmetic differed by up to +9.536743e-7. Thus CPU forward parity alone did not establish CUDA training parity. + +The corrected path retains the original scripted quaternion/body-transform +execution and flattened contiguous layout. Target construction and decoding +reuse the current unnormalized embedding, avoiding a second frame transform. +The reference adapter transports raw data to the model device before coordinate +math. The loss uses the original literal reciprocal-sqrt and native MSE +reduction. Cartpole wraps each position before differencing, as the source does. + +CUDA preprocessing also specializes its scripted graph after initial calls. +Both trainers now perform three real-layout preprocessing warmups before any +optimizer update or restoration of RNG state. Checkpoints bind this execution +contract; optimizer resume from an older incompatible contract is rejected. +Direct checkpoint comparisons also match optimizer state and the Torch RNG +streams used by training; unused Python/NumPy global states are recorded separately. + +| Actual diagnostic | Maximum parameter/buffer difference | +|---|---:| +| Ant, first CUDA optimizer update after correction | 0 | +| Ant, ten uninterrupted CUDA updates, original vs migrated trainer | 4.432135e-6 | +| Original Ant trainer, uninterrupted vs resumed ten updates | 0 | +| Migrated Ant trainer, uninterrupted vs resumed ten updates | 0 | +| Cartpole, ten CUDA updates, original vs migrated trainer | 0 | + +The unchanged absolute update tolerance is 1e-5. These are implementation and +resume diagnostics, **not** trained-model accuracy results. The full data, +initial tensors, normalizers, sampling streams, update counts and acceptance +thresholds were not reduced to obtain them. + +## Full training and reference scope + +The immutable `implementation-v2` includes the corrections above. Two robot +plans each contain two teacher domains × two trainers × five seeds. Every cell +retains **4,995,000 updates**, B512/H10, six layers, twelve heads and the declared +192/384 embedding widths. All 40 cells remain assigned exactly once. Experiments +use physical GPUs **0–6** only. The interrupted GPU7 cell resumed from its actual +80,000-update checkpoint on GPU6; its four queued cells were also reassigned. + +Reference training data remain the actual original-repository outputs generated +in Python3.8.20 / Torch2.2.2 / NumPy1.24.4, with Warp1.8.0 for Cartpole and +Warp1.5.1 for Ant. Separate train/validation/test raw files and their measured +physical arrays and hashes are retained. The original sampler actually produced +1,024,000 training transitions per robot. Current-Newton training uses the +separate frozen corpus, with 1,005,600 training transitions per robot. + +Both trainers execute under Python3.11 / Torch2.7.0+cu128 in the resolved +EmbodiChain environment, with the unmodified numerical `VanillaTrainer` methods +as the reference. The original-source data environment is distinct from this +controlled trainer-comparison environment. This does not establish a historical +Torch2.2 training run or literal paper-recipe reproduction. + +Fresh readiness reports find no full-profile discrepancies, verify both teacher +corpora and accept measured pinned-main numerical evidence. They correctly mark +training and rollout qualification **pending**. No partial checkpoint is used by +the full evaluation driver. + +Actual partial validation curves are retained at +`/workspace/exp_data/embodichain-pr822-accuracy-20261020/qualification-v2/training-validation-progress.png`, +with the exact plotted records in the adjacent JSON. These seed0 curves show +training progress; they are not the five-seed final test comparison. + +## Frozen evaluation data + +Repeated GPU Newton teacher runs were not byte-identical even at equal seeds. +Therefore the native evaluation now generates one common corpus per robot and +motion scenario, then reuses its ordered state/control content for every seed +and both trainers. All 2,048 requested trajectories are retained, including +physical failures and nonfinite states. For fixed PD feedback, every recorded +teacher action is independently checked against its own teacher state and +sealed policy identity. The neural simulator still computes each action from +its own evolving state. + +All six corpora have now completed: **12,288 trajectories / 12,288,000 macro +steps**. Every shard hash was rechecked, initial states agree across the three +motions for each robot, all60 native-case teacher guards accept them, and their +physics signatures match the training corpora despite different batch sizes. + +The superseded evaluation seed500 overlaps the current training candidates' +per-trajectory seeds0–12,499. Before observing full-model evaluation results, +the v2 protocol changes evaluation to **seed1,000,000**, with trajectories seeded +through1,002,047. The audit reads actual recorded trajectory seeds rather than +only the dataset's top-level seed. It also reads the initial q/qd of every row +in all four train/validation/test corpora: each robot's 2,048 new initial states +has **zero exact canonical-state overlap** with either domain. Signed zeros and +Ant quaternion sign are normalized for this check. This is a duplicate audit, +not a claim of statistical independence. + +The frozen evaluation has **120 cases**: 40 validation-selected full-budget +models × passive/random-effort/fixed-PD scenarios, each using 2,048 trajectories +and horizons1/10/100/1000. Historical models run through the original isolated +Warp integrator; current-domain models run the installed native NeRD provider. +Historical paired runs must also identify identical ordered teacher content. +The fixed PD baseline remains explicitly distinct from the paper's policy. + +The driver verifies training completion receipts, full budgets, seeds, dataset +and sampling identities, model hashes, and native common-corpus hashes before +attaching measurements. Actual guard probes reject a real 1,000-update pilot, +an incorrect corpus seed/population/domain and physical GPU7. Final paired +error, failure, nonfinite and policy-return gates retain their original margins. + +Separate native diagnostics also executed all four horizons through1,000steps +using these new PD traces and the earlier1,000-update workflow models. Their +selected teacher-content hashes match the exact eight-trajectory prefixes, and +the neural actions differ from teacher actions because each is evaluated from +its own state. These are pipeline diagnostics, not formal full-model evaluations: + +| Pilot diagnostic at1,000steps | Teacher failed | Neural failed | q MAE | qd MAE | +|---|---:|---:|---:|---:| +| Cartpole,8trajectories | 8/8 | 8/8 | 1.96837 | 16.84159 | +| Ant,8trajectories | 6/8 | 7/8 | 0.64865 | 1.17569 | + +Neither diagnostic produced nonfinite states. The poor physical outcomes remain +visible; completion of inference is not evidence of useful model accuracy. + +The stable cache root is +`/workspace/exp_data/embodichain-pr822-accuracy-20261020`. Its overlay filesystem +had only46GiB free when checked. The still-empty v2 rollout directories were +hash-verified and linked to separate local NVMe storage with about1.8TiB free; +all120 specification bytes and workspace access paths remain unchanged. Active +training and teacher-generation files were not moved. + +## Validation and remaining work + +The final focused affected-code suite passed **276 tests, with 8 optional +reference tests skipped** in193.14s. Actual trusted-source update/resume tests +and CUDA loss/transport checks were exercised separately; overlapping counts +are not added. An earlier broader unrelated URDF integration crashed inside the +SDK, so this scoped result must not be presented as an all-repository CI pass. + +The remaining work is completion of all40 full-budget training cells, all120 +full-population evaluations, trace verification, paired five-seed confidence +bounds, learning/error curves and the final human-readable qualification report. +Unfinished, failed or undefined results remain visible and cannot be treated as +acceptance. This report supplements the +[earlier trained-model/workflow evidence](neural-dynamics-trained-comparison-20261020.md). + +[Machine-readable evidence](neural-dynamics-cuda-alignment-20261020.json) binds the +actual CUDA diagnostics, dataset audits, immutable plans, guard checks and +observed processes. Local evidence paths are not download URLs. diff --git a/design/reports/neural-dynamics-full-budget-20261021.md b/design/reports/neural-dynamics-full-budget-20261021.md new file mode 100644 index 000000000..c7a455379 --- /dev/null +++ b/design/reports/neural-dynamics-full-budget-20261021.md @@ -0,0 +1,150 @@ +# NeRD 完整训练预算推进 + +阶段记录：2026-10-21（Asia/Shanghai）。**状态：Cartpole 完整预算已验收，Ant 继续训练；两机器人整体目标尚未完成。** + +本轮在[小预算归因结果](neural-dynamics-numerical-attribution-20261020.md)上落地数据保真、显式数值策略和官方采样重放。Cartpole、Ant 各一个种子，从真实官方零步初值开始，以 `total_steps=4,995,000` 运行。前 5,000 步为同一正式任务的检查前缀，通过后恢复原 Adam、RNG 和采样游标继续；没有另设短训练学习率，也没有恢复旧 40 组矩阵。 + +## 已通过的节点 + +| 检查 | 结果 | +|---|---| +| 官方数据保真 | 两机器人 × train／validation／test，6 个实际 B512/H10 批次的输入、目标与支持环境逐字节一致 | +| 初始化 | 真实官方 `initial.pt` 的 45 项参数及所有 FP32 归一化 mean／var 逐位保留，实际更新数为 0 | +| 官方采样 | 原 Torch2.2 CPU DataLoader 的 999 个训练 epoch、2,997 次完整洗牌在 Torch2.7 逐一验证 SHA，一致 | +| GPU 恢复 | 连续 100 步与 50 步后新进程恢复到 100 步的模型、Adam、RNG、采样及配置逐字节一致 | +| 5,000 步门禁 | 两模型的归档 checkpoint、实际完整固定测试集、生产导出及 portable 重载推理均通过 | +| 正式续训 | 两模型实际 500,000 步 checkpoint 均经独立审计；预算未变，Adam、RMS、游标、学习率及导出一致 | + +Cartpole 使用物理 GPU1，Ant 使用 GPU6。Ant 最初在 GPU3 验证；该卡随后出现其他任务，因此保留原证据并在空闲 GPU6 重做完整恢复验证，没有清理外部进程。本任务不使用 GPU7。 + +## 冻结的训练合同 + +Cartpole 为 2,713,668 参数、embedding192；Ant 为 10,891,037 参数、embedding384。两者均为 6 层、12 头、B512/H10、FP32、无 AMP，Adam `(0.9, 0.999)`／`eps=1e-8`、梯度裁剪 1。原始 1,024,000 条训练 transition 形成 931,840 个窗口，另有独立 validation／test 各 512 条轨迹。官方发布 LFS 模型的配置指向 Cartpole 5M、Ant 10M 数据，需与本轮已有 1M 数据明确区分；该 pickle 不证明实际训练的末次更新数。 + +原 recipe 有 1,000 个 epoch，epoch0 仅评估，所以完整实际训练预算为 `999 × 5,000 = 4,995,000` 次更新。学习率按 5,000 步间隔、offset1，从配置端点 `0.001→0.0001` 调度：首 5,000 步为 `0.0009991`，第 5,001 步为 `0.0009982`，最后一个训练 epoch 为 `0.0001009`。validation／checkpoint 每 5,000 步执行，日志每 500 步记录；test 不参与选模。 + +原始数据来自 Python3.8.20／Torch2.2.2 支持环境，Cartpole Warp1.8.0、Ant Warp1.5.1。新增 `preserve_legacy_velocity=True` 导入选项保留未旋转 source-world 的原始 FP32 spatial velocity，避免 Ant 经 COM 往返产生的目标舍入；默认 native 导入行为保留。新增数据指纹与旧实验不同，旧检查点不能混入此正式任务。 + +当前训练环境为 Python3.11.14／Torch2.7.0+cu128／Warp1.17.0／Newton1.7.0，SDK 为 DexSim MR!1483 的 b61e920 构建、版本0.5.1。选择显式 `numerical_profile=reference`：math SDPA，以及固定顺序 FP64 累加 FP32 各参数范数。该策略连同 attention backend 进入 checkpoint／portable metadata；生产默认策略和旧官方默认基线未被替换。 + +原始采样重放保留 DataLoader 每 epoch 丢弃未消耗尾部的行为。每机器人约 219KiB 的 seed／hash descriptor 和单次约 7.1MiB permutation 缓存替代约 20GB 全索引矩阵。原始数据驻留预算为 4GiB，包含新增保真字段。 + +运行使用独立只读源码快照，共 1,480 个文件。协议校验源码、数据、初值、采样、版本和 GPU UUID；训练目录绑定协议 SHA，继承的文件锁防止重复 worker。完成判定要求真实 `last.pt.global_step==4,995,000`，分别记录终点与 validation 选出的 best，提前出现 best 不算预算完成。 + +## 5,000 步阶段结果 + +两个模型均在完整 46,592 个 test 窗口、465,920 个有效位置上评估： + +| 模型 | 固定 test 归一化 MSE | 5,000 步耗时，含初始化／验证／保存 | 峰值 Torch 分配显存 | +|---|---:|---:|---:| +| Cartpole | 0.001019071375 | 62.43 秒 | 1.29GiB | +| Ant | 0.008243989798 | 114.77 秒 | 2.39GiB | + +短测稳态吞吐约为 Cartpole 98、Ant 50 更新／秒。全预算初估约 14、28 小时，实际耗时受后续运行资源与验证开销影响。当前进度以[运行记录](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/progress.json)和实际 checkpoint 为准。 + +下图快照生成于 2026-10-21 11:07:31 UTC，通过各自单个稳定文件描述符读取实际已提交 Cartpole 4,995,000、Ant 3,100,000 步。Cartpole 达到 4,995,000 步，Ant 仍以 4,995,000 为目标；总图保持 interim，`whole_budget_complete=False`。两者原 5,000 步 test 点保留，Cartpole 新增 GPU1 的终点 last（4,995,000）与验证集选定 best（4,940,000）固定 test 点，局部放大明确区分二者；这些点直接绑定 parent completion 的实际评估收据。Ant 仍只有 5,000 步 GPU test 点，未把当前未评估步数画成 test；两者 500,000 步 CPU fixed-test 在下文单列。刷新前 PNG、PDF、逐点快照及原绘图脚本已按原字节、大小和修改时间另存于 [Cartpole 4,995,000／Ant 2,915,000 归档](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/figures/archive/committed-cart4995000-ant2915000/archive-manifest.json)；[上一轮 4,545,000／2,310,000 归档](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/figures/archive/committed-cart4545000-ant2310000/archive-manifest.json)；[更早的 4,190,000／2,130,000 快照](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/figures/archive/committed-cart4190000-ant2130000/metrics-snapshot.json)和此前历史归档保留。 + +![正式训练的中间曲线](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/figures/training-progress.png) + +[PDF 图](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/figures/training-progress.pdf)与[逐点快照](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/figures/metrics-snapshot.json)保留完整预算及实际已提交步数。 + +Cartpole 已于 2026-10-21 09:52 UTC 通过[完整预算正式验收](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/readiness/cartpole-full-budget-acceptance-20261021T095225Z.json)。[终点预算与导出审计](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/readiness/cartpole-4995000-budget-export-audit-20261021T083759Z.json)确认 last 实际 4,995,000、验证选定 best 实际 4,940,000，核对 45 项 Adam、51 项导出 tensor、冻结 RMS、999 次 validation 选模、源码／环境以及原 5,000 步前缀连续恢复证据；best 不替代实际预算。该 Cartpole 审计的 Ant 中途计数为 2,660,000。最新 [Ant 独立 CPU 审计](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/readiness/ant-live-budget-audit-20261021T105916Z.json)于 2026-10-21 10:59:16 UTC 保留实际 3,075,000 步稳定副本，学习率为 `0.00044650000000000007`，341 项合同检查全部通过，`findings=[]`；45 项 Adam、冻结 RMS、采样／RNG 游标及源码／环境绑定均保持。上图于 2026-10-21 11:07:31 UTC 独立读到 Ant 3,100,000，审计和图快照按各自时间、计数记录；Ant 目标仍为实际 4,995,000 步。此前 [06:33 UTC 中途审计](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/readiness/live-budget-audit-20261021T063357Z.json)的 Cartpole 4,515,000／Ant 2,295,000 记录继续保留；整体目标等待 Ant 原 GPU6 任务实际完成 4,995,000 步及终点验收。 + +同一官方支持环境、每场景固定 16 条轨迹的 100 步关节位置 qMAE： + +| 模型／条件 | passive | random_effort | fixed_policy | +|---|---:|---:|---:| +| Cartpole 官方 LFS／原默认推理 | 0.021108 | 0.061312 | 0.017129 | +| Cartpole 本轮 5,000 步／显式 math | 0.588071 | 0.647311 | 0.484932 | +| Ant 官方 LFS／原默认推理 | 0.195290 | 0.204721 | 0.140644 | +| Ant 本轮 5,000 步／显式 math | 0.409015 | 0.373054 | 0.417351 | + +这批长预测仍明显落后于官方发布模型。5,000 步仅为目标预算约 0.1%；公开配置的数据量及数值路径也不同，不能据此将差距归因于迁移实现或运行环境。teacher／raw action 前缀逐项哈希一致，全部失败轨迹保留；新模型明确使用 math，官方 LFS baseline 保持其原默认推理。完整 H1／10／100 指标见实验目录下的 evaluation 报告。 + +后续第一个 50,000 步观察阈值已复评。Cartpole 采用首次捕获的实际 60,000 步模型，Ant 为实际 50,000 步模型，未将两者误标为相同步数： + +| 模型／实际步数 | passive H100 qMAE | random_effort | fixed_policy | +|---|---:|---:|---:| +| Cartpole／60,000 | 0.327882 | 0.353821 | 0.261394 | +| Ant／50,000 | 0.166831 | 0.277809 | 0.169040 | + +两个模型的长预测总体继续改善。Ant 被动场景已低于官方 LFS 的同项误差，但另外两个场景仍偏高；Cartpole 差距仍较大。不能将一个场景较低的误差扩展为全部精度验收通过。完整指标及哈希见[阶段复评](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/evaluation/milestones/RESULTS.md)。 + +Cartpole 的实际 500,000 步节点随后完成原支持环境复评，H100 qMAE 进一步降至 passive `0.0760825`、random_effort `0.1179903`、fixed_policy `0.1098365`。其 H1 qMAE 分别为 `0.00016520`／`0.00018086`／`0.00017016`；H10 为 `0.00182072`／`0.00201143`／`0.00184519`。相同固定语料上仍高于官方 LFS，完整 4,995,000 步尚未完成。 + +同一 500,000 步 portable export 另经当前 Python3.11／Torch2.7 **CPU** 生产离线 fixed-test 评估：46,592 窗口、465,920 位置的 normalized MSE 为 `2.041487108986725e-5`。该 CPU 评估与原支持环境 H1／10／100 分开记录；test 不参与选模，未使用额外 GPU。[离线结果及模型 SHA](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/evaluation/offline-tests/cartpole-000500000-current-cpu.json)。 + +[500,000 步独立审计](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/readiness/cartpole-500000-audit.json)核对了 45 项 Adam 游标、分段学习率、冻结 RMS、官方采样及数据指纹、源码和导出位模式，并从 NPZ 复算三模式指标与失败数。观察器只记账一次，全部 16 条轨迹保留；该节点不表示完整预算结束。 + +Ant 的实际 500,000 步节点也已在 Python3.8.20／Torch2.2.2／Warp1.5.1 官方支持环境完成固定 16 条轨迹复评： + +| Ant 500,000 步端点 qMAE | passive | random_effort | fixed_policy | +|---|---:|---:|---:| +| H1 | 0.00052551 | 0.00055877 | 0.00053560 | +| H10 | 0.0137785 | 0.0126133 | 0.0151508 | +| H100 | 0.137212 | 0.259807 | 0.123302 | +| 官方 LFS 原默认 H100 | 0.195290 | 0.204721 | 0.140644 | + +三个模式的 H100 qMAE 均较 50,000 步继续下降；其中 passive、fixed_policy 的该项均值已低于官方 LFS，random_effort 仍高约 26.9%。random_effort 的 H100 qdMAE 仍为 `2.03222`，高于官方 LFS 的 `1.53084`。三个模式的神经模型物理失败数分别为 8／8／6，teacher 为 8／9／6，所有失败轨迹均保留在 16 条均值中。这些结果限于当前固定语料与指标；公开模型数据量不同，无法据此单独归因训练实现、环境或物理参数。fixed_policy 保持同一 PD 规则，各模型按自身预测状态生成反馈动作。 + +逐轨迹比较进一步限制了均值结论：三个模式 H100 qMAE 优于 LFS 的轨迹数为 7／16、5／16、8／16；fixed_policy 的 qMAE 中位数为 `0.10015rad`，高于 LFS 的 `0.06886rad`。random_effort 的 H100 qdMAE 从 50,000 步 `1.95453` 回升至 `2.03222`，而同一前缀平均从 `1.08781` 降至 `0.87030`；端点与前缀分别记录。其 H100 逐轨迹 qdMAE 的 P90 为 `5.03614`，LFS 为 `3.25634`，速度长尾仍待改善。qdMAE 沿用 14 维混合速度的原定义，包含基座线速度和角速度，未统一标为 rad/s。完整分位数、配对差值、离群轨迹及首次失败步数见[Ant 500,000 步分析](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/evaluation/ant-500000-analysis.md)与[逐项证据](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/evaluation/ant-500000-analysis.json)。 + +random_effort 的 16 条轨迹在 H100 端点的最大单关节速度绝对误差为 `39.37205rad/s`，对应 seed1000008 的 hip_1；teacher 与模型在此前均未触发协议的物理失败规则。该极值与“逐轨迹平均速度误差最高的轨迹”分别统计，数值有限或未触发失败规则均不能代表速度精度已经达标。 + +[分布统计独立复核](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/readiness/ant-500000-analysis-review.json)从原 NPZ 复算配对计数、分位数及 72 项同单位关节速度极值，29 份输入 SHA 均匹配。最差均值轨迹与全部轨迹最大分量的范围澄清已落实，旧版分析归档，原模型和评估结果保持不变。 + +同一 Ant 导出模型在当前 Python3.11／Torch2.7 **CPU** 的完整固定 test 上得到 normalized MSE `0.0011545916441009268`，仍为 46,592 窗口、465,920 位置。相较 5,000 步记录明显下降，但两次 test 分别运行在 CPU 与 GPU，不构成同设备逐位比较。test 不参与选模，完整预算继续。[离线 test](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/evaluation/offline-tests/ant-000500000-current-cpu.json)和[Ant 500,000 步独立审计](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/readiness/ant-500000-audit.json)保留模型 SHA、45 项 Adam 更新数、初始 RMS、采样、1,480 份源码及三个模式的完整轨迹核对。 + +## Cartpole 4,995,000 步终点与验证选定 best + +Cartpole 的 last 达到 `999 × 5,000 = 4,995,000` 次实际更新，best 由固定 validation 选在 4,940,000 步；test 和 rollout 均未参与重选。[Parent completion](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/training/cartpole/completion.json)绑定的当前训练环境 GPU1 生产离线 fixed-test 如下，两次均覆盖 46,592 窗口、465,920 个有效位置，使用显式 `reference`／math： + +| Cartpole 权重 | 实际更新数 | 固定 test normalized MSE | +|---|---:|---:| +| 终点 last | 4,995,000 | `1.518426341927072e-07` | +| validation-selected best | 4,940,000 | `1.4898648293093266e-07` | + +以下为原 Python3.8.20／Torch2.2.2 CPU／Warp1.8.0 支持环境、每模式固定 16 条 held-out 轨迹的 H100 端点均值。**q 总 MAE 是 m 与 rad 的混合数值平均，qd 总 MAE 是 m/s 与 rad/s 的混合数值平均，均没有单一物理单位。** 滑块位置／杆角／滑块速度／杆角速度的四分量单位表，以及完整 H1／10／100 端点、前缀指标，见[Cartpole 完整预算分析](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/evaluation/cartpole-full-budget-analysis-20261021T083836Z.md)和[逐项 JSON](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/evaluation/cartpole-full-budget-analysis-20261021T083836Z.json)。 + +| 模式 | LFS q 总 MAE | last q 总 MAE | best q 总 MAE | LFS qd 总 MAE | last qd 总 MAE | best qd 总 MAE | +|---|---:|---:|---:|---:|---:|---:| +| passive | 0.021108404 | 0.0028777127 | 0.0028534411 | 0.077126942 | 0.0083973707 | 0.011800601 | +| random_effort | 0.061312363 | 0.0043903147 | 0.0043696053 | 0.24482719 | 0.011709001 | 0.018219153 | +| fixed_policy | 0.017129084 | 0.0027469154 | 0.0030762274 | 0.040733926 | 0.0063099302 | 0.010844845 | + +last 与 best 在三个模式的 H1／10／100 端点及前缀 q／qd 聚合均值均低于原默认官方 LFS；**该优势不覆盖每条轨迹**。[独立原始轨迹复核](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/readiness/cartpole-full-budget-analysis-review-20261021T094012Z.json)确认 81 份输入 SHA、15 份 NPZ 和 540 项分布重算，最大数值差 `3.608224830031759e-16`。每模式全部 16 条轨迹均保留，包括 H100 累计物理失败；passive／random_effort／PD 的 teacher、last、best 失败数均分别为 13／14／12。数值有限和失败数相同不表示所有轨迹稳定，逐 seed 的首次失败时间也并非全部相同。 + +best 的 H100 qd 总 MAE 均值比 last 高 40.5%／55.6%／71.9%（passive／random_effort／PD），速度尾部也更大：random_effort 最差逐轨迹 qd 总 MAE 为 best `0.1458285`、last `0.0447483`；PD 为 best `0.0999441`、last `0.0233717`。best 的 fixed-test nMSE 较低不能保证长 rollout 速度更好；保留验证集选择和终点 last，不用这些结果重选模型。 + +passive／random_effort 使用相同记录动作；PD 只保持同一反馈规则及 teacher 动作，各模型依据自身预测状态计算实际动作，动作可以分叉。seed1000006 的 t49→t50 中，last 与 LFS 位于 ±π 两侧，环绕后的物理杆角差约 `0.000608rad`，实际滑块力却相差约 `314.17N`；这是边界附近 2π×50 增益跳变，不能将其称为逐位相同输入动作，也不能将表示差当作约 6.28rad 的真实角误差。 + +公开 Cartpole LFS 配置指向 5M 数据，本轮约 1M，公开权重实际完成的更新数未知；本次不是同数据复现，也不能从数据、预算、数值策略和设备均不同的比较作单因素因果推断。官方默认 LFS baseline 与旧 40 组产物保留。Cartpole 满预算只证明上述训练／恢复／导出／离线推理合同通过，manifest 仍为 `unqualified`／`offline_only`，不赋予现代 Newton 部署资格；Ant 及两机器人整体目标仍未完成。 + +## 官方固定权重的注意力推理对照 + +另在同一支持环境对两份真实 LFS 权重增加六项 math-only 推理，保持 45 项参数、全部 RMS（含 count）、19 类物理字段及固定 teacher 前缀。此项只改变推理后端，没有将公开模型重新训练或标为 reference 训练；原默认 baseline 和既有观察器的 82 个受保护文件保持原 SHA。 + +| 官方权重 | passive H100 qMAE 差 | random_effort | fixed_policy | +|---|---:|---:|---:| +| Cartpole：math − default | −1.0682e−6 | −1.3225e−7 | −5.7835e−7 | +| Ant：math − default | +7.5128e−5 | −5.8886e−5 | +1.4384e−5 | + +六项 H1 预测均为零差；18 个 H1／10／100 聚合及 288 条逐轨迹记录已独立从 NPZ 复算。这里 H100 端点平均 qMAE 的后端差值很小，将官方模型也置于 math 后，本轮模型与它的该项平均差距基本不变。该结论不涵盖所有指标或整条轨迹的一致性。 + +Ant random_effort 的 seed1000012 轨迹在第 80 步，原 Y-up spatial ν_y 分量差达 5.8726m/s，native `qd[8]`（hip_2 速度）差达 3.3945rad/s；同一前缀的关节角最大差约 0.03066rad。两个速度峰值并非 2π 绕回。两分支在第 64 步均已触发物理失败，失败后的轨迹仍完整保留。passive／random_effort 的实际动作相同；fixed_policy 只保持同一 PD 规则，两分支按各自预测状态计算反馈，实际动作可以分叉。 + +完整结果见[推理后端对照](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/evaluation/lfs-attention-control/REPORT.md)及[独立复核](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/readiness/lfs-attention-control-review.json)。 + +## 后续节点和证据 + +Cartpole 已完成 4,995,000 步及独立终点验收，Ant 沿原配置和原 GPU6 任务继续完成实际 4,995,000 步；独立 CPU 观察器按既定节点保存实际已提交检查点并复评固定轨迹。固定 test、长预测和训练预算分别记录，Ant 仍需终点 last、验证选定 best 的导出和原支持环境验收。该历史物理训练不自动取得当前 Newton 部署资格，原始物理 metadata 不会被替换。 + +本轮 focused CPU／上下文检查为 170 passed、6 skipped、5 deselected；Black26.3.1 全树检查、2748/2748 API 覆盖和 context-map 检查通过。另有真实 GPU 恢复和 CPU launcher 负测，拒绝错误协议／锁、分数或非有限 Adam 游标及错误 moment。项目上下文选择 revise：保真导入、采样格式和持久化数值策略改变了跨模块合同，已更新 neural-dynamics overview。 + +旧产物另经只读保留审计：40 个旧计划 cell 的冻结 identity／artifacts，以及 8 个已停止检查点的原始 SHA 全部匹配；共检查 175 个文件、2,211,814,225 字节。原官方 1,000 步默认基线的选模权重、portable 权重、源码和原始训练环境标识保持一致，未发现旧队列／旧 spec 训练进程。该审计只证明保留与停止状态，不将旧检查点计入本轮完整预算。 + +- [冻结协议](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/protocol.json) +- [正式 GPU 验证](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/validation/completion.json) +- [5,000 步及恢复后的独立审计](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/readiness/live-prefix-audit.json) +- [旧基线与 40 组产物保留审计](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/readiness/legacy-preservation-audit.json) +- [Cartpole 配置](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/configs/cartpole.json)／[Ant 配置](/workspace/exp_data/embodichain-pr822-accuracy-20261021/full-budget-v1/configs/ant.json) diff --git a/design/reports/neural-dynamics-numerical-attribution-20261020.json b/design/reports/neural-dynamics-numerical-attribution-20261020.json new file mode 100644 index 000000000..462d84d6d --- /dev/null +++ b/design/reports/neural-dynamics-numerical-attribution-20261020.json @@ -0,0 +1,195 @@ +{ + "completed_local_time": "2026-10-21T00:04:59.850802+08:00", + "scope": "Follow-up numerical attribution, original full-size models, seed0, CPU. Original default baseline/frozen small-budget artifacts preserved.", + "ant_current_environment": { + "total_updates": 1000, + "new_updates": 990, + "model_max_abs_difference": 0.0, + "adam_max_abs_difference": 0.0, + "test_normalized_mse": 0.025661787018179893, + "grad_norm_log_max_abs_difference": 1.4901161193847656e-06, + "scope": "Current Torch2.7 environment original numerical methods versus migrated implementation on shared canonical data. Not the official-supported environment baseline. Does not isolate Torch environment versus legacy/canonical roundtrip effects." + }, + "cartpole_operator_findings": { + "first_default_sdpa_output_max_abs_difference": 1.7881393432617188e-07, + "first_default_step_parameter_max_abs_difference": 4.614586941897869e-06, + "same_gradient_adam_replay_max_abs_difference": 0.0, + "math_only_first_norm_split_step": 4, + "total_norm22": 1.6617134809494019, + "total_norm27": 1.6617133617401123, + "policy": "joint math SDPA plus fixed-order FP64 sum of squared FP32 per-parameter norms, cast back toFP32" + }, + "cartpole_common_policy": { + "training_updates_each": 1000, + "model_max_abs_difference": 0.0, + "adam_tensor_max_abs_difference": 0.0, + "common_test_normalized_mse": 0.007527741757068006, + "production_test_normalized_mse": 0.007527743803930806, + "test_reduction_difference": 2.046862800780813e-09, + "observed_checkpoints": [ + 250, + 500, + 750, + 1000 + ], + "model_and_adam_bytes_independently_verified": true, + "default_baseline_replaced": false + }, + "ant_transport": { + "runtime_canonical_q_qd_controls_exact": true, + "gradient_results_corrected": true, + "maximum_sampled_training_target_difference": 9.5367431640625e-07, + "sampled_training_batches_matching_canonical": 7, + "all_complete_test_batches_matching_canonical": 91, + "common_raw_test": { + "A_source22_raw": 0.025272413924500182, + "B_source22_roundtrip": 0.02508594701578329, + "D_current27_canonical": 0.02566179076393882 + }, + "common_roundtrip_test": { + "A_source22_raw": 0.025272413924500182, + "B_source22_roundtrip": 0.02508594701578329, + "D_current27_canonical": 0.02566179076393882 + }, + "common_test_contrasts": { + "raw": { + "A_to_B_absolute": -0.00018646690871689311, + "A_to_B_relative_percent": -0.7378278516407377, + "B_to_D_absolute": 0.0005758437481555301, + "B_to_D_relative_percent": 2.295483394719877, + "interpretation": "Conditional contrasts along one fixed-seed path. Not global additive causal fractions." + }, + "roundtrip": { + "A_to_B_absolute": -0.00018646690871689311, + "A_to_B_relative_percent": -0.7378278516407377, + "B_to_D_absolute": 0.0005758437481555301, + "B_to_D_relative_percent": 2.295483394719877, + "interpretation": "Conditional contrasts along one fixed-seed path. Not global additive causal fractions." + } + }, + "limits": [ + "One seed and 1000 updates only. No global additive causal percentages or final-convergence claim.", + "A_to_B changes training spatial/COM transport in an otherwise reproduced supported-runtime source loop; B_to_D is the remaining current-runtime contrast subject to verified full-batch encoding and recorded double-LR metadata differences.", + "Same normalization statistics, initial weights, true captured sample order, full architecture, CPU thread count and source raw data. No source physics, asset, data-generation or Newton resimulation difference is introduced.", + "All raw and roundtrip test values here use the same original supported model runtime. They are new controlled-comparison values, and do not overwrite the historical main-report numbers.", + "Canonical tensor equivalence is directly checked at seven complete real training batches and every complete fixed-test batch. This is not an all-1000-batch numerical invariance proof." + ] + }, + "observation_fix": { + "status": "passed", + "original_entry_updates": 1, + "all_model_parameter_max_abs_difference_vs_frozen_step1": 0.0, + "normalizers_unchanged": true, + "physical_before_forward_sidecar_verified": true, + "normalized_features_match_frozen_legacy": true, + "target_unchanged": true, + "source_tracked_files_unchanged": true, + "entry_exit_code": 0, + "actual_full_batch_consumer_checks": { + "new": { + "stage": "physical_before_model_forward", + "max_abs_differences": { + "states_embedding": 0.0, + "joint_f": 0.0, + "target": 0.0 + } + }, + "legacy": { + "stage": "legacy_unlabeled_normalized_after_model_forward", + "max_abs_differences": { + "states_embedding": 0.0, + "joint_f": 0.0, + "target": 0.0 + } + } + }, + "first_loss_unchanged": true + }, + "validation": { + "focused_tests": 10, + "black": "26.3.1", + "python38_compile": true, + "context_map_check": true, + "shared_production_defaults_unchanged": true + }, + "evidence": { + "ant-same-environment/completion.json": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/followup-alignment-v1/ant-same-environment/completion.json", + "sha256": "675b3d8dcf2fe6631dffcc251f8efbbc15cb17a524cef58353c84fb6c71ae0c6" + }, + "ant-same-environment/comparison.json": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/followup-alignment-v1/ant-same-environment/comparison.json", + "sha256": "a22ae41e87344594426e27ec0552c3cb7df8c78666240e4cf050fad867e7f3b6" + }, + "cartpole-numerics/completion.json": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/followup-alignment-v1/cartpole-numerics/completion.json", + "sha256": "5c15879dd405fb0b5893c6488125f2da07ff092c170b7bcdb7c4ee19373b29d6" + }, + "cartpole-numerics/numerical-attribution-summary.json": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/followup-alignment-v1/cartpole-numerics/numerical-attribution-summary.json", + "sha256": "07e9a316c1554a14e5d8e85e4c74b4343d2da09768a060d246f30dfe798f7f17" + }, + "cartpole-compatible-1000/completion.json": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/followup-alignment-v1/cartpole-compatible-1000/completion.json", + "sha256": "1b1c644d6d5a23dbcd33f673a8c346ec326ccfcc2caaf8109cd8ca46ee66ead5" + }, + "cartpole-compatible-1000/comparison.json": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/followup-alignment-v1/cartpole-compatible-1000/comparison.json", + "sha256": "eb112fd7edf247cd7e262a325f646ffe252581978962bfec2964ee92ba0cb978" + }, + "cartpole-compatible-migrated-1000/completion.json": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/followup-alignment-v1/cartpole-compatible-migrated-1000/completion.json", + "sha256": "3b21edfb28fe770f45bb49a8918218616072c2073ff10da6a3371f1a4dc29cfd" + }, + "cartpole-compatible-migrated-1000/comparison.json": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/followup-alignment-v1/cartpole-compatible-migrated-1000/comparison.json", + "sha256": "4809c830cb0ac77833975de449347dc065f09d06053797ae7c5910dba41bf63f" + }, + "ant-transport/completion.json": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/followup-alignment-v1/ant-transport/completion.json", + "sha256": "996292db1f19cb941f59c4a3e122b62416c992430a08de7fd447e50d2261e17e" + }, + "ant-transport/summary.json": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/followup-alignment-v1/ant-transport/summary.json", + "sha256": "a90e74fe7ca0a97bf3a4da4fdaee41a6b55d7959005ec650bb433d9a946f2698" + }, + "ant-transport-training-1000/completion.json": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/followup-alignment-v1/ant-transport-training-1000/completion.json", + "sha256": "e9cd0c706dfb71825abe2314d0d9f53dd4575d400e73d2081fbbf4e356659629" + }, + "ant-transport-training-1000/summary.json": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/followup-alignment-v1/ant-transport-training-1000/summary.json", + "sha256": "f683325587a608937bde491a36204fc6694dce5fe77b021c5266d2cece720325" + }, + "observation-official-entry-one-step/regression.json": { + "path": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/followup-alignment-v1/observation-official-entry-one-step/regression.json", + "sha256": "d1680d33b463d39a93890b99338b24aadd291aeba19e44b789669ba203476564" + } + }, + "resources": { + "full40_queue": "stopped; do not restart", + "new5000_steps_or_model_seeds": 0, + "all_followup_training_cpu": true, + "allowed_physical_gpus": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6 + ] + }, + "remaining": "Explicit numerical policy is validated for Cartpole CPU under a shared profile; preserve official raw targets/velocity for exact Ant source-recipe data fidelity before any new across-seed claim. No original-default full convergence, GPU or Newton accuracy claim.", + "verification": "/workspace/exp_data/embodichain-pr822-accuracy-20261020/followup-alignment-v1/verification.json", + "peer_review": { + "ant_runtime_transport": "passed after correcting double-normalized gradient probe; baseline/sham/content hashes verified", + "cartpole_operator_attribution": "passed; actual ten batches/SDPA/norm/Adam verified independently", + "cartpole_production_1000": "passed; parameter/Adam tensor byte identity and exact test-reduction replay verified independently", + "ant_transport_training_1000": "passed; raw selfcheck, same-runtime evaluator,7 training batches/91 test batches, and causal limits reviewed" + }, + "figures": { + "figures/ant-training-attribution.png": "69d0f0be06506a1050568b4a7187154e91bbaee19ff7830e4b33aef31a738c75", + "figures/ant-training-attribution.svg": "5aad38b21c92846fb1f04d9bdc8d88fb8f01dacca85fb1b00dfa10ebc24c8b71" + } +} diff --git a/design/reports/neural-dynamics-numerical-attribution-20261020.md b/design/reports/neural-dynamics-numerical-attribution-20261020.md new file mode 100644 index 000000000..52bdd174d --- /dev/null +++ b/design/reports/neural-dynamics-numerical-attribution-20261020.md @@ -0,0 +1,86 @@ +# NeRD 数值差异归因与 1,000 步复核 + +完成时间：2026-10-21 00:05（Asia/Shanghai）。实验目录沿用起始日 20261020。 + +三个子 agent 完成并交叉审查了训练实现、Torch 数值路径和 Ant 数据运输三条链路。**相同当前环境、相同 canonical 数据下，两机器人迁移训练更新均已验证到 1,000 步完全一致。Cartpole 在一套明确的共同数值策略下，也完成了原环境与真实迁移训练器的 1,000 步对齐。Ant 另外证实了速度运输舍入对训练后长预测的影响。** + +本报告接续[原始入口小预算比较](neural-dynamics-small-budget-20261020.md)。旧官方默认基线、原主实验数据／权重／日志全部保留；共同策略是独立受控变体，未替换旧基线、修改生产默认数值策略或恢复 40 组队列。全部新增训练使用 CPU，模型规模、B512／H10、种子 0 和实际官方采样索引均保持原约定；没有增加 5,000 步训练或模型种子。 + +## 同环境训练实现 + +Ant 将独立复制的第 10 步 source-method 控制精确恢复到 1,000 步。全部 1,000 条 loss／学习率一致；10／250／500／750／1,000 步的 59 项模型 state、135 项 Adam 张量与参数组完全一致，最大差为 0，固定测试归一化 MSE 为 `0.025661787018179893`。Cartpole 的同环境 1,000 步参数与 Adam 一致性已在上一轮完成。 + +该结论针对同一 canonical 数据协议：源方法也通过 transport 后读取共同数据，不能据此把 Ant 原始 legacy 数据与 canonical 数据的差异全归到 Torch。Ant 的日志 `grad_norm` 仍有约 `1.49e-6` 差异，来自报告字段的不同范数归约；实际模型和优化器更新一致。 + +## Cartpole：找到两处数值分叉并跑通共同策略 + +默认 CPU 路径的首个前向差异位于第一层 SDPA 输出，最大约 `1.79e-7`；此前线性层、LayerNorm 和 q／k／v 均一致。完全相同梯度跨两版重放裁剪和 Adam 时，更新一致，因此首步 Adam 实现不是分叉源。接近 `eps=1e-8` 的梯度扰动会被同一 Adam 公式放大，复现了首步约 `4.61e-6` 的参数差。 + +共同指定 math SDPA 后，第 4 步仍出现归约差异：对同一个 45 元素 FP32 参数范数向量，Torch 2.2／2.7 的总范数分别为 `1.6617134809494019`／`1.6617133617401123`，相差 1 ULP；关闭 foreach 无效。共同将这 45 个 FP32 范数的总平方和以固定顺序 FP64 累加，再转回 FP32 裁剪，消除了已观察的第二处差异。 + +在这个新共同策略下，独立运行原方法 Torch 2.2、原方法 Torch 2.7，以及当前环境生产 `DynamicsTrainer.step` 三个 1,000 步控制。真实迁移训练器没有重写，策略只在独立 worker 的上下文中注入。 + +| 检查 | 结果 | +|---|---| +| 1,000 条 loss／裁剪总范数／LR | 逐值一致 | +| 250／500／750／1,000 步模型与 Adam 张量 | 45 参数、135 Adam 张量，原始字节均一致 | +| 冻结归一化 mean／var | 6 项张量一致 | +| 同源原始固定测试归一化 MSE | 三者均 `0.007527741757068006` | +| 生产离线评估 MSE | `0.007527743803930806` | + +测试末两项相差约 `2.05e-9`。独立复核对同一组 91 批均值分别执行原 FP32 累加和生产 Python 加权累加，准确复现了该差异。Torch 2.7 参数组额外存在 `decoupled_weight_decay=False`，已单列为版本元数据；这里的 Adam 一致性指张量状态和共同有效超参数，而非跨版本字典结构完全相同。 + +**共同策略不复现原官方默认 flash 轨迹。** 默认官方 Cartpole 测试 MSE `0.006280350161122752` 仍保留；不能拿新的 `0.00752774` 替换它并宣称默认精度通过。本证据限于一个种子、CPU、1,000 步和已直接检查的阶段，不涉及任意设备或最终收敛。 + +## Ant：区分推理运输与训练运输 + +原始 spatial velocity 与 native COM velocity 的减后加在 FP32 下不能保证逐位恢复。首个真实训练批的目标最大差 `9.536743164e-7` 已精确复现；EC target 与往返后的原始 target 一致。固定权重时，在两模型、三场景各 16 条轨迹中，每宏步增加速度往返后，完整 101 个状态点的 q、COM qd 和控制均不变，1／10／100 步指标和失败计数不变。纯状态回写 sham 也与原轨迹完全一致。 + +梯度探针最初误用了原模型 forward 后已归一化的冻结快照，导致再次归一化。独立审查发现后，旧脚本／结果已完整归档并标 INVALID；更正版本从原 HDF5 重建真实输入，经一次预处理和一次模型归一化。更正后的首批目标扰动使初始化／官方 1,000 步模型／迁移 1,000 步模型的最大梯度差分别约 `9.31e-10`／`1.86e-9`／`1.40e-9`，没有执行优化器更新。原目标和 rollout 证据未改。 + +随后只新增一项官方支持环境 Torch 2.2 的 1,000 步运输控制 B；旧原始训练 A 和旧当前环境迁移 D 直接复用。新 loop 首先在无运输条件下精确复现旧官方第 1／10 步模型和 loss；B 再仅在原预处理前对 states／next_states 做 spatial→COM→spatial。原初值、RMS、记录索引、默认注意力／裁剪、Adam、FP32／8 线程不变。 + +7 个实际完整训练批的输入／目标与 canonical 逐字节相同，在这 7 批内 A／B 输入相同、差异位于目标；全部 91 个完整测试批也完成 canonical 一致性核验。训练批检查是抽查，不能表述为全部 1,000 批输入逐值证明。 + +三套权重统一在同一个支持环境原始 evaluator 上评估 46,592 个固定测试窗口： + +| 训练条件 | 共同 raw-test MSE | 共同 roundtrip-test MSE | +|---|---:|---:| +| A：官方 2.2／原始 raw | 0.025272413925 | 0.025272413925 | +| B：官方 2.2／速度往返 | 0.025085947016 | 0.025085947016 | +| D：当前 2.7／迁移 canonical | 0.025661790764 | 0.025661790764 | + +运输控制 B 的单步测试略改善，但相同 held-out 轨迹上的 100 步关节位置误差增加： + +| 场景 | A qMAE | B qMAE | D qMAE | A／B／D 失败数 | +|---|---:|---:|---:|---| +| passive | 0.472263 | 0.578930 | 0.568879 | 12／14／15（各 16 条） | +| random_effort | 0.334289 | 0.420823 | 0.455908 | 11／12／12（各 16 条） | +| fixed_policy | 0.459634 | 0.488172 | 0.500939 | 12／14／11（各 16 条） | + +![Ant 单步测试与长预测的不同响应](/workspace/exp_data/embodichain-pr822-accuracy-20261020/followup-alignment-v1/figures/ant-training-attribution.png) + +这证明运输舍入经训练后可对本种子的长预测产生可见影响；不能把 Ant 的全部长预测差异归到 Torch，也不能把 A→B、B→D 变成全局可加性的因果百分比。这些对照不是当前 Newton 与旧 Warp 物理差异的测量，没有重新生成 teacher 或调整物理参数。所有失败轨迹均保留。 + +D 日志第 14／19 步的 LR 有一个 float64 ULP 差，但全部有效 FP32 LR 和带偏差修正的 Adam 步长一致；D 恢复为原模型格式时，RMS count 为默认先验而非旧估计器样本数，但 mean／var 全等，冻结 normalize 不读 count。这两项元数据差异均显式记录，未改检查点。 + +## 观测修复与验证 + +本轮只修改实验观测入口和相应测试：首批物理特征在原 loss／forward 前保存，使用独立 CPU clone，写入包含数据阶段和 tensor SHA 的 sidecar；消费者按新 physical stage 或旧的无 sidecar normalized stage 比较。未知格式／阶段和错误 SHA 会拒绝。原冻结快照未重写。 + +10 个 focused tests 全部通过，覆盖回调在归一化前捕获、旧／新数据兼容、真实差异检测和错误 metadata。另在 Python 3.8／Torch 2.2 实际运行官方原入口 1 步，模型参数、统计量和 loss 与冻结第 1 步完全一致；实际完整首批的新／旧消费者比较也均为 0 差。Black 26.3.1 全树检查、Python3.8 编译、whitespace 和 context-map 检查通过。生产训练／物理／codec 默认值未改，原 source tracked files 无变动。 + +项目上下文选择不更新：本轮改变的是独立实验观察数据的阶段标识和复现说明，生产层级、配置路由与物理契约仍准确；实验 README 已更新。全部原始比较、当前同环境控制、共同数值策略控制与运输控制均分开记录。 + +## 后续建议 + +下一步应将已验证的共同数值策略整理为显式实验配置，并为官方训练复现保留原始 spatial velocity／target 的保真通道，减少导入后的训练舍入变量。确认这两项后再考虑增加一个配对种子；本轮没有启动 5,000 步或恢复 40 组。当前证据支持“迁移更新逻辑对齐”，不支持“原默认官方训练、GPU 与当前 Newton 长预测全部精度验收通过”。 + +详细证据： +- [Ant 同环境 1,000 步](/workspace/exp_data/embodichain-pr822-accuracy-20261020/followup-alignment-v1/ant-same-environment/README.md) +- [Cartpole 算子诊断](/workspace/exp_data/embodichain-pr822-accuracy-20261020/followup-alignment-v1/cartpole-numerics/REPORT.md) +- [Cartpole 跨环境共同策略](/workspace/exp_data/embodichain-pr822-accuracy-20261020/followup-alignment-v1/cartpole-compatible-1000/REPORT.md) +- [真实迁移共同策略闭环](/workspace/exp_data/embodichain-pr822-accuracy-20261020/followup-alignment-v1/cartpole-compatible-migrated-1000/REPORT.md) +- [Ant 推理运输／更正说明](/workspace/exp_data/embodichain-pr822-accuracy-20261020/followup-alignment-v1/ant-transport/report.md) +- [Ant 训练运输与三模型共同评估](/workspace/exp_data/embodichain-pr822-accuracy-20261020/followup-alignment-v1/ant-transport-training-1000/REPORT.md) +- [本轮机器汇总](/workspace/exp_data/embodichain-pr822-accuracy-20261020/followup-alignment-v1/summary.json) diff --git a/design/reports/neural-dynamics-official-model-alignment-20261020.json b/design/reports/neural-dynamics-official-model-alignment-20261020.json new file mode 100644 index 000000000..8047b8c04 --- /dev/null +++ b/design/reports/neural-dynamics-official-model-alignment-20261020.json @@ -0,0 +1,507 @@ +{ + "scope": "Official NeRD original model is the primary reference; Newton teacher experiments are auxiliary only.", + "official_source_commit": "370e89d2b8ce68aa3578aa7812859c681b7502bc", + "status": "model_port_verified_with_matching_torch; full_pipeline_and_training_accuracy_pending", + "robots": { + "cartpole": { + "matched-torch": { + "scope": "Original pretrained model versus migrated copy of the same weights; official supported environment generates comparison data. No Newton teacher.", + "reference": { + "source_commit": "370e89d2b8ce68aa3578aa7812859c681b7502bc", + "robot": "cartpole", + "python": "3.8.20", + "torch": "2.2.2+cpu", + "warp": "1.8.0", + "numpy": "1.24.4", + "device": "cpu", + "dtype": "float32", + "count": 32, + "horizon": 100, + "checkpoint_sha256": "287dde2134eb8f5277c17a8f7cc7abbe64063e97cb481bdb6e8d76f9536e5366", + "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": "ad34247ffa9b0d12ef9b0c56a86cb0d9bf91f19743dd3d2fa7e170954dbba43e", + "test.raw.h5": "a27e18397a13f4feb5d48945776b80e02af6ab33c27f27cc8f1692c3b0481c4c", + "official-rollout.npz": "e226ca2a8e9e4872a8edce0d9dcd1cbefebacaa4fb23bed6c7860617e200d2d0" + }, + "model_inputs": [ + "states_embedding", + "joint_acts" + ], + "state_frame": "world", + "body_com": [ + 0.0, + 0.0, + 0.0 + ] + }, + "migrated_environment": { + "python": "3.10.12", + "torch": "2.2.2+cpu", + "warp": "1.8.0" + }, + "model_sha256": "9908eab6ce6312f5e8cc77dcc7272b76e76178ba01907bc1f78b2468d719e88d", + "same_original_checkpoint": true, + "histories": { + "1": { + "windows": 64, + "direct_model_max_abs": 0.0, + "codec_and_model_max_abs": 4.76837158203125e-07, + "target_max_abs": 2.384185791015625e-07, + "feature_max_abs": { + "states_embedding": 2.384185791015625e-07, + "joint_f": 0.0 + }, + "official_normalized_mse": 2.3746521037537605e-05, + "migrated_normalized_mse": 2.3746797523926944e-05, + "relative_mse_change": 1.1682510375976562e-05, + "forward_tolerance_passed": true + }, + "9": { + "windows": 64, + "direct_model_max_abs": 0.0, + "codec_and_model_max_abs": 1.3113021850585938e-06, + "target_max_abs": 2.384185791015625e-07, + "feature_max_abs": { + "states_embedding": 2.384185791015625e-07, + "joint_f": 0.0 + }, + "official_normalized_mse": 1.5774938219692558e-05, + "migrated_normalized_mse": 1.577491821080912e-05, + "relative_mse_change": -1.2516975402832031e-06, + "forward_tolerance_passed": true + }, + "10": { + "windows": 64, + "direct_model_max_abs": 0.0, + "codec_and_model_max_abs": 9.5367431640625e-07, + "target_max_abs": 2.384185791015625e-07, + "feature_max_abs": { + "states_embedding": 2.384185791015625e-07, + "joint_f": 0.0 + }, + "official_normalized_mse": 0.001277043716982007, + "migrated_normalized_mse": 0.0012770441826432943, + "relative_mse_change": 3.5762786865234375e-07, + "forward_tolerance_passed": true + } + }, + "rollouts": { + "1": { + "max_abs_state_difference": 0.0, + "mean_abs_state_difference": 0.0, + "all_finite": true + }, + "10": { + "max_abs_state_difference": 0.0, + "mean_abs_state_difference": 0.0, + "all_finite": true + }, + "100": { + "max_abs_state_difference": 0.0, + "mean_abs_state_difference": 0.0, + "all_finite": true + } + }, + "current_simulator_imports": [], + "qualification": "bounded_real_data_model_port_comparison_not_full_training_quality" + }, + "current-torch": { + "scope": "Original pretrained model versus migrated copy of the same weights; official supported environment generates comparison data. No Newton teacher.", + "reference": { + "source_commit": "370e89d2b8ce68aa3578aa7812859c681b7502bc", + "robot": "cartpole", + "python": "3.8.20", + "torch": "2.2.2+cpu", + "warp": "1.8.0", + "numpy": "1.24.4", + "device": "cpu", + "dtype": "float32", + "count": 32, + "horizon": 100, + "checkpoint_sha256": "287dde2134eb8f5277c17a8f7cc7abbe64063e97cb481bdb6e8d76f9536e5366", + "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": "ad34247ffa9b0d12ef9b0c56a86cb0d9bf91f19743dd3d2fa7e170954dbba43e", + "test.raw.h5": "a27e18397a13f4feb5d48945776b80e02af6ab33c27f27cc8f1692c3b0481c4c", + "official-rollout.npz": "e226ca2a8e9e4872a8edce0d9dcd1cbefebacaa4fb23bed6c7860617e200d2d0" + }, + "model_inputs": [ + "states_embedding", + "joint_acts" + ], + "state_frame": "world", + "body_com": [ + 0.0, + 0.0, + 0.0 + ] + }, + "migrated_environment": { + "python": "3.11.14", + "torch": "2.7.0+cu128", + "warp": "1.8.0" + }, + "model_sha256": "9908eab6ce6312f5e8cc77dcc7272b76e76178ba01907bc1f78b2468d719e88d", + "same_original_checkpoint": true, + "histories": { + "1": { + "windows": 64, + "direct_model_max_abs": 0.0, + "codec_and_model_max_abs": 4.76837158203125e-07, + "target_max_abs": 2.384185791015625e-07, + "feature_max_abs": { + "states_embedding": 2.384185791015625e-07, + "joint_f": 0.0 + }, + "official_normalized_mse": 2.3746521037537605e-05, + "migrated_normalized_mse": 2.3746797523926944e-05, + "relative_mse_change": 1.1682510375976562e-05, + "forward_tolerance_passed": true + }, + "9": { + "windows": 64, + "direct_model_max_abs": 5.960464477539062e-07, + "codec_and_model_max_abs": 1.5497207641601562e-06, + "target_max_abs": 2.384185791015625e-07, + "feature_max_abs": { + "states_embedding": 2.384185791015625e-07, + "joint_f": 0.0 + }, + "official_normalized_mse": 1.5774938219692558e-05, + "migrated_normalized_mse": 1.577476541569922e-05, + "relative_mse_change": -1.0967254638671875e-05, + "forward_tolerance_passed": true + }, + "10": { + "windows": 64, + "direct_model_max_abs": 9.5367431640625e-07, + "codec_and_model_max_abs": 9.5367431640625e-07, + "target_max_abs": 2.384185791015625e-07, + "feature_max_abs": { + "states_embedding": 2.384185791015625e-07, + "joint_f": 0.0 + }, + "official_normalized_mse": 0.001277043716982007, + "migrated_normalized_mse": 0.0012770409230142832, + "relative_mse_change": -2.205371856689453e-06, + "forward_tolerance_passed": true + } + }, + "rollouts": { + "1": { + "max_abs_state_difference": 0.0, + "mean_abs_state_difference": 0.0, + "all_finite": true + }, + "10": { + "max_abs_state_difference": 2.384185791015625e-06, + "mean_abs_state_difference": 7.98434456328323e-08, + "all_finite": true + }, + "100": { + "max_abs_state_difference": 8.96453857421875e-05, + "mean_abs_state_difference": 2.113152049787459e-06, + "all_finite": true + } + }, + "current_simulator_imports": [], + "qualification": "bounded_real_data_model_port_comparison_not_full_training_quality" + }, + "largest_current_runtime_difference": { + "step": 100, + "trajectory": 31, + "raw_state_channel": 3, + "difference": 8.96453857421875e-05 + } + }, + "ant": { + "matched-torch": { + "scope": "Original pretrained model versus migrated copy of the same weights; official supported environment generates comparison data. No Newton teacher.", + "reference": { + "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 + ] + }, + "migrated_environment": { + "python": "3.10.12", + "torch": "2.2.2+cpu", + "warp": "1.5.1" + }, + "model_sha256": "c8d8c46cf3672a1e38fcc6544651e8a33e0238dda1090a04b1b49434784cb09a", + "same_original_checkpoint": true, + "histories": { + "1": { + "windows": 64, + "direct_model_max_abs": 0.0, + "codec_and_model_max_abs": 0.00040435791015625, + "target_max_abs": 6.794929504394531e-06, + "feature_max_abs": { + "states_embedding": 1.9073486328125e-06, + "contact_normals": 2.980232238769531e-07, + "contact_points_1": 1.7881393432617188e-06, + "contact_depths": 0.0, + "joint_f": 0.0, + "gravity_dir": 2.980232238769531e-07 + }, + "official_normalized_mse": 0.000704414036590606, + "migrated_normalized_mse": 0.0007043941295705736, + "relative_mse_change": -2.8252601623535156e-05, + "forward_tolerance_passed": false + }, + "9": { + "windows": 64, + "direct_model_max_abs": 0.0, + "codec_and_model_max_abs": 0.00033402442932128906, + "target_max_abs": 8.52346420288086e-06, + "feature_max_abs": { + "states_embedding": 2.1457672119140625e-06, + "contact_normals": 4.172325134277344e-07, + "contact_points_1": 1.6093254089355469e-06, + "contact_depths": 0.0, + "joint_f": 0.0, + "gravity_dir": 4.172325134277344e-07 + }, + "official_normalized_mse": 0.007895898073911667, + "migrated_normalized_mse": 0.007895909249782562, + "relative_mse_change": 1.430511474609375e-06, + "forward_tolerance_passed": false + }, + "10": { + "windows": 64, + "direct_model_max_abs": 0.0, + "codec_and_model_max_abs": 0.000457763671875, + "target_max_abs": 1.4722347259521484e-05, + "feature_max_abs": { + "states_embedding": 2.1457672119140625e-06, + "contact_normals": 3.5762786865234375e-07, + "contact_points_1": 1.7285346984863281e-06, + "contact_depths": 0.0, + "joint_f": 0.0, + "gravity_dir": 3.5762786865234375e-07 + }, + "official_normalized_mse": 0.007553629111498594, + "migrated_normalized_mse": 0.007553651463240385, + "relative_mse_change": 2.9802322387695312e-06, + "forward_tolerance_passed": false + } + }, + "rollouts": { + "1": { + "max_abs_state_difference": 0.0, + "mean_abs_state_difference": 0.0, + "all_finite": true + }, + "10": { + "max_abs_state_difference": 0.0, + "mean_abs_state_difference": 0.0, + "all_finite": true + }, + "100": { + "max_abs_state_difference": 0.0, + "mean_abs_state_difference": 0.0, + "all_finite": true + } + }, + "current_simulator_imports": [], + "qualification": "bounded_real_data_model_port_comparison_not_full_training_quality" + }, + "current-torch": { + "scope": "Original pretrained model versus migrated copy of the same weights; official supported environment generates comparison data. No Newton teacher.", + "reference": { + "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 + ] + }, + "migrated_environment": { + "python": "3.11.14", + "torch": "2.7.0+cu128", + "warp": "1.5.1" + }, + "model_sha256": "c8d8c46cf3672a1e38fcc6544651e8a33e0238dda1090a04b1b49434784cb09a", + "same_original_checkpoint": true, + "histories": { + "1": { + "windows": 64, + "direct_model_max_abs": 0.0, + "codec_and_model_max_abs": 0.00040435791015625, + "target_max_abs": 6.794929504394531e-06, + "feature_max_abs": { + "states_embedding": 1.9073486328125e-06, + "contact_normals": 2.980232238769531e-07, + "contact_points_1": 1.7881393432617188e-06, + "contact_depths": 0.0, + "joint_f": 0.0, + "gravity_dir": 2.980232238769531e-07 + }, + "official_normalized_mse": 0.000704414036590606, + "migrated_normalized_mse": 0.0007043941295705736, + "relative_mse_change": -2.8252601623535156e-05, + "forward_tolerance_passed": false + }, + "9": { + "windows": 64, + "direct_model_max_abs": 3.8623809814453125e-05, + "codec_and_model_max_abs": 0.00033473968505859375, + "target_max_abs": 8.52346420288086e-06, + "feature_max_abs": { + "states_embedding": 2.1457672119140625e-06, + "contact_normals": 4.172325134277344e-07, + "contact_points_1": 1.6093254089355469e-06, + "contact_depths": 0.0, + "joint_f": 0.0, + "gravity_dir": 4.172325134277344e-07 + }, + "official_normalized_mse": 0.007895898073911667, + "migrated_normalized_mse": 0.007895910181105137, + "relative_mse_change": 1.5497207641601562e-06, + "forward_tolerance_passed": false + }, + "10": { + "windows": 64, + "direct_model_max_abs": 1.33514404296875e-05, + "codec_and_model_max_abs": 0.0004630088806152344, + "target_max_abs": 1.4722347259521484e-05, + "feature_max_abs": { + "states_embedding": 2.1457672119140625e-06, + "contact_normals": 3.5762786865234375e-07, + "contact_points_1": 1.7285346984863281e-06, + "contact_depths": 0.0, + "joint_f": 0.0, + "gravity_dir": 3.5762786865234375e-07 + }, + "official_normalized_mse": 0.007553629111498594, + "migrated_normalized_mse": 0.007553649600595236, + "relative_mse_change": 2.7418136596679688e-06, + "forward_tolerance_passed": false + } + }, + "rollouts": { + "1": { + "max_abs_state_difference": 1.1920928955078125e-07, + "mean_abs_state_difference": 6.101768534882979e-10, + "all_finite": true + }, + "10": { + "max_abs_state_difference": 0.0019707679748535156, + "mean_abs_state_difference": 1.9508488549035974e-05, + "all_finite": true + }, + "100": { + "max_abs_state_difference": 7.820811748504639, + "mean_abs_state_difference": 0.005333449225872755, + "all_finite": true + } + }, + "current_simulator_imports": [], + "qualification": "bounded_real_data_model_port_comparison_not_full_training_quality" + }, + "largest_current_runtime_difference": { + "step": 80, + "trajectory": 10, + "raw_state_channel": 20, + "difference": 7.820811748504639 + } + } + } +} diff --git a/design/reports/neural-dynamics-official-model-alignment-20261020.md b/design/reports/neural-dynamics-official-model-alignment-20261020.md new file mode 100644 index 000000000..78dd4b6de --- /dev/null +++ b/design/reports/neural-dynamics-official-model-alignment-20261020.md @@ -0,0 +1,113 @@ +# Official NeRD model alignment — 2026-10-20 + +The reference for this follow-up is the actual pretrained NeRD model executed +from the unchanged official repository in its supported version environment. +Earlier Newton-teacher loss, parameter reconstruction and fine-tuning runs are +auxiliary diagnostics. They do not satisfy this model-comparison requirement. + +## Frozen reference + +- Repository: `NVlabs/neural-robot-dynamics`, commit + `370e89d2b8ce68aa3578aa7812859c681b7502bc`; tracked source verified unchanged. +- Official execution: Python **3.8.20**, Torch **2.2.2+cpu**, NumPy **1.24.4**. + Cartpole uses Warp **1.8.0**; Ant uses Warp **1.5.1**, as required by the + repository's pretrained-model reproduction note. +- The official objects loaded from the provided `model.pt` run directly. + Neither EmbodiChain nor DexSim/Newton is imported into reference execution. +- Each robot generates **32 × 100** transitions through the original sampler, + with seed 500. The original model produces frozen predictions for 64 real + windows at each history length 1/9/10, plus a 32-trajectory 100-step neural + rollout with history rollover. Both sides use identical weights, normalizers, + initial states, commands and feature order. +- CPU execution maps the original integrator's CPU CUDA-stream request to + Torch's no-op `None` stream. Numerical source methods remain unchanged. +- Raw HDF5 retains the upstream sampler's uninitialized non-up-axis gravity + feature. Both comparison sides explicitly replace that transport feature + with the measured physical gravity direction; the replacement is recorded. + This is not an undocumented repair of the original dataset. + +## Results + +The migrated reference run uses Python 3.10.12, Torch 2.2.2+cpu, NumPy 1.24.4 +and the same robot-specific old Warp. Python 3.10 satisfies EmbodiChain's +minimum Python version while keeping the numerical Torch/Warp dependencies +matched to the official oracle. + +| Robot | Same-feature forward maximum difference, H1/9/10 | Original-runtime neural rollout maximum difference, 32 × 100 | Status | +|---|---:|---:|---| +| Cartpole | 0 | 0 | Bounded model/history port agrees exactly | +| Ant | 0 | 0 | Bounded model/history port agrees exactly | + +This verifies migration of the provided pretrained weights. It does **not** +verify a newly trained model against the official trained model or complete +full-budget training noninferiority. + +The complete canonical-state conversion path is measured separately: + +| Robot | Maximum feature difference | Maximum prediction difference | Existing elementwise tolerance | +|---|---:|---:|---| +| Cartpole | 2.38e-7 | 1.31e-6 | Passed | +| Ant | 2.15e-6 | 4.58e-4 | **Not passed** | + +Ant's normalized-MSE relative changes are only −0.002825% through +0.000298% +in these windows, but that average does not override the failed elementwise +check (`atol=1e-5`, `rtol=1e-4`). Frame transport and twist/COM conversion remain +a precision investigation target. No tolerance was relaxed. + +A separate run retains the current environment, Python 3.11.14 / Torch +2.7.0+cu128, while still executing CPU old-Warp physics and the same reference +inputs. Its maximum same-feature output differences are 9.54e-7 for Cartpole +and 3.86e-5 for Ant. Over 100 original-runtime neural steps, raw-state maximum +component differences are 8.96e-5 and 7.82 respectively; mean absolute +component differences are 2.11e-6 and 0.00533. These are mixed-coordinate raw +state diagnostics, not meter-valued position errors. All compared trajectories +remain finite. Environment-dependent accumulation must be explained before +asserting full current-runtime equivalence; the matched-version zero result +cannot replace this check. + +## Reproduction and artifacts + +The independent official worker is +`scripts/experiments/neural_dynamics/official_model_reference.py`; the migrated +worker is `scripts/experiments/neural_dynamics/compare_official_model.py`. +Use explicit trusted source/checkpoint paths and a fresh output directory. +The official worker enforces the exact Python/Torch/Warp versions above. + +```bash +# OFFICIAL_PY points to Python 3.8.20 with Torch 2.2.2 and the correct Warp. +# For Ant, put the isolated Warp 1.5.1 installation first on PYTHONPATH. +"$OFFICIAL_PY" scripts/experiments/neural_dynamics/official_model_reference.py \ + --source "$NERD_SOURCE" --checkpoint "$OFFICIAL_MODEL_PT" \ + --robot ant --output "$REFERENCE_OUTPUT" --trust-external-code + +# PORT_PY points to supported EmbodiChain Python with the chosen Torch version. +"$PORT_PY" scripts/experiments/neural_dynamics/compare_official_model.py \ + --source "$NERD_SOURCE" --reference "$REFERENCE_OUTPUT" \ + --model "$PORTABLE_MODEL" --output "$COMPARISON_JSON" --trust-external-code +``` + +Reference receipts include source commit, checkpoint and data/prediction hashes, +actual numerical dependencies, transport decisions and an assertion that the +current simulators were not imported. The comparison verifies those hashes +before reading inputs. The original model normalizes by replacing dictionary +entries; fixture generation preserves physical features before model execution +so they are not normalized twice. + +Executed artifacts are retained under +`/root/.cache/embodichain/pr822-accuracy/strict-official-{cartpole,ant}-v3/` and +`*-official-model-{matched,current}-torch.{json,npz}`. The aggregate results are +[checked in as JSON](neural-dynamics-official-model-alignment-20261020.json). +The local paths are not published download URLs. + +## Remaining acceptance work + +1. Diagnose the measured Ant canonical conversion residual without weakening + the existing numerical gate; preserve the exact official oracle. +2. Separate Torch/runtime numerical accumulation from data/parameter and model + migration effects. Compare physical state metrics as well as raw components. +3. For **training quality**, train/evaluate the original and migrated trainers + on the same official-environment datasets, normalizers, seed matrix, update + budget and checkpoint-selection protocol. The earlier Newton-domain fine-tune + is not evidence for that requirement. +4. Complete the declared five-seed, full-data/full-budget and 2048-trajectory + long-horizon qualification before changing that acceptance status. diff --git a/design/reports/neural-dynamics-parameter-alignment-20261020.json b/design/reports/neural-dynamics-parameter-alignment-20261020.json new file mode 100644 index 000000000..2cc802827 --- /dev/null +++ b/design/reports/neural-dynamics-parameter-alignment-20261020.json @@ -0,0 +1,12367 @@ +{ + "status": "partial_not_qualified", + "identity": { + "embodichain_head": "660121cb46185703555665a53a1767ef963d84da", + "dexsim_mr": "http://192.168.3.16/Engine/dexsim/-/merge_requests/1483", + "dexsim_head": "b61e920fbb222e50213574ad57d885e7b96148b9", + "sdk_artifact_job": 302950, + "reference_head": "370e89d2b8ce68aa3578aa7812859c681b7502bc", + "sdk_wheel": "/root/.cache/embodichain/pr822-accuracy/sdk-artifacts/build_Release/lib/python_package/pip_package/dexsim_engine-0.5.1-cp311-cp311-manylinux_2_35_x86_64.whl", + "sdk_wheel_sha256": "a7f9e55584627f5809f76dcef91075674e3fe49b22602e68e24b7e02bebc608b", + "newton": "1.7.0", + "warp": "1.17.0" + }, + "parameters": { + "cartpole": { + "legacy_body_mass": [ + 50.0, + 2.5 + ], + "native_body_mass": [ + 1.0, + 10.0 + ], + "legacy_body_inertia": [ + [ + [ + 1.21833336353302, + 0.0, + 0.0 + ], + [ + 0.0, + 1.21833336353302, + 0.0 + ], + [ + 0.0, + 0.0, + 2.0933332443237305 + ] + ], + [ + [ + 0.2188541740179062, + 0.0, + 0.0 + ], + [ + 0.0, + 0.2188541740179062, + 0.0 + ], + [ + 0.0, + 0.0, + 0.011041666381061077 + ] + ] + ], + "native_body_inertia": [ + [ + 1.0, + 0.0, + 0.0, + 0.0, + 1.0, + 0.0, + 0.0, + 0.0, + 1.0 + ], + [ + 1.0, + 0.0, + 0.0, + 0.0, + 1.0, + 0.0, + 0.0, + 0.0, + 1.0 + ] + ], + "legacy_joint_armature": [ + 0.009999999776482582, + 0.009999999776482582 + ], + "native_joint_armature": [ + 0.01, + 0.01 + ], + "legacy_shape_mu": [ + 0.25, + 0.25, + 0.25, + 0.25, + 0.25, + 0.5 + ], + "native_shape_mu": [ + 1.0, + 1.0, + 1.0, + 1.0, + 1.0 + ] + }, + "ant": { + "legacy_body_mass": [ + 31.653451919555664, + 7.831550598144531, + 13.518441200256348, + 7.831550598144531, + 13.518441200256348, + 7.831550598144531, + 13.518441200256348, + 7.831550598144531, + 13.518441200256348 + ], + "native_body_mass": [ + 0.48388024694610887, + 0.03915775304929302, + 0.06759220389338635, + 0.03915775304929302, + 0.06759220389338635, + 0.03915775304929302, + 0.06759220389338635, + 0.03915775304929302, + 0.06759220389338635 + ], + "legacy_body_inertia": [ + [ + [ + 0.5960065126419067, + 0.0, + 0.0 + ], + [ + 0.0, + 0.5960065126419067, + 0.0 + ], + [ + 0.0, + 0.0, + 1.0890781879425049 + ] + ], + [ + [ + 0.07864078879356384, + -0.04495241865515709, + 0.0 + ], + [ + -0.04495241865515709, + 0.07864080369472504, + 0.0 + ], + [ + 0.0, + 0.0, + 0.12359321862459183 + ] + ], + [ + [ + 0.298422247171402, + -0.24653588235378265, + 0.0 + ], + [ + -0.24653588235378265, + 0.29842230677604675, + 0.0 + ], + [ + 0.0, + 0.0, + 0.5449581742286682 + ] + ], + [ + [ + 0.07864078879356384, + 0.04495241865515709, + 0.0 + ], + [ + 0.04495241865515709, + 0.07864080369472504, + 0.0 + ], + [ + 0.0, + 0.0, + 0.12359321862459183 + ] + ], + [ + [ + 0.298422247171402, + 0.24653588235378265, + 0.0 + ], + [ + 0.24653588235378265, + 0.29842230677604675, + 0.0 + ], + [ + 0.0, + 0.0, + 0.5449581742286682 + ] + ], + [ + [ + 0.07864080369472504, + -0.044952426105737686, + 0.0 + ], + [ + -0.04495241865515709, + 0.07864079624414444, + 0.0 + ], + [ + 0.0, + 0.0, + 0.12359318882226944 + ] + ], + [ + [ + 0.29842233657836914, + -0.24653591215610504, + 0.0 + ], + [ + -0.24653591215610504, + 0.29842230677604675, + 0.0 + ], + [ + 0.0, + 0.0, + 0.5449580550193787 + ] + ], + [ + [ + 0.07864080369472504, + 0.044952426105737686, + 0.0 + ], + [ + 0.04495241865515709, + 0.07864079624414444, + 0.0 + ], + [ + 0.0, + 0.0, + 0.12359318882226944 + ] + ], + [ + [ + 0.29842233657836914, + 0.24653591215610504, + 0.0 + ], + [ + 0.24653591215610504, + 0.29842230677604675, + 0.0 + ], + [ + 0.0, + 0.0, + 0.5449580550193787 + ] + ] + ], + "native_body_inertia": [ + [ + 0.011120358482003212, + 0.0, + 0.0, + 0.0, + 0.011120358482003212, + 0.0, + 0.0, + 0.0, + 0.013585715554654598 + ], + [ + 0.0003432039520703256, + -0.00022476207232102752, + -9.245874899033169e-12, + -0.00022476207232102752, + 0.0003432039520703256, + -1.590164978859665e-11, + -2.4111877425236905e-11, + -1.894602703889081e-11, + 0.000567966082599014 + ], + [ + 0.0014421113301068544, + -0.001232679234817624, + -9.025284686270396e-11, + -0.001232679234817624, + 0.0014421113301068544, + -9.082672808302661e-11, + -1.0758841734181956e-10, + -1.0390709820340405e-10, + 0.002674790797755122 + ], + [ + 0.0003432039520703256, + 0.00022476207232102752, + 9.245874899033169e-12, + 0.00022476207232102752, + 0.0003432039520703256, + -1.590164978859665e-11, + 2.4111877425236905e-11, + -1.894602703889081e-11, + 0.000567966082599014 + ], + [ + 0.0014421113301068544, + 0.001232679234817624, + 9.025284686270396e-11, + 0.001232679234817624, + 0.0014421113301068544, + -9.082672808302661e-11, + 1.0758841734181956e-10, + -1.0390709820340405e-10, + 0.002674790797755122 + ], + [ + 0.0003432039520703256, + -0.00022476207232102752, + 9.245874899033169e-12, + -0.00022476207232102752, + 0.0003432039520703256, + 1.590164978859665e-11, + 2.4111877425236905e-11, + 1.894602703889081e-11, + 0.000567966082599014 + ], + [ + 0.0014421113301068544, + -0.001232679234817624, + 9.025284686270396e-11, + -0.001232679234817624, + 0.0014421113301068544, + 9.082672808302661e-11, + 1.0758841734181956e-10, + 1.0390709820340405e-10, + 0.002674790797755122 + ], + [ + 0.0003432039520703256, + 0.00022476207232102752, + -9.245874899033169e-12, + 0.00022476207232102752, + 0.0003432039520703256, + 1.590164978859665e-11, + -2.4111877425236905e-11, + 1.894602703889081e-11, + 0.000567966082599014 + ], + [ + 0.0014421113301068544, + 0.001232679234817624, + -9.025284686270396e-11, + 0.001232679234817624, + 0.0014421113301068544, + 9.082672808302661e-11, + -1.0758841734181956e-10, + 1.0390709820340405e-10, + 0.002674790797755122 + ] + ], + "legacy_joint_armature": [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.009999999776482582, + 0.009999999776482582, + 0.009999999776482582, + 0.009999999776482582, + 0.009999999776482582, + 0.009999999776482582, + 0.009999999776482582, + 0.009999999776482582 + ], + "native_joint_armature": [ + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.01, + 0.01, + 0.01, + 0.01, + 0.01, + 0.01, + 0.01, + 0.01 + ], + "legacy_shape_mu": [ + 0.75, + 0.75, + 0.75, + 0.75, + 0.75, + 0.75, + 0.75, + 0.75, + 0.75, + 0.75, + 0.75, + 0.75, + 0.75, + 0.5 + ], + "native_shape_mu": [ + 1.5, + 1.5, + 1.5, + 1.5, + 1.5, + 1.5, + 1.5, + 1.5, + 1.5, + 1.5, + 1.5, + 1.5, + 1.5 + ] + } + }, + "physics_comparison": { + "cartpole": { + "physics": { + "robot": "cartpole", + "reference": { + "source": { + "root": "/root/.cache/embodichain/pr822-accuracy/official-pinned", + "commit": "370e89d2b8ce68aa3578aa7812859c681b7502bc", + "files": { + "algorithms/vanilla_trainer.py": "7027fbc0d543479b7b4a558114b6ae3a89b93729049f9edcc031299398cbc3fd", + "integrators/integrator_neural.py": "2c42402a702f04a7b9d2a7880f1c6cd5ba57488e04eb78828d786713d1208a6b", + "models/models.py": "5956cadc0390830376c8744f66c80dec94b29465dbd08a5d3c8c4c2f1f14b8f6", + "models/model_transformer.py": "8d8fad21b296c7b54e7e42052fdca7363b5040534481f6d60f415cc073e5820d", + "models/base_models.py": "07313429e3b107db602d19e513796a5d37c4dbf418ba275d50f9e48c27cb3414", + "utils/torch_utils.py": "98492259e1ee2268a02cbd71596283502c1ef2b72672960c80a13da2fb00d89f", + "utils/running_mean_std.py": "c7229216dd77dd8573c2db238b6e3df51f85a05286ad85b6f905b2c0e4bde651" + }, + "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" + }, + "warp": "1.8.0", + "substeps": 5, + "mass_update_interval": 5, + "gravity": [ + -0.0, + -9.81, + -0.0 + ], + "engine": "legacy", + "profile": "warp_reference", + "robot": "cartpole", + "inputs_sha256": "2118cb8143b912a5861a453c0f17c7c6d2ab57e9cc817e0787e64ad509d4b36d", + "trace_sha256": "37e86180a312e3b53a0d5303dd129794a40689ddd597274970bb0e230401d6e0", + "horizon": 1000, + "count": 8 + }, + "actual": { + "robot": "cartpole", + "dt": 0.016666666666666666, + "source_commit": "476dc67fccc6a99be293bf611f897b1c06f13e30", + "asset_sha256": "f31276531c03050a3ac3a3e24443899f38a4424586db17140978383433dec237", + "joint_names": [ + "slider_to_cart", + "cart_to_pole" + ], + "q_dim": 2, + "qd_dim": 2, + "default_q": [ + 0.0, + 0.0 + ], + "default_qd": [ + 0.0, + 0.0 + ], + "effort_limits": [ + 1500.0, + 0.0 + ], + "codec": { + "robot": "cartpole", + "base_com": [ + 0.0, + 0.0, + 0.0 + ], + "continuous_q_indices": [ + 1 + ] + }, + "teacher_profile": "current_newton_import_generalized_effort", + "source_actuators_imported": false, + "source_solver_options_imported": false, + "physics_profile": "warp_reference", + "physics_signature": "b44acc2eceab2d5c9f301bbfac5dd53f32d90977ef4d6adb75ad4c1e0862f073", + "physics_signature_schema": "nerd-reference-physics-v2", + "gravity": [ + 0.0, + 0.0, + -9.8100004196167 + ], + "contact_schema": "nerd-static-plane-v1", + "teacher": { + "solver": { + "solver_type": "featherstone", + "update_mass_matrix_interval": 5 + }, + "num_substeps": 5, + "physics_profile": "warp_reference", + "versions": { + "newton": "1.7.0", + "warp-lang": "1.17.0", + "torch": "2.7.0+cu128", + "dexsim-engine": "0.5.1" + } + }, + "teacher_origin": { + "kind": "current_newton", + "device": "cpu", + "solver": { + "solver_type": "featherstone", + "update_mass_matrix_interval": 5 + }, + "num_substeps": 5, + "versions": { + "newton": "1.7.0", + "warp-lang": "1.17.0", + "torch": "2.7.0+cu128", + "dexsim-engine": "0.5.1" + }, + "dt": 0.016666666666666666, + "physics_signature": "b44acc2eceab2d5c9f301bbfac5dd53f32d90977ef4d6adb75ad4c1e0862f073", + "physics_signature_schema": "nerd-reference-physics-v2" + }, + "dexsim_commit": "b61e920fbb222e50213574ad57d885e7b96148b9", + "engine": "newton", + "profile": "warp_reference", + "inputs_sha256": "2118cb8143b912a5861a453c0f17c7c6d2ab57e9cc817e0787e64ad509d4b36d", + "trace_sha256": "13c2770aea58bb1bbdb03f43bcbf322896a1129b356b732f08c2625f3eefb934", + "horizon": 1000, + "count": 8 + }, + "horizons": { + "1": { + "finite_trajectories": 8, + "nonfinite_trajectories": 0, + "errors": { + "q_mae": 0.0, + "qd_mae": 1.1641532182693481e-08, + "base_position_error": 0.0, + "base_orientation_error_rad": 0.0 + } + }, + "10": { + "finite_trajectories": 8, + "nonfinite_trajectories": 0, + "errors": { + "q_mae": 3.914465196430683e-09, + "qd_mae": 4.459725460037589e-08, + "base_position_error": 0.0, + "base_orientation_error_rad": 0.0 + } + }, + "100": { + "finite_trajectories": 8, + "nonfinite_trajectories": 0, + "errors": { + "q_mae": 3.026152262464166e-07, + "qd_mae": 9.274093827116303e-07, + "base_position_error": 0.0, + "base_orientation_error_rad": 0.0 + } + }, + "1000": { + "finite_trajectories": 8, + "nonfinite_trajectories": 0, + "errors": { + "q_mae": 0.004342447319500934, + "qd_mae": 0.010397311032855214, + "base_position_error": 0.0, + "base_orientation_error_rad": 0.0 + } + } + }, + "population": "All requested trajectories; errors remain undefined when any trajectory becomes nonfinite." + }, + "native": { + "robot": "cartpole", + "reference": { + "source": { + "root": "/root/.cache/embodichain/pr822-accuracy/official-pinned", + "commit": "370e89d2b8ce68aa3578aa7812859c681b7502bc", + "files": { + "algorithms/vanilla_trainer.py": "7027fbc0d543479b7b4a558114b6ae3a89b93729049f9edcc031299398cbc3fd", + "integrators/integrator_neural.py": "2c42402a702f04a7b9d2a7880f1c6cd5ba57488e04eb78828d786713d1208a6b", + "models/models.py": "5956cadc0390830376c8744f66c80dec94b29465dbd08a5d3c8c4c2f1f14b8f6", + "models/model_transformer.py": "8d8fad21b296c7b54e7e42052fdca7363b5040534481f6d60f415cc073e5820d", + "models/base_models.py": "07313429e3b107db602d19e513796a5d37c4dbf418ba275d50f9e48c27cb3414", + "utils/torch_utils.py": "98492259e1ee2268a02cbd71596283502c1ef2b72672960c80a13da2fb00d89f", + "utils/running_mean_std.py": "c7229216dd77dd8573c2db238b6e3df51f85a05286ad85b6f905b2c0e4bde651" + }, + "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" + }, + "warp": "1.8.0", + "substeps": 5, + "mass_update_interval": 5, + "gravity": [ + -0.0, + -9.81, + -0.0 + ], + "engine": "legacy", + "profile": "warp_reference", + "robot": "cartpole", + "inputs_sha256": "2118cb8143b912a5861a453c0f17c7c6d2ab57e9cc817e0787e64ad509d4b36d", + "trace_sha256": "37e86180a312e3b53a0d5303dd129794a40689ddd597274970bb0e230401d6e0", + "horizon": 1000, + "count": 8 + }, + "actual": { + "robot": "cartpole", + "dt": 0.016666666666666666, + "source_commit": "476dc67fccc6a99be293bf611f897b1c06f13e30", + "asset_sha256": "f31276531c03050a3ac3a3e24443899f38a4424586db17140978383433dec237", + "joint_names": [ + "slider_to_cart", + "cart_to_pole" + ], + "q_dim": 2, + "qd_dim": 2, + "default_q": [ + 0.0, + 0.0 + ], + "default_qd": [ + 0.0, + 0.0 + ], + "effort_limits": [ + 1500.0, + 0.0 + ], + "codec": { + "robot": "cartpole", + "base_com": [ + 0.0, + 0.0, + 0.0 + ], + "continuous_q_indices": [ + 1 + ] + }, + "teacher_profile": "current_newton_import_generalized_effort", + "source_actuators_imported": false, + "source_solver_options_imported": false, + "physics_profile": "native", + "physics_signature": "045e56ec566450721d6e132adc5b6572adbe426d34d6da6f04ada51f670a0df7", + "physics_signature_schema": "nerd-reference-physics-v2", + "gravity": [ + 0.0, + 0.0, + -9.8100004196167 + ], + "contact_schema": "nerd-static-plane-v1", + "teacher": { + "solver": { + "solver_type": "featherstone", + "update_mass_matrix_interval": 5 + }, + "num_substeps": 5, + "physics_profile": "native", + "versions": { + "newton": "1.7.0", + "warp-lang": "1.17.0", + "torch": "2.7.0+cu128", + "dexsim-engine": "0.5.1" + } + }, + "teacher_origin": { + "kind": "current_newton", + "device": "cpu", + "solver": { + "solver_type": "featherstone", + "update_mass_matrix_interval": 5 + }, + "num_substeps": 5, + "versions": { + "newton": "1.7.0", + "warp-lang": "1.17.0", + "torch": "2.7.0+cu128", + "dexsim-engine": "0.5.1" + }, + "dt": 0.016666666666666666, + "physics_signature": "045e56ec566450721d6e132adc5b6572adbe426d34d6da6f04ada51f670a0df7", + "physics_signature_schema": "nerd-reference-physics-v2" + }, + "dexsim_commit": "b61e920fbb222e50213574ad57d885e7b96148b9", + "engine": "newton", + "profile": "native", + "inputs_sha256": "2118cb8143b912a5861a453c0f17c7c6d2ab57e9cc817e0787e64ad509d4b36d", + "trace_sha256": "90246fa6b4b67c7fa273b184e24e44e5fe4ffd3599a2e1fb8ab089c4b3158a7d", + "horizon": 1000, + "count": 8 + }, + "horizons": { + "1": { + "finite_trajectories": 8, + "nonfinite_trajectories": 0, + "errors": { + "q_mae": 0.0008549786871299148, + "qd_mae": 0.08555843774229288, + "base_position_error": 0.0, + "base_orientation_error_rad": 0.0 + } + }, + "10": { + "finite_trajectories": 8, + "nonfinite_trajectories": 0, + "errors": { + "q_mae": 0.020638723357842538, + "qd_mae": 0.347311472625006, + "base_position_error": 0.0, + "base_orientation_error_rad": 0.0 + } + }, + "100": { + "finite_trajectories": 8, + "nonfinite_trajectories": 0, + "errors": { + "q_mae": 0.6030942024626555, + "qd_mae": 2.2402471829054047, + "base_position_error": 0.0, + "base_orientation_error_rad": 0.0 + } + }, + "1000": { + "finite_trajectories": 8, + "nonfinite_trajectories": 0, + "errors": { + "q_mae": 1.7770247821814062, + "qd_mae": 2.7704254585227135, + "base_position_error": 0.0, + "base_orientation_error_rad": 0.0 + } + } + }, + "population": "All requested trajectories; errors remain undefined when any trajectory becomes nonfinite." + } + }, + "ant": { + "physics": { + "robot": "ant", + "reference": { + "source": { + "root": "/root/.cache/embodichain/pr822-accuracy/official-pinned", + "commit": "370e89d2b8ce68aa3578aa7812859c681b7502bc", + "files": { + "algorithms/vanilla_trainer.py": "7027fbc0d543479b7b4a558114b6ae3a89b93729049f9edcc031299398cbc3fd", + "integrators/integrator_neural.py": "2c42402a702f04a7b9d2a7880f1c6cd5ba57488e04eb78828d786713d1208a6b", + "models/models.py": "5956cadc0390830376c8744f66c80dec94b29465dbd08a5d3c8c4c2f1f14b8f6", + "models/model_transformer.py": "8d8fad21b296c7b54e7e42052fdca7363b5040534481f6d60f415cc073e5820d", + "models/base_models.py": "07313429e3b107db602d19e513796a5d37c4dbf418ba275d50f9e48c27cb3414", + "utils/torch_utils.py": "98492259e1ee2268a02cbd71596283502c1ef2b72672960c80a13da2fb00d89f", + "utils/running_mean_std.py": "c7229216dd77dd8573c2db238b6e3df51f85a05286ad85b6f905b2c0e4bde651" + }, + "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" + }, + "warp": "1.5.1", + "substeps": 16, + "mass_update_interval": 16, + "gravity": [ + -0.0, + -9.81, + -0.0 + ], + "engine": "legacy", + "profile": "warp_reference", + "robot": "ant", + "inputs_sha256": "be321297ba576f39a02497d1c9baef37b5a1b2df371983077ebac8d8baee2fee", + "trace_sha256": "3394fe34c740987dfd38b8a19001e6f2176744361dd2088d01b63063bd63ca44", + "horizon": 1000, + "count": 8 + }, + "actual": { + "robot": "ant", + "dt": 0.016666666666666666, + "source_commit": "476dc67fccc6a99be293bf611f897b1c06f13e30", + "asset_sha256": "23e9a720465a71fde3fd9fd2d5a18f8c11ec642f3dca8921c66cb4e169421fa7", + "joint_names": [ + "hip_1", + "ankle_1", + "hip_2", + "ankle_2", + "hip_3", + "ankle_3", + "hip_4", + "ankle_4" + ], + "q_dim": 15, + "qd_dim": 14, + "default_q": [ + 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, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ], + "effort_limits": [ + 400.0, + 400.0, + 400.0, + 400.0, + 400.0, + 400.0, + 400.0, + 400.0 + ], + "codec": { + "robot": "ant", + "base_com": [ + 0.0, + 0.0, + 0.0 + ], + "continuous_q_indices": [] + }, + "teacher_profile": "current_newton_import_generalized_effort", + "source_actuators_imported": false, + "source_solver_options_imported": false, + "physics_profile": "warp_reference", + "physics_signature": "f718639389563ca930e60d55e1c622c2943aa7c29a1444d8520b9f6a1afd811c", + "physics_signature_schema": "nerd-reference-physics-v2", + "gravity": [ + 0.0, + 0.0, + -9.8100004196167 + ], + "contact_schema": "nerd-static-plane-v1", + "teacher": { + "solver": { + "solver_type": "featherstone", + "update_mass_matrix_interval": 16 + }, + "num_substeps": 16, + "physics_profile": "warp_reference", + "versions": { + "newton": "1.7.0", + "warp-lang": "1.17.0", + "torch": "2.7.0+cu128", + "dexsim-engine": "0.5.1" + } + }, + "teacher_origin": { + "kind": "current_newton", + "device": "cpu", + "solver": { + "solver_type": "featherstone", + "update_mass_matrix_interval": 16 + }, + "num_substeps": 16, + "versions": { + "newton": "1.7.0", + "warp-lang": "1.17.0", + "torch": "2.7.0+cu128", + "dexsim-engine": "0.5.1" + }, + "dt": 0.016666666666666666, + "physics_signature": "f718639389563ca930e60d55e1c622c2943aa7c29a1444d8520b9f6a1afd811c", + "physics_signature_schema": "nerd-reference-physics-v2" + }, + "dexsim_commit": "b61e920fbb222e50213574ad57d885e7b96148b9", + "engine": "newton", + "profile": "warp_reference", + "inputs_sha256": "be321297ba576f39a02497d1c9baef37b5a1b2df371983077ebac8d8baee2fee", + "trace_sha256": "b8d3250fb2ca1142b7f4f28773244974874fa63b0ba5e8035d733b35e33f44dc", + "horizon": 1000, + "count": 8 + }, + "horizons": { + "1": { + "finite_trajectories": 8, + "nonfinite_trajectories": 0, + "errors": { + "q_mae": 1.3220931577961892e-05, + "qd_mae": 0.00013303273382397104, + "base_position_error": 1.8774212876823425e-06, + "base_orientation_error_rad": 8.402926348215195e-06 + } + }, + "10": { + "finite_trajectories": 8, + "nonfinite_trajectories": 0, + "errors": { + "q_mae": 0.00010817542306540417, + "qd_mae": 0.0009031878210797525, + "base_position_error": 1.1629142254828825e-05, + "base_orientation_error_rad": 7.080164585794712e-05 + } + }, + "100": { + "finite_trajectories": 8, + "nonfinite_trajectories": 0, + "errors": { + "q_mae": 0.04863060728172968, + "qd_mae": 0.13723946905850556, + "base_position_error": 0.041775816515958894, + "base_orientation_error_rad": 0.028861641871272083 + } + }, + "1000": { + "finite_trajectories": 8, + "nonfinite_trajectories": 0, + "errors": { + "q_mae": 0.04225928523123275, + "qd_mae": 0.02386986058258808, + "base_position_error": 0.052635303684157406, + "base_orientation_error_rad": 0.06033432495269444 + } + } + }, + "population": "All requested trajectories; errors remain undefined when any trajectory becomes nonfinite." + }, + "native": { + "robot": "ant", + "reference": { + "source": { + "root": "/root/.cache/embodichain/pr822-accuracy/official-pinned", + "commit": "370e89d2b8ce68aa3578aa7812859c681b7502bc", + "files": { + "algorithms/vanilla_trainer.py": "7027fbc0d543479b7b4a558114b6ae3a89b93729049f9edcc031299398cbc3fd", + "integrators/integrator_neural.py": "2c42402a702f04a7b9d2a7880f1c6cd5ba57488e04eb78828d786713d1208a6b", + "models/models.py": "5956cadc0390830376c8744f66c80dec94b29465dbd08a5d3c8c4c2f1f14b8f6", + "models/model_transformer.py": "8d8fad21b296c7b54e7e42052fdca7363b5040534481f6d60f415cc073e5820d", + "models/base_models.py": "07313429e3b107db602d19e513796a5d37c4dbf418ba275d50f9e48c27cb3414", + "utils/torch_utils.py": "98492259e1ee2268a02cbd71596283502c1ef2b72672960c80a13da2fb00d89f", + "utils/running_mean_std.py": "c7229216dd77dd8573c2db238b6e3df51f85a05286ad85b6f905b2c0e4bde651" + }, + "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" + }, + "warp": "1.5.1", + "substeps": 16, + "mass_update_interval": 16, + "gravity": [ + -0.0, + -9.81, + -0.0 + ], + "engine": "legacy", + "profile": "warp_reference", + "robot": "ant", + "inputs_sha256": "be321297ba576f39a02497d1c9baef37b5a1b2df371983077ebac8d8baee2fee", + "trace_sha256": "3394fe34c740987dfd38b8a19001e6f2176744361dd2088d01b63063bd63ca44", + "horizon": 1000, + "count": 8 + }, + "actual": { + "robot": "ant", + "dt": 0.016666666666666666, + "source_commit": "476dc67fccc6a99be293bf611f897b1c06f13e30", + "asset_sha256": "23e9a720465a71fde3fd9fd2d5a18f8c11ec642f3dca8921c66cb4e169421fa7", + "joint_names": [ + "hip_1", + "ankle_1", + "hip_2", + "ankle_2", + "hip_3", + "ankle_3", + "hip_4", + "ankle_4" + ], + "q_dim": 15, + "qd_dim": 14, + "default_q": [ + 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, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ], + "effort_limits": [ + 400.0, + 400.0, + 400.0, + 400.0, + 400.0, + 400.0, + 400.0, + 400.0 + ], + "codec": { + "robot": "ant", + "base_com": [ + 0.0, + 0.0, + 0.0 + ], + "continuous_q_indices": [] + }, + "teacher_profile": "current_newton_import_generalized_effort", + "source_actuators_imported": false, + "source_solver_options_imported": false, + "physics_profile": "native", + "physics_signature": "b49fb353b96a8768112d2d066d9e75c745b72084265e9c4853f9b905ca3b1867", + "physics_signature_schema": "nerd-reference-physics-v2", + "gravity": [ + 0.0, + 0.0, + -9.8100004196167 + ], + "contact_schema": "nerd-static-plane-v1", + "teacher": { + "solver": { + "solver_type": "featherstone", + "update_mass_matrix_interval": 16 + }, + "num_substeps": 16, + "physics_profile": "native", + "versions": { + "newton": "1.7.0", + "warp-lang": "1.17.0", + "torch": "2.7.0+cu128", + "dexsim-engine": "0.5.1" + } + }, + "teacher_origin": { + "kind": "current_newton", + "device": "cpu", + "solver": { + "solver_type": "featherstone", + "update_mass_matrix_interval": 16 + }, + "num_substeps": 16, + "versions": { + "newton": "1.7.0", + "warp-lang": "1.17.0", + "torch": "2.7.0+cu128", + "dexsim-engine": "0.5.1" + }, + "dt": 0.016666666666666666, + "physics_signature": "b49fb353b96a8768112d2d066d9e75c745b72084265e9c4853f9b905ca3b1867", + "physics_signature_schema": "nerd-reference-physics-v2" + }, + "dexsim_commit": "b61e920fbb222e50213574ad57d885e7b96148b9", + "engine": "newton", + "profile": "native", + "inputs_sha256": "be321297ba576f39a02497d1c9baef37b5a1b2df371983077ebac8d8baee2fee", + "trace_sha256": "0bf89a0cf94267148585d6dccc52dded34bf62d43c75512834b5dfd01b7a5b22", + "horizon": 1000, + "count": 8 + }, + "horizons": { + "1": { + "finite_trajectories": 8, + "nonfinite_trajectories": 0, + "errors": { + "q_mae": 0.04324831923622696, + "qd_mae": 2.8115949213138913, + "base_position_error": 0.016508535030398546, + "base_orientation_error_rad": 0.02473029776161908 + } + }, + "10": { + "finite_trajectories": 0, + "nonfinite_trajectories": 8, + "errors": null + }, + "100": { + "finite_trajectories": 0, + "nonfinite_trajectories": 8, + "errors": null + }, + "1000": { + "finite_trajectories": 0, + "nonfinite_trajectories": 8, + "errors": null + } + }, + "population": "All requested trajectories; errors remain undefined when any trajectory becomes nonfinite." + } + } + }, + "training": { + "cartpole": { + "result": { + "global_step": 1000, + "budget": 1000, + "status": "completed", + "best_validation": 0.0012673945393544273, + "last": { + "step": 1000, + "loss": 0.0025523770600557327, + "learning_rate": 1.009e-06, + "grad_norm": 0.03288434445858002, + "validation": { + "normalized_mse": 0.0012673945393544273, + "target_mse": 5.796554771109186e-05, + "q_error_norm": 0.0007195487560238689, + "qd_error_norm": 0.0034404500154778363, + "valid_positions": 40960, + "windows": 4096, + "split": "validation" + } + }, + "runtime_identity": { + "schema_version": 1, + "python": "3.11.14", + "python_implementation": "CPython", + "platform": "Linux-5.15.0-119-generic-x86_64-with-glibc2.35", + "torch_version": "2.7.0+cu128", + "parameter_placements": [ + { + "device": "cuda:0", + "dtype": "torch.float32" + } + ], + "parameter_placement_status": "available", + "cuda": { + "used": true, + "runtime_version": "12.8", + "devices": [ + { + "device": "cuda:0", + "status": "available", + "name": "NVIDIA RTX PRO 5000 72GB Blackwell" + } + ], + "driver_inventory": { + "status": "available", + "physical_gpus": [ + { + "index": "0", + "uuid": "GPU-5fbdbdd3-4b93-066a-cf17-bd62a5261b2c", + "name": "NVIDIA RTX PRO 5000 72GB Blackwell", + "driver_version": "580.126.20" + }, + { + "index": "1", + "uuid": "GPU-3b85e500-27ff-35d9-eb8d-d54f9ff0898d", + "name": "NVIDIA RTX PRO 5000 72GB Blackwell", + "driver_version": "580.126.20" + }, + { + "index": "2", + "uuid": "GPU-72b782e4-97d4-c0fa-c0f8-78956448d7b4", + "name": "NVIDIA RTX PRO 5000 72GB Blackwell", + "driver_version": "580.126.20" + }, + { + "index": "3", + "uuid": "GPU-afe5dd63-2838-06db-b0fa-5d98d4b1eec5", + "name": "NVIDIA RTX PRO 5000 72GB Blackwell", + "driver_version": "580.126.20" + }, + { + "index": "4", + "uuid": "GPU-a5477f57-58a6-4f89-6e7a-338d8fcfc6f1", + "name": "NVIDIA RTX PRO 5000 72GB Blackwell", + "driver_version": "580.126.20" + }, + { + "index": "5", + "uuid": "GPU-569c23a7-9291-45ef-c727-ba74d63708db", + "name": "NVIDIA RTX PRO 5000 72GB Blackwell", + "driver_version": "580.126.20" + }, + { + "index": "6", + "uuid": "GPU-624342e0-6d34-01ea-c202-1d15bf69f63f", + "name": "NVIDIA RTX PRO 5000 72GB Blackwell", + "driver_version": "580.126.20" + }, + { + "index": "7", + "uuid": "GPU-0059f32f-0c76-b67a-ddd1-f3b2c3bb1508", + "name": "NVIDIA RTX PRO 5000 72GB Blackwell", + "driver_version": "580.126.20" + } + ] + } + }, + "warp": { + "status": "not_loaded" + } + } + }, + "config": { + "dataset": "/root/.cache/embodichain/pr822-accuracy/cartpole-aligned-data", + "output_dir": "/root/.cache/embodichain/pr822-accuracy/cartpole-fine-tune", + "fine_tune": "/root/.cache/embodichain/pr822-accuracy/cartpole-upstream-export", + "train": { + "total_steps": 1000, + "batch_size": 512, + "learning_rate": 1e-05, + "final_learning_rate": 1e-06, + "device": "cuda:0", + "seed": 0, + "validation_interval": 200, + "validation_batches": 8, + "checkpoint_interval": 200 + } + }, + "dataset_manifest": { + "fields": { + "states": { + "q": { + "dtype": "

+ +
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 +mkdir -p outputs/dynamics/cartpole/pretrained +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. + +For a live rollout, install a DexSim build exposing named Newton solver +extensions and the Scene snapshot/functional APIs. Use the matching reference +scene and physics profile: + +```bash +embodichain dynamics eval --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.rollout.pretrained.yaml +``` + +This command advances both the released model and a matching Newton teacher at +the same macro timestep. The preset evaluates 16 trajectories at horizons 1, 10 +and 100. It saves the report, per-trajectory metrics and raw state evidence. +Select a fresh output path with `--output` when repeating an evaluation. + +The released models were trained for `physics_profile: warp_reference`, which +reconstructs a specific reference asset and material configuration using current +Newton. The analytical teacher uses five substeps for Cartpole or sixteen for +Ant, while NeRD predicts the entire 1/60-second interval in one neural substep. +The default `native`/MuJoCo workflow configurations describe a different physical +system and do not match these released exports. + +## 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. + +## Deploy NeRD in a matching scene + +An installed solver entry point makes `nerd` available through the same Newton +configuration used for analytical solvers: + +```yaml +physics: newton +physics_config: + physics_dt: 0.016666666666666666 + num_substeps: 1 + solver_cfg: + solver_type: nerd + options: + checkpoint: outputs/models/neural_dynamics/nerd/cartpole/v1 +``` + +The scene must match the export's ordered joints, gravity, contact schema, +physical signature and timestep. The pretrained rollout presets construct that +reference scene. A model cannot be attached to a different importer/material +profile by changing its manifest or scaling its predicted state delta. + +NeRD consumes generalized efforts. It rejects unsupported free-base and external +forces and does not provide physical contact impulse or friction-force reports. +The current provider runs eager FP32 inference and does not support CUDA Graph +capture. Compare performance with the analytical solver's supported Graph mode +enabled, keeping latency and throughput scopes explicit. + +## Functional dynamics and gradients + +{class}`~embodichain.lab.gym.envs.differentiable_dynamics_env.DifferentiableDynamicsEnv` +and {class}`~embodichain.lab.sim.diff.dynamics.FunctionalDynamics` expose +graph-connected physical state, per-world history and valid history lengths. +The ordinary native step and functional step share the same learned transition; +publishing a functional result does not advance the world again. + +The Torch/Warp FK boundary supplies first-order derivatives. Shared tensor +storage alone does not establish a gradient path. Continuous state and contact +geometry have local derivatives, while contact selection, angle wrapping and +quaternion signs contain discrete branches. Physical parameters and contact +topology are not differentiable through this provider. + +Reset clears the selected worlds' histories while preserving other rows and +already-computed terminal rewards. Call `detach_state()` at an intentional +rollout boundary, and update trainable model parameters only after backward. + +## Add another neural solver + +Implement the DexSim solver-provider contract and expose it through an installed +entry point: + +```toml +[project.entry-points."dexsim.newton_solvers"] +my_solver = "my_package.solver:MySolverProvider" +``` + +The provider declares builder attributes and capabilities, then creates its +solver after the scene model is finalized. Mutable histories belong to each +World. The configured `solver_type` is the entry-point name; auto-selection does +not choose a learned model implicitly. Duplicate or unknown names fail. + +Keep state transforms in {mod}`embodichain.compute.dynamics`, model/checkpoint +code in {mod}`embodichain.learning.models.nerd`, offline training in +{mod}`embodichain.learning.dynamics.workflow`, and scene adaptation in +{mod}`embodichain.lab.sim.physics.solvers`. Dataset storage belongs to +{mod}`embodichain.data_pipeline.dynamics.shards`. diff --git a/docs/source/guides/index.rst b/docs/source/guides/index.rst index b892a7a30..ba6a9518c 100644 --- a/docs/source/guides/index.rst +++ b/docs/source/guides/index.rst @@ -19,6 +19,5 @@ Practical guides for common tasks in EmbodiChain. task_catalog robot_catalog policy_evaluation - neural_dynamics rlinf cli diff --git a/docs/source/guides/neural_dynamics.md b/docs/source/guides/neural_dynamics.md index 1c950722a..74fcac64e 100644 --- a/docs/source/guides/neural_dynamics.md +++ b/docs/source/guides/neural_dynamics.md @@ -1,170 +1,8 @@ -# Named dynamics solvers and NeRD +--- +orphan: true +--- -NeRD integrates a robot-specific learned state transition into the DexSim Newton -scene lifecycle. Cartpole uses contact-free dynamics; Ant uses fixed candidate -points projected onto a static plane. A model is bound to its ordered joints, -physical signature, gravity, feature schema and time interval. +# Neural dynamics -This feature requires a DexSim build exporting `RegisteredNewtonSolverCfg`, the -`dexsim.newton_solvers` registry and the Scene functional/snapshot APIs. The -development JIT package and a released source-free package are different artifact -types; use the actual tested build identity rather than assuming every 0.5.x wheel -contains these interfaces. - -The repository's existing `dexsim_engine==0.5.1rc1` dependency pin does not identify -the extension-enabled development artifact. Until a matching release and pin are -available, deployment requires the tested local DexSim wheel and this EmbodiChain -checkout installed as a pair. Install the ordinary dependencies first, then install -the local wheel and editable checkout with `--no-deps` so the older pin does not -replace it. Record the wheel hash and check the required APIs before running. -This development procedure does not qualify the existing rc1 distribution. - -## Register an implementation - -Declare one installed entry point whose name becomes the configuration value: - -```toml -[project.entry-points."dexsim.newton_solvers"] -my_solver = "my_package.solver:MySolverProvider" -``` - -The provider registers builder attributes before model finalization and creates -its extension after model, control and contacts exist. It returns a Newton -`SolverBase` implementation, explicit capabilities, and optional functional -services. Provider instances and mutable histories are per World. Install local -development packages with an editable install; a source file alone does not -register an entry point. - -```yaml -physics: newton -physics_config: - solver_cfg: - solver_type: my_solver - options: {} -``` - -There is no separate direct-provider configuration or imperative registration -route. Names must be unique and cannot override built-in aliases. AutoSolver -retains its scene-aware selection; it does not guess which neural model applies. - -## Generate, train and export - -The Cartpole and Ant workflow configurations live under -`extensions/embodichain_tasks/configs/tasks/classic_control//`. -These `dynamics.*.yaml` files configure commands, not Gym 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 -``` - -Use `embodichain dynamics --help` for path/device/seed overrides. -Configuration-file paths resolve relative to the configuration file; explicit CLI -path overrides resolve relative to the working directory. The `full` profiles -carry a much larger training budget. Smoke profiles establish executable flow, -not model accuracy or official benchmark qualification. - -Teacher collection owns an ordinary `SimulationManager` and uses its public -articulation/Scene APIs. Cartpole sets `enable_ground_plane: false`: hiding a -ground mesh would leave collisions active. CPU MuJoCo batches use -`separate_worlds: true`. Record the actual teacher configuration, including -ignored solver-specific material settings, with every dataset. - -HDF5 shards contain complete trajectories: T+1 states and T controls/contact -features. Only committed shards appear in the manifest. Sampling windows cannot -cross trajectory/reset/split boundaries. Collection and reading use bounded -buffers, and each data-loader worker owns its file handles. Training, export and -offline evaluation do not initialize DexSim. - -Training checkpoints retain optimizer, scheduler, RNG and committed sample cursor -for resume. Exported models contain tensor state, normalization buffers and JSON -metadata with checksums. They can be loaded directly by `solver_type: nerd`. -Legacy official Python-object checkpoints require the separate -`convert-checkpoint --trust-pickle` operation and an explicit deployment manifest. - -## Functional dynamics and model gradients - -`DifferentiableDynamicsEnv` uses graph-connected state and output hooks. Its -functional state contains physical tensors, per-world history and valid history -lengths. Native state publication is detached and does not own autograd storage. -`FunctionalDynamics` advances these states through SimulationManager; do not also -call `World.update` for the same interval. - -Torch network operations preserve gradients in evaluation mode. Newton FK uses -an explicit first-order VJP bridge. Tensor sharing between Warp and Torch alone -does not establish a gradient path. Ant's continuous contact geometry has local -derivatives; contact masks, wrap cuts and quaternion signs remain discrete branch -choices. Physical parameters and contact topology are not differentiable in v1. - -Resetting selected environments clears their histories without breaking other -rows' graphs or already-computed terminal rewards. Call `detach_state()` at an -intentional rollout boundary, and update model weights only after backward has -completed. Resident model weights are fixed by default; explicitly request -trainable weights or supply parameter tensors for a functional optimization. - -NeRD consumes effort controls and rejects unsupported free-base/external forces. -It provides no physical contact impulse/friction-force report. Its trained time -interval is validated; changing `dt` cannot be implemented by scaling a predicted -state delta. Both reference exports use 1/60 s and one neural substep. - -## Correctness and official accuracy evidence - -Implementation parity, training-step parity, matched-budget training accuracy, -and deployment rollout accuracy are separate gates. A completed short training -run or a finite gradient is not an official accuracy result. - -Comparisons pin source commits, assets, data, initialization, batch order, -normalization statistics, optimizer settings and checkpoint-selection criteria. -The official trainer adapter executes original methods in a separate process. -Current-Newton teacher data and official-Warp teacher data occupy separate rows -of the matched-data matrix. Do not attribute teacher changes to the model port. - -The historical importers and current Newton produce different masses, inertias -and materials even from identical asset bytes. Official checkpoint conversion -does not prove deployment compatibility. Preserve the historical physical profile -or retrain for the current teacher. The reference data adapter also records an -explicit repair of the official sampler's uninitialized gravity components; -raw arrays remain available and corrected inputs are labeled in the report. - -Formal reports require five training seeds, matched declared update budgets, -fixed evaluation trajectories and error-versus-time metrics. Test data never -selects checkpoints. Missing experiments remain explicitly pending. The complete -acceptance contract is maintained in the repository design document -`design/neural-dynamics-solvers.md`. - -## Historical Warp comparison - -Use `embodichain.learning.dynamics.reference_rollout` in a fresh process to -evaluate either trainer's portable weights with the original environment and -Transformer integrator. Cartpole requires Warp 1.8.0; Ant requires Warp 1.5.1. -Keep these dependencies in separate directories rather than replacing the -installed Newton/Warp runtime. The source checkout must match the pinned main -commit checked by the reference tools. - -```bash -PYTHONPATH=/path/to/isolated-warp:/path/to/pinned-nerd-source:$PWD \ - python -m embodichain.learning.dynamics.reference_rollout \ - --spec historical-rollout.json --trust-external-code -``` - -The JSON spec provides `source`, `model`, `robot`, `output`, `device`, `num_envs`, -`trajectories`, `horizons`, `seed`, `motion_mode` and `initial_states`. -`fixed_policy` currently accepts `policy: {"kind": "pd"}`. Teacher and learned -simulators evaluate the policy on their own states. CPU execution uses a recorded -stream compatibility adapter while retaining original simulation/network math. - -Historical physical parameters must match the model's recorded training source. -Older manifests without measured gravity remain pending; do not add a new label -to make them pass. Original/portable network equivalence, teacher-content hashes, -source/trace/model identities, error curves and every failed trajectory are -retained. This path does not make historical weights compatible with the current -Newton asset import. - -The training comparison is explicitly a controlled matched-data experiment. -Shared initialization, statistics, sample order and validation selection permit -port comparisons; they do not reproduce every upstream paper-recipe detail. -Full-profile gates inspect the actual immutable model's feature order, FP32 -weights, architecture and optimizer settings. Short runs remain short-run -evidence even when all paired metrics pass their numerical thresholds. +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 f28f75812..ae1451427 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -47,6 +47,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/lab/sim/physics/solvers/_reference_assets.py b/embodichain/lab/sim/physics/solvers/_reference_assets.py index fb1730426..ee91d14d3 100644 --- a/embodichain/lab/sim/physics/solvers/_reference_assets.py +++ b/embodichain/lab/sim/physics/solvers/_reference_assets.py @@ -58,6 +58,32 @@ } +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}") @@ -233,8 +259,11 @@ def make_reference_articulation( ) -> ArticulationDesc: """Build an independent public descriptor with imported physical properties. - ``warp_reference`` explicitly reproduces measured old-Warp asset parameters; - it does not select the old simulator or assert identical contact algorithms. + ``warp_reference`` reconstructs measured old-Warp asset parameters. The + host applies a solver-specific Ant friction conversion to fresh descriptors + when using Newton Featherstone; other solvers retain these source material + coefficients. This does not select the old simulator or assert identical + contact algorithms. Body/joint names retain their asset spelling, with only the importer's path prefix removed. The free root is represented by ``fixed_base=False`` and is created by Spawn. Default nonzero joint positions are in ``reference_profile`` diff --git a/embodichain/lab/sim/physics/solvers/collection.py b/embodichain/lab/sim/physics/solvers/collection.py index 727d375c9..c1c5bd3f9 100644 --- a/embodichain/lab/sim/physics/solvers/collection.py +++ b/embodichain/lab/sim/physics/solvers/collection.py @@ -50,7 +50,11 @@ from embodichain.lab.sim.cfg import NewtonPhysicsCfg, RenderCfg from ._nerd_binding import NeRDBinding -from ._reference_assets import make_reference_articulation, reference_profile +from ._reference_assets import ( + _ant_reference_friction, + make_reference_articulation, + reference_profile, +) __all__ = [ "DynamicsGenerationCfg", @@ -149,6 +153,34 @@ def __init__( ) neural = solver_cfg.get("solver_type") == "nerd" self._featherstone_teacher = solver_cfg.get("solver_type") == "featherstone" + friction_mapping = None + if robot == "ant" and physics_profile == "warp_reference": + if self._featherstone_teacher: + friction_mapping = _ant_reference_friction() + delta = friction_mapping["friction_smoothing"] + requested = solver_cfg.setdefault("friction_smoothing", delta) + if isinstance(requested, bool) or requested != delta: + raise ValueError( + "Ant warp_reference requires the coupled contact mapping: " + f"friction_smoothing={delta} with scaled shape/ground mu." + ) + elif neural: + # Older exports predate the Featherstone friction conversion. + # Their original material domain must still be reconstructed; + # the provider then validates the actual physics signature. + options = solver_cfg.get("options", {}) + checkpoint = Path(options["checkpoint"]) + exported = json.loads((checkpoint / "manifest.json").read_text()) + recorded = exported.get("reference_friction_mapping") + if recorded is not None: + expected = _ant_reference_friction() + if not isinstance(recorded, Mapping) or any( + recorded.get(key) != value + for key, value in expected.items() + if key != "interpretation" + ): + raise ValueError("Unsupported Ant reference friction mapping") + friction_mapping = dict(recorded) if self._featherstone_teacher: interval = solver_cfg.get("update_mass_matrix_interval", 1) if type(interval) is not int or interval < 1 or num_substeps % interval: @@ -203,7 +235,7 @@ def __init__( ke=1e5, kd=1e3, kf=1e3, - mu=0.5, + mu=(friction_mapping["ground_mu"] if friction_mapping else 0.5), margin=1e-5, gap=0.05, mu_torsional=0.0, @@ -217,9 +249,14 @@ def __init__( per_env=False, ) ) - self.articulation = self.sim.add_articulation_descriptor( - make_reference_articulation(robot, physics_profile) - ) + descriptor = make_reference_articulation(robot, physics_profile) + if friction_mapping is not None: + # Scale every freshly authored Ant shape, including the shared + # plane above, before Spawn finalizes the physical model. + for link in descriptor.links: + for collision in link.collisions: + collision.newton.mu *= friction_mapping["material_mu_scale"] + self.articulation = self.sim.add_articulation_descriptor(descriptor) self.sim.prepare() self.binding = NeRDBinding.from_metadata( dict(self.sim.spawn_result.dynamics_metadata()), robot @@ -250,6 +287,11 @@ def __init__( "physics_signature_schema" ], } + if friction_mapping is not None: + self.manifest["reference_friction_mapping"] = friction_mapping + self.manifest["teacher_origin"]["reference_friction_mapping"] = dict( + friction_mapping + ) self.reset() except Exception: self.close() diff --git a/embodichain/learning/dynamics/workflow.py b/embodichain/learning/dynamics/workflow.py index d06fc13ce..231e8433b 100644 --- a/embodichain/learning/dynamics/workflow.py +++ b/embodichain/learning/dynamics/workflow.py @@ -45,6 +45,7 @@ "gravity_transport_override", "source_physical_arrays", "legacy_velocity", + "reference_friction_mapping", } ) diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/README.md b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/README.md index 309550b30..0bdcc55a6 100644 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/README.md +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/README.md @@ -1,126 +1,67 @@ -# Ant neural dynamics workflow +# Ant neural dynamics + +These files configure the `embodichain dynamics` workflow. They are separate +from runnable Gym deployments. See the +[Neural Dynamics Models feature](../../../../../../docs/source/features/neural_dynamics.md) +for the architecture, published datasets and models, comparison videos and +supported deployment boundaries. + +## Configuration profiles + +| Profile | Purpose | +| --- | --- | +| `dynamics.*.smoke.yaml` | Small CPU collection/training/evaluation workflow | +| `dynamics.*.full.yaml` | Reference network scale and long training budget for a newly generated native dataset | +| `dynamics.generate.warp-reference.yaml` | Current-Newton reference-asset collection with Featherstone | +| `dynamics.eval.pretrained.yaml` | Offline test evaluation of the published Hugging Face dataset/model pair | +| `dynamics.rollout.pretrained.yaml` | Released NeRD model versus its matching Newton reference teacher | +| `env.nerd.yaml` | Registered Gym scene using an export from the matching native workflow | +| `agents/apg.nerd.yaml` | Policy learning through the dynamics environment | + +The published reference models use `warp_reference` physics. They do not match +the default native/MuJoCo-generated dataset or the native Gym preset merely +because the robot has the same name. The pretrained rollout preset reconstructs +the asset, materials and solver configuration expected by the released model. + +## Use published data and models + +Download `neural_dynamics/nerd/ant/v1` from the dataset repository +[DexForceAI/embodichain_datasets](https://huggingface.co/datasets/DexForceAI/embodichain_datasets) +into `outputs/datasets`, and from the model repository +[DexForceAI/embodichain_model](https://huggingface.co/DexForceAI/embodichain_model) +into `outputs/models`. Keep the path prefix when downloading; the feature page +provides a `snapshot_download` example. + +Run from the repository root: -These files configure `embodichain dynamics`; they are not Gym deployments and -have no runnable environment ID. The simulator host constructs the pinned NeRD -reference asset through EmbodiChain and DexSim Spawn. The generated manifest -provides the actual physics, asset and codec identity; training inherits it. +```bash +mkdir -p outputs/dynamics/ant/pretrained +embodichain dynamics eval --config extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.pretrained.yaml +embodichain dynamics eval --config extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.rollout.pretrained.yaml +``` -Run from the repository root (replace `smoke` with `full` for the single-GPU profile): +Offline evaluation needs no simulation. Live rollouts require an installed +DexSim build with named solver extensions and Scene functional/snapshot APIs. +Use `--device cpu` or `--device cuda:0` to select the device, and `--output` with a +new path to retain repeated rollout results. The reference models are research +artifacts; consult their model cards for measured error and supported versions. + +## Train and export a new model ```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 -embodichain dynamics eval --config extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.rollout.smoke.yaml ``` -Paths inside YAML resolve relative to the YAML file. CLI path overrides resolve -relative to the current directory. Output directories are under repository -`outputs/`, which is ignored by Git. Exports require a new destination directory. -Generation resumes with `--resume`; training resumes with -`--resume outputs/dynamics/ant/smoke/train/last.pt`. Exact resume preserves -budget, seed, device, configuration, immutable data fingerprint, optimizer, RNG -and the committed sampler cursor. Use `--fine-tune path/to/export` for a new run. - -The smoke profile uses CPU, 64 short trajectories, a small Transformer and -100 optimizer steps. It tests the complete plumbing and does not qualify model -accuracy. Ant smoke uses effort_scale=0.05 and default initial states; this is deliberately a different distribution from the official-range full profile. - -The full profile uses one CUDA device, the official six-layer, twelve-head -Transformer with embedding width 384, history 10, batch 512, and -4,995,000 Adam updates. This matches 999 trained epochs × 5000 updates in the -saved official configuration (epoch zero only evaluates). The learning rate is -held for 5000 updates, with offset one and a 1000-interval denominator, matching -the official epoch timing. Gradient clipping is a norm limit of 1.0; it does not -truncate temporal gradients. Validation uses the first 100 complete/partial -batches; final offline evaluation reads the entire test split. The expected -training budget is substantial and is not run automatically by tests. - -Default generation samples the full recorded force range: cart force ±1500 -with no pole actuation, or Ant scalar joints ±400. Teachers use MuJoCo-Warp with -10 substeps per 1/60-second macro step. CPU fixtures explicitly request separate -worlds. Source MJCF motor objects and solver XML options are not installed into -the teacher: the host applies actual generalized forces and records its solver -configuration. This profile is a current Newton teacher, not Warp 1.5.1 paper -reproduction. - -## Official comparison controls - -Default training uses deterministic sampling with replacement and streamed -population moments over overlapping training windows. These choices are explicit; -they do not reproduce the official dataloader's permutation or its running -normalizer initialization automatically. - -For a paired official comparison, pass the **same** immutable data, trajectory -split, initialized weights, batch order, normalizer artifact and update budget to -both trainers. `--statistics stats.pt` loads a tensor-only artifact containing -`schema_version: 1`, `inputs: {name: {mean, var}}`, and -`output: {mean, var}`. No validation/test samples may enter these statistics. -`--sample-order order.npy` reads a memory-mapped integer array of shape -`[total_steps, batch_size]`; each row indexes the full-window training dataset. -A `.pt` tensor is also accepted for short fixtures. The array hash and committed -cursor are checked on resume; large `.npy` orders are never loaded in full. - -Use `--seed 0` through `--seed 4` with distinct `--output-dir` values for the -five paired training seeds. Export each validation-selected checkpoint separately -and use evaluation's `--model` and `--output` overrides to retain every result. -Full rollout configurations request 2048 trajectories at horizons 1, 10, 100, -1000; repeat with `motion_mode: passive` and with the qualified fixed policy. -A common seed alone does not guarantee identical minibatches across trainers. - -No committed configuration is evidence that the full run passed. Required reports -include short numerical/Adam parity, both teachers × both trainers, five seed -pairs, confidence bounds, non-finite/failure rates, fixed-policy return and -software/asset/checkpoint hashes. `qualification.paired_noninferiority` implements -the seed-paired one-sided 95% bootstrap gate; fewer than five seeds is pending. -A 100-step smoke cannot replace the full budget or paper reproduction. - -## Registered differentiable task - -After exporting a matching model, `env.nerd.yaml` creates the registered -`AntDynamics-v1` task. Its checkpoint path resolves from the -working directory, unlike dynamics workflow file paths (config-relative). - -```bash -embodichain run-task --gym_config extensions/embodichain_tasks/configs/tasks/classic_control/ant/env.nerd.yaml --preview -embodichain train-rl --config extensions/embodichain_tasks/configs/tasks/classic_control/ant/agents/apg.nerd.yaml -``` - -Use preview for the Gym scene: this task does not implement an expert trajectory -collector. The APG configuration uses the existing learning-environment factory -and trainer, with graph-connected state/reward hooks. Its policy checkpoint is -separate from the NeRD dynamics-model export. - -## Legacy Warp parameter comparison - -`dynamics.generate.warp-reference.yaml` selects `physics_profile: warp_reference` -and Featherstone with 16 substeps per 1/60-second frame. It preserves measured -legacy asset masses/inertia, body armature, force-mode damping and contact-slot -order. Ant also uses an explicitly authored shared ground material. The simulator -is still current Newton; contact algorithms can differ from old Warp. - -Generate this bounded 512-trajectory dataset with: - -```bash -embodichain dynamics generate --config extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.warp-reference.yaml -``` - -Keep the same `physics_profile` in rollout evaluation. Existing exports remain -bound to their recorded physical signature; generate new data and train or -fine-tune a new export for this profile. The default `native` profile retains -modern importer semantics. Both profiles need an installed DexSim with named -solver extensions; the measured SDK branch is MR !1483 at `b61e920f`. - -The isolated historical/current replay tool is -`scripts/experiments/neural_dynamics/physics_alignment.py`. Its shared input and -trace hashes, plus finite/failure counts, are required when comparing trajectories. - - -## Supported-version reference data +Smoke profiles verify executable flow, not accuracy. For a different dataset, +copy the appropriate training config and set `dataset` and `output_dir` to its +location and a new training directory. For a different physical system, collect +new teacher data and train a matching model. -The [original data and paired evaluation workflow](../../../../../../scripts/experiments/neural_dynamics/README.md) -uses isolated Python3.8/Torch2.2 generation, immutable raw/canonical teacher -traces and independently evolving model feedback. Training and evaluation -runtimes are recorded separately; a current-Newton corpus cannot be relabeled -as original-Warp data. +Configuration paths resolve relative to the YAML file; CLI path overrides +resolve relative to the working directory. Exports require a new destination. +Resume an identical training run with `train --resume PATH/last.pt`, or initialize +a new run from a portable export with `train --fine-tune PATH`. Validation data +selects the checkpoint; the test split remains held out. 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 000000000..9a1c92879 --- /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.generate.warp-reference.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.warp-reference.yaml index c0f0f5453..15ff72bec 100644 --- 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 @@ -10,6 +10,7 @@ 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 diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.rollout.pretrained.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.rollout.pretrained.yaml new file mode 100644 index 000000000..a2cfa23ac --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.rollout.pretrained.yaml @@ -0,0 +1,22 @@ +# Published model versus its matching current-Newton analytical teacher. +mode: rollout +robot: ant +model: ../../../../../../outputs/models/neural_dynamics/nerd/ant/v1 +device: cuda:0 +num_envs: 16 +trajectories: 16 +horizons: +- 1 +- 10 +- 100 +seed: 1000000 +physics_profile: warp_reference +teacher_solver: + solver_type: featherstone + update_mass_matrix_interval: 16 + friction_smoothing: 0.001 +num_substeps: 16 +motion_mode: random_effort +initial_states: sample +effort_scale: 1.0 +output: ../../../../../../outputs/dynamics/ant/pretrained/rollout.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 index 0c2a73b08..882a52eb2 100644 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/README.md +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/README.md @@ -1,126 +1,67 @@ -# Cartpole neural dynamics workflow +# Cartpole neural dynamics + +These files configure the `embodichain dynamics` workflow. They are separate +from runnable Gym deployments. See the +[Neural Dynamics Models feature](../../../../../../docs/source/features/neural_dynamics.md) +for the architecture, published datasets and models, comparison videos and +supported deployment boundaries. + +## Configuration profiles + +| Profile | Purpose | +| --- | --- | +| `dynamics.*.smoke.yaml` | Small CPU collection/training/evaluation workflow | +| `dynamics.*.full.yaml` | Reference network scale and long training budget for a newly generated native dataset | +| `dynamics.generate.warp-reference.yaml` | Current-Newton reference-asset collection with Featherstone | +| `dynamics.eval.pretrained.yaml` | Offline test evaluation of the published Hugging Face dataset/model pair | +| `dynamics.rollout.pretrained.yaml` | Released NeRD model versus its matching Newton reference teacher | +| `env.nerd.yaml` | Registered Gym scene using an export from the matching native workflow | +| `agents/apg.nerd.yaml` | Policy learning through the dynamics environment | + +The published reference models use `warp_reference` physics. They do not match +the default native/MuJoCo-generated dataset or the native Gym preset merely +because the robot has the same name. The pretrained rollout preset reconstructs +the asset, materials and solver configuration expected by the released model. + +## Use published data and models + +Download `neural_dynamics/nerd/cartpole/v1` from the dataset repository +[DexForceAI/embodichain_datasets](https://huggingface.co/datasets/DexForceAI/embodichain_datasets) +into `outputs/datasets`, and from the model repository +[DexForceAI/embodichain_model](https://huggingface.co/DexForceAI/embodichain_model) +into `outputs/models`. Keep the path prefix when downloading; the feature page +provides a `snapshot_download` example. + +Run from the repository root: -These files configure `embodichain dynamics`; they are not Gym deployments and -have no runnable environment ID. The simulator host constructs the pinned NeRD -reference asset through EmbodiChain and DexSim Spawn. The generated manifest -provides the actual physics, asset and codec identity; training inherits it. +```bash +mkdir -p outputs/dynamics/cartpole/pretrained +embodichain dynamics eval --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.pretrained.yaml +embodichain dynamics eval --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.rollout.pretrained.yaml +``` -Run from the repository root (replace `smoke` with `full` for the single-GPU profile): +Offline evaluation needs no simulation. Live rollouts require an installed +DexSim build with named solver extensions and Scene functional/snapshot APIs. +Use `--device cpu` or `--device cuda:0` to select the device, and `--output` with a +new path to retain repeated rollout results. The reference models are research +artifacts; consult their model cards for measured error and supported versions. + +## Train and export a new model ```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 dynamics eval --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.rollout.smoke.yaml ``` -Paths inside YAML resolve relative to the YAML file. CLI path overrides resolve -relative to the current directory. Output directories are under repository -`outputs/`, which is ignored by Git. Exports require a new destination directory. -Generation resumes with `--resume`; training resumes with -`--resume outputs/dynamics/cartpole/smoke/train/last.pt`. Exact resume preserves -budget, seed, device, configuration, immutable data fingerprint, optimizer, RNG -and the committed sampler cursor. Use `--fine-tune path/to/export` for a new run. - -The smoke profile uses CPU, 64 short trajectories, a small Transformer and -100 optimizer steps. It tests the complete plumbing and does not qualify model -accuracy. Cartpole smoke uses default initial states; the full profile samples initial states. - -The full profile uses one CUDA device, the official six-layer, twelve-head -Transformer with embedding width 192, history 10, batch 512, and -4,995,000 Adam updates. This matches 999 trained epochs × 5000 updates in the -saved official configuration (epoch zero only evaluates). The learning rate is -held for 5000 updates, with offset one and a 1000-interval denominator, matching -the official epoch timing. Gradient clipping is a norm limit of 1.0; it does not -truncate temporal gradients. Validation uses the first 100 complete/partial -batches; final offline evaluation reads the entire test split. The expected -training budget is substantial and is not run automatically by tests. - -Default generation samples the full recorded force range: cart force ±1500 -with no pole actuation, or Ant scalar joints ±400. Teachers use MuJoCo-Warp with -10 substeps per 1/60-second macro step. CPU fixtures explicitly request separate -worlds. Source MJCF motor objects and solver XML options are not installed into -the teacher: the host applies actual generalized forces and records its solver -configuration. This profile is a current Newton teacher, not Warp 1.5.1 paper -reproduction. - -## Official comparison controls - -Default training uses deterministic sampling with replacement and streamed -population moments over overlapping training windows. These choices are explicit; -they do not reproduce the official dataloader's permutation or its running -normalizer initialization automatically. - -For a paired official comparison, pass the **same** immutable data, trajectory -split, initialized weights, batch order, normalizer artifact and update budget to -both trainers. `--statistics stats.pt` loads a tensor-only artifact containing -`schema_version: 1`, `inputs: {name: {mean, var}}`, and -`output: {mean, var}`. No validation/test samples may enter these statistics. -`--sample-order order.npy` reads a memory-mapped integer array of shape -`[total_steps, batch_size]`; each row indexes the full-window training dataset. -A `.pt` tensor is also accepted for short fixtures. The array hash and committed -cursor are checked on resume; large `.npy` orders are never loaded in full. - -Use `--seed 0` through `--seed 4` with distinct `--output-dir` values for the -five paired training seeds. Export each validation-selected checkpoint separately -and use evaluation's `--model` and `--output` overrides to retain every result. -Full rollout configurations request 2048 trajectories at horizons 1, 10, 100, -1000; repeat with `motion_mode: passive` and with the qualified fixed policy. -A common seed alone does not guarantee identical minibatches across trainers. - -No committed configuration is evidence that the full run passed. Required reports -include short numerical/Adam parity, both teachers × both trainers, five seed -pairs, confidence bounds, non-finite/failure rates, fixed-policy return and -software/asset/checkpoint hashes. `qualification.paired_noninferiority` implements -the seed-paired one-sided 95% bootstrap gate; fewer than five seeds is pending. -A 100-step smoke cannot replace the full budget or paper reproduction. - -## Registered differentiable task - -After exporting a matching model, `env.nerd.yaml` creates the registered -`CartpoleDynamics-v1` task. Its checkpoint path resolves from the -working directory, unlike dynamics workflow file paths (config-relative). - -```bash -embodichain run-task --gym_config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/env.nerd.yaml --preview -embodichain train-rl --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/agents/apg.nerd.yaml -``` - -Use preview for the Gym scene: this task does not implement an expert trajectory -collector. The APG configuration uses the existing learning-environment factory -and trainer, with graph-connected state/reward hooks. Its policy checkpoint is -separate from the NeRD dynamics-model export. - -## Legacy Warp parameter comparison - -`dynamics.generate.warp-reference.yaml` selects `physics_profile: warp_reference` -and Featherstone with 5 substeps per 1/60-second frame. It preserves measured -legacy asset masses/inertia, body armature, force-mode damping and contact-slot -order. Ant also uses an explicitly authored shared ground material. The simulator -is still current Newton; contact algorithms can differ from old Warp. - -Generate this bounded 512-trajectory dataset with: - -```bash -embodichain dynamics generate --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.warp-reference.yaml -``` - -Keep the same `physics_profile` in rollout evaluation. Existing exports remain -bound to their recorded physical signature; generate new data and train or -fine-tune a new export for this profile. The default `native` profile retains -modern importer semantics. Both profiles need an installed DexSim with named -solver extensions; the measured SDK branch is MR !1483 at `b61e920f`. - -The isolated historical/current replay tool is -`scripts/experiments/neural_dynamics/physics_alignment.py`. Its shared input and -trace hashes, plus finite/failure counts, are required when comparing trajectories. - - -## Supported-version reference data +Smoke profiles verify executable flow, not accuracy. For a different dataset, +copy the appropriate training config and set `dataset` and `output_dir` to its +location and a new training directory. For a different physical system, collect +new teacher data and train a matching model. -The [original data and paired evaluation workflow](../../../../../../scripts/experiments/neural_dynamics/README.md) -uses isolated Python3.8/Torch2.2 generation, immutable raw/canonical teacher -traces and independently evolving model feedback. Training and evaluation -runtimes are recorded separately; a current-Newton corpus cannot be relabeled -as original-Warp data. +Configuration paths resolve relative to the YAML file; CLI path overrides +resolve relative to the working directory. Exports require a new destination. +Resume an identical training run with `train --resume PATH/last.pt`, or initialize +a new run from a portable export with `train --fine-tune PATH`. Validation data +selects the checkpoint; the test split remains held out. 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 000000000..73b9c4926 --- /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.rollout.pretrained.yaml b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.rollout.pretrained.yaml new file mode 100644 index 000000000..8a33af90a --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.rollout.pretrained.yaml @@ -0,0 +1,21 @@ +# Published model versus its matching current-Newton analytical teacher. +mode: rollout +robot: cartpole +model: ../../../../../../outputs/models/neural_dynamics/nerd/cartpole/v1 +device: cuda:0 +num_envs: 16 +trajectories: 16 +horizons: +- 1 +- 10 +- 100 +seed: 1000000 +physics_profile: warp_reference +teacher_solver: + solver_type: featherstone + update_mass_matrix_interval: 5 +num_substeps: 5 +motion_mode: random_effort +initial_states: sample +effort_scale: 1.0 +output: ../../../../../../outputs/dynamics/cartpole/pretrained/rollout.json diff --git a/scripts/experiments/neural_dynamics/README.md b/scripts/experiments/neural_dynamics/README.md index 0e863c6fd..4900ad802 100644 --- a/scripts/experiments/neural_dynamics/README.md +++ b/scripts/experiments/neural_dynamics/README.md @@ -1,262 +1,64 @@ -# Original NeRD data and matched training - -Use two isolated environments. The original data worker uses Python 3.8.20, -Torch 2.2.2, NumPy 1.24.4 and Warp 1.8.0 for Cartpole or Warp 1.5.1 for Ant. The -EmbodiChain environment owns current-Newton collection, controlled training, -export and native inference. Use the matching SDK build for the named solver -and functional Scene contracts. Keep the unchanged original checkout at commit -`370e89d2b8ce68aa3578aa7812859c681b7502bc`. - -## Original-entry comparison with a small budget - -`official_small_budget.py` runs the actual upstream `train/train.py` through -`runpy`, including its constructor, DataLoaders, FP32 RunningMeanStd, loss, -epoch-zero validation, rollout evaluation and independent checkpoint selectors. -It observes the emitted batches rather than generating an approximate shuffle. -Use the supported original environment above. CPU execution maps the upstream -CUDA-stream request to `None`; the original numerical methods remain unchanged. - -Provide a JSON plan with `source`, `robot`, `output_dir`, `training_config`, -`seed: 0`, `steps: 1000` and `cpu_threads: 8`. Start from the original robot YAML, -change the existing dataset paths and set `algorithm.num_iters_per_epoch: 1000`. -Keep `num_epochs: 1000` and the original optimizer/model fields. The observer -stops before the second trained epoch, after the original validation, simulation -evaluation and checkpoint selection complete. Thus all 1,000 updates retain the -original first-trained-epoch rate, 0.0009991. Existing output directories are -rejected. The original source checkout stays unmodified. - -```bash -CUDA_VISIBLE_DEVICES="" PYTHONPATH="$NERD_WARP_ROOT" "$NERD_DATA_PYTHON" scripts/experiments/neural_dynamics/official_small_budget.py --config original-small-plan.json --trust-external-code -``` - -Import the same raw split files for the migrated trainer. For Ant, preserve the -original Y-up world coordinates in this experiment (`world_rotation_xyzw=None`) -so extra world rotation does not enter the comparison. This does not change the -native Newton deployment coordinate contract. The official run supplies the -actual initial weights, normalization statistics and consumed sample indices. - -```bash -CUDA_VISIBLE_DEVICES="" "$EMBODICHAIN_PYTHON" scripts/experiments/neural_dynamics/matched_small_budget.py --official-dir /workspace/exp_data/ant/official --dataset /workspace/exp_data/ant/original-frame-data --output-dir /workspace/exp_data/ant/migrated --device cpu -``` - -The migrated experiment calls production `DynamicsTrainer.step` and export APIs; -its outer loop reproduces the captured original sampling and validation selector. -This differs from the default CLI sampling/statistics protocol. Compare both CPU -runs first, and record any separate CUDA run as an environment diagnostic. - -New first-batch observations are saved **before** the original model forward, -with `first-preprocessed-batch.json` identifying physical features and binding -the tensor file's SHA256. The original model normalizes its input dictionary in -place. Earlier observations without that sidecar contain already-normalized -features; the matched reader preserves that legacy interpretation and records -the observed stage. Rebuild physical inputs from raw data for gradient probes, -rather than passing a legacy normalized snapshot through the model again. - -`official_small_budget_evaluate.py` evaluates the selected original checkpoint on -its untouched test split, records fixed validation curves and produces forward -golden tensors under Python3.8/Torch2.2. `official_small_budget_rollout.py` reuses -fixed prefixes of the supported-environment evaluation corpora below, evaluates -both trained weight sets through the original integrator in the same supported -runtime, and records endpoint/prefix errors at 1/10/100 steps. A migrated export -must first be restored into the original model format, preserving all weights -and normalizers. All failed trajectories remain in the measurements, and each -feedback trajectory uses its own predicted state. - -Historical-data model exports and native SDK models have different deployment -provenance. Verify the current SDK generate/train/export/infer workflow with -native-data models; never replace a historical model's missing native physics -signature to bypass deployment checks. - -## Reference/math exported-model inference - -`official_math_rollout.py` evaluates an explicitly `reference`/math portable -export through the unchanged original model and Transformer integrator. This is -the reusable entry used for the full-budget model's supported-environment -accuracy checks. The original-default baseline remains separately evaluated by -`official_small_budget_rollout.py` with its native official checkpoint. - -Use the original Python 3.8.20/Torch 2.2.2 environment and the robot's pinned Warp. -Provide fixed official teacher trajectories, the portable export and the exact -pure transport modules whose hashes were recorded when the teachers were -created. Source, asset, raw-state/action, physics and transport identities are -checked; changing a teacher or transport requires its own experiment identity. - -```bash -CUDA_VISIBLE_DEVICES="" PYTHONDONTWRITEBYTECODE=1 PYTHONPATH="$NERD_WARP_ROOT" "$NERD_DATA_PYTHON" scripts/experiments/neural_dynamics/official_math_rollout.py --source "$NERD_REFERENCE_SOURCE" --portable "$NERD_PORTABLE_MODEL" --scope checkpoint_accuracy --teacher "$NERD_TEACHER_DATA" --transport-root "$NERD_FROZEN_TRANSPORT" --output "$NERD_NEW_RESULT_JSON" --count 16 --trust-external-code -``` - -The output contains H1/10/100 endpoint and prefix metrics, complete NPZ traces, -source/runtime/physics provenance and the actual attention operator. Passive and -random-effort runs replay the same controls; PD recomputes each neural branch's -controls from its own predicted state. Failed trajectories remain included. -`--scope smoke` is reserved for zero-update pipeline checks. This tool neither -starts training nor qualifies historical weights for current-Newton deployment. - -`_portable_reference.py` and `_historical_physics.py` are standalone helpers for -this old Python environment. They do not import the installed EmbodiChain package -or SDK, and do not modify the original source or default baseline. - -## Cross-version numerical attribution - -Treat training implementation, runtime kernels and data transport as separate -controls. On shared canonical data, the current-environment original numerical -methods and migrated trainer match through 1,000 updates for both robots. This -does not remove the Ant raw-spatial/native-COM target-rounding distinction. - -CPU SDPA dispatch and the final FP32 gradient-norm reduction can differ between -Torch 2.2 and 2.7. A jointly specified math-attention and fixed-order total-norm -policy is a separate controlled variant. It changes the original default -numerical path and must not replace the unchanged official default baseline or -silently change production settings. Record optimizer tensor equality separately -from version-specific metadata such as `decoupled_weight_decay=False`. - -The [numerical attribution report](../../../design/reports/neural-dynamics-numerical-attribution-20261020.md) -records bounded operator probes, the 1,000-update controls and their limits. - -## Full-budget continuation with an explicit numerical policy - -Production `DynamicsTrainCfg.numerical_profile="reference"` selects math SDPA -and ordered FP64 accumulation of FP32 per-parameter gradient norms. Portable -exports persist `attention_backend="math"`; changing the policy requires a new -fine-tuning run. The original-method `run_reference` adapter supports only its -default policy and rejects math initial exports rather than silently ignoring -their execution metadata. Keep the original default baseline separate. - -For source-fidelity imports, `import_legacy_datasets` accepts -`preserve_legacy_velocity=True` with `world_rotation_xyzw=None`. It preserves -original FP32 velocity before COM conversion, verifies bit-exact transition -continuity, and supplies current/next raw velocity directly to production -encoding. This option does not change native deployment physics qualification. - -Keep `train.total_steps=4995000` for both initial and resumed phases. Use only -top-level `stop_after=5000` for the first checkpoint; then remove `fine_tune` and -`statistics`, provide `resume`, and retain the entire training configuration. -The exact source-shuffle descriptor includes the original DataLoader's epoch -tail dropping and source permutation hashes. Independent seeded permutations -or random-with-replacement orders do not reproduce that stream. - -The [full-budget progress report](../../../design/reports/neural-dynamics-full-budget-20261021.md) -binds the two formal runs, official zero-step initialization, GPU recovery gates, -fixed-test metrics and original-supported-runtime rollout comparisons. Both -actual committed cursors must reach the full budget before completion is claimed. - -## Training data and paired training - -`official_model_reference.py --generate-only` executes the unmodified original -sampler without a model. `NERD_DATA_PYTHON` is the original environment's -interpreter; `NERD_WARP_ROOT` is the robot's pinned Warp dependency directory: - -```bash -CUDA_VISIBLE_DEVICES="" PYTHONPATH="$NERD_WARP_ROOT" "$NERD_DATA_PYTHON" scripts/experiments/neural_dynamics/official_model_reference.py --source /path/to/neural-robot-dynamics --robot ant --output /workspace/exp_data/ant/reference-train --count 10000 --num-envs 512 --horizon 100 --seed 0 --generate-only --trust-external-code -``` - -Generate validation/test separately with different seeds and output paths. The -original sampler rounds counts to replica batches; inspect the measured count. -`import_legacy_datasets` in `embodichain.learning.dynamics.reference` imports raw -split files with canonical transport and explicit gravity-transport repair. -Preserve unchanged raw files and actual split identities, versions and hashes. - -Prepare the controlled matrix with -`embodichain.learning.dynamics.comparison.prepare_comparison`, supplying both -immutable teacher datasets, original source/dependency paths, five seeds and -complete model/train configuration. `run_comparison` executes frozen commands. -Both trainers share initialized tensors, normalizers, sample indices, updates -and the validation selector. Use distinct output directories and explicitly -resume incomplete work from its own optimizer/RNG checkpoint. This controlled -protocol does not establish literal paper-recipe reproduction. - -## Freeze supported-environment evaluation data - -`official_evaluation_data.py` retains all requested outcomes, including physical -failures and NaN/Inf, with no rejection/resampling. It runs original physics and -sampling unchanged. Pure Torch coordinate/observation files load directly under -an isolated namespace; their paths/hashes are recorded separately. The worker -does not import the EmbodiChain package, DexSim or Newton. - -Example `original-eval-data.json`: - -```json -{ - "source": "/path/to/neural-robot-dynamics", - "robot": "ant", - "output_dir": "/workspace/exp_data/ant/original-eval-passive", - "num_envs": 64, - "trajectories": 2048, - "horizon": 1000, - "seed": 1000000, - "initial_states": "sample", - "motion_mode": "passive" -} -``` - -```bash -CUDA_VISIBLE_DEVICES="" PYTHONPATH="$NERD_WARP_ROOT" "$NERD_DATA_PYTHON" scripts/experiments/neural_dynamics/official_evaluation_data.py --config original-eval-data.json --trust-external-code -``` - -Use new directories for the other scenarios. Set `motion_mode: random_effort`, -or set `motion_mode: fixed_policy` with Ant policy -`{"kind": "pd", "kp": 20.0, "kd": 1.0}`. Cartpole PD uses -`{"kind": "pd", "gains": [5.0, 50.0, 5.0, 20.0]}`. These are declared baselines, -not the paper's learned policy. Original raw states/efforts are preserved beside -canonical q/qd/efforts and measured physics. - -Choose seeds outside every stored training/validation/test candidate range and -audit actual initial states. Current collection uses generation seed plus -candidate index, so comparing top-level seeds alone is insufficient. Reuse one -immutable corpus per robot/scenario for both trainers and all model seeds. - -## Evaluate both exports against one corpus - -Run the historical evaluator in a fresh EmbodiChain Python process, with the -old Warp dependency directory prepended to `PYTHONPATH`. Do not replace installed -current Warp or import DexSim/Newton into that process. Example spec: - -```json -{ - "source": "/path/to/neural-robot-dynamics", - "robot": "ant", - "model": "/workspace/exp_data/ant/official-seed0/best-export", - "teacher_dataset": "/workspace/exp_data/ant/original-eval-passive", - "output": "/workspace/exp_data/ant/official-seed0/passive.json", - "device": "cpu", - "num_envs": 64, - "trajectories": 2048, - "horizons": [1, 10, 100, 1000], - "seed": 1000000, - "initial_states": "sample", - "motion_mode": "passive", - "cpu_threads": 4 -} -``` - -```bash -CUDA_VISIBLE_DEVICES="" PYTHONPATH="$NERD_WARP_ROOT" "$EMBODICHAIN_PYTHON" -m embodichain.learning.dynamics.reference_rollout --spec historical-eval.json --trust-external-code -``` - -The evaluator verifies generation versions/source/physics, raw/canonical state -transport, ordered teacher contents and feedback actions. The neural simulator -starts from the recorded initial state, evolves independently, and computes -feedback from its own state. Later teacher states are never copied into it. -Model-runtime versions and teacher-generation versions are recorded separately. -Omitting `teacher_dataset` retains live teacher generation for diagnostics. - -For the migrated export change only model/output. Repeat for both robots, all -three scenarios and five seeds, retaining complete budgets. Use -`embodichain dynamics` for separate current-Newton deployment evaluation; it also -accepts a verified `teacher_dataset`, including fixed-PD feedback traces. - -Attach results with `annotate_rollout_evidence`, then call -`write_comparison_report`. Training completion, numerical parity, hashes and all -error/failure/nonfinite/return gates are separate checks. Completed commands and -short pilots cannot substitute for full trained-model accuracy qualification. - -## Optional reference tests - -For the opt-in checkpoint tests, set `NERD_REFERENCE_ROOT` to a pinned checkout -with real LFS weights. When that checkout is original Warp main, also set -`NERD_REFERENCE_DEPS` to the matching isolated Warp dependency directory. A -checkout containing LFS pointer files cannot supply the pretrained checkpoint. -For the fixed-teacher tests, `NERD_SUPPORTED_EVAL_DATA` contains directories named -`cartpole-passive`, `ant-passive`, and the corresponding `random_effort` and -`fixed_policy` directories. These tests inspect original generation provenance -separately from the model-inference runtime. +# NeRD reference experiments + +These command-line tools compare an upstream NeRD implementation with the +EmbodiChain implementation. The normal dataset, training, export and solver +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 + +| Tool | Purpose | +| --- | --- | +| `official_model_reference.py` | Generate upstream reference data and model observations | +| `official_evaluation_data.py` | Freeze upstream teacher trajectories and provenance | +| `compare_official_model.py` | Compare a portable copy with an upstream checkpoint | +| `official_small_budget.py` | Observe the original training entry point and consumed batches | +| `matched_small_budget.py` | Replay captured initialization, statistics and sampling in the production trainer | +| `official_small_budget_evaluate.py` | Evaluate an observed original training run | +| `official_small_budget_rollout.py` | Evaluate checkpoints through the original integrator | +| `official_math_rollout.py` | Evaluate explicit reference/math exports in the original runtime | +| `physics_alignment.py` | Capture matched-input physics/version comparisons | +| `trajectory_closeness.py` | Measure kinematic motion differences in physical units | + +`_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/physics_alignment.py b/scripts/experiments/neural_dynamics/physics_alignment.py index 1c95bf845..b72d6f005 100644 --- a/scripts/experiments/neural_dynamics/physics_alignment.py +++ b/scripts/experiments/neural_dynamics/physics_alignment.py @@ -25,8 +25,11 @@ import argparse import hashlib import json +import os from pathlib import Path +import platform import sys +from typing import Any import numpy as np @@ -55,7 +58,16 @@ def _compare(args: argparse.Namespace) -> dict: ] if q.shape != rq.shape or qd.shape != rqd.shape: raise ValueError("Trace shapes differ") - torch.testing.assert_close(q[0], rq[0], rtol=1e-5, atol=1e-6) + if args.robot == "ant": + torch.testing.assert_close(q[0, :, :3], rq[0, :, :3], rtol=1e-5, atol=1e-6) + torch.testing.assert_close(q[0, :, 7:], rq[0, :, 7:], rtol=1e-5, atol=1e-6) + a = q[0, :, 3:7].double() + b = rq[0, :, 3:7].double() + a, b = a / a.norm(dim=-1, keepdim=True), b / b.norm(dim=-1, keepdim=True) + if not bool(torch.all(torch.abs((a * b).sum(-1)) > 1 - 1e-10)): + raise ValueError("Initial root orientations differ") + else: + torch.testing.assert_close(q[0], rq[0], rtol=1e-5, atol=1e-6) torch.testing.assert_close(qd[0], rqd[0], rtol=1e-5, atol=1e-6) errors = _error_values(args.robot, q, qd, rq, rqd) finite = torch.isfinite(torch.cat((q, qd, rq, rqd), -1)).all(-1) @@ -92,15 +104,31 @@ def _compare(args: argparse.Namespace) -> dict: } -def _inputs(robot: str, count: int, horizon: int) -> dict[str, np.ndarray]: +def _inputs( + robot: str, + count: int, + horizon: int, + motion_mode: str = "random_effort", + initial_states: str = "near_upright", + effort_scale: float = 0.02, +) -> dict[str, np.ndarray]: from embodichain.lab.sim.physics.solvers.collection import _sample rows = [ - _sample(robot, 42 + i, horizon, "random_effort", "default", 0.02) + _sample( + robot, + 42 + i, + horizon, + motion_mode, + "sample" if initial_states == "sample" else "default", + effort_scale, + ) for i in range(count) ] q, qd, efforts = [np.stack([row[j].numpy() for row in rows]) for j in range(3)] rng = np.random.default_rng(42) + if initial_states == "sample": + return {"q": q, "qd": qd, "efforts": efforts} if robot == "cartpole": q[:, 0] = rng.uniform(-0.25, 0.25, count) q[:, 1] = rng.uniform(-0.4, 0.4, count) @@ -112,9 +140,35 @@ def _inputs(robot: str, count: int, horizon: int) -> dict[str, np.ndarray]: return {"q": q, "qd": qd, "efforts": efforts} +def _numpy(value: Any) -> np.ndarray: + if hasattr(value, "detach"): + return value.detach().cpu().numpy() + return value.numpy() if hasattr(value, "numpy") else np.asarray(value) + + +def _physical_arrays(model: Any) -> dict[str, Any]: + """Record measured parameters without requiring cross-version equality.""" + result = {} + for name in ( + "body_mass", + "body_com", + "body_inertia", + "joint_armature", + "joint_target_ke", + "joint_target_kd", + "joint_damping", + ): + value = ( + model.get(name) if isinstance(model, dict) else getattr(model, name, None) + ) + if value is not None: + result[name] = _numpy(value).tolist() + return result + + def _legacy( args: argparse.Namespace, data: dict[str, np.ndarray] -) -> tuple[np.ndarray, np.ndarray, dict]: +) -> tuple[dict[str, np.ndarray], dict]: from embodichain.learning.dynamics.reference import reference_identity source = reference_identity(args.source) @@ -145,7 +199,11 @@ def _legacy( wp.sim.eval_fk( env.model, env.state_0.joint_q, env.state_0.joint_qd, None, env.state_0 ) - qs, vs = [], [] + qs, vs, bodies = [], [], [] + body_count = env.model.body_count // args.count + names = [str(name).rsplit("/", 1)[-1] for name in env.model.body_name[:body_count]] + if len(names) != len(set(names)): + raise ValueError("Legacy body names are not unique within one world") for step in range(args.horizon + 1): a = env.state_0.joint_q.numpy().reshape(args.count, -1).copy() v = env.state_0.joint_qd.numpy().reshape(args.count, -1).copy() @@ -161,6 +219,15 @@ def _legacy( a[:, 3:7] = (rotation * Rotation.from_quat(a[:, 3:7])).as_quat() qs.append(a) vs.append(v) + body = env.state_0.body_q.numpy().reshape(args.count, body_count, 7).copy() + flat = body.reshape(-1, 7) + flat[:, :3] = rotation.apply(flat[:, :3]) + # Invalid simulator quaternions remain in the trace and finite gate. + valid = np.isfinite(flat[:, 3:]).all(-1) & ( + np.linalg.norm(flat[:, 3:], axis=-1) > 0 + ) + flat[valid, 3:] = (rotation * Rotation.from_quat(flat[valid, 3:])).as_quat() + bodies.append(body) if step < args.horizon: force = data["efforts"][:, step] env.control.joint_act.assign( @@ -168,21 +235,23 @@ def _legacy( ) env.update() return ( - np.stack(qs), - np.stack(vs), + {"q": np.stack(qs), "qd": np.stack(vs), "body_q": np.stack(bodies)}, { "source": source, "warp": wp.__version__, "substeps": env.sim_substeps, "mass_update_interval": env.integrator.update_mass_matrix_every, "gravity": list(env.model.gravity), + "body_names": names, + "body_pose_frame": "Z-up world; legacy Y-up poses rotated +pi/2 around X", + "physical_arrays": _physical_arrays(env.model), }, ) def _current( args: argparse.Namespace, data: dict[str, np.ndarray] -) -> tuple[np.ndarray, np.ndarray, dict]: +) -> tuple[dict[str, np.ndarray], dict]: import torch import dexsim from embodichain.lab.sim.physics.solvers.collection import DynamicsSimulation @@ -191,7 +260,7 @@ def _current( solver = ( None if args.solver == "featherstone" - else {"solver_type": "mujoco_warp", "use_mujoco_cpu": True} + else {"solver_type": "mujoco_warp", "use_mujoco_cpu": args.device == "cpu"} ) if args.profile == "native" and solver is None: solver = { @@ -200,21 +269,55 @@ def _current( } with DynamicsSimulation( args.robot, + device=args.device, num_envs=args.count, physics_profile=args.profile, num_substeps=substeps, solver=solver, ) as env: - env.reset(torch.from_numpy(data["q"]), torch.from_numpy(data["qd"])) - qs, vs = [], [] + env.reset( + torch.from_numpy(data["q"]).to(env.device), + torch.from_numpy(data["qd"]).to(env.device), + ) + qs, vs, bodies = [], [], [] + model = dict(env.sim.spawn_result.dynamics_metadata()) + worlds = _numpy(model["body_world"]) + body_ids = np.flatnonzero(worlds == (0 if int(model["world_count"]) else -1)) + names = [str(model["body_label"][i]).rsplit("/", 1)[-1] for i in body_ids] + if len(names) != len(set(names)): + raise ValueError("Current body names are not unique within one world") + all_ids = [] + for world in range(args.count): + selected = np.flatnonzero( + worlds == (world if int(model["world_count"]) else -1) + ) + labels = { + str(model["body_label"][i]).rsplit("/", 1)[-1]: int(i) for i in selected + } + all_ids.append([labels[name] for name in names]) + body_indices = torch.tensor(all_ids, device=env.device, dtype=torch.long) for step in range(args.horizon + 1): state, _ = env.snapshot() - qs.append(state["q"].numpy()) - vs.append(state["qd"].numpy()) + qs.append(_numpy(state["q"]).copy()) + vs.append(_numpy(state["qd"]).copy()) + physical = env.sim.spawn_result.dynamics_snapshot().physical + bodies.append(physical["body_q"][body_indices].cpu().numpy().copy()) if step < args.horizon: - env.step(torch.from_numpy(data["efforts"][:, step])) + env.step(torch.from_numpy(data["efforts"][:, step]).to(env.device)) metadata = {**env.manifest, "dexsim_commit": dexsim.__commit_id__} - return np.stack(qs), np.stack(vs), metadata + metadata.update( + body_names=names, + body_pose_frame="Z-up world", + physical_arrays=_physical_arrays(model), + actual_device=str(env.device), + cuda_visible_devices=os.environ.get("CUDA_VISIBLE_DEVICES"), + actual_gpu_uuid=( + str(torch.cuda.get_device_properties(env.device).uuid) + if env.device.type == "cuda" + else None + ), + ) + return {"q": np.stack(qs), "qd": np.stack(vs), "body_q": np.stack(bodies)}, metadata def main() -> None: @@ -224,6 +327,7 @@ def main() -> None: ) parser.add_argument("--robot", choices=["cartpole", "ant"], required=True) parser.add_argument("--source", type=Path) + parser.add_argument("--device", choices=["cpu", "cuda:0"], default="cpu") parser.add_argument("--trust-external-code", action="store_true") parser.add_argument( "--profile", choices=["native", "warp_reference"], default="warp_reference" @@ -232,12 +336,21 @@ def main() -> None: "--solver", choices=["featherstone", "mujoco_warp"], default="featherstone" ) parser.add_argument("--substeps", type=int) + parser.add_argument( + "--motion-mode", choices=["passive", "random_effort"], default="random_effort" + ) + parser.add_argument( + "--initial-states", choices=["near_upright", "sample"], default="near_upright" + ) + parser.add_argument("--effort-scale", type=float, default=0.02) parser.add_argument("--count", type=int, default=8) parser.add_argument("--horizon", type=int, default=1000) parser.add_argument("--inputs", type=Path, required=True) parser.add_argument("--compare-with", type=Path) parser.add_argument("--output", type=Path, required=True) args = parser.parse_args() + if args.engine == "legacy" and args.device != "cpu": + parser.error("legacy execution remains pinned to --device cpu") if args.engine == "compare": if args.output.exists(): raise FileExistsError(args.output) @@ -258,14 +371,21 @@ def main() -> None: else: if args.engine != "newton": parser.error("generate shared inputs with the current worker first") - data = _inputs(args.robot, args.count, args.horizon) + data = _inputs( + args.robot, + args.count, + args.horizon, + args.motion_mode, + args.initial_states, + args.effort_scale, + ) args.inputs.parent.mkdir(parents=True, exist_ok=True) np.savez(args.inputs, **data) if data["efforts"].shape[:2] != (args.count, args.horizon): raise ValueError("input shape differs from count/horizon") - q, qd, meta = (_legacy if args.engine == "legacy" else _current)(args, data) + traces, meta = (_legacy if args.engine == "legacy" else _current)(args, data) args.output.parent.mkdir(parents=True, exist_ok=True) - np.savez(args.output, q=q, qd=qd) + np.savez(args.output, **traces) meta.update( engine=args.engine, profile=args.profile, @@ -274,6 +394,11 @@ def main() -> None: trace_sha256=hashlib.sha256(args.output.read_bytes()).hexdigest(), horizon=args.horizon, count=args.count, + motion_mode=args.motion_mode, + initial_states=args.initial_states, + effort_scale=args.effort_scale, + python=platform.python_version(), + runner_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), ) args.output.with_suffix(".json").write_text(json.dumps(meta, indent=2, default=str)) print(args.output) diff --git a/scripts/experiments/neural_dynamics/trajectory_closeness.py b/scripts/experiments/neural_dynamics/trajectory_closeness.py new file mode 100644 index 000000000..b35661554 --- /dev/null +++ b/scripts/experiments/neural_dynamics/trajectory_closeness.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. +# ---------------------------------------------------------------------------- + +"""Compare isolated physics traces in physical units using only NumPy. + +This experimental gate measures motion closeness. It does not require identical +internal model trees, physical-array bytes, velocities or solver implementations. +""" + +from __future__ import annotations + +import argparse +from datetime import datetime, timezone +import hashlib +import json +from pathlib import Path +from typing import Any + +import numpy as np + +__all__: list[str] = [] + + +def _hash(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _orientation(a: np.ndarray, b: np.ndarray) -> np.ndarray: + a = a / np.linalg.norm(a, axis=-1, keepdims=True) + b = b / np.linalg.norm(b, axis=-1, keepdims=True) + return 2 * np.arccos(np.clip(np.abs(np.sum(a * b, axis=-1)), 0, 1)) + + +def _rotate(q: np.ndarray, v: np.ndarray) -> np.ndarray: + q = q / np.linalg.norm(q, axis=-1, keepdims=True) + t = 2 * np.cross(q[..., :3], v) + return v + q[..., 3:] * t + np.cross(q[..., :3], t) + + +def _summary(values: np.ndarray) -> dict[str, Any]: + values = np.asarray(values, dtype=np.float64) + finite = np.isfinite(values) + result: dict[str, Any] = { + "count": int(values.size), + "nonfinite_count": int((~finite).sum()), + "per_trajectory": [float(x) if np.isfinite(x) else None for x in values], + } + for key in ("mean", "median", "p90", "p95", "max"): + result[key] = None + if bool(finite.all()): + result.update( + mean=float(values.mean()), + median=float(np.median(values)), + p90=float(np.quantile(values, 0.9)), + p95=float(np.quantile(values, 0.95)), + max=float(values.max()), + ) + return result + + +def _first(mask: np.ndarray) -> list[int | None]: + return [int(np.flatnonzero(row)[0]) if bool(row.any()) else None for row in mask.T] + + +def _load(path: Path) -> tuple[dict[str, np.ndarray], dict[str, Any]]: + receipt = json.loads(path.with_suffix(".json").read_text()) + if _hash(path) != receipt["trace_sha256"]: + raise ValueError(f"Trace hash differs from receipt: {path}") + with np.load(path, allow_pickle=False) as arrays: + data = {name: arrays[name].astype(np.float64) for name in arrays.files} + if data["q"].shape[:2] != data["qd"].shape[:2]: + raise ValueError("Position/velocity time and trajectory dimensions differ") + if data["q"].ndim != 3 or data["q"].shape[0] < 101: + raise ValueError("Expected time-major traces with at least 100 transitions") + return data, receipt + + +def _metrics( + robot: str, + ref: dict[str, np.ndarray], + actual: dict[str, np.ndarray], + metadata: tuple[dict[str, Any], dict[str, Any]], +) -> tuple[dict[str, np.ndarray], dict[str, Any]]: + rq, q = ref["q"], actual["q"] + rv, v = ref["qd"], actual["qd"] + if rq.shape != q.shape or rv.shape != v.shape: + raise ValueError("Trace state shapes differ") + metrics = {} + if robot == "cartpole": + delta = q[..., 1] - rq[..., 1] + metrics.update( + slider_position_abs_m=np.abs(q[..., 0] - rq[..., 0]), + pole_angle_abs_rad=np.abs(np.arctan2(np.sin(delta), np.cos(delta))), + slider_velocity_abs_m_per_s=np.abs(v[..., 0] - rv[..., 0]), + pole_velocity_abs_rad_per_s=np.abs(v[..., 1] - rv[..., 1]), + ) + # A fixed massless root is not an independently moving free base. + interpretation = {"root_pose": "not applicable: fixed Cartpole root"} + else: + metrics.update( + joint_angle_MAE_rad=np.abs(q[..., 7:] - rq[..., 7:]).mean(-1), + root_position_L2_m=np.linalg.norm(q[..., :3] - rq[..., :3], axis=-1), + root_orientation_geodesic_rad=_orientation(q[..., 3:7], rq[..., 3:7]), + joint_velocity_MAE_rad_per_s=np.abs(v[..., 6:] - rv[..., 6:]).mean(-1), + root_linear_velocity_L2_m_per_s=np.linalg.norm( + v[..., :3] - rv[..., :3], axis=-1 + ), + root_angular_velocity_L2_rad_per_s=np.linalg.norm( + v[..., 3:6] - rv[..., 3:6], axis=-1 + ), + ) + for index in range(8): + metrics[f"joint_{index}_angle_abs_rad"] = np.abs( + q[..., 7 + index] - rq[..., 7 + index] + ) + interpretation = {"joint_angles": "eight limited scalar joints; no wrapping"} + common_names = [] + if "body_q" in ref and "body_q" in actual: + names_r, names_a = (row["body_names"] for row in metadata) + common_names = [name for name in names_r if name in names_a] + if not common_names: + raise ValueError("No common named links in body traces") + errors = [] + for name in common_names: + rbody = ref["body_q"][..., names_r.index(name), :] + body = actual["body_q"][..., names_a.index(name), :] + error = np.linalg.norm(body[..., :3] - rbody[..., :3], axis=-1) + metrics[f"link_{name}_position_L2_m"] = error + errors.append(error) + if robot == "cartpole" and name == "pole": + tip = body[..., :3] + _rotate(body[..., 3:], np.array([0.0, 0.0, 1.0])) + rtip = rbody[..., :3] + _rotate( + rbody[..., 3:], np.array([0.0, 0.0, 1.0]) + ) + metrics["pole_tip_position_L2_m"] = np.linalg.norm(tip - rtip, axis=-1) + errors.append(metrics["pole_tip_position_L2_m"]) + metrics["key_link_max_position_L2_m"] = np.stack(errors).max(0) + interpretation.update( + common_named_links=common_names, + link_matching="Semantic name intersection; extra fixed or collapsed bodies do not require identical trees", + prefix="Within each trajectory mean over steps 1..H; distributions across all requested trajectories", + endpoint="At step H; distributions across all requested trajectories", + selection="No failed or nonfinite trajectories are dropped", + velocity="Descriptive diagnostics; not an automatic motion gate", + ) + return metrics, interpretation + + +def _compare(reference: Path, actual: Path, thresholds: Path) -> dict[str, Any]: + ref, rm = _load(reference) + data, am = _load(actual) + if rm["robot"] != am["robot"] or rm["inputs_sha256"] != am["inputs_sha256"]: + raise ValueError("Traces must have the same robot and input SHA256") + robot = rm["robot"] + limits = json.loads(thresholds.read_text()) + metrics, interpretation = _metrics(robot, ref, data, (rm, am)) + initial_keys = ( + ["slider_position_abs_m", "pole_angle_abs_rad"] + if robot == "cartpole" + else [ + "joint_angle_MAE_rad", + "root_position_L2_m", + "root_orientation_geodesic_rad", + ] + ) + if any(np.max(metrics[name][0]) > 1e-5 for name in initial_keys): + raise ValueError("Canonical initial states differ beyond 1e-5 physical units") + if not np.allclose(data["qd"][0], ref["qd"][0], rtol=1e-5, atol=1e-6): + raise ValueError("Canonical initial velocities differ") + finite_masks = [] + for trace in (ref, data): + mask = np.isfinite(trace["q"]).all(-1) & np.isfinite(trace["qd"]).all(-1) + if robot == "ant": + mask &= np.linalg.norm(trace["q"][..., 3:7], axis=-1) > 0 + if "body_q" in trace: + mask &= np.isfinite(trace["body_q"]).all((-1, -2)) + mask &= (np.linalg.norm(trace["body_q"][..., 3:], axis=-1) > 0).all(-1) + finite_masks.append(mask) + horizon_results, checks = {}, [] + for h in (1, 10, 100): + valid = finite_masks[0][: h + 1].all(0) & finite_masks[1][: h + 1].all(0) + row = { + "finite_trajectories": int(valid.sum()), + "nonfinite_trajectories": int((~valid).sum()), + "endpoint": {key: _summary(value[h]) for key, value in metrics.items()}, + "prefix_mean": { + key: _summary(value[1 : h + 1].mean(0)) + for key, value in metrics.items() + }, + } + checks.append( + { + "horizon": h, + "check": "all requested trajectories finite", + "passed": bool(valid.all()), + } + ) + for key, bounds in limits["robots"][robot]["horizons"][str(h)].items(): + for kind in ("endpoint", "prefix_mean"): + for statistic, bound in bounds.items(): + value = row[kind].get(key, {}).get(statistic) + checks.append( + { + "horizon": h, + "metric": key, + "aggregation": kind, + "statistic": statistic, + "value": value, + "upper_bound": bound, + "passed": value is not None and value <= bound, + } + ) + horizon_results[str(h)] = row + failures = {} + for label, trace in (("reference", ref), ("actual", data)): + if robot == "cartpole": + task = (np.abs(trace["q"][..., 0]) > 4) | (np.abs(trace["qd"]) > 10).any(-1) + criterion = "official Cartpole: abs(slider)>4m or abs(either velocity)>10 (respective units)" + else: + task = trace["q"][..., 2] < 0.2 + criterion = "descriptive torso/root height below 0.2m; falling is retained, not automatic rejection" + failures[label] = { + "first_nonfinite_step": _first( + ~finite_masks[0 if label == "reference" else 1] + ), + "first_task_failure_step": _first(task), + "task_failure_criterion": criterion, + "max_abs_q": ( + float(np.nanmax(np.abs(trace["q"]))) + if np.isfinite(trace["q"]).all() + else None + ), + "max_abs_qd": ( + float(np.nanmax(np.abs(trace["qd"]))) + if np.isfinite(trace["qd"]).all() + else None + ), + } + passed = all(row["passed"] for row in checks) + return { + "schema": "embodichain.physics_trajectory_closeness.v1", + "created_utc": datetime.now(timezone.utc).isoformat(), + "status": ( + "motion_closeness_passed" + if passed + else "motion_closeness_requires_analysis" + ), + "motion_gate_passed": passed, + "robot": robot, + "requested_trajectories": int(data["q"].shape[1]), + "horizons": horizon_results, + "checks": checks, + "failures": failures, + "interpretation": interpretation, + "reference_receipt": rm, + "actual_receipt": am, + "bound_files_sha256": { + str(path): _hash(path) + for path in ( + reference, + reference.with_suffix(".json"), + actual, + actual.with_suffix(".json"), + thresholds, + Path(__file__), + ) + }, + "qualification": "Small fixed-input teacher motion gate only; not a neural accuracy, same-data convergence or native inference qualification", + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--reference", type=Path, required=True) + parser.add_argument("--actual", type=Path, required=True) + parser.add_argument("--thresholds", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if args.output.exists(): + raise FileExistsError(args.output) + result = _compare(args.reference, args.actual, args.thresholds) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(result, indent=2, allow_nan=False) + "\n") + print( + json.dumps( + { + "status": result["status"], + "output": str(args.output), + "sha256": _hash(args.output), + } + ) + ) + + +if __name__ == "__main__": + main() diff --git a/tests/sim/physics/solvers/test_collection.py b/tests/sim/physics/solvers/test_collection.py index 05efac3b1..6349a6d4e 100644 --- a/tests/sim/physics/solvers/test_collection.py +++ b/tests/sim/physics/solvers/test_collection.py @@ -92,6 +92,89 @@ def test_warp_parameter_profile_preserves_reset_replay(robot, substeps) -> None: assert host.manifest["physics_profile"] == "warp_reference" assert host.manifest["teacher_origin"]["kind"] == "current_newton" assert host.manifest["teacher"]["solver"]["solver_type"] == "featherstone" + if robot == "ant": + # Validate the actual Spawn-built scene, including its shared plane. + metadata = host.sim.spawn_result.dynamics_metadata() + body = metadata["shape_body"].cpu().numpy() + mu = metadata["shape_material_mu"].cpu().numpy() + np.testing.assert_allclose(mu[body >= 0], 0.00075) + np.testing.assert_allclose(mu[body < 0], 0.0005) + assert host.manifest["teacher"]["solver"]["friction_smoothing"] == 0.001 + semantics = host.manifest["reference_friction_mapping"] + assert semantics["material_mu_scale"] == 0.001 + assert semantics["reference_shape_mu"] == 0.75 + assert semantics["reference_ground_mu"] == 0.5 + assert ( + host.manifest["teacher_origin"]["reference_friction_mapping"] + == semantics + ) + else: + assert "friction_smoothing" not in host.manifest["teacher"]["solver"] + assert "reference_friction_mapping" not in host.manifest + + +@pytest.mark.parametrize("smoothing", [1.0, 0.0, 1e-5, True, float("nan")]) +def test_ant_warp_contact_mapping_rejects_inconsistent_solver_before_scene( + monkeypatch, smoothing +) -> None: + from embodichain.lab.sim.physics.solvers import collection + + def forbidden_scene(*args, **kwargs): + pytest.fail("Invalid coupled parameters must fail before allocating a scene") + + monkeypatch.setattr(collection, "SimulationManager", forbidden_scene) + with pytest.raises(ValueError, match="coupled contact mapping"): + collection.DynamicsSimulation( + "ant", + physics_profile="warp_reference", + num_substeps=16, + solver={ + "solver_type": "featherstone", + "update_mass_matrix_interval": 16, + "friction_smoothing": smoothing, + }, + ) + + +def test_ant_warp_mujoco_keeps_original_material_domain() -> None: + from embodichain.lab.sim.physics.solvers.collection import DynamicsSimulation + + with DynamicsSimulation( + "ant", + physics_profile="warp_reference", + solver={"solver_type": "mujoco_warp", "use_mujoco_cpu": True}, + ) as host: + metadata = host.sim.spawn_result.dynamics_metadata() + body = metadata["shape_body"].cpu().numpy() + mu = metadata["shape_material_mu"].cpu().numpy() + np.testing.assert_allclose(mu[body >= 0], 0.75) + np.testing.assert_allclose(mu[body < 0], 0.5) + assert "friction_smoothing" not in host.manifest["teacher"]["solver"] + assert "reference_friction_mapping" not in host.manifest + + +def test_ant_warp_neural_mapping_rejects_unknown_marker_before_scene( + tmp_path: Path, monkeypatch +) -> None: + from embodichain.lab.sim.physics.solvers import collection + + checkpoint = tmp_path / "unknown-mapping" + checkpoint.mkdir() + (checkpoint / "manifest.json").write_text( + json.dumps({"reference_friction_mapping": {"schema": "unknown"}}) + ) + + def forbidden_scene(*args, **kwargs): + pytest.fail("Unknown mapping must fail before allocating a scene") + + monkeypatch.setattr(collection, "SimulationManager", forbidden_scene) + with pytest.raises(ValueError, match="Unsupported Ant reference friction mapping"): + collection.DynamicsSimulation( + "ant", + physics_profile="warp_reference", + num_substeps=1, + solver={"solver_type": "nerd", "options": {"checkpoint": str(checkpoint)}}, + ) def test_fixed_policy_adapter_uses_own_state_and_records_hash() -> None: @@ -112,6 +195,51 @@ def _helper(): return module +@pytest.mark.parametrize("mapped", [False, True]) +def test_ant_warp_contact_mapping_survives_nerd_scene_rebuild( + tmp_path: Path, mapped: bool +) -> None: + """Exports and neural scenes must bind the same actual material domain.""" + from embodichain.lab.sim.physics.solvers.collection import DynamicsSimulation + + from embodichain.learning.models.nerd import export_model, load_model + + solver = None if mapped else {"solver_type": "mujoco_warp", "use_mujoco_cpu": True} + with DynamicsSimulation( + "ant", num_substeps=16, physics_profile="warp_reference", solver=solver + ) as teacher: + expected_signature = teacher.binding.physics_signature + prototype = _helper()._small_export( + teacher.sim.spawn_result, tmp_path / "prototype-ant", "ant" + ) + network, _ = load_model(prototype) + checkpoint = tmp_path / "exported-ant" + export_model(checkpoint, network, teacher.manifest) + with DynamicsSimulation( + "ant", + physics_profile="warp_reference", + num_substeps=1, + solver={"solver_type": "nerd", "options": {"checkpoint": str(checkpoint)}}, + ) as neural: + metadata = neural.sim.spawn_result.dynamics_metadata() + body = metadata["shape_body"].cpu().numpy() + mu = metadata["shape_material_mu"].cpu().numpy() + np.testing.assert_allclose(mu[body >= 0], 0.00075 if mapped else 0.75) + np.testing.assert_allclose(mu[body < 0], 0.0005 if mapped else 0.5) + assert neural.binding.physics_signature == expected_signature + if mapped: + assert ( + neural.manifest["reference_friction_mapping"]["material_mu_scale"] + == 0.001 + ) + else: + assert "reference_friction_mapping" not in neural.manifest + neural.step(torch.zeros((1, 14))) + state, _ = neural.snapshot() + assert torch.isfinite(state["q"]).all() + assert torch.isfinite(state["qd"]).all() + + def _check_pipeline(tmp_path: Path, robot: str) -> None: from embodichain.lab.sim.physics.solvers.collection import DynamicsSimulation diff --git a/tests/sim/physics/solvers/test_reference_assets.py b/tests/sim/physics/solvers/test_reference_assets.py index b9a061079..4409c1870 100644 --- a/tests/sim/physics/solvers/test_reference_assets.py +++ b/tests/sim/physics/solvers/test_reference_assets.py @@ -204,7 +204,7 @@ def test_warp_cartpole_import_preserves_measured_mass_and_body_armature(): def test_warp_ant_import_preserves_geometry_density_and_force_mode(): - """Check measured Warp 1.5.1 masses, inertia regularization and friction.""" + """Check measured Warp 1.5.1 mass, body armature and source materials.""" builder = _import_reference("ant", "warp_reference") np.testing.assert_allclose( builder.body_mass, @@ -219,6 +219,7 @@ def test_warp_ant_import_preserves_geometry_density_and_force_mode(): ) np.testing.assert_allclose(builder.joint_damping, 0.0) np.testing.assert_allclose(builder.joint_target_kd, 0.0) + # The host applies any solver-specific material conversion after import. np.testing.assert_allclose(builder.shape_material_mu, 0.75) np.testing.assert_allclose(builder.shape_margin, 0.0) transform = builder.shape_transform[1] @@ -231,6 +232,19 @@ def test_warp_ant_import_preserves_geometry_density_and_force_mode(): ) +def test_native_ant_import_keeps_original_modern_material_and_mass_defaults(): + """The explicit legacy-motion recipe must not change the native domain.""" + builder = _import_reference("ant") + # The modern MJCF parser honors the asset's authored friction=1.5. + np.testing.assert_allclose(builder.shape_material_mu, 1.5) + np.testing.assert_allclose(builder.joint_damping[6:], 0.1) + np.testing.assert_allclose( + builder.body_mass, + [0.48388024694610887, *[0.039157753856, 0.067592205994] * 4], + rtol=1e-6, + ) + + def test_unknown_reference_physics_profile_rejected(): with pytest.raises(ValueError, match="physics profile"): _import_reference("ant", "unregistered") From 8290fc49493a24f0ae8b0b45a79c4685e3ca6d91 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 24 Oct 2026 09:28:20 +0000 Subject: [PATCH 5/7] fix(dynamics): create fresh evaluation output directories --- docs/source/features/neural_dynamics.md | 1 - embodichain/learning/dynamics/workflow.py | 6 +++--- .../configs/tasks/classic_control/ant/README.md | 1 - .../configs/tasks/classic_control/cartpole/README.md | 1 - tests/cli/test_dynamics.py | 6 ++++-- 5 files changed, 7 insertions(+), 8 deletions(-) diff --git a/docs/source/features/neural_dynamics.md b/docs/source/features/neural_dynamics.md index ab09a67e3..b7f1120e5 100644 --- a/docs/source/features/neural_dynamics.md +++ b/docs/source/features/neural_dynamics.md @@ -94,7 +94,6 @@ snapshot_download( Evaluate the saved test transitions without starting a simulator: ```bash -mkdir -p outputs/dynamics/cartpole/pretrained embodichain dynamics eval --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.pretrained.yaml ``` diff --git a/embodichain/learning/dynamics/workflow.py b/embodichain/learning/dynamics/workflow.py index 231e8433b..4e152fe4b 100644 --- a/embodichain/learning/dynamics/workflow.py +++ b/embodichain/learning/dynamics/workflow.py @@ -215,9 +215,9 @@ def run_eval(config: dict[str, Any]) -> dict[str, Any]: "attention_backend": model.cfg.attention_backend, } if config.get("output"): - Path(config["output"]).write_text( - json.dumps(result, indent=2, allow_nan=False) - ) + 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() diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/README.md b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/README.md index 0bdcc55a6..0bbb893de 100644 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/ant/README.md +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/README.md @@ -35,7 +35,6 @@ provides a `snapshot_download` example. Run from the repository root: ```bash -mkdir -p outputs/dynamics/ant/pretrained embodichain dynamics eval --config extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.eval.pretrained.yaml embodichain dynamics eval --config extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.rollout.pretrained.yaml ``` diff --git a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/README.md b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/README.md index 882a52eb2..e59309414 100644 --- a/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/README.md +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/README.md @@ -35,7 +35,6 @@ provides a `snapshot_download` example. Run from the repository root: ```bash -mkdir -p outputs/dynamics/cartpole/pretrained embodichain dynamics eval --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.eval.pretrained.yaml embodichain dynamics eval --config extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.rollout.pretrained.yaml ``` diff --git a/tests/cli/test_dynamics.py b/tests/cli/test_dynamics.py index 2170f6016..9736ed9f0 100644 --- a/tests/cli/test_dynamics.py +++ b/tests/cli/test_dynamics.py @@ -184,7 +184,7 @@ def test_real_offline_cli_train_export_eval(tmp_path): "mode": "offline", "dataset": "data", "model": "export", - "output": "result.json", + "output": "reports/evaluation/result.json", "device": "cpu", } eval_cfg = tmp_path / "eval.yaml" @@ -212,7 +212,7 @@ def test_real_offline_cli_train_export_eval(tmp_path): env=env, ) assert result.returncode == 0, result.stdout + result.stderr - report = json.loads((tmp_path / "result.json").read_text()) + report = json.loads((tmp_path / "reports/evaluation/result.json").read_text()) assert report["split"] == "test" and report["windows"] == 3 @@ -252,8 +252,10 @@ def train(config): "dynamics.train.full.yaml", "dynamics.eval.smoke.yaml", "dynamics.eval.full.yaml", + "dynamics.eval.pretrained.yaml", "dynamics.rollout.smoke.yaml", "dynamics.rollout.full.yaml", + "dynamics.rollout.pretrained.yaml", } for robot in ("cartpole", "ant"): robot_configs = list( From 63beff70553546ed3d6694647113d93e59a4aa96 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 24 Oct 2026 14:33:45 +0000 Subject: [PATCH 6/7] 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 b4ac4f4c0..2f3a5ff94 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 f795d24b6..5db12b3fa 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 fd46bc648..7e609fb2f 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 000000000..6403f7b44 --- /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 bfe2e2bf7..52476fc8e 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 10f9c879a..f8c60d7ed 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 095a621c3..98eb57935 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 000000000..c5b1c79d1 --- /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 000000000..74fcac64e --- /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 993bb22d0..813d0c8a4 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 000000000..6f98492a1 --- /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 8a7210013..315ec493a 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 000000000..ffc3c550f --- /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 000000000..e0aceb53b --- /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 000000000..57a7236f0 --- /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 000000000..0b332e2b6 --- /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 000000000..e198e2787 --- /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 000000000..2bb9ad240 --- /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 000000000..04c237b0d --- /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 000000000..3f5981a49 --- /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 8e3ea1615..9597ede6b 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 000000000..b1f8e18a5 --- /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 000000000..c09a147d2 --- /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 000000000..692844139 --- /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 000000000..9663efcab --- /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 000000000..49e2f8b1b --- /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 000000000..061778784 --- /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 373314c61..b83c68b8a 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 000000000..b1f8e18a5 --- /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 000000000..85d127da0 --- /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 000000000..def714dcf --- /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 000000000..ff0c3d164 --- /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 000000000..593ae1b98 --- /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 000000000..74b5f5414 --- /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 000000000..76ad6a6f1 --- /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 000000000..6cfdff3db --- /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 000000000..3226dd71f --- /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 000000000..406fe3378 --- /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 000000000..e31c13ec5 --- /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 000000000..4e152fe4b --- /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 000000000..e30e79e71 --- /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 000000000..1d3ba033e --- /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 000000000..f161a55d1 --- /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 000000000..3abd86c07 --- /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 000000000..d7d18c550 --- /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 000000000..2e4a7e59b --- /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 000000000..9a1c92879 --- /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 000000000..0caca1bfb --- /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 000000000..da72a9f08 --- /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 000000000..c134f5c0d --- /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 000000000..15ff72bec --- /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 000000000..212b3e951 --- /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 000000000..853394306 --- /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 000000000..5e0bac733 --- /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 000000000..7cd4eef56 --- /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 000000000..73b9c4926 --- /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 000000000..6df395f4c --- /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 000000000..8ee89e2b3 --- /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 000000000..c5e4d1431 --- /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 000000000..a7fd7e8b3 --- /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 000000000..cb287d7f1 --- /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 000000000..cf1815fde --- /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 000000000..89772bd6f --- /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 000000000..62360ee3a --- /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 000000000..69dd5e9c6 --- /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 000000000..0cf9daa0e --- /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 000000000..f75718a72 --- /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 000000000..c6cffc09c --- /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 000000000..81404b197 --- /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 000000000..ec4d2d8d3 --- /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 000000000..b40933776 --- /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 000000000..0f28e8e23 --- /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 000000000..7d04d4080 --- /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 000000000..25b7af7b3 --- /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 000000000..8ab2ac998 --- /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 000000000..f44fe077e --- /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 367259147..c0922dd29 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 000000000..cf534b2da --- /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 000000000..8737a5f7c --- /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 000000000..f2dbb6cdb --- /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 000000000..ba93266d3 --- /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 000000000..6f65154d5 --- /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 000000000..c04d60fd1 --- /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 000000000..380a3401b --- /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 000000000..c7da56531 --- /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 000000000..15042ae47 --- /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 000000000..470e14e1b --- /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 000000000..98a62860c --- /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 000000000..7257b0c7b --- /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 000000000..788bca3ab --- /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 000000000..7886d24b0 --- /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 000000000..9f711f329 --- /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 000000000..195fb75f1 --- /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 000000000..5c31b58f2 --- /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 000000000..77ce6d093 --- /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 6e5120608..1b9ac04d6 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 7/7] 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 00aa7240a..0d05ee881 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 e198e2787..b79134308 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.