Add planner-oriented atomic skill benchmark

agent_context/topics/motion-planning/motion-planning.md

@@ -9,7 +9,6 @@
 | TOPPRA planner | `embodichain/lab/sim/planners/toppra_planner.py` → `ToppraPlanner`, `ToppraPlannerCfg`, `ToppraPlanOptions` |
 | Neural planner | `embodichain/lab/sim/planners/neural_planner.py` → `NeuralPlanner`, `NeuralPlannerCfg`, `NeuralPlanOptions` |
 | cuRobo planner | `embodichain/lab/sim/planners/curobo/curobo_planner.py` → `CuroboPlanner`, `CuroboPlannerCfg`, `CuroboWorldCfg`, `CuroboPlanOptions` |
-| Planner assets | `embodichain/data/assets/planner_assets.py` → `download_neural_planner_checkpoint()` |
 | Motion generator | `embodichain/lab/sim/planners/motion_generator.py` → `MotionGenerator`, `MotionGenCfg`, `MotionGenOptions` |
 | Planner utilities & data types | `embodichain/lab/sim/planners/utils.py` → `PlanState`, `PlanResult`, `MoveType`, `MovePart`, `TrajectorySampleMethod`, `interpolate_xpos_batched` |
 
@@ -40,7 +39,7 @@ Config hierarchy:
 ```
 BasePlannerCfg            robot_uid (MISSING), planner_type
   ├─ ToppraPlannerCfg     planner_type = "toppra", max_workers, mp_context
-  └─ NeuralPlannerCfg     planner_type = "neural", checkpoint_path (MISSING)
+  └─ NeuralPlannerCfg     planner_type = "neural", onnx_model_path (MISSING)
 
 MotionGenCfg              planner_cfg (MISSING — must be a BasePlannerCfg subclass)
 
@@ -89,11 +88,11 @@ Worker details:
 
 Learning-based EEF waypoint planner. Franka Panda only.
 
-- Checkpoint: `download_neural_planner_checkpoint()` from HuggingFace (gated, needs `HF_TOKEN`)
+- Runtime: standalone `.onnx` policy with normalization embedded; install the `nmg` extra
 - Use via `MotionGenerator` with `planner_type="neural"` and `plan_opts=NeuralPlanOptions(...)`
 - Input: `EEF_MOVE` `PlanState` list with batched `xpos:(B, 4, 4)`
-- Key cfg: `checkpoint_path` (from download), `control_part`
-- Natively batched: transformer forward, reach checks, and convergence holds all operate on `(B, ...)`.
+- Key cfg: `onnx_model_path`, `control_part`, `num_waypoints`, `policy_frame_from_world`, `runtime_tcp_from_policy_tcp`
+- The NMG exporter produces a dynamic-batch ONNX policy; NeuralPlanner rolls out all environments together.
 
 ### CuroboPlanner collision worlds
 

docs/source/overview/sim/planners/index.rst

@@ -27,8 +27,8 @@ The `embodichain` project provides a unified interface for robot trajectory plan
 
 These tools can be used to generate smooth and dynamically feasible robot trajectories. Install NVIDIA's CUDA-matched cuRobo source package separately when collision-aware planning against an explicit cuRobo world is required.
 
-Use NeuralPlanner (experimental) when you have a trained APG checkpoint and need
-learned EEF waypoint rollout on Franka Panda.
+Use NeuralPlanner (experimental) when you have a standalone NMG ONNX policy and
+need learned EEF waypoint rollout on Franka Panda.
 
 See also
 --------

docs/source/overview/sim/planners/neural_planner.md

@@ -3,21 +3,24 @@
 ````{admonition} Experimental
 :class: warning
 
-`NeuralPlanner` is an **experimental** feature. The API, checkpoint format,
+`NeuralPlanner` is an **experimental** feature. The API, ONNX model contract,
 and default parameters may change without a deprecation cycle. It is currently
 only validated on the **Franka Panda** robot.
 ````
 
 `NeuralPlanner` is a learning-based EEF waypoint planner. It rolls out a
-trained APG checkpoint through `MotionGenerator` to reach Cartesian targets.
+standalone NMG ONNX policy through `MotionGenerator` to reach Cartesian targets.
+The ONNX graph must include raw-observation normalization.
 
 ## Configuration
 
-Pre-trained checkpoints are hosted on HuggingFace and can be downloaded with
-`download_neural_planner_checkpoint()` (requires `HF_TOKEN` environment variable).
+Install the optional runtime and export the trained policy to ONNX before use:
+
+```bash
+pip install -e '.[nmg]'
+```
 
 ```python
-from embodichain.data.assets.planner_assets import download_neural_planner_checkpoint
 from embodichain.lab.sim.planners import (
     MotionGenCfg,
     MotionGenOptions,
@@ -28,13 +31,13 @@ from embodichain.lab.sim.planners import (
 )
 from embodichain.lab.sim.planners.neural_planner import NeuralPlanOptions
 
-checkpoint_path = download_neural_planner_checkpoint()
+onnx_model_path = "/path/to/best_mean.onnx"
 
 motion_generator = MotionGenerator(
     cfg=MotionGenCfg(
         planner_cfg=NeuralPlannerCfg(
             robot_uid=robot.uid,
-            checkpoint_path=checkpoint_path,
+            onnx_model_path=onnx_model_path,
             control_part="main_arm",
         )
     )
@@ -57,7 +60,18 @@ result = motion_generator.generate(
 ## Example
 
 ```bash
-python examples/sim/planners/neural_planner.py --headless --device cuda
+python examples/sim/planners/neural_planner.py \
+  --headless --device cuda \
+  --onnx-model-path /path/to/best_mean.onnx
 ```
 
-The example downloads the checkpoint automatically on first run.
+The NMG exporter produces a dynamic-batch ONNX policy, so one graph can serve
+single-env and env-batched rollout. If the runtime robot base frame or TCP differs
+from training, configure `policy_frame_from_world` and
+`runtime_tcp_from_policy_tcp` as explicit homogeneous transforms. The conversion is
+
+```text
+policy_T_policy_tcp = policy_T_world
+                    @ world_T_runtime_tcp
+                    @ runtime_tcp_T_policy_tcp
+```

embodichain/lab/sim/atomic_actions/core.py

@@ -545,7 +545,16 @@ def build_plan(
             raise ValueError("Trajectory and planning context batch sizes must match.")
         if timed.robot_dof != context.robot.robot_dof:
             raise ValueError("Trajectory robot_dof must match the planning context.")
-        timed = timed.hold_rows(success_mask, context.robot.qpos)
+        planner = (
+            None
+            if self._planning_services is None
+            else self._planning_services.motion_generator.planner
+        )
+        preserve_failed_positions = (
+            getattr(planner, "preserve_failed_plan_positions", False) is True
+        )
+        if not preserve_failed_positions:
+            timed = timed.hold_rows(success_mask, context.robot.qpos)
 
         commands = self._joint_command_sequence(
             request,

embodichain/lab/sim/atomic_actions/engine.py

@@ -491,9 +491,18 @@ def compile(
                     f"Skill {plan.skill_id!r} emits non-joint runtime commands and "
                     "cannot be used with offline joint-trajectory compilation."
                 )
-            trajectory = plan.joint_trajectory.hold_rows(
-                step_success,
-                previous_qpos,
+            preserve_failed_positions = (
+                getattr(
+                    self.motion_generator.planner,
+                    "preserve_failed_plan_positions",
+                    False,
+                )
+                is True
+            )
+            trajectory = (
+                plan.joint_trajectory
+                if preserve_failed_positions
+                else plan.joint_trajectory.hold_rows(step_success, previous_qpos)
             )
             plans.append(plan)
             trajectories.append(trajectory)

embodichain/lab/sim/planners/base_planner.py

@@ -224,6 +224,9 @@ def __init__(self, cfg: BasePlannerCfg):
     waypoints for a joint-only backend.
     """
 
+    supports_heterogeneous_waypoints: bool = False
+    """Whether one plan may contain an ordered mixture of movement types."""
+
     preserve_plan_samples: bool = False
     """Whether callers must retain this planner's returned sample points exactly.
 

embodichain/lab/sim/planners/motion_generator.py

@@ -212,6 +212,11 @@ def supports_joint_trajectory_validation(self) -> bool:
             is True
         )
 
+    @property
+    def supports_heterogeneous_waypoints(self) -> bool:
+        """Whether the selected backend accepts mixed waypoint movement types."""
+        return getattr(self.planner, "supports_heterogeneous_waypoints", False) is True
+
     @property
     def dynamic_collision_entity_ids(self) -> tuple[str, ...]:
         """Return canonical dynamic-obstacle IDs declared by the planner."""
@@ -488,11 +493,19 @@ def generate(
             )
 
         move_types = {state.move_type for state in target_states}
-        if len(move_types) != 1:
+        heterogeneous = len(move_types) > 1
+        if heterogeneous and not self.supports_heterogeneous_waypoints:
             names = sorted(move_type.name for move_type in move_types)
-            raise ValueError(f"All target states must share move_type; got {names}.")
+            raise ValueError(
+                f"{type(self.planner).__name__} does not support heterogeneous "
+                f"waypoints; got {names}."
+            )
+        if heterogeneous and options.strategy == "ik_interp":
+            raise ValueError(
+                "strategy='ik_interp' does not support heterogeneous waypoints."
+            )
         move_type = target_states[0].move_type
-        use_interpolation = (
+        use_interpolation = not heterogeneous and (
             options.preserve_cartesian_samples
             or options.strategy == "ik_interp"
             or (
@@ -512,9 +525,11 @@ def _generate_with_planner(
         options: MotionGenOptions,
     ) -> PlanResult:
         """Dispatch batched targets through the configured planner backend."""
+        move_types = {state.move_type for state in target_states}
         move_type = target_states[0].move_type
         should_preinterpolate = (
-            options.is_interpolate
+            len(move_types) == 1
+            and options.is_interpolate
             and not self.planner.supports_move_type(MoveType.EEF_MOVE)
             and self.planner.supports_move_type(MoveType.JOINT_MOVE)
         )
@@ -566,9 +581,11 @@ def _generate_with_planner(
         else:
             target_plan_states = target_states
 
-        unsupported_move_types = (
-            set() if self.planner.supports_move_type(move_type) else {move_type}
-        )
+        unsupported_move_types = {
+            candidate
+            for candidate in move_types
+            if not self.planner.supports_move_type(candidate)
+        }
         if not should_preinterpolate and unsupported_move_types:
             unsupported_names = sorted(
                 move_type.name for move_type in unsupported_move_types
@@ -868,7 +885,14 @@ def normalize_derivative(
 
         velocities = normalize_derivative(result.velocities, "velocities")
         accelerations = normalize_derivative(result.accelerations, "accelerations")
-        if start_qpos is not None and not success.all():
+        preserve_failed_positions = (
+            getattr(self.planner, "preserve_failed_plan_positions", False) is True
+        )
+        if (
+            start_qpos is not None
+            and not success.all()
+            and not preserve_failed_positions
+        ):
             held = (
                 start_qpos.to(dtype=positions.dtype).unsqueeze(1).expand_as(positions)
             )