diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index b4ac4f4c0..7344ef423 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -268,6 +268,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 f795d24b6..306016837 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` | @@ -23,6 +27,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/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 6c614c6ec..8d237b0f6 100644 --- a/agent_context/topics/env-framework/configuration.md +++ b/agent_context/topics/env-framework/configuration.md @@ -77,6 +77,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 the selected environment 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..01c9bf54c 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` @@ -9,7 +13,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..bdd8a5fa2 --- /dev/null +++ b/agent_context/topics/neural-dynamics/neural-dynamics.md @@ -0,0 +1,180 @@ +# 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/` | +| Matched Scene 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` | + +The independent model/data/training layer lives below the solver integration. +`data_pipeline/dynamics/generation.py` collects directly with built-in Newton +solvers and remains the default `dynamics generate` path. An explicit +`collector: scene` selects this layer's matched Scene collector; its presets use +that setting so exported models bind the actual deployment scene. Direct-Newton +data uses `nerd-direct-newton-builder-v1` and must not be relabeled as a matching +Scene signature. Both collectors share asset interpretation and shard formats. + +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 + +`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. +For Ant Featherstone, Huber smoothing and both materials' Coulomb coefficients +are coupled to approximate legacy Euclidean friction. NeRD reconstructs that +material domain from the export's `reference_friction_mapping`; missing markers +retain the older domain, and the provider still validates its physical signature. +Other analytical solvers keep their ordinary material coefficients. The profile +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 +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. +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. 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 + +`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 [Neural Dynamics Models](../../../docs/source/features/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 +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. + +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/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 0dbd10542..490d6e2a2 100644 --- a/agent_context/topics/simulation-system/simulation-system.md +++ b/agent_context/topics/simulation-system/simulation-system.md @@ -26,6 +26,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. `VisualObject` owns render-only root poses, visibility and material assignments, with no rigid-body buffers or collision API. Its config lives in @@ -104,8 +107,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..685fb9b7d --- /dev/null +++ b/design/neural-dynamics-solvers.md @@ -0,0 +1,120 @@ +# Neural dynamics architecture + +Neural dynamics replaces the state-transition solver while retaining scene +construction, controls, collision geometry and simulation lifecycle. The user +workflow, downloadable artifacts and examples are documented in +[Neural Dynamics Models](../docs/source/features/neural_dynamics.md). + +## Ownership + +| Layer | Owner | Responsibility | +| --- | --- | --- | +| Solver discovery and native lifecycle | DexSim `dexsim.newton_solvers` | Installed provider discovery, construction, stepping, reset and publication | +| Tensor mathematics | `embodichain/compute/dynamics/` | State/velocity transport, history operations and physical-unit errors | +| Model implementation | `embodichain/learning/models/nerd/` | Network, normalizers and portable weights | +| Offline data | `embodichain/data_pipeline/dynamics/` | Committed HDF5 shards, trajectory splits and sequence windows | +| Offline learning | `embodichain/learning/dynamics/` | Training, resume, evaluation, export and optional reference adapters | +| Scene integration | `embodichain/lab/sim/physics/solvers/` | Physical binding, contact adaptation and native/functional transitions | +| Task configuration | `extensions/embodichain_tasks/configs/tasks/classic_control/` | Reference robot scenes and runnable workflow examples | +| Research tools | `scripts/experiments/neural_dynamics/` | Isolated upstream-environment and cross-version comparisons | + +Compute does not import simulation or learning. Model loading and offline +training do not initialize DexSim or a simulator. Generic discovery belongs to +DexSim; adding a provider does not require an EmbodiChain dispatch branch. +The current model is `NeRDModel`; generic extension points do not require an +additional abstract neural-model hierarchy. + +## Solver deployment + +The installed package registers `NeRDProvider` through the +`dexsim.newton_solvers` entry-point group. A deployment selects +`solver_type: nerd` and supplies a portable export directory. Unknown names, +duplicate providers and incompatible interfaces fail instead of falling back. +There is no provider-path import or imperative registration API. + +Each World owns its solver, histories and buffers. A solver handles the complete +scene. The initial NeRD implementation supports homogeneous Cartpole or Ant +articulations, generalized joint efforts, and its declared static-plane contact +schema. Capability metadata distinguishes reset, first-order differentiability, +contact reports, external forces and CUDA Graph support. Unsupported quantities +are unavailable rather than fabricated as zero. + +A compatible DexSim build must expose the named-solver and functional Scene +contracts; a package version string alone does not establish that capability. + +## Physical and numerical contracts + +Portable exports bind robot topology, joint order, timestep, gravity, contact +schema, base COM and a measured physics signature. The runtime reconstructs and +validates that physical domain before stepping; users cannot relabel historical +weights to bypass the checks. + +`physics_profile=warp_reference` reconstructs the reference assets. Ant with +current Featherstone also uses a coupled conversion from legacy Euclidean +friction to Newton Huber/Coulomb parameters. Both material coefficients and the +smoothing scale form one mapping. The export retains that mapping so inference +recreates the trained physical signature. Other solver profiles retain their +own semantics. Equal numeric parameter trees do not imply equal physics across +versions. + +The codec owns conversions between Newton COM-linear-first velocities and +NeRD angular-first spatial twists, including reference-point transport. Relative +state targets use the frame at the current timestep. Training and inference use +the same codec and frozen input/output normalizers. Exported attention policy is +part of the numerical contract; changing it is an explicit model variant. + +## Functional execution + +Native and functional NeRD calls share one transition implementation. Functional +state consists of owned physical tensors, history tensors and per-world valid +history lengths. Native ping-pong buffers are publication storage, not tensors +retained for backward. + +Warp FK has an explicit first-order Torch VJP bridge; shared storage alone does +not connect autograd. Contact selection, angle wrapping and quaternion signs +have branch-local derivatives. Functional publication advances time once and +adopts detached state/history without executing a second transition. Selected +resets preserve unselected trajectories and previous reward graphs. +`detach_state()` defines the rollout differentiation boundary. + +## Data, training and exports + +Each committed trajectory stores T+1 states and T efforts/contact inputs. +Training, validation and test splits are trajectory-level; sampled windows do +not cross split or episode boundaries. Worker processes own their file handles. +Optional resident caches have explicit byte budgets and preserve tensor values. + +Training supports scratch initialization, fine-tuning and exact checkpoint +resume with optimizer, scheduler, RNG, normalization and consumed-sample state. +Reference DataLoader descriptors preserve shuffled pass order and dropped tails. +Offline training checkpoints remain local training artifacts. Published inference +packages contain tensor weights and JSON deployment metadata without optimizer +states, process logs or workstation-specific paths. + +Reusable datasets are published under +`DexForceAI/embodichain_datasets/neural_dynamics/nerd//v1`; inference +models are published under +`DexForceAI/embodichain_model/neural_dynamics/nerd//v1`. +Cards state the actual data scale, supported physics and measured limitations. +The source repository contains the implementation, configuration, tests and +feature documentation; chronological training reports remain outside the PR. + +## Validation boundaries + +Validate codecs/history in `tests/compute/dynamics`, data persistence/windows in +`tests/data_pipeline/dynamics`, and model/trainer contracts in +`tests/learning/models` and `tests/learning/dynamics`. Solver lifecycle, physical +signatures and contact adaptation are covered in `tests/sim/physics/solvers`; +functional environments are covered in `tests/sim/diff` and `tests/gym/envs`. + +Published artifacts must load through the production readers and match their +metadata, checksums and physical domain. Model accuracy is measured against the +matching analytical solver using common initial states and actual action tapes. +Own-state feedback-policy comparisons are labelled closed-loop behavior. + +Network parity, trainer parity, solver deployment, task success and reproduction +of a paper's final accuracy are separate claims. Accuracy reports distinguish +endpoint error from temporal mean, retain failed trajectories, and record the +sample population and horizon. Performance comparisons use synchronized warmed +calls, equivalent macro timesteps and the analytical solver's supported CUDA +Graph path; throughput does not establish an accuracy acceptance claim. 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/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..65b002520 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 ----------------------- @@ -649,6 +730,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 ------------------------------------- @@ -659,6 +753,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.sim.scene.expansion.layout ------------------------------------------ @@ -1359,8 +1496,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 @@ -1369,6 +1507,12 @@ autograd. It does not advance the Newton dynamics solver. NewtonStepFunc tape_context + +.. autosummary:: + + FunctionalDynamics + detach_dynamics_state + embodichain.lab.sim.diff.bridge ------------------------------- @@ -1655,6 +1799,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 -------------------------------------------------- @@ -2639,6 +2939,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 -------------------------------------------- @@ -2650,6 +2963,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/features/neural_dynamics.md b/docs/source/features/neural_dynamics.md new file mode 100644 index 000000000..711b406df --- /dev/null +++ b/docs/source/features/neural_dynamics.md @@ -0,0 +1,229 @@ +# Neural Dynamics Models + +A neural dynamics model predicts how a robot's physical state changes under +applied joint efforts and contact. EmbodiChain supports these models as named +Newton solver extensions, so the simulation world, scene objects, state access, +and environment lifecycle remain shared with analytical physics. + +[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. + +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. + +These Scene-matched generation presets declare `collector: scene` and require +the DexSim solver/Scene extension. The independent model/data/training PR also +provides `collector: newton` (the CLI default), which collects directly with +Newton and needs no DexSim imports. Its separate physical signature describes +a different model-building path; it cannot be relabeled to qualify a Scene +export. Existing published datasets remain available for offline training. + +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/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..009ed4f9e --- /dev/null +++ b/embodichain/cli/dynamics.py @@ -0,0 +1,262 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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": + collector = config.pop("collector", "newton") + if collector == "newton": + from embodichain.data_pipeline.dynamics.generation import ( + generate_dataset, + ) + elif collector == "scene": + from embodichain.lab.sim.physics.solvers.collection import ( + generate_dataset, + ) + else: + raise ValueError("collector must be newton or scene") + 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 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..b79134308 --- /dev/null +++ b/embodichain/data/nerd/ATTRIBUTION.md @@ -0,0 +1,21 @@ +# NeRD reference robot assets + +`cartpole_single.urdf` and `ant.xml` are unmodified copies from +NVlabs/neural-robot-dynamics, commit +`476dc67fccc6a99be293bf611f897b1c06f13e30`, under +`envs/newton_envs/assets/`. The upstream repository license is reproduced in +`LICENSE.md` (Apache License 2.0; NVIDIA Corporation and affiliates, 2025). + +Source: https://github.com/NVlabs/neural-robot-dynamics/tree/476dc67fccc6a99be293bf611f897b1c06f13e30/envs/newton_envs/assets + +EmbodiChain interprets these assets using the installed Newton public importer. +The independent data collector finalizes an analytical Newton model directly; +Scene-based neural solver integration is a separate consumer of the same asset +interpretation. Neither path claims to reproduce legacy Warp 1.5.1 physics. + +The host supplies generalized forces and its own solver configuration rather +than installing the source motor actuators. Dataset manifests identify the +importer profile, asset hash, actual parameter hashes and solver versions. +`warp_reference` reconstructs measured legacy mass, inertia and contact-parameter +semantics; it does not select the historical simulator or guarantee identical +trajectories across solver versions. diff --git a/embodichain/data/nerd/LICENSE.md b/embodichain/data/nerd/LICENSE.md 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/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 9f885ccc3..a62204a61 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -729,6 +729,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"), @@ -746,10 +747,13 @@ class ComponentCfg: ), ) - env_cfg.robot = _robot_cfg_from_dict(config["robot"]) - env_cfg.robot.enable_sensor = env_cfg.enable_sensor - if env_cfg.robot.source is not None: - resolve_asset_paths(env_cfg.robot) + 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 + if env_cfg.robot.source is not None: + resolve_asset_paths(env_cfg.robot) # 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 902e3296e..7118e57e2 100644 --- a/embodichain/lab/sim/objects/articulation.py +++ b/embodichain/lab/sim/objects/articulation.py @@ -2846,8 +2846,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..b1d7ab2a5 --- /dev/null +++ b/embodichain/lab/sim/physics/solvers/_reference_assets.py @@ -0,0 +1,249 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""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 +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", +] + +from embodichain.data_pipeline.dynamics._reference_assets import ( + _asset, + _ant_reference_friction, + _import_reference, + reference_asset_hash, + reference_joint_names, + reference_profile, +) + + +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, physics_profile: str = "native" +) -> ArticulationDesc: + """Build an independent public descriptor with imported physical properties. + + ``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`` + 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, 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") + 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..c1c5bd3f9 --- /dev/null +++ b/embodichain/lab/sim/physics/solvers/collection.py @@ -0,0 +1,1484 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 ( + _ant_reference_friction, + 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, + 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": "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" + 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: + 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" and physics_profile == "native", + 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 ( + {"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=(friction_mapping["ground_mu"] if friction_mapping else 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, + ) + ) + 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 + ) + 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, + "physics_profile": physics_profile, + "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" + ], + } + 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() + 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.""" + # 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) + .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 + physics_profile: str = "native" + + 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.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 DynamicsSimulation( + cfg.robot, + device=cfg.device, + 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"] = { + "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 + ) + self.identity["target"] = self.target.detach().cpu().tolist() + 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), + physics_profile=config.get("physics_profile", "native"), + ) 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 _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], + 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. 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 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. + + 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", + "physics_profile", + } + 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 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 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" + ) + 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, + physics_profile=config.get("physics_profile", "native"), + ) 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 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 + ): + 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, + "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, + "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..e262b845b --- /dev/null +++ b/embodichain/lab/sim/physics/solvers/nerd.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. +# ---------------------------------------------------------------------------- + +"""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, 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() + ) + 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 6a4eb01ce..b7dbacbcc 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -147,6 +147,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 @@ -248,7 +249,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 @@ -300,6 +305,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.""" @@ -981,7 +993,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())) @@ -1735,7 +1747,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 @@ -1961,6 +1973,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. @@ -2165,6 +2287,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, @@ -2714,6 +2839,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. @@ -3341,6 +3504,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/lab/sim/spawn/scene.py b/embodichain/lab/sim/spawn/scene.py index ed4bea6ad..d035e431c 100644 --- a/embodichain/lab/sim/spawn/scene.py +++ b/embodichain/lab/sim/spawn/scene.py @@ -159,7 +159,7 @@ def declare( if uid in self._assets: raise ValueError(f"Spawn asset uid is already declared: {uid!r}.") self._bind_resource_materials(descriptor) - self._initialize_facade_declaration(facade) + self._initialize_facade_declaration(facade, descriptor) declaration = _AssetDeclaration( kind=kind, descriptor=descriptor, @@ -248,7 +248,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) @@ -439,7 +439,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 @@ -450,7 +452,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/__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..0bbb893de --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/README.md @@ -0,0 +1,66 @@ +# 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: + +```bash +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 +``` + +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 +``` + +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. + +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/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.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..7970df21a --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.full.yaml @@ -0,0 +1,19 @@ +# Dynamics workflow configuration; not a runnable Gym deployment. +robot: ant +collector: scene +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..2149adeb6 --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.smoke.yaml @@ -0,0 +1,20 @@ +# Dynamics workflow configuration; not a runnable Gym deployment. +robot: ant +collector: scene +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..59645bfdf --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/ant/dynamics.generate.warp-reference.yaml @@ -0,0 +1,20 @@ +# Bounded parameter-alignment collection; not full accuracy qualification. +robot: ant +collector: scene +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.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.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/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..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/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..e59309414 --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/README.md @@ -0,0 +1,66 @@ +# 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: + +```bash +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 +``` + +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 +``` + +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. + +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/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.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..929ffe7b6 --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.full.yaml @@ -0,0 +1,19 @@ +# Dynamics workflow configuration; not a runnable Gym deployment. +robot: cartpole +collector: scene +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..f4ff461d5 --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.smoke.yaml @@ -0,0 +1,20 @@ +# Dynamics workflow configuration; not a runnable Gym deployment. +robot: cartpole +collector: scene +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..5c33bb3c7 --- /dev/null +++ b/extensions/embodichain_tasks/configs/tasks/classic_control/cartpole/dynamics.generate.warp-reference.yaml @@ -0,0 +1,19 @@ +# Bounded parameter-alignment collection; not full accuracy qualification. +robot: cartpole +collector: scene +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.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.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/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..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/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 9812e2926..22f2a1608 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", @@ -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_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..4900ad802 --- /dev/null +++ b/scripts/experiments/neural_dynamics/README.md @@ -0,0 +1,64 @@ +# 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/_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..b72d6f005 --- /dev/null +++ b/scripts/experiments/neural_dynamics/physics_alignment.py @@ -0,0 +1,408 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 +import os +from pathlib import Path +import platform +import sys +from typing import Any + +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") + 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) + 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, + 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, + 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) + 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 _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[dict[str, 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, 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() + 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) + 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( + (force[:, 6:] if args.robot == "ant" else force).ravel() + ) + env.update() + return ( + {"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[dict[str, 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": args.device == "cpu"} + ) + if args.profile == "native" and solver is None: + solver = { + "solver_type": "featherstone", + "update_mass_matrix_interval": substeps, + } + 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"]).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(_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]).to(env.device)) + metadata = {**env.manifest, "dexsim_commit": dexsim.__commit_id__} + 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: + 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("--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" + ) + parser.add_argument( + "--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) + 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.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") + traces, meta = (_legacy if args.engine == "legacy" else _current)(args, data) + args.output.parent.mkdir(parents=True, exist_ok=True) + np.savez(args.output, **traces) + 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, + 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) + + +if __name__ == "__main__": + main() 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/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..ff3d743f2 --- /dev/null +++ b/tests/cli/test_dynamics.py @@ -0,0 +1,348 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Lazy 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\ncollector: scene\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": "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.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 = [] + 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", + "dynamics.rollout.smoke.yaml", + "dynamics.rollout.full.yaml", + "dynamics.rollout.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(["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/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 000000000..9c39fbd51 Binary files /dev/null and b/tests/compute/dynamics/data/ant-official-frame.npz differ 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/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..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/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/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 new file mode 100644 index 000000000..6349a6d4e --- /dev/null +++ b/tests/sim/physics/solvers/test_collection.py @@ -0,0 +1,644 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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")), + ("physics_profile", "unregistered"), + ], +) +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() + + +@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" + 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: + 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 + + +@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 + + 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 + ) + 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: + 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..4409c1870 --- /dev/null +++ b/tests/sim/physics/solvers/test_reference_assets.py @@ -0,0 +1,279 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 json +from pathlib import Path + +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"]) +@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): + 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") + + +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 mass, body armature and source materials.""" + 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) + # 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] + 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_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") + + +@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 2dc2f62c7..67e4aecb7 100644 --- a/tests/sim/spawn/test_scene.py +++ b/tests/sim/spawn/test_scene.py @@ -198,6 +198,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 4b655fcd2..ae3846642 100644 --- a/tests/sim/test_sim_manager.py +++ b/tests/sim/test_sim_manager.py @@ -66,6 +66,75 @@ (0.0, 0.0, 1.0), ) + +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( + 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 +2105,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 6e5120608..ca6d09ed3 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", @@ -74,6 +75,7 @@ def test_map_registers_the_supported_context_domains() -> None: "scene-expansion", "scene-services", "differentiable-env", + "neural-dynamics", } assert set(topics) == expected_topic_ids