From bb6da087aea3abe09d5327dae01e3d2013c5276d Mon Sep 17 00:00:00 2001 From: Xinyi YUAN Date: Tue, 29 Sep 2026 17:29:58 +0900 Subject: [PATCH 1/4] feat(motion): let the neural planner retime its rollout under dynamic limits `NeuralPlanner` integrates `action * action_scale` per step and reports `NeuralPlannerCfg.dt` as a constant, so its timing is bookkeeping rather than something any dynamic limit informed: at the default scale a saturated step implies 20 rad/s against 2.62 rad/s for the slowest FR3 joint. The rollout is a usable geometric path, but nothing turned it into an executable trajectory. - Add `NeuralPlannerCfg.constraints`, default `None`, which preserves the current nominal-timing behavior exactly. When set, the rollout samples are re-parameterized under those velocity and acceleration limits and the returned timing, derivatives and poses come from the solved trajectory. - Expose `toppra_planner.retime_joint_paths()` for this: the time parameterization `ToppraPlanner` already owns, as a batched entry point for planners that emit geometry without timing. It selects backends as the planner's `auto` mode does -- Warp on CUDA or when a gradient is needed, NumPy otherwise -- so an NMG rollout on the GPU is retimed there, batched, and differentiably. - Recompute poses from the retimed positions. Retiming resamples along the fitted spline, so carrying `xpos_list` over from the rollout would describe samples that are no longer the ones being returned. - Report an environment whose path cannot be parameterized as failed. The caller asked for limit-respecting output and none exists for that row. - Share one result assembly between the planner and the new entry point rather than duplicating the tail-padding of unequal-length rows. Retiming bounds the derivatives but cannot smooth the path: jerk still comes from the policy, and the cost of respecting the limits appears as duration. Refs #684 Co-Authored-By: Claude Opus 5.5 --- .../topics/motion-planning/planner-details.md | 16 ++ docs/source/api_reference/public_api.rst | 6 + .../lab/sim/motion/planners/neural_planner.py | 69 +++++- .../lab/sim/motion/planners/toppra_planner.py | 213 ++++++++++++++---- .../motion/planners/test_neural_batched.py | 2 + .../motion/planners/test_neural_planner.py | 83 +++++++ 6 files changed, 342 insertions(+), 47 deletions(-) diff --git a/agent_context/topics/motion-planning/planner-details.md b/agent_context/topics/motion-planning/planner-details.md index b1ea0b513..ab7889b70 100644 --- a/agent_context/topics/motion-planning/planner-details.md +++ b/agent_context/topics/motion-planning/planner-details.md @@ -104,6 +104,14 @@ the graph. Validate complete observation-layout metadata/fingerprint before rollout; model dimensions and waypoint capacity belong to the export contract. Inspect the `policy-deploy` package extra for runtime dependencies. +`NeuralPlannerCfg.dt` is bookkeeping, not solved timing: the rollout integrates +a joint delta per step, so a saturated step implies `action_scale / dt` rad/s +and overshoots a typical arm's velocity limit several times over. Setting +`constraints` swaps that nominal timing for a TOPPRA parameterization through +the shared entry point below; it resamples the path, so positions and poses +change with the timing, and it cannot reduce the path's own jerk. Default +remains `None`, which preserves the nominal-timing behavior. + Runtime/training frame differences must be expressed through explicit policy frame and TCP transforms. Target/FK quaternions follow the shared [simulation convention](../simulation-system/simulation-system.md); do not @@ -123,6 +131,14 @@ semantics are owned by [collision worlds](collision-worlds.md). ## Retiming and playback +`toppra_planner.py::retime_joint_paths()` is the shared entry point for +re-parameterizing an already-planned joint path under velocity and acceleration +limits, used by planners that emit geometry without executable timing. It is +the planner's own time parameterization with the same backend selection, so +CUDA paths and gradient requests run on Warp; it resamples along the fitted +spline, so recompute derived poses from its output rather than reusing the +input samples. + `compute/trajectory/timing.py` owns time-domain differentiation/resampling and `retime_to_control_grid()`. Zero-time position changes are invalid; repeated-time unchanged samples may represent padding/junctions. Time resampling preserves diff --git a/docs/source/api_reference/public_api.rst b/docs/source/api_reference/public_api.rst index d4b037867..ea1e1c533 100644 --- a/docs/source/api_reference/public_api.rst +++ b/docs/source/api_reference/public_api.rst @@ -1176,11 +1176,17 @@ embodichain.lab.sim.motion.planners.toppra_planner .. currentmodule:: embodichain.lab.sim.motion.planners.toppra_planner +Time-optimal joint-path parameterization. The planner solves timing for +waypoint input; :func:`retime_joint_paths` applies the same parameterization to +paths another planner already produced, such as a closed-loop policy rollout +whose samples carry only nominal timing. + .. autosummary:: ToppraPlanner ToppraPlannerCfg ToppraPlanOptions + retime_joint_paths embodichain.lab.sim.motion.planners.se3 --------------------------------------- diff --git a/embodichain/lab/sim/motion/planners/neural_planner.py b/embodichain/lab/sim/motion/planners/neural_planner.py index 3caf83b98..b4ef02d9d 100644 --- a/embodichain/lab/sim/motion/planners/neural_planner.py +++ b/embodichain/lab/sim/motion/planners/neural_planner.py @@ -316,7 +316,26 @@ class NeuralPlannerCfg(BasePlannerCfg): """ dt: float = 0.01 - """Nominal timestep reported in PlanResult.""" + """Nominal timestep reported in PlanResult when no retiming is requested. + + The rollout integrates a joint delta per step and never solves a duration, + so this value is bookkeeping rather than executable timing: at the default + ``action_scale`` a saturated step implies ``action_scale / dt`` rad/s, + which exceeds a typical arm's joint velocity limit several times over. Set + :attr:`constraints` to replace it with a solved time parameterization. + """ + + constraints: dict | None = None + """Optional ``velocity`` and ``acceleration`` limits for output retiming. + + ``None`` keeps the rollout's nominal timing and current behavior. When set, + the closed-loop samples are treated as a geometric path and re-parameterized + under these limits, so the returned timing and derivatives are executable. + Keys and value shapes match :class:`ToppraPlanOptions`: a scalar or one + value per controlled joint. Retiming resamples the path, so it changes the + reported positions and poses as well as the timing, and it cannot reduce + the path's own jerk -- only the duration over which it is traversed. + """ @configclass @@ -583,10 +602,12 @@ def plan( dt[:, 0] = 0.0 positions_t = positions_t.permute(1, 0, 2) xpos_t = xpos_t.permute(1, 0, 2, 3) + success = active_idx >= episode_k + if self.cfg.constraints is not None: + return self._retimed_result(success, positions_t, control_part) velocities_t, accelerations_t = self._compute_vel_acc_via_finite_diff( positions_t, dt ) - success = active_idx >= episode_k return PlanResult( success=success, positions=positions_t, @@ -596,6 +617,50 @@ def plan( dt=dt, ) + def _retimed_result( + self, + success: torch.Tensor, + positions: torch.Tensor, + control_part: str, + ) -> PlanResult: + """Re-parameterize the rollout path under the configured limits. + + Retiming resamples the path, so poses are recomputed from the returned + positions rather than carried over from the rollout. An environment + whose path cannot be parameterized is reported as failed: the caller + asked for limit-respecting output and none is available for that row. + + Args: + success: Per-env rollout convergence of shape ``(B,)``. + positions: Rollout joint samples of shape ``(B, N, DOF)``. + control_part: Robot control part used for forward kinematics. + + Returns: + PlanResult with solved timing, derivatives and recomputed poses. + """ + from .toppra_planner import retime_joint_paths + + retimed = retime_joint_paths( + positions, + constraints=dict(self.cfg.constraints), + device=self.device, + ) + poses = torch.stack( + [ + self._fk_matrix(retimed.positions[:, index], control_part) + for index in range(retimed.positions.shape[1]) + ], + dim=1, + ) + return PlanResult( + success=success & retimed.success.to(success.device), + positions=retimed.positions, + velocities=retimed.velocities, + accelerations=retimed.accelerations, + xpos_list=poses, + dt=retimed.dt, + ) + def _parse_waypoints(self, target_states: list[PlanState]) -> tuple[ torch.Tensor, torch.Tensor, diff --git a/embodichain/lab/sim/motion/planners/toppra_planner.py b/embodichain/lab/sim/motion/planners/toppra_planner.py index 1cf21e21c..49bad508b 100644 --- a/embodichain/lab/sim/motion/planners/toppra_planner.py +++ b/embodichain/lab/sim/motion/planners/toppra_planner.py @@ -19,6 +19,7 @@ from copy import deepcopy import os +from collections.abc import Mapping, Sequence from typing import Literal import torch @@ -174,7 +175,171 @@ def _worker_init() -> None: _set_parent_death_signal() -__all__ = ["ToppraPlanner", "ToppraPlannerCfg", "ToppraPlanOptions"] +def _assemble_batched_results( + results: list[dict], + dofs: int, + device: torch.device, +) -> PlanResult: + """Stack per-env solves into one batched :class:`PlanResult`. + + Env trajectories may differ in length. Shorter rows are tail-padded by + repeating their final waypoint with zero velocity and acceleration, so + every output tensor shares the same ``(B, N, DOF)`` or ``(B, N)`` shape. + + Args: + results: Per-env dicts as returned by :func:`_toppra_solve_one_env`. + dofs: Per-env degrees of freedom. + device: Device for the assembled tensors. + + Returns: + PlanResult with env-batched positions, derivatives, timing and success. + """ + b = len(results) + max_n = max(r["n"] for r in results) + positions = np.zeros((b, max_n, dofs), dtype=np.float32) + velocities = np.zeros((b, max_n, dofs), dtype=np.float32) + accelerations = np.zeros((b, max_n, dofs), dtype=np.float32) + dt = np.zeros((b, max_n), dtype=np.float32) + success = np.zeros((b,), dtype=bool) + for i, r in enumerate(results): + n = r["n"] + positions[i, :n] = r["positions"] + velocities[i, :n] = r["velocities"] + accelerations[i, :n] = r["accelerations"] + dt[i, :n] = r["dt"] + success[i] = r["success"] + # tail-pad: repeat final waypoint for held-pose rows + if n < max_n: + positions[i, n:] = r["positions"][-1] + velocities[i, n:] = 0.0 + accelerations[i, n:] = 0.0 + return PlanResult( + success=torch.as_tensor(success, device=device), + positions=torch.as_tensor(positions, device=device), + velocities=torch.as_tensor(velocities, device=device), + accelerations=torch.as_tensor(accelerations, device=device), + dt=torch.as_tensor(dt, device=device), + ) + + +def retime_joint_paths( + positions: torch.Tensor, + *, + constraints: Mapping[str, float | Sequence[float] | torch.Tensor], + sample_method: TrajectorySampleMethod = TrajectorySampleMethod.QUANTITY, + sample_interval: float | int | None = None, + grid_size: int = 100, + device: torch.device | None = None, +) -> PlanResult: + """Re-parameterize already-planned joint paths under dynamic limits. + + This is the time parameterization of :class:`ToppraPlanner`, exposed for + planners that produce geometry without executable timing. A closed-loop + policy rollout, for example, emits one joint sample per control step on a + nominal ``dt`` that no dynamic limit informed; feeding those samples here + replaces that timing with a solved one that respects the limits. + + Backend selection follows :class:`ToppraPlanner`'s ``auto`` mode: Warp on + CUDA inputs or when a gradient is required, NumPy otherwise. Both enforce + the limits over every grid interval rather than only at grid points. + + The path geometry is resampled along the fitted spline, so poses derived + from the input samples must be recomputed from the returned positions. + + Args: + positions: Planned joint paths of shape ``(B, N, DOF)``. + constraints: ``velocity`` and ``acceleration`` limits, each a scalar, + one value per joint, or signed ``(DOF, 2)`` bounds, matching + :class:`ToppraPlanOptions`. + sample_method: Fixed output quantity or approximately fixed time step. + sample_interval: Output count for ``QUANTITY`` or seconds for ``TIME``. + ``None`` keeps the input sample count and is only meaningful for + ``QUANTITY``, which leaves array shapes unchanged while the timing + behind them changes. + grid_size: Minimum parameterization grid size, as for + :class:`ToppraPlannerCfg`. + device: Device for the assembled tensors; defaults to the input's. + + Returns: + PlanResult whose ``success`` is per-env and false where the path could + not be parameterized. + + Raises: + ValueError: If ``positions`` is not a ``(B, N, DOF)`` tensor, the + constraints omit ``velocity`` or ``acceleration``, or ``TIME`` + sampling is requested without an explicit interval in seconds. + """ + if positions.dim() != 3: + raise ValueError( + f"positions must have shape (B, N, DOF), got {tuple(positions.shape)}." + ) + missing = sorted({"velocity", "acceleration"}.difference(constraints)) + if missing: + raise ValueError(f"constraints is missing required keys: {missing}.") + if sample_interval is None and sample_method is TrajectorySampleMethod.TIME: + # The default is a sample count; reading it as seconds would silently + # collapse a dense path to two samples whenever its duration is shorter. + raise ValueError( + "sample_interval is required for TIME sampling; it is a duration in " + "seconds, not the sample count the default supplies." + ) + resolved_device = positions.device if device is None else device + batch, count, dofs = positions.shape + resolved_interval = count if sample_interval is None else sample_interval + velocity, acceleration = constraints["velocity"], constraints["acceleration"] + needs_grad = torch.is_grad_enabled() and ( + positions.requires_grad + or any( + isinstance(limit, torch.Tensor) and limit.requires_grad + for limit in (velocity, acceleration) + ) + ) + if positions.is_cuda or needs_grad: + from embodichain.compute.trajectory._toppra_warp import _retime_toppra_warp + + quantity = sample_method == TrajectorySampleMethod.QUANTITY + waypoints = positions if positions.is_floating_point() else positions.double() + try: + result = _retime_toppra_warp( + waypoints, + velocity, + acceleration, + sample_count=int(resolved_interval) if quantity else None, + sample_dt=None if quantity else float(resolved_interval), + grid_size=grid_size, + ) + except (ValueError, FloatingPointError, OverflowError): + return _assemble_batched_results( + [_empty_failure(dofs) for _ in range(batch)], dofs, resolved_device + ) + return PlanResult( + **{key: value.to(resolved_device) for key, value in result.items()} + ) + + samples = positions.detach().cpu().numpy().astype(np.float64) + to_numpy = lambda limit: ( + limit.detach().cpu().numpy() if isinstance(limit, torch.Tensor) else limit + ) + results = [ + _toppra_solve_one_env( + samples[index], + to_numpy(velocity), + to_numpy(acceleration), + sample_method, + resolved_interval, + grid_size, + ) + for index in range(batch) + ] + return _assemble_batched_results(results, dofs, resolved_device) + + +__all__ = [ + "ToppraPlanner", + "ToppraPlannerCfg", + "ToppraPlanOptions", + "retime_joint_paths", +] @configclass @@ -535,47 +700,5 @@ def plan( return self._assemble_batched_result(results, dofs) def _assemble_batched_result(self, results: list[dict], dofs: int) -> PlanResult: - """Stack per-env TOPPRA results into a batched :class:`PlanResult`. - - Each entry of ``results`` is the dict returned by - :func:`_toppra_solve_one_env`. Env trajectories may have different - lengths (``n``); this method pads shorter trajectories out to the - longest by repeating their final waypoint (held pose) with zero - velocity and acceleration, so every output tensor shares the same - ``(B, N, DOF)`` / ``(B, N)`` shape. - - Args: - results: list of per-env result dicts (length ``B``). - dofs: per-env degrees of freedom. - - Returns: - PlanResult with env-batched tensors (``success`` ``(B,)``, - ``positions``/``velocities``/``accelerations`` ``(B, N, DOF)``, - ``dt`` ``(B, N)``, ``duration`` ``(B,)``). - """ - b = len(results) - max_n = max(r["n"] for r in results) - positions = np.zeros((b, max_n, dofs), dtype=np.float32) - velocities = np.zeros((b, max_n, dofs), dtype=np.float32) - accelerations = np.zeros((b, max_n, dofs), dtype=np.float32) - dt = np.zeros((b, max_n), dtype=np.float32) - success = np.zeros((b,), dtype=bool) - for i, r in enumerate(results): - n = r["n"] - positions[i, :n] = r["positions"] - velocities[i, :n] = r["velocities"] - accelerations[i, :n] = r["accelerations"] - dt[i, :n] = r["dt"] - success[i] = r["success"] - # tail-pad: repeat final waypoint for held-pose rows - if n < max_n: - positions[i, n:] = r["positions"][-1] - velocities[i, n:] = 0.0 - accelerations[i, n:] = 0.0 - return PlanResult( - success=torch.as_tensor(success, device=self.device), - positions=torch.as_tensor(positions, device=self.device), - velocities=torch.as_tensor(velocities, device=self.device), - accelerations=torch.as_tensor(accelerations, device=self.device), - dt=torch.as_tensor(dt, device=self.device), - ) + """Stack per-env TOPPRA results into a batched :class:`PlanResult`.""" + return _assemble_batched_results(results, dofs, self.device) diff --git a/tests/sim/motion/planners/test_neural_batched.py b/tests/sim/motion/planners/test_neural_batched.py index e2ebc039c..df263aeee 100644 --- a/tests/sim/motion/planners/test_neural_batched.py +++ b/tests/sim/motion/planners/test_neural_batched.py @@ -87,6 +87,7 @@ def test_plan_returns_batched_success(self, monkeypatch): "dt": 0.01, "control_part": "arm", "num_arm_joints": 7, + "constraints": None, }, )() @@ -158,6 +159,7 @@ def test_converged_env_holds_qpos(self): "dt": 0.01, "control_part": "arm", "num_arm_joints": 7, + "constraints": None, }, )() diff --git a/tests/sim/motion/planners/test_neural_planner.py b/tests/sim/motion/planners/test_neural_planner.py index 73f91caa7..52a27446a 100644 --- a/tests/sim/motion/planners/test_neural_planner.py +++ b/tests/sim/motion/planners/test_neural_planner.py @@ -31,6 +31,7 @@ ) from embodichain.lab.sim.motion.planners import neural_planner as neural_planner_module from embodichain.lab.sim.motion.planners.neural_planner import NeuralPlanOptions +from embodichain.lab.sim.motion.planners.toppra_planner import retime_joint_paths from embodichain.lab.sim.sim_manager import SimulationManager NUM_ARM_JOINTS = 7 @@ -402,6 +403,88 @@ def test_neural_planner_returns_velocities_and_accelerations(tmp_path, monkeypat assert torch.isfinite(result.accelerations).all() +class SaturatingOnnxPolicy(FakeOnnxPolicy): + """Policy that always saturates, the worst case for implied velocity.""" + + def __call__(self, obs: torch.Tensor) -> torch.Tensor: + self.last_obs = obs.clone() + return torch.ones(obs.shape[0], NUM_ARM_JOINTS, device=obs.device) + + +def _rollout(tmp_path, monkeypatch, *, policy=None, **cfg_overrides): + """Roll the neural planner out toward one identity pose target.""" + model_path = _create_fake_onnx_model(tmp_path) + fake_sim = FakeSimulationManager() + monkeypatch.setattr( + SimulationManager, + "get_instance", + classmethod(lambda cls, instance_id=0: fake_sim), + ) + if policy is not None: + monkeypatch.setattr(neural_planner_module, "_OnnxPolicy", policy) + motion_generator = MotionGenerator( + cfg=MotionGenCfg( + planner_cfg=NeuralPlannerCfg( + robot_uid="fake_robot", + onnx_model_path=model_path, + control_part="main_arm", + **cfg_overrides, + ) + ) + ) + return motion_generator.generate( + target_states=[ + PlanState.single(move_type=MoveType.EEF_MOVE, xpos=torch.eye(4)) + ], + options=MotionGenOptions( + plan_opts=NeuralPlanOptions( + control_part="main_arm", + start_qpos=torch.zeros(NUM_ARM_JOINTS), + ), + ), + ) + + +def test_neural_planner_keeps_nominal_timing_without_constraints(tmp_path, monkeypatch): + result = _rollout(tmp_path, monkeypatch, policy=SaturatingOnnxPolicy) + + # The default path is unchanged: one nominal interval per rollout step. + intervals = result.dt[0, 1:] + assert torch.allclose(intervals, torch.full_like(intervals, 0.01)) + # 0.2 rad per 0.01 s is 20 rad/s, which no arm can execute. + assert float(result.velocities.abs().max()) == pytest.approx(20.0, rel=1e-3) + + +def test_neural_planner_constraints_replace_timing_and_bound_dynamics( + tmp_path, monkeypatch +): + nominal = _rollout(tmp_path, monkeypatch, policy=SaturatingOnnxPolicy) + retimed = _rollout( + tmp_path, + monkeypatch, + policy=SaturatingOnnxPolicy, + constraints={"velocity": 1.0, "acceleration": 2.0}, + ) + + assert bool(retimed.success.all()) + assert float(retimed.velocities.abs().max()) <= 1.0 + 1e-6 + assert float(retimed.accelerations.abs().max()) <= 2.0 + 1e-6 + # Respecting the limits costs time; the nominal timing understated it. + assert float(retimed.dt.sum()) > float(nominal.dt.sum()) + # Retiming resamples the path, so poses must track the returned positions. + assert retimed.xpos_list.shape[:2] == retimed.positions.shape[:2] + + +def test_retime_joint_paths_rejects_malformed_input(): + positions = torch.zeros(1, 4, NUM_ARM_JOINTS) + with pytest.raises(ValueError, match="velocity"): + retime_joint_paths(positions, constraints={"acceleration": 1.0}) + with pytest.raises(ValueError, match=r"\(B, N, DOF\)"): + retime_joint_paths( + positions[0], constraints={"velocity": 1.0, "acceleration": 1.0} + ) + + def test_neural_planner_disables_grad_for_all_fk_calls(tmp_path, monkeypatch): model_path = _create_fake_onnx_model(tmp_path) fake_sim = FakeSimulationManager() From e13601cf347a155694705be3f6a00aeddc7bcbcf Mon Sep 17 00:00:00 2001 From: Xinyi YUAN Date: Tue, 29 Sep 2026 17:29:43 +0900 Subject: [PATCH 2/4] fix(motion): derive retimed success from the trajectory returned Retiming fits a spline through the rollout samples and resamples it, so the returned trajectory is not the one the rollout verified. Carrying the rollout's success flag across that boundary lets failures through: - The resampled grid can step past a waypoint the rollout stopped on, so a multi-waypoint plan could report success while its returned samples never come within tolerance of a required target. Arrival is now re-checked on the returned grid with the rollout's own test. - A spline through samples the rollout clamped at a joint limit can overshoot that limit between them, since the fit sees only the sampled values. The parameterization bounds velocity and acceleration, not position, so retimed positions are now checked against the joint limits. - A row that changes pose in zero total time is reported as failed. The current parameterization gives tiny moves a positive duration, so this is a guard against a regression rather than a fix for present behavior. Also annotate `NeuralPlannerCfg.constraints` as `dict[str, float | list[float]]` rather than a bare `dict`, per the repository's public-API requirement. Refs #684 Co-Authored-By: Claude Opus 5.5 --- .../lab/sim/motion/planners/neural_planner.py | 118 ++++++++++++++++-- .../motion/planners/test_neural_planner.py | 29 ++++- 2 files changed, 136 insertions(+), 11 deletions(-) diff --git a/embodichain/lab/sim/motion/planners/neural_planner.py b/embodichain/lab/sim/motion/planners/neural_planner.py index b4ef02d9d..1b6d4e63e 100644 --- a/embodichain/lab/sim/motion/planners/neural_planner.py +++ b/embodichain/lab/sim/motion/planners/neural_planner.py @@ -325,7 +325,7 @@ class NeuralPlannerCfg(BasePlannerCfg): :attr:`constraints` to replace it with a solved time parameterization. """ - constraints: dict | None = None + constraints: dict[str, float | list[float]] | None = None """Optional ``velocity`` and ``acceleration`` limits for output retiming. ``None`` keeps the rollout's nominal timing and current behavior. When set, @@ -604,7 +604,21 @@ def plan( xpos_t = xpos_t.permute(1, 0, 2, 3) success = active_idx >= episode_k if self.cfg.constraints is not None: - return self._retimed_result(success, positions_t, control_part) + return self._retimed_result( + success, + positions_t, + control_part, + waypoints=( + waypoints_pos, + waypoints_quat, + waypoints_joint, + pos_mask, + rot_mask, + joint_mask, + ), + episode_k=episode_k, + qpos_limits=(lower, upper), + ) velocities_t, accelerations_t = self._compute_vel_acc_via_finite_diff( positions_t, dt ) @@ -622,18 +636,30 @@ def _retimed_result( success: torch.Tensor, positions: torch.Tensor, control_part: str, + *, + waypoints: tuple[torch.Tensor, ...], + episode_k: int, + qpos_limits: tuple[torch.Tensor, torch.Tensor], ) -> PlanResult: """Re-parameterize the rollout path under the configured limits. - Retiming resamples the path, so poses are recomputed from the returned - positions rather than carried over from the rollout. An environment - whose path cannot be parameterized is reported as failed: the caller - asked for limit-respecting output and none is available for that row. + Retiming fits a spline through the rollout samples and resamples it, so + the returned trajectory is not the one the rollout verified. Success is + therefore re-derived from the samples actually returned rather than + carried over: the resampled grid can step past a waypoint the rollout + stopped on, and a spline through samples clamped at a joint limit can + overshoot that limit between them. A row is reported as successful only + when its returned trajectory parameterizes, moves in nonzero time, + stays inside the joint limits, and still reaches every waypoint. Args: success: Per-env rollout convergence of shape ``(B,)``. positions: Rollout joint samples of shape ``(B, N, DOF)``. control_part: Robot control part used for forward kinematics. + waypoints: Parsed waypoint targets and masks from + :meth:`_parse_waypoints`. + episode_k: Number of waypoints this episode must reach. + qpos_limits: Lower and upper joint-position limits. Returns: PlanResult with solved timing, derivatives and recomputed poses. @@ -645,22 +671,94 @@ def _retimed_result( constraints=dict(self.cfg.constraints), device=self.device, ) + retimed_positions = retimed.positions poses = torch.stack( [ - self._fk_matrix(retimed.positions[:, index], control_part) - for index in range(retimed.positions.shape[1]) + self._fk_matrix(retimed_positions[:, index], control_part) + for index in range(retimed_positions.shape[1]) ], dim=1, ) + success = success & retimed.success.to(success.device) + success = success & self._retimed_is_executable( + retimed_positions, retimed.dt, qpos_limits + ) + success = success & self._retimed_reaches_waypoints( + retimed_positions, control_part, waypoints, episode_k + ) return PlanResult( - success=success & retimed.success.to(success.device), - positions=retimed.positions, + success=success, + positions=retimed_positions, velocities=retimed.velocities, accelerations=retimed.accelerations, xpos_list=poses, dt=retimed.dt, ) + def _retimed_is_executable( + self, + positions: torch.Tensor, + dt: torch.Tensor, + qpos_limits: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + """Return whether each retimed row has real timing and legal positions. + + The shared retiming kernel returns a zero-duration result for a path + whose endpoints nearly coincide, which would claim a move with no time + to execute it. Spline interpolation through samples the rollout clamped + at a joint limit can also overshoot that limit between them, since the + fit sees only the sampled values. + """ + lower, upper = qpos_limits + arm = positions[..., : self._action_dim] + moves = (arm - arm[:, :1]).abs().amax(dim=(1, 2)) > 1.0e-9 + timed = dt.sum(dim=-1) > 0.0 + within = ((arm >= lower) & (arm <= upper)).all(dim=-1).all(dim=-1) + return within & (timed | ~moves) + + def _retimed_reaches_waypoints( + self, + positions: torch.Tensor, + control_part: str, + waypoints: tuple[torch.Tensor, ...], + episode_k: int, + ) -> torch.Tensor: + """Return whether each retimed row still reaches every waypoint in order. + + Applies the rollout's own arrival test to the resampled grid, so a + trajectory that passes a waypoint only between output samples is not + reported as reaching it. + """ + ( + waypoints_pos, + waypoints_quat, + waypoints_joint, + pos_mask, + rot_mask, + joint_mask, + ) = waypoints + batch = positions.shape[0] + active_idx = torch.zeros(batch, dtype=torch.long, device=self.device) + for index in range(positions.shape[1]): + qpos = positions[:, index] + reached = self._is_active_reached( + qpos[:, : self._action_dim], + self._fk_pose_xyzw(qpos, control_part), + waypoints_pos, + waypoints_quat, + waypoints_joint, + pos_mask, + rot_mask, + joint_mask, + active_idx, + ) + active_idx = torch.where( + reached & (active_idx < episode_k), active_idx + 1, active_idx + ) + if bool((active_idx >= episode_k).all()): + break + return active_idx >= episode_k + def _parse_waypoints(self, target_states: list[PlanState]) -> tuple[ torch.Tensor, torch.Tensor, diff --git a/tests/sim/motion/planners/test_neural_planner.py b/tests/sim/motion/planners/test_neural_planner.py index 52a27446a..eaee82870 100644 --- a/tests/sim/motion/planners/test_neural_planner.py +++ b/tests/sim/motion/planners/test_neural_planner.py @@ -32,6 +32,7 @@ from embodichain.lab.sim.motion.planners import neural_planner as neural_planner_module from embodichain.lab.sim.motion.planners.neural_planner import NeuralPlanOptions from embodichain.lab.sim.motion.planners.toppra_planner import retime_joint_paths +from embodichain.lab.sim.motion.planners.utils import TrajectorySampleMethod from embodichain.lab.sim.sim_manager import SimulationManager NUM_ARM_JOINTS = 7 @@ -477,14 +478,40 @@ def test_neural_planner_constraints_replace_timing_and_bound_dynamics( def test_retime_joint_paths_rejects_malformed_input(): positions = torch.zeros(1, 4, NUM_ARM_JOINTS) + limits = {"velocity": 1.0, "acceleration": 1.0} with pytest.raises(ValueError, match="velocity"): retime_joint_paths(positions, constraints={"acceleration": 1.0}) with pytest.raises(ValueError, match=r"\(B, N, DOF\)"): + retime_joint_paths(positions[0], constraints=limits) + # The default interval is a sample count; reading it as seconds would + # silently collapse a dense path to two samples. + with pytest.raises(ValueError, match="TIME sampling"): retime_joint_paths( - positions[0], constraints={"velocity": 1.0, "acceleration": 1.0} + positions, + constraints=limits, + sample_method=TrajectorySampleMethod.TIME, ) +def test_retimed_success_reflects_the_trajectory_actually_returned( + tmp_path, monkeypatch +): + # A joint limit the rollout clamps against is also a limit the spline must + # respect, and the resampled grid must still reach the goal. Success is + # re-derived from the returned samples rather than inherited. + retimed = _rollout( + tmp_path, + monkeypatch, + policy=SaturatingOnnxPolicy, + constraints={"velocity": 1.0, "acceleration": 2.0}, + ) + arm = retimed.positions[..., :NUM_ARM_JOINTS] + # FakeRobot declares +-2.0 rad on every joint. + assert float(arm.max()) <= 2.0 + 1e-6 + assert float(arm.min()) >= -2.0 - 1e-6 + assert float(retimed.dt.sum()) > 0.0 + + def test_neural_planner_disables_grad_for_all_fk_calls(tmp_path, monkeypatch): model_path = _create_fake_onnx_model(tmp_path) fake_sim = FakeSimulationManager() From 97ebc8c037096e6d6b264aca0a9d887f96df608a Mon Sep 17 00:00:00 2001 From: Xinyi YUAN Date: Tue, 29 Sep 2026 17:29:43 +0900 Subject: [PATCH 3/4] fix(motion): sample the retimed trajectory on the control period The retimed output sample count defaulted to the rollout's own step count, which is however many steps the policy happened to take. Sampling a trajectory stretched by more than an order of magnitude on so few points leaves a waypoint the rollout stopped on falling between two output samples, and the arrival re-check then reports it missed. Measured on a five-waypoint case against a 10 mm tolerance, worst waypoint error by output sample count: rollout (13 samples) 8.77 mm within tolerance retimed, N = 13 10.55 mm outside retimed, N = 64 8.64 mm within retimed, N = 1024 8.64 mm unchanged from N = 64 This was measured with the previous parameterization backend, but the effect is a property of output sample density, not of the parameterization: between 64 and 1024 samples the error does not move. The margin is what makes it fragile: the policy stops as soon as it is barely inside tolerance, leaving about 1 mm of slack for any resampling to consume. Sample the solved trajectory on `NeuralPlannerCfg.dt` instead. That field was already the assumed control period, and using it here makes the output an execution-ready grid rather than a rollout-shaped one. Refs #684 Co-Authored-By: Claude Opus 5.5 --- .../lab/sim/motion/planners/neural_planner.py | 22 ++++++++++++++++--- .../motion/planners/test_neural_planner.py | 20 +++++++++++++++++ 2 files changed, 39 insertions(+), 3 deletions(-) diff --git a/embodichain/lab/sim/motion/planners/neural_planner.py b/embodichain/lab/sim/motion/planners/neural_planner.py index 1b6d4e63e..33d1fc240 100644 --- a/embodichain/lab/sim/motion/planners/neural_planner.py +++ b/embodichain/lab/sim/motion/planners/neural_planner.py @@ -30,7 +30,12 @@ _infer_batch_size, validate_plan_options, ) -from embodichain.lab.sim.motion.planners.utils import MoveType, PlanResult, PlanState +from embodichain.lab.sim.motion.planners.utils import ( + MoveType, + PlanResult, + PlanState, + TrajectorySampleMethod, +) from embodichain.utils import configclass, logger from embodichain.utils.math import quat_error_magnitude, quat_from_matrix @@ -316,13 +321,17 @@ class NeuralPlannerCfg(BasePlannerCfg): """ dt: float = 0.01 - """Nominal timestep reported in PlanResult when no retiming is requested. + """Output sampling period in seconds. + + Without :attr:`constraints` this is bookkeeping rather than executable + timing. The rollout integrates a joint delta per step and never solves a duration, so this value is bookkeeping rather than executable timing: at the default ``action_scale`` a saturated step implies ``action_scale / dt`` rad/s, which exceeds a typical arm's joint velocity limit several times over. Set - :attr:`constraints` to replace it with a solved time parameterization. + :attr:`constraints` to replace it with a solved time parameterization, in + which case this becomes the period that trajectory is sampled on. """ constraints: dict[str, float | list[float]] | None = None @@ -666,9 +675,16 @@ def _retimed_result( """ from .toppra_planner import retime_joint_paths + # Sample the solved trajectory on the configured control period rather + # than on the rollout's step count. The rollout emits however many + # steps it happened to take, which is far too coarse once the duration + # stretches: a waypoint the rollout stopped on then falls between two + # output samples and is reported as missed. retimed = retime_joint_paths( positions, constraints=dict(self.cfg.constraints), + sample_method=TrajectorySampleMethod.TIME, + sample_interval=float(self.cfg.dt), device=self.device, ) retimed_positions = retimed.positions diff --git a/tests/sim/motion/planners/test_neural_planner.py b/tests/sim/motion/planners/test_neural_planner.py index eaee82870..bd2c9e35f 100644 --- a/tests/sim/motion/planners/test_neural_planner.py +++ b/tests/sim/motion/planners/test_neural_planner.py @@ -493,6 +493,26 @@ def test_retime_joint_paths_rejects_malformed_input(): ) +def test_retiming_samples_on_the_control_period_not_the_rollout_step_count( + tmp_path, monkeypatch +): + # The rollout emits however many steps it happened to take. Inheriting that + # count leaves a stretched trajectory sampled far too coarsely, so a + # waypoint the rollout stopped on falls between two output samples. + nominal = _rollout(tmp_path, monkeypatch, policy=SaturatingOnnxPolicy) + retimed = _rollout( + tmp_path, + monkeypatch, + policy=SaturatingOnnxPolicy, + constraints={"velocity": 1.0, "acceleration": 2.0}, + ) + duration = float(retimed.dt.sum()) + assert duration > float(nominal.dt.sum()) + # One sample per nominal dt over the solved duration, not one per step. + assert retimed.positions.shape[1] > nominal.positions.shape[1] + assert retimed.positions.shape[1] == pytest.approx(duration / 0.01, rel=0.05) + + def test_retimed_success_reflects_the_trajectory_actually_returned( tmp_path, monkeypatch ): From 20250151e7fb1df1cdd9e351723fc48550dd8f5e Mon Sep 17 00:00:00 2001 From: Xinyi YUAN Date: Tue, 29 Sep 2026 17:29:43 +0900 Subject: [PATCH 4/4] perf(motion): run post-retime forward kinematics as one batched call After retiming, the planner recomputed every sample's pose with one FK call per time step, then did it a second time inside the waypoint re-check. For a batch of eight environments that was 618 ms of an 1188 ms plan. Compute the poses once with `Robot.compute_batch_fk`, which takes the whole (B, K, DOF) trajectory, and reuse them for the waypoint re-check, which now only maps them to the policy frame. On the Franka FR3 benchmark robot the batched result is identical to the per-step one (max difference 0.0); the post-retime checks fall from 618 ms to about 100 ms. The FK cost depends on the output sample count, not on which parameterization backend produced it. Refs #684 Co-Authored-By: Claude Opus 5.5 --- .../lab/sim/motion/planners/neural_planner.py | 28 +++++++++++-------- .../motion/planners/test_neural_planner.py | 7 +++++ 2 files changed, 23 insertions(+), 12 deletions(-) diff --git a/embodichain/lab/sim/motion/planners/neural_planner.py b/embodichain/lab/sim/motion/planners/neural_planner.py index 33d1fc240..bf9a91dd0 100644 --- a/embodichain/lab/sim/motion/planners/neural_planner.py +++ b/embodichain/lab/sim/motion/planners/neural_planner.py @@ -688,19 +688,16 @@ def _retimed_result( device=self.device, ) retimed_positions = retimed.positions - poses = torch.stack( - [ - self._fk_matrix(retimed_positions[:, index], control_part) - for index in range(retimed_positions.shape[1]) - ], - dim=1, + # One batched FK over every sample; the waypoint re-check reuses it. + poses = self.robot.compute_batch_fk( + qpos=retimed_positions, name=control_part, to_matrix=True ) success = success & retimed.success.to(success.device) success = success & self._retimed_is_executable( retimed_positions, retimed.dt, qpos_limits ) success = success & self._retimed_reaches_waypoints( - retimed_positions, control_part, waypoints, episode_k + retimed_positions, poses, waypoints, episode_k ) return PlanResult( success=success, @@ -735,7 +732,7 @@ def _retimed_is_executable( def _retimed_reaches_waypoints( self, positions: torch.Tensor, - control_part: str, + poses: torch.Tensor, waypoints: tuple[torch.Tensor, ...], episode_k: int, ) -> torch.Tensor: @@ -753,13 +750,16 @@ def _retimed_reaches_waypoints( rot_mask, joint_mask, ) = waypoints - batch = positions.shape[0] + batch, samples = positions.shape[:2] + policy_poses = self._policy_pose_xyzw(poses.flatten(0, 1)).view( + batch, samples, -1 + ) active_idx = torch.zeros(batch, dtype=torch.long, device=self.device) - for index in range(positions.shape[1]): + for index in range(samples): qpos = positions[:, index] reached = self._is_active_reached( qpos[:, : self._action_dim], - self._fk_pose_xyzw(qpos, control_part), + policy_poses[:, index], waypoints_pos, waypoints_quat, waypoints_joint, @@ -881,7 +881,11 @@ def _fk_matrix(self, qpos: torch.Tensor, control_part: str) -> torch.Tensor: def _fk_pose_xyzw(self, qpos: torch.Tensor, control_part: str) -> torch.Tensor: """Return the policy-frame FK pose as ``xyz + xyzw``.""" - fk = self._to_policy_frame(self._fk_matrix(qpos, control_part)) + return self._policy_pose_xyzw(self._fk_matrix(qpos, control_part)) + + def _policy_pose_xyzw(self, fk_matrix: torch.Tensor) -> torch.Tensor: + """Map ``(M, 4, 4)`` runtime TCP poses to policy-frame ``xyz + xyzw``.""" + fk = self._to_policy_frame(fk_matrix) pos = fk[:, :3, 3] # ``quat_from_matrix`` is an EmbodiChain ``xyzw`` producer; converting # it again would turn a valid pose into a different rotation. diff --git a/tests/sim/motion/planners/test_neural_planner.py b/tests/sim/motion/planners/test_neural_planner.py index bd2c9e35f..82db1f152 100644 --- a/tests/sim/motion/planners/test_neural_planner.py +++ b/tests/sim/motion/planners/test_neural_planner.py @@ -118,6 +118,13 @@ def compute_fk( return torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]]).repeat(batch, 1) +def _fake_batch_fk(self, qpos, name=None, env_ids=None, to_matrix=False): + return torch.eye(4).repeat(qpos.shape[0], qpos.shape[1], 1, 1) + + +FakeRobot.compute_batch_fk = _fake_batch_fk + + class FakeSimulationManager: def __init__(self): self.robot = FakeRobot()