[yuecideng]

Description

Add NeRD as an installable model extension and runnable SL showcase, keeping model-specific collection, reference experiments and raw assets out of the core. embodichain.learning.sl owns shared supervised updates, training and recovery; the NeRD recipe uses that trainer by default.

The branch includes current main at 8cb2b5353 (#846), with no conflicts or changes to upstream motion code. #857 was merged into this PR's feature branch, rather than main; its SL foundation remains included here. The #860 NeRD-to-SL integration is also consolidated here as requested.

Refs #822, #848 and #860.

Ownership and entry points

Area	Final owner
Shared supervised batches, objectives, updates, offline loop and run ownership	embodichain.learning.sl; embodichain train-sl
Existing shared MLPs and normalizers used by RL	embodichain.learning.models
NeRD model, portable checkpoints, physical codec, shard/window adapter, objective and SL factory	embodichain_models.dynamics.nerd, from extensions/embodichain_models/embodichain_models/
Teacher collection, historical training compatibility and independent reference/acceptance workflows	examples/learning/nerd/, with private _support helpers
Presets and user-facing preparation/train/export/evaluation workflow	examples/learning/nerd/configs/; embodichain run-example learning/nerd
Reference robot files and their upstream license/attribution	Pinned Hugging Face resources downloaded by the showcase

Remove the NeRD-only compute/dynamics, data_pipeline/dynamics, learning/dynamics, learning/models/nerd and specialized core dynamics CLI. Preserve released general MLP/RL imports. The model extension is bundled in the same wheel/version; it has no separate distribution or registry and does not import the example, Newton, Warp or the simulator for inference/SL training.

The extension has two first-level domains defined by prediction contract: policy predicts actions; dynamics predicts future state, geometry, observations or latent transitions. Canonical packages under embodichain_models are policy.bc, dynamics.nerd, dynamics.rwm and dynamics.rigidformer. Only NeRD is implemented here; the other three are migration targets for their follow-ups. RWM belongs to dynamics because it predicts next robot state from state/action/context; world_models is not a parallel domain or alternate namespace. Modalities, network architecture and training methods do not determine the domain or imply gradient/solver capabilities. The model catalog contract owns this naming.

Downloadable assets

The four former packaged reference files now live at reference_assets/v1 on Hugging Face. Their bytes, attribution and upstream Apache license are preserved; all 116 pre-existing remote files remain unchanged. Raw XML/URDF files are absent from the core and wheel.

Preparation pins the asset revision above, dataset revision b9ec58f36da99cc99a93f97adfe74ebc9fd66586 and portable-model revision 78e581526570069410bf269b38c0c52c800b24a7. It verifies manifests and payload/member hashes, uses caller-owned caches, and supports --assets-only and --local-files-only. Imports do not download resources.

embodichain run-example learning/nerd prepare --robot cartpole --resource-root ./outputs/nerd-resources
embodichain run-example learning/nerd train --robot cartpole --resource-root ./outputs/nerd-resources --device cuda:0
embodichain run-example learning/nerd export --checkpoint ./outputs/nerd/cartpole/train/last.pt --output ./outputs/nerd/cartpole/export
embodichain run-example learning/nerd eval --robot cartpole --resource-root ./outputs/nerd-resources --model ./outputs/nerd/cartpole/export --device cpu

The default recipe uses common SL checkpoints. Both common SL and historical NeRD checkpoints export through the showcase, preserving floating precision and ordered feature columns. Relative default outputs resolve from the caller's directory after wheel installation. See the showcase README for custom recipes, resume and historical workflows.

Review fixes

Greptile finding	Resolution
Fresh runs erase previous results	Shared durable run admission rejects foreign destinations before updates/writes; history is append-only.
Recovery loses selected best weights	last.pt preserves an owned CPU best snapshot and selection boundary for recovery into a fresh directory.
Full Ant resident cache cannot fit	Full/native Ant and shared SL presets default to bounded non-resident reads.
Feature order changes on resume	Ordered feature layout participates in recovery identity and portable export.
SL resume overwrites another run	Destination run identity is checked independently of its logged update count.
Export drops floating precision	Load, resume and both export formats restore saved floating dtype before loading tensors.
Rejected candidates disappear after interruption	Cursor/rejection evidence commits with shard-manifest transactions without changing data fingerprints.
FP64 requires input normalization	Features follow model dtype independently of normalization; loss uses the operands' common dtype without quantizing FP32 targets during BF16 autocast.
Another dataset's statistics are accepted	Statistics fingerprint preflight precedes normalization/output/RNG mutation, including disabled initialization.
FP64 Ant evaluation fails	Physical states follow prediction dtype before Ant decoding. Prior source-bound checks exercised FP32/FP64 validation, portable export/load and held-out evaluation.
Completed comparison runs appear pending	Keep the source import root separate from the experiment root. Two-cell execution checks verified receipt/summary placement and reuse without invoking workers again.

Independent migration review also found and fixed the implicit evaluation preset's output-directory resolution. Installed-showcase execution checks verified that outputs stay in the caller's directory.

Validation

At the requesting user's direction, remove all NeRD-specific unit tests from this PR: 19 test modules and their helper/package/fixture files, 27 files in total. Shared SL, generic packaging, documentation and CI checks remain. Clear the deleted test paths from the context map and route NeRD validation to runnable showcase/reference workflows.

• 75 remaining focused checks pass, covering shared SL, context routing, ordinary-example packaging and release metadata.
• With the CI setting HF_ENDPOINT=https://hf-mirror.com and an empty NeRD resource cache, a real download of the pinned reference bundle succeeds from the canonical Hub. Checksum verification and offline cache reuse also pass.
• Actual SDK-blocked CPU showcase passes: generate 64 Cartpole trajectories/256 transitions, train through common SL for two updates, export last.pt, and run held-out inference over 12 test windows. No NeRD unit-test files are added or retained.
• Black 26.10.1 black ., whitespace checks, context validation and 3168/3168 API documentation coverage pass.

Earlier long numerical qualification remains bound to its recorded source revisions and the frozen acceptance protocol. Historical test counts describe those older revisions rather than the current test inventory.

CI follow-up

The failing test job had four NeRD asset-download failures and 10,117 passing checks. Its global mirror returned missing Hub commit metadata for the newly published asset revision. Fetch this small pinned reference bundle explicitly from https://huggingface.co; retain the configured mirrors for other data/model downloads and preserve all pinned revisions, hashes, cache and offline semantics.

Commit 2c9ee8875 includes the download fix and the requested removal of NeRD unit tests. The push triggers a new CI run; local validation is reported separately from its eventual result.

Contracts, dependencies and review

last.pt is resumable state; historical new best.pt files are model-only selected snapshots. The default common SL recipe exports its committed last weights. Older complete NeRD checkpoints and portable models remain readable; trusted legacy Python-object conversion still requires explicit --trust-pickle. Physical metadata/provenance is preserved rather than relabeled, and collector resume still requires the recorded source identity.

DDP #864, RWM #859, RigidFormer #858 and draft BC #861 remain separate follow-ups and will need to follow the new extension/import layout. This PR does not claim full-budget retraining, paper-level reproduction, native learned-solver deployment or completed ODS integration.

Keep this PR open for the requesting user's review; auto-merge is disabled.

Type of change

• New feature, enhancement, bug fix and documentation update.

Checklist

• Run black . with Black 26.10.1.
• Review affected docs/context and revise the owning guidance.
• Document public APIs and pass the read-only coverage gate.
• Retain shared SL/general checks and verify NeRD through its actual showcase, without dedicated NeRD unit tests.
• Validate the actual installed wheel and preserve optional dependency boundaries.
• Move raw reference assets to pinned Hugging Face downloads with license/attribution.
• Synchronize current main and leave merge to the user's review.

[greptile-apps]

Confidence Score: 4/5

[Medium impact] The changes reach a shared training workflow, a bundled model package, and the NeRD showcase.

Fix CPU history lengths in padded CUDA inference before merging.

Findings

1.  Padded GPU inference fails ▶

Fix with agent prompt

### Issue 1
extensions/embodichain_models/embodichain_models/dynamics/nerd/model.py:323-324
`NeRDModel.evaluate()` accepts and checks CPU `valid_lengths`, but leaves the gather index on CPU. With a CUDA model and CUDA inputs, passing CPU lengths-such as the `length` tensor returned by `DynamicsWindowDataset.get_batch()`-then fails at `value.gather()` instead of returning predictions.

Move the checked lengths or the resulting index to `value.device` before gathering.

```suggestion
        index = (valid_lengths.to(value.device) - 1)[:, None, None].expand(
            -1, 1, value.shape[-1]
        )
        return value.gather(1, index)
```

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

Summary

Adds the bundled embodichain_models.dynamics.nerd package and shared supervised training through embodichain train-sl.

• Model recipes can train through one shared supervised-learning command.
• NeRD now ships as a reusable dynamics model in the main wheel.
• The NeRD showcase now guides people from pinned resources to offline evaluation.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  Prepare[Showcase prepare] --> Resources[Pinned datasets and models]
  Resources --> Recipe[NeRD SL recipe]
  Recipe --> Trainer[Shared SL trainer]
  Trainer --> Checkpoint[last.pt]
  Checkpoint --> Export[Showcase export]
  Export --> Model[Portable NeRD model]
  Model --> Evaluate[Offline evaluation]

Loading

_{Reviews (12) · Last reviewed commit: "Merge latest main and reconcile model an..." · Reviewed by Greptile}

[greptile-apps]

Padded GPU inference fails

NeRDModel.evaluate() accepts and checks CPU valid_lengths, but leaves the gather index on CPU. With a CUDA model and CUDA inputs, passing CPU lengths—such as the length tensor returned by DynamicsWindowDataset.get_batch()—then fails at value.gather() instead of returning predictions.

Move the checked lengths or the resulting index to value.device before gathering.

Suggested change

	index = (valid_lengths - 1)[:, None, None].expand(-1, 1, value.shape[-1])
	return value.gather(1, index)
	index = (valid_lengths.to(value.device) - 1)[:, None, None].expand(
	-1, 1, value.shape[-1]
	)
	return value.gather(1, index)

Prompt To Fix With AI

This is a comment left during a code review.
Path: extensions/embodichain_models/embodichain_models/dynamics/nerd/model.py
Line: 323-324

Comment:
**Padded GPU inference fails**

`NeRDModel.evaluate()` accepts and checks CPU `valid_lengths`, but leaves the gather index on CPU. With a CUDA model and CUDA inputs, passing CPU lengths—such as the `length` tensor returned by `DynamicsWindowDataset.get_batch()`—then fails at `value.gather()` instead of returning predictions.

Move the checked lengths or the resulting index to `value.device` before gathering.

```suggestion
        index = (valid_lengths.to(value.device) - 1)[:, None, None].expand(
            -1, 1, value.shape[-1]
        )
        return value.gather(1, index)
```

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.