Add experimental NeuralPlanner for learned waypoint pass through

[yangchen73]

Description

This PR adds an experimental NeuralPlanner that uses an APG transformer checkpoint to follow EEF waypoints via MotionGenerator.

• Add NeuralPlanner, NeuralPlannerCfg, and NeuralPlanOptions
• Add Franka example and unit tests
• Add docs

Type of change

• New feature (non-breaking change which adds functionality)
• Documentation update

Screenshots

neural_planner.mp4

Checklist

• I have run the black . command to format the code base.
• I have made corresponding changes to the documentation
• I have added tests that prove my fix is effective or that my feature works
• Dependencies have been updated, if applicable.

[Copilot]

Pull request overview

This PR introduces an experimental learning-based NeuralPlanner that rolls out an APG transformer policy to follow EEF waypoints via the existing MotionGenerator planning interface, along with supporting assets download utilities, documentation, an example, and unit tests.

Changes:

• Added NeuralPlanner / NeuralPlannerCfg / NeuralPlanOptions and registered the new planner type in MotionGenerator.
• Added HuggingFace checkpoint download helper for neural planner checkpoints.
• Added Franka waypoint example, tests, and documentation pages for NeuralPlanner.

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
embodichain/lab/sim/planners/neural_planner.py	New NeuralPlanner implementation + APG transformer inference actor and checkpoint loading.
embodichain/lab/sim/planners/motion_generator.py	Registers "neural" planner type; updates plan option defaulting logic inside generate().
embodichain/lab/sim/planners/__init__.py	Exposes NeuralPlanner symbols via package imports.
embodichain/data/assets/planner_assets.py	Adds download_neural_planner_checkpoint() helper for gated HF checkpoint download.
embodichain/data/assets/__init__.py	Re-exports planner assets module.
examples/sim/planners/neural_planner.py	New Franka Panda waypoint rollout example using NeuralPlanner.
tests/sim/planners/test_neural_planner.py	Adds unit tests for planner registration and basic generation behavior with a fake checkpoint.
docs/source/overview/sim/planners/neural_planner.md	New NeuralPlanner documentation page (experimental warning + usage snippet).
docs/source/overview/sim/planners/motion_generator.md	Updates MotionGenerator docs to mention NeuralPlanner support.
docs/source/overview/sim/planners/index.rst	Adds NeuralPlanner to planners overview and toctree.
agent_context/topics/motion-planning/motion-planning.md	Updates internal motion-planning context to include NeuralPlanner and assets entry points.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[yangchen73]

checked in the test

[yangchen73]

Planner options go through plan_opts (same as ToppraPlanOptions)

[yuecideng]

Mention that the policy of Neural Planner is specified for a robot.

[yuecideng]

Code Review: Neural Planner

Automated review pass against commit 6168d0b (current PR head). Focus: (1) the Neural Planner implementation, and (2) the structure of its integration into MotionGenerator.

Verdict: Merge with fixes

No Critical issues — the code is functional, tests pass, and the planner-registry integration is structurally sound. Several Important issues remain, most clustered around the MotionGenerator integration boundary, which is the primary architectural question for this PR.

---

Strengths

• Clean planner-registry integration. embodichain/lab/sim/planners/motion_generator.py:96-99 registers "neural": (NeuralPlanner, NeuralPlannerCfg) via the same _support_planner_dict / register_planner_type pattern as "toppra". NeuralPlanner subclasses BasePlanner exactly as ToppraPlanner does — the plug-in boundary at the planner level is correct and consistent with the established pattern.
• Defensive checkpoint loading. embodichain/lab/sim/planners/neural_planner.py:223-291 validates every expected key (agent, obs_normalizer.mean/var), checks policy_arch == "transformer", derives architecture dims from the checkpoint, and emits clear error messages listing available keys. More defensive than ToppraPlanner.
• Correct quaternion conventions. Traced the full path: robot.compute_fk(to_matrix=False) → [pos, quat_wxyz]; _fk_pose_xyzw converts wxyz→xyzw for the obs; _is_active_reached converts back to wxyz for quat_error_magnitude. All consistent.
• Correct qpos-limits indexing (neural_planner.py:330-332): get_qpos_limits returns (N, dof, 2), correctly indexed to (dof,2) then [:action_dim, 0/1].
• Good asset-download UX (embodichain/data/assets/planner_assets.py:81-91): actionable error with the three fix steps; proxy env-var normalization is a thoughtful touch.
• Tests run and pass (4 passed, CPU-only via FakeRobot). Conventions largely honored: copyright headers present on all new files; from __future__ import annotations + __all__; @configclass with MISSING for required checkpoint_path.

---

Issues

Important (should fix before merge)

1. is_interpolate=True is incompatible with the neural planner and is unguarded
motion_generator.py:154-221 + neural_planner.py:402-406

MotionGenerator.generate()'s interpolation path converts EEF waypoints into PlanState(move_type=JOINT_MOVE, qpos=...) targets (lines 212-219) and passes them to self.planner.plan(). But NeuralPlanner._parse_waypoints rejects anything but EEF_MOVE with xpos. So a user who naively sets is_interpolate=True (a core MotionGenerator capability, used by every Toppra test) gets a ValueError from inside the planner rather than an upfront rejection.

Why it matters: This is the central integration finding. It shows the neural planner does not cleanly fit the "interpolate-then-plan" pipeline MotionGenerator is built around — TOPPRA consumes joint waypoints; the neural planner consumes raw EEF targets and runs its own rollout. The mismatch should be surfaced at the MotionGenerator boundary, not deep in _parse_waypoints.

Fix: Add a guard in MotionGenerator.generate() (or __init__) that rejects is_interpolate=True when the planner is NeuralPlanner, with a message explaining the neural planner runs its own waypoint rollout. Alternatively, document the contract on MotionGenOptions.is_interpolate.

2. control_part / start_qpos contract change is a usability hazard
motion_generator.py:223-231

The old line options.plan_opts.control_part = options.control_part was replaced by a comment saying planner-specific options "must be set on plan_opts explicitly." This means MotionGenOptions.control_part and MotionGenOptions.start_qpos — both real fields on MotionGenOptions (motion_generator.py:55-59) that Toppra users set routinely — are silently ignored on the neural path unless duplicated into NeuralPlanOptions. The example only sets them on NeuralPlanOptions, so a user porting the Toppra pattern will be surprised.

Why it matters: Two fields with the same name on two option objects, where only one is read depending on the planner, is a leaky abstraction. The removal of propagation is a behavior change with no deprecation path and no validation that NeuralPlanOptions.control_part is set until plan() runs (neural_planner.py:319-324).

Fix: Either propagate MotionGenOptions.control_part/start_qpos into plan_opts when the planner-specific options leave them None, or document the contract prominently and validate control_part at MotionGenerator.generate() entry rather than mid-rollout.

3. torch.load(..., weights_only=False) is a security smell
neural_planner.py:229

weights_only=False permits arbitrary code execution from the checkpoint. The checkpoint comes from a gated HuggingFace repo, which mitigates but does not eliminate the risk; torch ≥2.6 defaults to weights_only=True for this reason. The payload here is dicts/tensors/strings, all of which are safe under weights_only=True.

Fix: Try weights_only=True first; fall back to False only if a UnpicklingError occurs, with a logged warning.

4. Public API docstrings missing — convention violation
neural_planner.py:205-206, 311-315

NeuralPlanner has no class docstring and plan() has no docstring, while ToppraPlanner and ToppraPlanner.plan() both have Google-style docstrings. AGENTS.md requires "Google-style docstrings with Sphinx directives" on all public APIs. _WaypointTransformerActor and _layer_init also lack type annotations.

Fix: Add docstrings to NeuralPlanner, plan, _load_checkpoint. (NeuralPlannerCfg/NeuralPlanOptions field docs are good.)

5. Leaky PlanResult contract: neural returns velocities=None/accelerations=None
neural_planner.py:379-385

TOPPRA returns positions, velocities, accelerations, and a real dt/duration from the time-optimal solution. The neural planner returns only positions + xpos_list, fabricates a constant dt (cfg.dt=0.01) and duration=(N-1)*dt. BasePlanner.is_satisfied_constraint (base_planner.py:140-191) dereferences vels.dtype unconditionally — feeding it a neural result would crash. Nothing in the type system or docs signals this.

Why it matters: Consumers that treat PlanResult uniformly (the whole point of MotionGenerator's abstraction) will break or silently get None. The fabricated dt/duration also misrepresent the trajectory's actual timing.

Fix: Either (a) document on NeuralPlanner.plan/the class that velocities/accelerations/dt are nominal-or-absent and guard is_satisfied_constraint against None, or (b) compute a defensible dt and leave velocities/accelerations None with a clear note.

6. Tests verify plumbing, not planning behavior
tests/sim/planners/test_neural_planner.py:92-102, 146

FakeRobot.compute_fk(to_matrix=False) returns a hardcoded [[0,0,0,1,0,0,0]] (origin + identity quat) regardless of qpos, and the test's waypoint is torch.eye(4) (also origin + identity). So _is_active_reached is trivially satisfied on step 0 — success is True no matter what actions the policy emits. assert result.success is True (line 146) proves the loop wiring, not that the policy drives the arm to a target. The shape/finiteness assertions (lines 147-151) are honest; the success assertion is not. Also, none of the three tests exercise the MotionGenerator integration failure modes (e.g. is_interpolate=True rejection, missing control_part).

Fix: Make FakeRobot.compute_fk depend on qpos (e.g. a trivial linear FK) so reaching a non-origin waypoint requires the policy to actually move, and add a test for the is_interpolate=True incompatibility once issue #1 is guarded.

Minor (nice to have)

• num_arm_joints cfg field is misleading (neural_planner.py:185): the actor hard-raises unless action_dim == 7 (neural_planner.py:74-78) and _parse_obs hardcodes 7-DoF slices throughout. The field is configurable but only 7 works. Either remove the field or gate it with a clear error.
• Redundant quaternion round-trip (neural_planner.py:436-440, 488-489): _fk_pose_xyzw builds xyzw, then _is_active_reached converts it back to wxyz. Keeping FK in wxyz and converting once for the obs would be cleaner.
• Double FK per loop step (neural_planner.py:342, 362): compute_fk called twice per iteration (obs + reach check). Cheap on CPU, but 2× solver calls on the GPU path.
• default_plan_options asymmetry (neural_planner.py:220 vs Toppra): only NeuralPlanner implements it, so MotionGenerator's hasattr fallback (motion_generator.py:224) is effectively neural-only. Adding it to ToppraPlanner/BasePlanner would make the plan_opts=None path uniform.
• Test file missing from __future__ import annotations (test_neural_planner.py): uses str | None annotations; works on 3.11 but violates the "every file" convention in AGENTS.md.
• _layer_init return type annotation missing (neural_planner.py:42): def _layer_init(layer: nn.Linear, ...) → add -> nn.Linear.

---

Recommendations

1. Decide explicitly whether MotionGenerator is the right host for a learned policy. The planner registry is clean, but NeuralPlanner is semantically a policy rollout, not a trajectory optimizer. The is_interpolate incompatibility ([WIP]: Init Project #1) and the duplicated control_part/start_qpos (Try to fix docs build #2) are symptoms of forcing a policy into a pipeline built for optimizers. For an experimental feature, staying in MotionGenerator is acceptable if the incompatibilities are guarded and documented; if neural planning becomes permanent, consider a sibling PolicyMotionGenerator that bypasses the interpolation path entirely.
2. Add an integration test that drives MotionGenerator with a real (or non-trivial fake) Franka and asserts the end-effector actually converges to a non-trivial waypoint, not just that the loop runs.
3. Run black . before the next commit — the branch is at "WIP" and formatting hasn't been verified against black==26.3.1.

---

Assessment

Ready to merge: With fixes.

The neural planner implementation itself is solid — checkpoint loading is defensive, quaternion/limit handling is correct, and tests pass. The integration into MotionGenerator is structurally sound at the planner-registry level (correct use of BasePlanner/_support_planner_dict), but it has real seams: the neural planner cannot consume the JOINT_MOVE targets that MotionGenerator's interpolation path produces (unguarded, #1), the control_part/start_qpos contract was changed in a way that duplicates fields and silently ignores MotionGenOptions versions (#2), and the PlanResult is missing velocities/accelerations with no documented contract (#5). The tests verify plumbing rather than planning behavior (#6). Issues #1, #2, #4, #6 should be resolved before merge; #3, #5 can land as follow-ups given the "experimental" labeling, but should not be dropped.

---

Automated code review. Findings reviewed against commit 6168d0b.

[yuecideng]

Rename to nmg_weights would be better. and place to a weights folder

[yuecideng]

Move solver as well

[yuecideng]

How we handle start qpos now?