From 3885b1fff388250a7dec41801b4eff30ab5290dd Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 17 Oct 2026 08:35:22 +0000 Subject: [PATCH 1/3] feat(motion): add differentiable Warp TOPPRA pipeline --- .../topics/motion-planning/planner-details.md | 62 +- .../embodichain.lab.sim.motion.planners.rst | 14 + embodichain/compute/trajectory/_toppra.py | 237 ++++++ .../compute/trajectory/_toppra_warp.py | 611 +++++++++++++++ .../compute/trajectory/_warp/toppra.py | 698 ++++++++++++++++++ .../lab/sim/motion/planners/toppra_planner.py | 467 ++++++++---- pyproject.toml | 2 +- .../motion_generation/planners/toppra.py | 6 +- .../test_toppra_dependency.py | 39 + tests/compute/test_toppra.py | 614 +++++++++++++++ tests/compute/test_toppra_cpu.py | 329 +++++++++ .../sim/motion/planners/test_joint_limits.py | 4 +- .../motion/planners/test_toppra_batched.py | 586 ++++++++++++++- tests/sim/motion/test_motion_imports.py | 39 + 14 files changed, 3533 insertions(+), 175 deletions(-) create mode 100644 embodichain/compute/trajectory/_toppra.py create mode 100644 embodichain/compute/trajectory/_toppra_warp.py create mode 100644 embodichain/compute/trajectory/_warp/toppra.py create mode 100644 tests/benchmark/motion_generation/test_toppra_dependency.py create mode 100644 tests/compute/test_toppra.py create mode 100644 tests/compute/test_toppra_cpu.py diff --git a/agent_context/topics/motion-planning/planner-details.md b/agent_context/topics/motion-planning/planner-details.md index bcdbd3a8f..d0c00e834 100644 --- a/agent_context/topics/motion-planning/planner-details.md +++ b/agent_context/topics/motion-planning/planner-details.md @@ -77,17 +77,54 @@ retain their existing timing contracts and do not consume these backend limits. ### TOPPRA -`toppra_planner.py` fans batched solves into independent CPU work. The module-level -NumPy worker must remain picklable and must not touch CUDA/Warp/simulation state. -Single-worker/single-row execution can run inline. The automatic multiprocessing -context uses spawn with GPU simulation to avoid forking initialized CUDA state. -Inspect the config for overrides; do not force fork merely to reduce overhead. +`toppra_planner.py` selects Warp on CPU and CUDA in automatic mode. NumPy +remains an explicit reference backend. Numerical ownership is in +`compute/trajectory/_toppra.py`, `_toppra_warp.py`, and `_warp/toppra.py`. +Robot-backed planning retains shared speed-cap and control-part resolution, +including gradients of uncapped requested limits. +Both preserve uniform not-a-knot cubic geometry, the existing waypoint +deduplication tolerance and final knot, and signed velocity/acceleration bounds +containing rest. Planner outputs retain float32 and quantity sampling retains +the existing integer coercion, independently of backend or input precision. +Interval bounds are conservative; timing need not match the former external +TOPPRA package. Tiny real moves must retain positive duration. + +Warp constructs splines, bounds, controllable sets and samples on the input +device. It honors the current Torch stream and reads a scalar to size output. +CPU inference amortizes sufficiently large batches over a shared bounded thread +pool; small requests and adjoint-recording calls stay inline. Row blocks own +their numerical buffers, inherit Torch grad/inference modes and merge using +the global endpoint/zero-derivative/zero-dt padding contract. Ordinary inference +launches must not record into Warp's process-global tape. CPU pool state is +discarded after fork; it is not owned by an individual planner's `close()`. +Explicit NumPy selection fans batched solves into independent CPU processes, +which must remain picklable and avoid CUDA/Warp/simulation state. +Single-worker/single-row NumPy execution can run inline. The automatic process +context uses spawn with GPU simulation; do not force fork to reduce overhead. + +Warp records first-order adjoints when Torch gradient tracking is enabled and +waypoints or tensor velocity/acceleration bounds require gradients. Both time +and quantity sampling propagate position, velocity, acceleration and duration +losses through the timing solve. Scalar, per-joint and signed tensor bounds keep +their graphs on either device; TOPPRA option copies preserve these tensor +references while copying mutable containers. Explicit NumPy selection rejects +gradient requests; `torch.no_grad()` retains ordinary forward planning. +Filtering, grid topology, sample counts and active-constraint choices are +discrete: derivatives apply locally away from their switches. Time sampling +keeps `ceil(duration / sample_dt) + 1` points uniformly spread across the full +duration, with the integer count fixed during backward. Failed rows contribute +zero gradients; stationary rows differentiate only their retained position. +Higher-order gradients are unsupported. Test finite-difference agreement, +generator option copying/resampling and CUDA backward stream behavior in the +TOPPRA tests. Variable-length outputs pad with the terminal position without adding duration. Per-row failures preserve other rows; a broken worker pool is discarded and rebuilt on a later call. Test these lifecycle boundaries in -`tests/sim/motion/planners/test_toppra_batched.py` and keep dependency versions -in the package declarations rather than this context. +`tests/sim/motion/planners/test_toppra_batched.py`; numerical limits, geometry, +CPU/Warp agreement and CUDA stream behavior are covered in +`tests/compute/test_toppra.py`; CPU projection, thread isolation, variable-length +assembly and pool inheritance are covered in `tests/compute/test_toppra_cpu.py`. ### Trapezoidal and Double-S @@ -133,6 +170,17 @@ 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. All CPU/CUDA rows run on Warp, +including single-row calls, so batch size does not change waypoint cleanup +precision or the helper's float32/float64 output dtype. It resamples +along the fitted spline, so recompute derived poses from its output rather +than reusing the input samples. Optional per-row path lengths exclude +rectangular tail holds from spline fitting; without them, an explicit repeated +final waypoint retains the historical final-knot geometry. + `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/embodichain/embodichain.lab.sim.motion.planners.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.planners.rst index 98e73e777..0b166f09b 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.planners.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.planners.rst @@ -144,3 +144,17 @@ Utilities :members: :undoc-members: :show-inheritance: + +Joint path retiming +------------------ + +Re-parameterize existing joint paths with the differentiable Warp TOPPRA +backend. The helper preserves float32/float64 input precision, supports valid +per-row waypoint prefixes, and returns explicit timing with isolated failures. + +.. autosummary:: + :nosignatures: + + embodichain.lab.sim.motion.planners.toppra_planner.retime_joint_paths + +.. autofunction:: embodichain.lab.sim.motion.planners.toppra_planner.retime_joint_paths diff --git a/embodichain/compute/trajectory/_toppra.py b/embodichain/compute/trajectory/_toppra.py new file mode 100644 index 000000000..bddb731b0 --- /dev/null +++ b/embodichain/compute/trajectory/_toppra.py @@ -0,0 +1,237 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Internal NumPy reference for rest-to-rest TOPPRA planning. + +The path remains the planner's uniform, not-a-knot cubic spline. Squared +path speeds at adjacent gridpoints obey half-planes a*x + b*y <= c. +Bernstein acceleration bounds and analytic derivative extrema constrain the +whole interval. Backward controllability and a greedy forward pass optimize +this conservative discretization; they do not claim the continuous optimum. +""" + +from __future__ import annotations + +from numbers import Integral + +import numpy as np +from scipy.interpolate import CubicSpline + +__all__ = [] + + +def _limits(value: object, dofs: int) -> np.ndarray: + """Canonicalize symmetric magnitudes or signed lower/upper bounds.""" + array = np.asarray(value, dtype=np.float64) + if array.ndim == 0: + array = np.full(dofs, float(array)) + if array.shape == (dofs,): + array = np.stack((-array, array), axis=-1) + if array.shape != (dofs, 2) or not np.isfinite(array).all(): + raise ValueError("limits must be finite scalar, (DOF,), or (DOF, 2)") + if np.any(array[:, 0] > 0) or np.any(array[:, 1] < 0): + raise ValueError("rest-to-rest limits must contain zero") + return array + + +def _validate_sampling(sample_count: int | None, sample_dt: float | None) -> None: + if (sample_count is None) == (sample_dt is None): + raise ValueError("specify exactly one of sample_count and sample_dt") + if sample_count is not None and ( + isinstance(sample_count, bool) + or not isinstance(sample_count, Integral) + or sample_count < 2 + ): + raise ValueError("sample_count must be an integer >= 2") + if sample_dt is not None and (not np.isfinite(sample_dt) or sample_dt <= 0): + raise ValueError("sample_dt must be finite and positive") + + +def _validate_grid_size(grid_size: int) -> None: + if ( + isinstance(grid_size, bool) + or not isinstance(grid_size, Integral) + or grid_size < 3 + ): + raise ValueError("grid_size must be an integer >= 3") + + +class _NumpyToppra: + """Numerical reference; construction raises on infeasible input rows.""" + + def __init__( + self, + points: np.ndarray, + velocity: object, + acceleration: object, + grid_size: int = 100, + ) -> None: + _validate_grid_size(grid_size) + points = np.asarray(points) + if points.ndim != 2 or min(points.shape) < 1 or points.shape[0] < 2: + raise ValueError("waypoints must have shape (N >= 2, DOF >= 1)") + if not np.isfinite(points).all(): + raise ValueError("waypoints must be finite") + self.vlim = _limits(velocity, points.shape[1]) + self.alim = _limits(acceleration, points.shape[1]) + frozen = np.all(self.vlim == 0, axis=1) + if np.any(points[:, frozen] != points[:1, frozen]): + raise ValueError("zero-speed joints must hold at every waypoint") + # Preserve the planner's historical tolerance and final knot, including + # a repeated endpoint: dropping it changes the fitted cubic geometry. + keep = [0] + for i in range(1, len(points)): + if np.max(np.abs(points[i] - points[keep[-1]])) >= 1e-6: + keep.append(i) + if keep[-1] != len(points) - 1: + keep.append(len(points) - 1) + self.points = points[keep].astype(np.float64) + self.duration = 0.0 + if np.all(self.points == self.points[0]): + return + knot_count = len(self.points) + base_intervals = max(grid_size - 1, 10 * knot_count - 1) + self.path = CubicSpline(np.linspace(0, 1, knot_count), self.points) + subdivisions = (base_intervals + knot_count - 2) // (knot_count - 1) + self.grid = np.linspace(0, 1, (knot_count - 1) * subdivisions + 1) + # Segment-aligned subdivision never creates almost-duplicate knots. + self.a, self.b, self.c = self._constraints() + self._solve() + + def _constraints(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + ds = np.diff(self.grid)[:, None] + first, last = self.path(self.grid[:-1], 1), self.path(self.grid[1:], 1) + second_first, second_last = self.path(self.grid[:-1], 2), self.path( + self.grid[1:], 2 + ) + derivative = np.stack((first, first + 0.5 * ds * second_first, last), axis=1) + zero = np.zeros_like(first) + ax = np.stack((second_first, 0.5 * second_last, zero), axis=1) - derivative / ( + 2 * ds[:, None] + ) + ay = np.stack((zero, 0.5 * second_first, second_last), axis=1) + derivative / ( + 2 * ds[:, None] + ) + # Normalize each sign independently, including zero one-sided bounds. + upper, lower = self.alim[:, 1], -self.alim[:, 0] + us, ls = np.where(upper > 0, upper, 1.0), np.where(lower > 0, lower, 1.0) + shape = (len(ds), -1) + a = np.concatenate( + ((ax / us).reshape(shape), (-ax / ls).reshape(shape)), axis=1 + ) + b = np.concatenate( + ((ay / us).reshape(shape), (-ay / ls).reshape(shape)), axis=1 + ) + c = np.broadcast_to( + np.r_[np.tile(upper / us, 3), np.tile(lower / ls, 3)], a.shape + ).copy() + p0, p1, p2 = np.moveaxis(derivative, 1, 0) + quadratic = p0 - 2 * p1 + p2 + ratio = np.clip( + np.divide(p0 - p1, quadratic, out=np.zeros_like(p0), where=quadratic != 0), + 0, + 1, + ) + extremum = (1 - ratio) ** 2 * p0 + 2 * ratio * (1 - ratio) * p1 + ratio**2 * p2 + high = np.maximum(np.maximum(p0, p2), extremum) + low = np.minimum(np.minimum(p0, p2), extremum) + cap = np.full_like(first, np.inf) + np.divide(self.vlim[:, 1], high, out=cap, where=high > 0) + negative_cap = np.full_like(first, np.inf) + np.divide(self.vlim[:, 0], low, out=negative_cap, where=low < 0) + cap = np.min(np.minimum(cap, negative_cap), axis=1) ** 2 + a = np.column_stack((a, np.ones(len(ds)), np.zeros(len(ds)))) + b = np.column_stack((b, np.zeros(len(ds)), np.ones(len(ds)))) + c = np.column_stack((c, cap, cap)) + return a, b, c + + def _solve(self) -> None: + count = len(self.grid) + controllable = np.zeros(count) + # All limits contain rest, so the lower controllable bound is zero. + for i in range(count - 2, -1, -1): + a, b, c = self.a[i], self.b[i], self.c[i] + positive, negative = b > 0, b < 0 + coefficients = ( + a[negative, None] * b[positive] - b[negative, None] * a[positive] + ) + rhs = c[negative, None] * b[positive] - b[negative, None] * c[positive] + mask = coefficients > 0 + cap = np.min(rhs[mask] / coefficients[mask], initial=np.inf) + bound = c.copy() + bound[negative] -= b[negative] * controllable[i + 1] + mask = a > 0 + cap = min(cap, np.min(bound[mask] / a[mask], initial=np.inf)) + if not np.isfinite(cap) or cap < 0: + raise ValueError("path has no finite controllable speed") + controllable[i] = cap + x = np.zeros(count) + for i in range(count - 1): + a, b, c = self.a[i], self.b[i], self.c[i] + residual = c - a * x[i] + rounding = ( + 16 * np.finfo(float).eps * np.maximum(np.abs(c), np.abs(a * x[i])) + ) + positive = b > 0 + high = min( + controllable[i + 1], + np.min( + (residual[positive] + rounding[positive]) / b[positive], + initial=np.inf, + ), + ) + x[i + 1] = max(0.0, high) + violation = a * x[i] + b * x[i + 1] - c + tolerance = 1e-10 * np.maximum( + np.maximum(np.abs(a * x[i]), np.abs(b * x[i + 1])), np.abs(c) + ) + if np.any(violation > tolerance): + raise ValueError("path forward pass is infeasible") + self.speeds = np.sqrt(x) + ds = np.diff(self.grid) + denominator = self.speeds[:-1] + self.speeds[1:] + if np.any(denominator <= 0): + raise ValueError("path contains a zero-speed interval") + self.accelerations = np.diff(x) / (2 * ds) + self.times = np.r_[0.0, np.cumsum(2 * ds / denominator)] + if not np.isfinite(self.times).all() or np.any(np.diff(self.times) <= 0): + raise ValueError("path timing is not representable") + self.duration = float(self.times[-1]) + + def sample(self, times: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + times = np.clip(np.asarray(times), 0.0, self.duration) + if self.duration == 0: + positions = np.broadcast_to( + self.points[0], (len(times), self.points.shape[1]) + ).copy() + return positions, np.zeros_like(positions), np.zeros_like(positions) + index = np.clip( + np.searchsorted(self.times, times, side="right") - 1, 0, len(self.times) - 2 + ) + elapsed = times - self.times[index] + accel = self.accelerations[index] + speed = np.maximum(0, self.speeds[index] + accel * elapsed) + s = self.grid[index] + self.speeds[index] * elapsed + 0.5 * accel * elapsed**2 + s = np.clip(s, self.grid[index], self.grid[index + 1]) + positions = self.path(s) + velocities = self.path(s, 1) * speed[:, None] + accelerations = ( + self.path(s, 2) * speed[:, None] ** 2 + self.path(s, 1) * accel[:, None] + ) + positions[times == 0] = self.points[0] + positions[times == self.duration] = self.points[-1] + velocities[(times == 0) | (times == self.duration)] = 0 + return positions, velocities, accelerations diff --git a/embodichain/compute/trajectory/_toppra_warp.py b/embodichain/compute/trajectory/_toppra_warp.py new file mode 100644 index 000000000..cb2d5dc30 --- /dev/null +++ b/embodichain/compute/trajectory/_toppra_warp.py @@ -0,0 +1,611 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Internal tensor adapter for the planner's Warp TOPPRA backend. + +Waypoint data and numerical work stay on the input device. Scalar reads +validate and size variable-length output; no path is solved through NumPy. +""" + +from __future__ import annotations + +from contextlib import nullcontext +from concurrent.futures import ThreadPoolExecutor, wait +import os +from threading import Lock +from typing import Any + +import numpy as np +import torch +import warp as wp +from torch.autograd.function import once_differentiable + +from ._toppra import _limits, _validate_grid_size, _validate_sampling +from ._warp import toppra as kernels + +__all__ = [] + +_CPU_POOL: ThreadPoolExecutor | None = None +_CPU_POOL_LOCK = Lock() + + +def _reset_cpu_pool() -> None: + """Discard inherited executor state; fork does not preserve its threads.""" + global _CPU_POOL, _CPU_POOL_LOCK + _CPU_POOL = None + _CPU_POOL_LOCK = Lock() + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_reset_cpu_pool) + + +def _cpu_pool() -> ThreadPoolExecutor: + """Reuse at most eight CPU workers across concurrent inference requests.""" + global _CPU_POOL + with _CPU_POOL_LOCK: + if _CPU_POOL is None: + _CPU_POOL = ThreadPoolExecutor(max_workers=8, thread_name_prefix="toppra") + return _CPU_POOL + + +def _cpu_workers(batch: int, count: int, dofs: int, grid_size: int) -> int: + """Amortize task overhead while respecting the caller's CPU affinity.""" + available = ( + len(os.sched_getaffinity(0)) + if hasattr(os, "sched_getaffinity") + else (os.cpu_count() or 1) + ) + # Long paths contain enough work to use smaller row blocks. Tiny batches + # and low-DoF paths stay inline instead of paying for thread submission. + row_grain = 4 if count >= 32 else 8 + work = batch * max(grid_size - 1, 10 * count - 1) * (6 * dofs + 2) ** 2 + return max(1, min(8, available, batch // row_grain, work // 1_000_000)) + + +def _cpu_chunk( + waypoints: torch.Tensor, + velocity: torch.Tensor, + acceleration: torch.Tensor, + path_lengths: torch.Tensor, + sample_count: int | None, + sample_dt: float | None, + grid_size: int, + grad_enabled: bool, + inference_mode: bool, +) -> dict[str, torch.Tensor]: + """Run a row block with the caller's Torch thread-local execution modes.""" + with torch.inference_mode(inference_mode), torch.set_grad_enabled(grad_enabled): + return _launch_toppra( + waypoints, + velocity, + acceleration, + path_lengths, + sample_count=sample_count, + sample_dt=sample_dt, + grid_size=grid_size, + ) + + +def _launch_cpu_parallel( + waypoints: torch.Tensor, + velocity: torch.Tensor, + acceleration: torch.Tensor, + path_lengths: torch.Tensor, + workers: int, + sample_count: int | None, + sample_dt: float | None, + grid_size: int, +) -> dict[str, torch.Tensor]: + """Compute independent row blocks and restore the global padding contract.""" + wp.init() + # Compile/load once on the submitting thread, before concurrent launches. + wp.load_module(module=kernels, device="cpu") + pool = _cpu_pool() + block = (len(waypoints) + workers - 1) // workers + futures = [ + pool.submit( + _cpu_chunk, + waypoints[start : start + block], + velocity, + acceleration, + path_lengths[start : start + block], + sample_count, + sample_dt, + grid_size, + torch.is_grad_enabled(), + torch.is_inference_mode_enabled(), + ) + for start in range(0, len(waypoints), block) + ] + try: + results = [future.result() for future in futures] + except BaseException: + for future in futures: + future.cancel() + wait(futures) + raise + samples = max(result["positions"].shape[1] for result in results) + merged = {"success": torch.cat([result["success"] for result in results])} + for key in ("positions", "velocities", "accelerations", "dt"): + parts = [] + for result in results: + value = result[key] + extra = samples - value.shape[1] + if extra: + shape = (value.shape[0], extra, *value.shape[2:]) + padding = ( + value[:, -1:].expand(shape) + if key == "positions" + else value.new_zeros(shape) + ) + value = torch.cat((value, padding), dim=1) + parts.append(value) + merged[key] = torch.cat(parts) + return merged + + +class _ToppraAutograd(torch.autograd.Function): + """Bridge Warp's first-order adjoints into the caller's Torch graph.""" + + @staticmethod + def forward( + ctx: Any, + waypoints: torch.Tensor, + velocity: torch.Tensor, + acceleration: torch.Tensor, + path_lengths: torch.Tensor, + sample_count: int | None, + sample_dt: float | None, + grid_size: int, + ) -> tuple[torch.Tensor, ...]: + state: dict[str, Any] = {} + result = _launch_toppra( + waypoints, + velocity, + acceleration, + path_lengths, + sample_count=sample_count, + sample_dt=sample_dt, + grid_size=grid_size, + state=state, + ) + ctx.state = state + ctx.save_for_backward(waypoints, velocity, acceleration) + ctx.set_materialize_grads(False) + ctx.mark_non_differentiable(result["success"]) + return tuple( + result[key] + for key in ("success", "positions", "velocities", "accelerations", "dt") + ) + + @staticmethod + @once_differentiable + def backward( + ctx: Any, success_grad: None, *output_grads: torch.Tensor | None + ) -> tuple: + saved = ctx.saved_tensors + waypoints = saved[0] + state = ctx.state + stream = ( + wp.ScopedStream( + wp.stream_from_torch(torch.cuda.current_stream(waypoints.device)) + ) + if waypoints.is_cuda + else nullcontext() + ) + with stream: + # A retained Torch graph may be differentiated repeatedly. Copy the + # returned input gradient before the next backward clears the tape. + state["tape"].zero() + seeds = { + array: wp.from_torch( + grad.to(torch.float64).contiguous(), dtype=wp.float64 + ) + for array, grad in zip(state["outputs"], output_grads) + if grad is not None + } + if seeds: + state["tape"].backward(grads=seeds) + gradients = tuple( + wp.to_torch(array.grad).to(tensor.dtype).clone() if needed else None + for array, tensor, needed in zip( + state["inputs"], saved, ctx.needs_input_grad + ) + ) + return *gradients, None, None, None, None + + +def _retime_toppra_warp( + waypoints: torch.Tensor, + velocity: object, + acceleration: object, + *, + path_lengths: torch.Tensor | None = None, + sample_count: int | None = None, + sample_dt: float | None = None, + grid_size: int = 100, +) -> dict[str, torch.Tensor]: + """Retime (B, N, DOF) paths, isolating invalid/infeasible rows. + + Failed rows contain zeros. Stationary rows have two zero-time samples. + ``path_lengths`` can exclude rectangular tail padding independently for + each row. Shorter successful outputs pad with their endpoint and zero + derivatives/dt. + The floating outputs use the input dtype and device. Computation uses + float64 on both Warp CPU and CUDA, and honors the current Torch stream. + CPU inference uses a shared, bounded thread pool for sufficiently large + batches; small requests and all adjoint-recording calls stay inline. + + Waypoints and tensor motion limits support first-order Warp adjoints on + CPU and CUDA with either sampling mode. Derivatives hold discrete filtering, + grid topology, sample counts and active constraints fixed, and are not + guaranteed at their switches. In time mode, sample_dt sets an upper bound + on spacing; samples remain uniform over the full solved duration. + Failed rows contribute zero gradients; stationary rows differentiate only + the retained position. Higher-order derivatives are unsupported. + """ + _validate_sampling(sample_count, sample_dt) + _validate_grid_size(grid_size) + if waypoints.ndim != 3 or min(waypoints.shape) < 1 or waypoints.shape[1] < 2: + raise ValueError("waypoints must have shape (B >= 1, N >= 2, DOF >= 1)") + if waypoints.dtype not in (torch.float32, torch.float64): + raise ValueError("waypoints must have float32 or float64 dtype") + velocity = _tensor_limits(velocity, waypoints.shape[2], waypoints.device) + acceleration = _tensor_limits(acceleration, waypoints.shape[2], waypoints.device) + path_lengths = _tensor_path_lengths( + path_lengths, waypoints.shape[0], waypoints.shape[1], waypoints.device + ) + if torch.is_grad_enabled() and any( + tensor.requires_grad for tensor in (waypoints, velocity, acceleration) + ): + values = _ToppraAutograd.apply( + waypoints, + velocity, + acceleration, + path_lengths, + sample_count, + sample_dt, + grid_size, + ) + return dict( + zip(("success", "positions", "velocities", "accelerations", "dt"), values) + ) + if waypoints.device.type == "cpu": + workers = _cpu_workers(*waypoints.shape, grid_size) + if workers > 1: + return _launch_cpu_parallel( + waypoints, + velocity, + acceleration, + path_lengths, + workers, + sample_count, + sample_dt, + grid_size, + ) + return _launch_toppra( + waypoints, + velocity, + acceleration, + path_lengths, + sample_count=sample_count, + sample_dt=sample_dt, + grid_size=grid_size, + ) + + +def _tensor_path_lengths( + value: torch.Tensor | None, + batch: int, + count: int, + device: torch.device, +) -> torch.Tensor: + """Validate per-row valid waypoint counts on the input device.""" + if value is None: + return torch.full((batch,), count, dtype=torch.int32, device=device) + if not isinstance(value, torch.Tensor) or value.shape != (batch,): + raise ValueError(f"path_lengths must have shape ({batch},)") + if value.dtype not in ( + torch.uint8, + torch.int8, + torch.int16, + torch.int32, + torch.int64, + ): + raise ValueError("path_lengths must contain integers") + lengths = value.to(device=device, dtype=torch.int32).contiguous() + if bool(((lengths < 2) | (lengths > count)).any().item()): + raise ValueError(f"path_lengths must be between 2 and {count}") + return lengths + + +def _tensor_limits(value: object, dofs: int, device: torch.device) -> torch.Tensor: + """Normalize tensor limits without detaching their gradient or visiting CPU.""" + if not isinstance(value, torch.Tensor): + # Keep Python/NumPy validation on the host; no extra CUDA round trips + # are needed for the existing constant-limit planning interface. + return torch.as_tensor(_limits(value, dofs), dtype=torch.float64, device=device) + bounds = value.to(device=device, dtype=torch.float64) + if bounds.ndim == 0: + bounds = bounds.expand(dofs) + if bounds.shape == (dofs,): + bounds = torch.stack((-bounds, bounds), dim=-1) + if bounds.shape != (dofs, 2): + raise ValueError("limits must be finite scalar, (DOF,), or (DOF, 2)") + if not bool( + ( + torch.isfinite(bounds).all() + & (bounds[:, 0] <= 0).all() + & (bounds[:, 1] >= 0).all() + ).item() + ): + raise ValueError("limits must be finite and contain zero") + return bounds.contiguous() + + +def _launch_toppra( + waypoints: torch.Tensor, + velocity: torch.Tensor, + acceleration: torch.Tensor, + path_lengths: torch.Tensor, + *, + sample_count: int | None = None, + sample_dt: float | None = None, + grid_size: int = 100, + state: dict[str, Any] | None = None, +) -> dict[str, torch.Tensor]: + """Run the shared forward kernels, recording a tape only for autograd.""" + batch, count, dofs = waypoints.shape + # NumPy 1.x promotes a float32 scalar comparison to float64; NumPy 2.x + # keeps float32. Match the installed CPU backend at the dedup boundary. + scalar_type = np.float32 if waypoints.dtype == torch.float32 else np.float64 + duplicate_tolerance = np.asarray( + 1e-6, dtype=np.result_type(scalar_type(0), 1e-6) + ).item() + wp.init() + device = wp.device_from_torch(waypoints.device) + stream = ( + wp.ScopedStream( + wp.stream_from_torch(torch.cuda.current_stream(waypoints.device)) + ) + if waypoints.is_cuda + else nullcontext() + ) + differentiable = state is not None + tape = wp.Tape() if differentiable else nullcontext() + with stream: + tensor = waypoints.detach().to(dtype=torch.float64).contiguous() + points = wp.from_torch(tensor, dtype=wp.float64, requires_grad=differentiable) + clean = wp.zeros((batch, count, dofs), dtype=wp.float64, device=device) + kept = wp.zeros((batch, count), dtype=int, device=device) + counts = wp.zeros(batch, dtype=int, device=device) + intervals = wp.zeros(batch, dtype=int, device=device) + success = wp.zeros(batch, dtype=int, device=device) + base = max(grid_size - 1, 10 * count - 1) + capacity = base + count - 1 + wp.launch( + kernels._prepare, + dim=batch, + inputs=[ + points, + wp.from_torch(path_lengths), + wp.from_torch(velocity.detach(), dtype=wp.float64, requires_grad=False), + clean, + kept, + counts, + intervals, + success, + grid_size, + int(waypoints.dtype == torch.float32), + wp.float64(duplicate_tolerance), + ], + device=device, + record_tape=False, + ) + with tape: + if differentiable: + clean = wp.zeros( + (batch, count, dofs), + dtype=wp.float64, + device=device, + requires_grad=True, + ) + wp.launch( + kernels._gather, + dim=clean.shape, + inputs=[points, kept, counts, success, clean], + device=device, + record_tape=differentiable, + ) + diagonal = wp.zeros_like(clean, requires_grad=False) + eliminated = wp.zeros_like(clean) + second = wp.zeros_like(clean) + coefficients = wp.zeros( + (batch, count - 1, dofs, 4), + dtype=wp.float64, + device=device, + requires_grad=differentiable, + ) + wp.launch( + kernels._spline, + dim=(batch, dofs), + inputs=[ + clean, + counts, + success, + diagonal, + eliminated, + second, + coefficients, + ], + device=device, + record_tape=differentiable, + ) + velocity_wp = wp.from_torch( + velocity.detach(), dtype=wp.float64, requires_grad=differentiable + ) + acceleration_wp = wp.from_torch( + acceleration.detach(), dtype=wp.float64, requires_grad=differentiable + ) + velocity_caps = wp.zeros( + (batch, capacity, dofs), + dtype=wp.float64, + device=device, + requires_grad=differentiable, + ) + a = wp.zeros( + (batch, capacity, 6 * dofs + 2), + dtype=wp.float64, + device=device, + requires_grad=differentiable, + ) + b, c = wp.zeros_like(a), wp.zeros_like(a) + wp.launch( + kernels._constraints, + dim=(batch, capacity), + inputs=[ + coefficients, + counts, + intervals, + success, + velocity_wp, + acceleration_wp, + velocity_caps, + a, + b, + c, + ], + device=device, + record_tape=differentiable, + ) + wp.launch( + kernels._velocity_caps, + dim=(batch, capacity), + inputs=[velocity_caps, intervals, success, a, b, c], + device=device, + record_tape=differentiable, + ) + caps = wp.zeros( + (batch, capacity), + dtype=wp.float64, + device=device, + requires_grad=differentiable, + ) + wp.launch( + ( + kernels._project_caps_cpu + if device.is_cpu and not differentiable + else kernels._project_caps + ), + dim=(batch, capacity), + inputs=[a, b, c, intervals, caps], + device=device, + record_tape=differentiable, + ) + controllable = wp.zeros( + (batch, capacity + 1), + dtype=wp.float64, + device=device, + requires_grad=differentiable, + ) + x, times = wp.zeros_like(controllable), wp.zeros_like(controllable) + durations = wp.zeros( + batch, dtype=wp.float64, device=device, requires_grad=differentiable + ) + wp.launch( + kernels._solve, + dim=batch, + inputs=[ + a, + b, + c, + intervals, + caps, + controllable, + x, + times, + durations, + success, + ], + device=device, + record_tape=differentiable, + ) + duration_tensor = wp.to_torch(durations) + ok = wp.to_torch(success).bool() + moving = ok & (wp.to_torch(counts) > 1) + if sample_count is not None: + sample_lengths = torch.full( + (batch,), sample_count, device=waypoints.device, dtype=torch.int32 + ) + else: + # Integer output lengths are frozen for this forward/backward + # pair; timing and sampling still differentiate through duration. + desired = torch.ceil(duration_tensor.detach() / sample_dt) + 1 + if bool((desired > torch.iinfo(torch.int32).max).any().item()): + raise ValueError("sample_dt requests too many samples") + sample_lengths = desired.to(torch.int32).clamp_min(2) + sample_lengths = torch.where(moving, sample_lengths, 2).contiguous() + samples = int(sample_lengths.max().item()) + positions = wp.zeros( + (batch, samples, dofs), + dtype=wp.float64, + device=device, + requires_grad=differentiable, + ) + velocities, accelerations = wp.zeros_like(positions), wp.zeros_like( + positions + ) + dt = wp.zeros( + (batch, samples), + dtype=wp.float64, + device=device, + requires_grad=differentiable, + ) + wp.launch( + kernels._sample, + dim=(batch, samples, dofs), + inputs=[ + clean, + coefficients, + counts, + intervals, + success, + x, + times, + durations, + wp.from_torch(sample_lengths), + positions, + velocities, + accelerations, + dt, + ], + device=device, + record_tape=differentiable, + ) + if state is not None: + state.update( + tape=tape, + inputs=(points, velocity_wp, acceleration_wp), + outputs=(positions, velocities, accelerations, dt), + ) + return { + "success": ok, + "positions": wp.to_torch(positions).to(waypoints.dtype), + "velocities": wp.to_torch(velocities).to(waypoints.dtype), + "accelerations": wp.to_torch(accelerations).to(waypoints.dtype), + "dt": wp.to_torch(dt).to(waypoints.dtype), + } diff --git a/embodichain/compute/trajectory/_warp/toppra.py b/embodichain/compute/trajectory/_warp/toppra.py new file mode 100644 index 000000000..3f0c58379 --- /dev/null +++ b/embodichain/compute/trajectory/_warp/toppra.py @@ -0,0 +1,698 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Double-precision TOPPRA kernels; parallelism is across paths and intervals. + +Integer searches live in separate functions so Warp replays their selected +indices before differentiating the active expression. Loop-carried min/max +accumulators do not retain the intermediate values needed by the adjoint. +""" + +from __future__ import annotations + +import warp as wp + +__all__ = [] + + +@wp.kernel +def _prepare( + points: wp.array3d(dtype=wp.float64), + path_lengths: wp.array(dtype=int), + velocity: wp.array2d(dtype=wp.float64), + clean: wp.array3d(dtype=wp.float64), + kept: wp.array2d(dtype=int), + counts: wp.array(dtype=int), + intervals: wp.array(dtype=int), + success: wp.array(dtype=int), + grid_size: int, + single_precision: int, + duplicate_tolerance: wp.float64, +): + row = wp.tid() + count = int(0) + last_kept = int(-1) + valid = int(1) + path_count = path_lengths[row] + for i in range(points.shape[1]): + if i < path_count: + different = int(0) + for j in range(points.shape[2]): + value = points[row, i, j] + if not wp.isfinite(value): + valid = 0 + # Preserve explicit runtime zero-speed joints before waypoint + # filtering can hide a sub-tolerance interior excursion. + if ( + velocity[j, 0] == wp.float64(0.0) + and velocity[j, 1] == wp.float64(0.0) + and value != points[row, 0, j] + ): + valid = 0 + if count == 0: + different = 1 + elif single_precision != 0: + # Match the legacy subtraction and threshold rounding before + # the numerical solve is promoted to double precision. + if ( + wp.float64( + wp.abs( + wp.float32(value) - wp.float32(clean[row, count - 1, j]) + ) + ) + >= duplicate_tolerance + ): + different = 1 + elif wp.abs(value - clean[row, count - 1, j]) >= duplicate_tolerance: + different = 1 + if different != 0: + kept[row, count] = i + for j in range(points.shape[2]): + clean[row, count, j] = points[row, i, j] + count += 1 + last_kept = i + last_index = path_count - 1 + if last_kept != last_index: + kept[row, count] = last_index + for j in range(points.shape[2]): + clean[row, count, j] = points[row, last_index, j] + count += 1 + if count == 2: + stationary = int(1) + for j in range(points.shape[2]): + if clean[row, 0, j] != clean[row, 1, j]: + stationary = 0 + if stationary != 0: + count = 1 + counts[row] = count + success[row] = valid + if count > 1 and valid != 0: + base = wp.max(grid_size - 1, 10 * count - 1) + subdivisions = (base + count - 2) // (count - 1) + intervals[row] = (count - 1) * subdivisions + + +@wp.kernel +def _gather( + points: wp.array3d(dtype=wp.float64), + kept: wp.array2d(dtype=int), + counts: wp.array(dtype=int), + success: wp.array(dtype=int), + clean: wp.array3d(dtype=wp.float64), +): + row, i, joint = wp.tid() + if success[row] != 0 and i < counts[row]: + clean[row, i, joint] = points[row, kept[row, i], joint] + + +@wp.func +def _speed(x: wp.float64): + # Zero speed at the fixed rest boundaries is constant. Avoid 0 * inf in + # the adjoint of sqrt; degenerate interior stops use the same convention. + value = wp.float64(0.0) + if x > wp.float64(0.0): + value = wp.sqrt(x) + return value + + +@wp.kernel +def _spline( + points: wp.array3d(dtype=wp.float64), + counts: wp.array(dtype=int), + success: wp.array(dtype=int), + diagonal: wp.array3d(dtype=wp.float64), + eliminated: wp.array3d(dtype=wp.float64), + second: wp.array3d(dtype=wp.float64), + coefficients: wp.array4d(dtype=wp.float64), +): + row, joint = wp.tid() + n = counts[row] + if success[row] == 0 or n < 2: + return + h = wp.float64(1.0) / wp.float64(n - 1) + if n == 3: + value = ( + points[row, 2, joint] + - wp.float64(2.0) * points[row, 1, joint] + + points[row, 0, joint] + ) / (h * h) + for i in range(3): + second[row, i, joint] = value + elif n > 3: + # Uniform not-a-knot boundaries eliminate the endpoint unknowns. + # Keep elimination RHS values separate from back substitution so the + # adjoint can read every forward intermediate without overwritten data. + for i in range(1, n - 1): + diag = wp.float64(4.0) + rhs = ( + wp.float64(6.0) + * ( + points[row, i + 1, joint] + - wp.float64(2.0) * points[row, i, joint] + + points[row, i - 1, joint] + ) + / (h * h) + ) + if i == 1 or i == n - 2: + diag = wp.float64(6.0) + if i > 1 and i < n - 2: + factor = wp.float64(1.0) / diagonal[row, i - 1, joint] + if i > 2: + diag -= factor + rhs -= factor * eliminated[row, i - 1, joint] + diagonal[row, i, joint] = diag + eliminated[row, i, joint] = rhs + for k in range(n - 2): + i = n - 2 - k + rhs = eliminated[row, i, joint] + wp.float64(0.0) + if i > 1 and i < n - 2: + rhs -= second[row, i + 1, joint] + second[row, i, joint] = rhs / diagonal[row, i, joint] + second[row, 0, joint] = ( + wp.float64(2.0) * second[row, 1, joint] - second[row, 2, joint] + ) + second[row, n - 1, joint] = ( + wp.float64(2.0) * second[row, n - 2, joint] - second[row, n - 3, joint] + ) + for i in range(n - 1): + left, right = second[row, i, joint], second[row, i + 1, joint] + coefficients[row, i, joint, 0] = points[row, i, joint] + coefficients[row, i, joint, 1] = ( + points[row, i + 1, joint] - points[row, i, joint] + ) / h - h * (wp.float64(2.0) * left + right) / wp.float64(6.0) + coefficients[row, i, joint, 2] = left / wp.float64(2.0) + coefficients[row, i, joint, 3] = (right - left) / (wp.float64(6.0) * h) + + +@wp.func +def _joint_constraints( + coefficients: wp.array4d(dtype=wp.float64), + velocity: wp.array2d(dtype=wp.float64), + acceleration: wp.array2d(dtype=wp.float64), + a: wp.array3d(dtype=wp.float64), + b: wp.array3d(dtype=wp.float64), + c: wp.array3d(dtype=wp.float64), + row: int, + i: int, + joint: int, + segment: int, + delta: wp.float64, + ds: wp.float64, +): + dofs = coefficients.shape[2] + cap = wp.float64(wp.inf) + cb = coefficients[row, segment, joint, 1] + cc = coefficients[row, segment, joint, 2] + cd = coefficients[row, segment, joint, 3] + end = delta + ds + first = cb + delta * (wp.float64(2.0) * cc + wp.float64(3.0) * delta * cd) + last = cb + end * (wp.float64(2.0) * cc + wp.float64(3.0) * end * cd) + second_first = wp.float64(2.0) * cc + wp.float64(6.0) * delta * cd + second_last = wp.float64(2.0) * cc + wp.float64(6.0) * end * cd + middle = first + wp.float64(0.5) * ds * second_first + quadratic = first - wp.float64(2.0) * middle + last + ratio = wp.float64(0.0) + if quadratic != wp.float64(0.0): + ratio = wp.clamp((first - middle) / quadratic, wp.float64(0.0), wp.float64(1.0)) + one = wp.float64(1.0) - ratio + extremum = ( + one * one * first + + wp.float64(2.0) * ratio * one * middle + + ratio * ratio * last + ) + high, low = wp.max(wp.max(first, last), extremum), wp.min( + wp.min(first, last), extremum + ) + if high > wp.float64(0.0): + limit = velocity[joint, 1] / high + cap = wp.min(cap, limit * limit) + if low < wp.float64(0.0): + limit = velocity[joint, 0] / low + cap = wp.min(cap, limit * limit) + upper, lower = acceleration[joint, 1], -acceleration[joint, 0] + us = wp.where(upper > wp.float64(0.0), upper, wp.float64(1.0)) + ls = wp.where(lower > wp.float64(0.0), lower, wp.float64(1.0)) + for k in range(3): + deriv, sx, sy = first, second_first, wp.float64(0.0) + if k == 1: + deriv, sx, sy = ( + middle, + wp.float64(0.5) * second_last, + wp.float64(0.5) * second_first, + ) + elif k == 2: + deriv, sx, sy = last, wp.float64(0.0), second_last + ax = sx - deriv / (wp.float64(2.0) * ds) + ay = sy + deriv / (wp.float64(2.0) * ds) + index = k * dofs + joint + a[row, i, index] = ax / us + b[row, i, index] = ay / us + c[row, i, index] = upper / us + index += 3 * dofs + a[row, i, index] = -ax / ls + b[row, i, index] = -ay / ls + c[row, i, index] = lower / ls + return cap + + +@wp.kernel +def _constraints( + coefficients: wp.array4d(dtype=wp.float64), + counts: wp.array(dtype=int), + intervals: wp.array(dtype=int), + success: wp.array(dtype=int), + velocity: wp.array2d(dtype=wp.float64), + acceleration: wp.array2d(dtype=wp.float64), + velocity_caps: wp.array3d(dtype=wp.float64), + a: wp.array3d(dtype=wp.float64), + b: wp.array3d(dtype=wp.float64), + c: wp.array3d(dtype=wp.float64), +): + row, i = wp.tid() + total = intervals[row] + if success[row] == 0 or i >= total: + return + dofs = coefficients.shape[2] + subdivisions = total // (counts[row] - 1) + segment = i // subdivisions + ds = wp.float64(1.0) / wp.float64(total) + delta = wp.float64(i % subdivisions) * ds + for joint in range(dofs): + velocity_caps[row, i, joint] = _joint_constraints( + coefficients, + velocity, + acceleration, + a, + b, + c, + row, + i, + joint, + segment, + delta, + ds, + ) + + +@wp.func +def _velocity_active(velocity_caps: wp.array3d(dtype=wp.float64), row: int, i: int): + active = int(0) + for joint in range(1, velocity_caps.shape[2]): + if velocity_caps[row, i, joint] < velocity_caps[row, i, active]: + active = joint + return active + + +@wp.kernel +def _velocity_caps( + velocity_caps: wp.array3d(dtype=wp.float64), + intervals: wp.array(dtype=int), + success: wp.array(dtype=int), + a: wp.array3d(dtype=wp.float64), + b: wp.array3d(dtype=wp.float64), + c: wp.array3d(dtype=wp.float64), +): + row, i = wp.tid() + if success[row] == 0 or i >= intervals[row]: + return + active = _velocity_active(velocity_caps, row, i) + cap = velocity_caps[row, i, active] + index = 6 * velocity_caps.shape[2] + a[row, i, index] = wp.float64(1.0) + b[row, i, index] = wp.float64(0.0) + c[row, i, index] = cap + a[row, i, index + 1] = wp.float64(0.0) + b[row, i, index + 1] = wp.float64(1.0) + c[row, i, index + 1] = cap + + +@wp.func +def _projection_active( + a: wp.array3d(dtype=wp.float64), + b: wp.array3d(dtype=wp.float64), + c: wp.array3d(dtype=wp.float64), + row: int, + i: int, +): + cap = wp.float64(wp.inf) + active_j, active_k = int(-1), int(-1) + for j in range(a.shape[2]): + aj, bj, cj = a[row, i, j], b[row, i, j], c[row, i, j] + if bj >= wp.float64(0.0) and aj > wp.float64(0.0): + if cj / aj < cap: + cap = cj / aj + active_j, active_k = j, -1 + if bj < wp.float64(0.0): + for k in range(a.shape[2]): + bk = b[row, i, k] + if bk > wp.float64(0.0): + coefficient = aj * bk - bj * a[row, i, k] + if coefficient > wp.float64(0.0): + rhs = cj * bk - bj * c[row, i, k] + if rhs / coefficient < cap: + cap = rhs / coefficient + active_j, active_k = j, k + return active_j, active_k + + +@wp.kernel +def _project_caps( + a: wp.array3d(dtype=wp.float64), + b: wp.array3d(dtype=wp.float64), + c: wp.array3d(dtype=wp.float64), + intervals: wp.array(dtype=int), + caps: wp.array2d(dtype=wp.float64), +): + row, i = wp.tid() + if i >= intervals[row]: + return + active_j, active_k = _projection_active(a, b, c, row, i) + value = wp.float64(wp.inf) + if active_j >= 0: + selected_aj, selected_bj, selected_cj = ( + a[row, i, active_j], + b[row, i, active_j], + c[row, i, active_j], + ) + if active_k < 0: + value = selected_cj / selected_aj + else: + selected_ak, selected_bk, selected_ck = ( + a[row, i, active_k], + b[row, i, active_k], + c[row, i, active_k], + ) + value = (selected_cj * selected_bk - selected_bj * selected_ck) / ( + selected_aj * selected_bk - selected_bj * selected_ak + ) + caps[row, i] = value + + +@wp.kernel(enable_backward=False) +def _project_caps_cpu( + a: wp.array3d(dtype=wp.float64), + b: wp.array3d(dtype=wp.float64), + c: wp.array3d(dtype=wp.float64), + intervals: wp.array(dtype=int), + caps: wp.array2d(dtype=wp.float64), +): + """Project the 2-D feasible polygon by descending active intersections. + + Rest is feasible, so each positive-x pair bound is an upper bound on the + optimum. Start at the independent upper bound, find the tightest lower + and upper y boundaries there, and descend to their intersection until + y is feasible. Each pass scans O(D) planes instead of enumerating all + O(D**2) pairs. A rounding-stalled or nonfinite intersection uses the + exhaustive reference projection; this also bounds the iteration count. + """ + row, i = wp.tid() + if i >= intervals[row]: + return + cap = wp.float64(wp.inf) + for j in range(a.shape[2]): + seed_a, seed_b = a[row, i, j], b[row, i, j] + if seed_a > wp.float64(0.0) and seed_b >= wp.float64(0.0): + cap = wp.min(cap, c[row, i, j] / seed_a) + feasible = int(0) + for iteration in range(a.shape[2]): + if not wp.isfinite(cap): + break + lower = wp.float64(0.0) + upper = wp.float64(wp.inf) + lower_index, upper_index = int(-1), int(-1) + for j in range(a.shape[2]): + plane_b = b[row, i, j] + if plane_b != wp.float64(0.0): + ordinate = (c[row, i, j] - a[row, i, j] * cap) / plane_b + if plane_b < wp.float64(0.0) and ordinate > lower: + lower, lower_index = ordinate, j + elif plane_b > wp.float64(0.0) and ordinate < upper: + upper, upper_index = ordinate, j + if lower <= upper: + feasible = 1 + break + if lower_index < 0 or upper_index < 0: + break + aj, bj, cj = ( + a[row, i, lower_index], + b[row, i, lower_index], + c[row, i, lower_index], + ) + ak, bk, ck = ( + a[row, i, upper_index], + b[row, i, upper_index], + c[row, i, upper_index], + ) + coefficient = aj * bk - bj * ak + if coefficient <= wp.float64(0.0): + break + bound = (cj * bk - bj * ck) / coefficient + if not wp.isfinite(bound) or bound < wp.float64(0.0) or bound >= cap: + break + cap = bound + if feasible == 0: + active_j, active_k = _projection_active(a, b, c, row, i) + cap = wp.float64(wp.inf) + if active_j >= 0: + fallback_a = a[row, i, active_j] + fallback_b = b[row, i, active_j] + fallback_c = c[row, i, active_j] + if active_k < 0: + cap = fallback_c / fallback_a + else: + other_a = a[row, i, active_k] + other_b = b[row, i, active_k] + other_c = c[row, i, active_k] + cap = (fallback_c * other_b - fallback_b * other_c) / ( + fallback_a * other_b - fallback_b * other_a + ) + caps[row, i] = cap + + +@wp.func +def _controllable_active( + a: wp.array3d(dtype=wp.float64), + b: wp.array3d(dtype=wp.float64), + c: wp.array3d(dtype=wp.float64), + caps: wp.array2d(dtype=wp.float64), + controllable: wp.array2d(dtype=wp.float64), + row: int, + i: int, +): + cap = caps[row, i] + wp.float64(0.0) + active = int(-1) + for j in range(a.shape[2]): + aj, bj = a[row, i, j], b[row, i, j] + if aj > wp.float64(0.0) and bj < wp.float64(0.0): + bound = (c[row, i, j] - bj * controllable[row, i + 1]) / aj + if bound < cap: + cap = bound + active = j + return active + + +@wp.func +def _controllable_step( + a: wp.array3d(dtype=wp.float64), + b: wp.array3d(dtype=wp.float64), + c: wp.array3d(dtype=wp.float64), + caps: wp.array2d(dtype=wp.float64), + controllable: wp.array2d(dtype=wp.float64), + row: int, + i: int, +): + active = _controllable_active(a, b, c, caps, controllable, row, i) + value = caps[row, i] + wp.float64(0.0) + if active >= 0: + value = (c[row, i, active] - b[row, i, active] * controllable[row, i + 1]) / a[ + row, i, active + ] + return value + + +@wp.func +def _forward_bound(a: wp.float64, b: wp.float64, c: wp.float64, x: wp.float64): + product = a * x + rounding = wp.float64(3.552713678800501e-15) * wp.max(wp.abs(product), wp.abs(c)) + return (c - product + rounding) / b + + +@wp.func +def _reachable_active( + a: wp.array3d(dtype=wp.float64), + b: wp.array3d(dtype=wp.float64), + c: wp.array3d(dtype=wp.float64), + controllable: wp.array2d(dtype=wp.float64), + x: wp.array2d(dtype=wp.float64), + row: int, + i: int, +): + high = controllable[row, i + 1] + wp.float64(0.0) + active = int(-1) + for j in range(a.shape[2]): + if b[row, i, j] > wp.float64(0.0): + bound = _forward_bound(a[row, i, j], b[row, i, j], c[row, i, j], x[row, i]) + if bound < high: + high = bound + active = j + return active + + +@wp.func +def _reachable_step( + a: wp.array3d(dtype=wp.float64), + b: wp.array3d(dtype=wp.float64), + c: wp.array3d(dtype=wp.float64), + controllable: wp.array2d(dtype=wp.float64), + x: wp.array2d(dtype=wp.float64), + row: int, + i: int, +): + active = _reachable_active(a, b, c, controllable, x, row, i) + value = controllable[row, i + 1] + wp.float64(0.0) + if active >= 0: + value = _forward_bound( + a[row, i, active], b[row, i, active], c[row, i, active], x[row, i] + ) + return wp.max(wp.float64(0.0), value) + + +@wp.kernel +def _solve( + a: wp.array3d(dtype=wp.float64), + b: wp.array3d(dtype=wp.float64), + c: wp.array3d(dtype=wp.float64), + intervals: wp.array(dtype=int), + caps: wp.array2d(dtype=wp.float64), + controllable: wp.array2d(dtype=wp.float64), + x: wp.array2d(dtype=wp.float64), + times: wp.array2d(dtype=wp.float64), + durations: wp.array(dtype=wp.float64), + success: wp.array(dtype=int), +): + row = wp.tid() + total = intervals[row] + if success[row] == 0 or total == 0: + return + for reverse in range(total): + i = total - 1 - reverse + cap = _controllable_step(a, b, c, caps, controllable, row, i) + if not wp.isfinite(cap) or cap < wp.float64(0.0): + success[row] = 0 + return + controllable[row, i] = cap + for i in range(total): + current = x[row, i] + high = _reachable_step(a, b, c, controllable, x, row, i) + for j in range(a.shape[2]): + left, right, check_bound = ( + a[row, i, j] * current, + b[row, i, j] * high, + c[row, i, j], + ) + tolerance = wp.float64(1e-10) * wp.max( + wp.max(wp.abs(left), wp.abs(right)), wp.abs(check_bound) + ) + if left + right - check_bound > tolerance: + success[row] = 0 + return + x[row, i + 1] = high + denominator = _speed(current) + _speed(high) + if denominator <= wp.float64(0.0): + success[row] = 0 + return + arrival = times[row, i] + wp.float64(2.0) / (wp.float64(total) * denominator) + if not wp.isfinite(arrival) or arrival <= times[row, i]: + success[row] = 0 + return + times[row, i + 1] = arrival + durations[row] = times[row, total] + + +@wp.func +def _sample_interval( + times: wp.array2d(dtype=wp.float64), row: int, total: int, query: wp.float64 +): + lo = int(0) + hi = total + 0 + while lo + 1 < hi: + mid = (lo + hi) // 2 + if times[row, mid] <= query: + lo = mid + else: + hi = mid + return wp.min(lo, total - 1) + + +@wp.kernel +def _sample( + points: wp.array3d(dtype=wp.float64), + coefficients: wp.array4d(dtype=wp.float64), + counts: wp.array(dtype=int), + intervals: wp.array(dtype=int), + success: wp.array(dtype=int), + x: wp.array2d(dtype=wp.float64), + times: wp.array2d(dtype=wp.float64), + durations: wp.array(dtype=wp.float64), + sample_counts: wp.array(dtype=int), + positions: wp.array3d(dtype=wp.float64), + velocities: wp.array3d(dtype=wp.float64), + accelerations: wp.array3d(dtype=wp.float64), + dt: wp.array2d(dtype=wp.float64), +): + row, sample, joint = wp.tid() + if success[row] == 0: + return + n = counts[row] + valid = sample_counts[row] + duration = durations[row] + if sample >= valid or n == 1: + positions[row, sample, joint] = points[row, n - 1, joint] + return + step = duration / wp.float64(valid - 1) + query = step * wp.float64(sample) + if joint == 0 and sample > 0: + dt[row, sample] = step + total = intervals[row] + i = _sample_interval(times, row, total, query) + elapsed = query - times[row, i] + accel = (x[row, i + 1] - x[row, i]) * wp.float64(total) / wp.float64(2.0) + initial = _speed(x[row, i]) + speed = wp.max(wp.float64(0.0), initial + accel * elapsed) + s = wp.clamp( + wp.float64(i) / wp.float64(total) + + initial * elapsed + + wp.float64(0.5) * accel * elapsed * elapsed, + wp.float64(i) / wp.float64(total), + wp.float64(i + 1) / wp.float64(total), + ) + segment = wp.min(int(s * wp.float64(n - 1)), n - 2) + delta = s - wp.float64(segment) / wp.float64(n - 1) + ca, cb = coefficients[row, segment, joint, 0], coefficients[row, segment, joint, 1] + cc, cd = coefficients[row, segment, joint, 2], coefficients[row, segment, joint, 3] + first = cb + delta * (wp.float64(2.0) * cc + wp.float64(3.0) * delta * cd) + second = wp.float64(2.0) * cc + wp.float64(6.0) * delta * cd + position = ca + delta * (cb + delta * (cc + delta * cd)) + velocity_value = first * speed + accelerations[row, sample, joint] = second * speed * speed + first * accel + if sample == 0: + position = points[row, 0, joint] + velocity_value = wp.float64(0.0) + elif sample == valid - 1: + position = points[row, n - 1, joint] + velocity_value = wp.float64(0.0) + positions[row, sample, joint] = position + velocities[row, sample, joint] = velocity_value diff --git a/embodichain/lab/sim/motion/planners/toppra_planner.py b/embodichain/lab/sim/motion/planners/toppra_planner.py index 408a43c62..5495668de 100644 --- a/embodichain/lab/sim/motion/planners/toppra_planner.py +++ b/embodichain/lab/sim/motion/planners/toppra_planner.py @@ -17,8 +17,11 @@ from __future__ import annotations from dataclasses import field +from copy import deepcopy import os +from collections.abc import Mapping, Sequence +from typing import Literal import torch import numpy as np @@ -40,32 +43,11 @@ _resolve_joint_constraints, ) -try: - import toppra as ta - import toppra.constraint as constraint -except ImportError: - logger.log_error( - "toppra not installed. Install with `pip install toppra==0.6.3`", ImportError - ) - -ta.setup_logging(level="WARN") - - -def _build_constraint_arrays(value, acc, dofs: int) -> tuple[np.ndarray, np.ndarray]: - """Expand limits, using compatibility defaults for standalone ``None`` inputs.""" - if value is None: - value = _DEFAULT_VELOCITY - if acc is None: - acc = _DEFAULT_ACCELERATION - if isinstance(value, (float, int)): - vlims = np.array([[-value, value] for _ in range(dofs)]) - else: - vlims = np.array(value) - if isinstance(acc, (float, int)): - alims = np.array([[-acc, acc] for _ in range(dofs)]) - else: - alims = np.array(acc) - return vlims, alims +from embodichain.compute.trajectory._toppra import ( + _NumpyToppra, + _validate_grid_size, + _validate_sampling, +) def _toppra_solve_one_env( @@ -74,6 +56,7 @@ def _toppra_solve_one_env( acc_constraint, sample_method: "TrajectorySampleMethod", sample_interval: float | int, + grid_size: int = 100, ) -> dict: """Solve a single-env TOPPRA trajectory. Pure numpy/scipy — picklable, no torch/robot. @@ -90,83 +73,37 @@ def _toppra_solve_one_env( ``dt`` ``(N_b,)``, ``success`` bool, and ``n`` int. """ dofs = waypoints.shape[1] - vlims, alims = _build_constraint_arrays(vel_constraint, acc_constraint, dofs) - - if vlims.shape in ((dofs,), (dofs, 2)): - frozen = vlims == 0 if vlims.ndim == 1 else (vlims == 0).all(axis=1) - if np.any(waypoints[:, frozen] != waypoints[:1, frozen]): - return _empty_failure(dofs) - # TOPPRA requires nonzero intervals. Proven-constant joints impose no - # timing constraint, so use a finite interval only inside the solver. - vlims[frozen] = 1.0 if vlims.ndim == 1 else [-1.0, 1.0] - - if sample_method == TrajectorySampleMethod.TIME and sample_interval <= 0: - return _empty_failure(dofs) - if sample_method == TrajectorySampleMethod.QUANTITY and sample_interval < 2: - return _empty_failure(dofs) - - # Remove consecutive duplicate waypoints. Long plateaus of identical points - # (e.g. when start_qpos equals the first target and joint-space interpolation - # is enabled) can make TOPPRA's controllable-set computation numerically - # ill-conditioned and fail with "Instance is not controllable". - dup_tol = 1e-6 - keep = [0] - for i in range(1, len(waypoints)): - if np.max(np.abs(waypoints[i] - waypoints[keep[-1]])) >= dup_tol: - keep.append(i) - if keep[-1] != len(waypoints) - 1: - keep.append(len(waypoints) - 1) - waypoints = waypoints[keep] - - # Trivial same-waypoint shortcut - if len(waypoints) == 2 and np.sum(np.abs(waypoints[1] - waypoints[0])) < 1e-3: - pos = np.stack([waypoints[0], waypoints[1]]) - return { - "positions": pos, - "velocities": np.zeros_like(pos), - "accelerations": np.zeros_like(pos), - "dt": np.array([0.0, 0.0], dtype=np.float32), - "success": True, - "n": 2, - } - - ss = np.linspace(0.0, 1.0, len(waypoints)) + if vel_constraint is None: + vel_constraint = _DEFAULT_VELOCITY + if acc_constraint is None: + acc_constraint = _DEFAULT_ACCELERATION try: - path = ta.SplineInterpolator(ss, waypoints) - pc_vel = constraint.JointVelocityConstraint(vlims) - pc_acc = constraint.JointAccelerationConstraint(alims) - instance = ta.algorithm.TOPPRA( - [pc_vel, pc_acc], - path, - parametrizer="ParametrizeConstAccel", - gridpt_min_nb_points=max(100, 10 * len(waypoints)), + sample_dt = ( + sample_interval if sample_method == TrajectorySampleMethod.TIME else None ) - jnt_traj = instance.compute_trajectory() - except Exception: - return _empty_failure(dofs) - - if jnt_traj is None: - return _empty_failure(dofs) - - duration = float(jnt_traj.duration) - if duration <= 0: + sample_count = ( + int(sample_interval) + if sample_method == TrajectorySampleMethod.QUANTITY + else None + ) + _validate_sampling(sample_count, sample_dt) + trajectory = _NumpyToppra(waypoints, vel_constraint, acc_constraint, grid_size) + duration = trajectory.duration + if duration == 0: + count = 2 + elif sample_dt is not None: + count = max(2, int(np.ceil(duration / sample_dt)) + 1) + else: + count = int(sample_count) + ts = np.linspace(0.0, duration, count) + positions, velocities, accelerations = trajectory.sample(ts) + except (ValueError, RuntimeError, FloatingPointError, OverflowError): return _empty_failure(dofs) - - if sample_method == TrajectorySampleMethod.TIME: - n_points = max(2, int(np.ceil(duration / sample_interval)) + 1) - ts = np.linspace(0.0, duration, n_points) - else: - ts = np.linspace(0.0, duration, num=int(sample_interval)) - - positions = np.array([jnt_traj.eval(t) for t in ts]) - velocities = np.array([jnt_traj.evald(t) for t in ts]) - accelerations = np.array([jnt_traj.evaldd(t) for t in ts]) - dt = np.diff(ts, prepend=0.0).astype(np.float32) return { "positions": positions, "velocities": velocities, "accelerations": accelerations, - "dt": dt, + "dt": np.diff(ts, prepend=0.0).astype(np.float32), "success": True, "n": len(ts), } @@ -250,17 +187,194 @@ 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], + path_lengths: torch.Tensor | None = None, + 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. + + Warp handles both CPU and CUDA inputs. Sufficiently large CPU inference + batches share a bounded thread pool; small requests and gradient recording + run inline. Interval constraints enforce limits between grid points. + Floating outputs retain float32/float64 input precision; other input dtypes + are promoted to float32 before solving. + + 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`. + path_lengths: Optional valid prefix length for each batch row. This + excludes rectangular tail padding from spline fitting while + preserving explicit repeated endpoints inside each prefix. + 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, + ``path_lengths`` is malformed, 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 + if path_lengths is not None: + if not isinstance(path_lengths, torch.Tensor) or path_lengths.shape != (batch,): + raise ValueError(f"path_lengths must have shape ({batch},)") + if path_lengths.dtype not in ( + torch.uint8, + torch.int8, + torch.int16, + torch.int32, + torch.int64, + ): + raise ValueError("path_lengths must contain integers") + path_lengths = path_lengths.to(device=positions.device, dtype=torch.int32) + if bool(((path_lengths < 2) | (path_lengths > count)).any().item()): + raise ValueError(f"path_lengths must be between 2 and {count}") + path_lengths = path_lengths.contiguous() + resolved_interval = count if sample_interval is None else sample_interval + velocity, acceleration = constraints["velocity"], constraints["acceleration"] + from embodichain.compute.trajectory._toppra_warp import _retime_toppra_warp + + quantity = sample_method == TrajectorySampleMethod.QUANTITY + waypoints = ( + positions + if positions.dtype in (torch.float32, torch.float64) + else positions.float() + ) + try: + result = _retime_toppra_warp( + waypoints, + velocity, + acceleration, + path_lengths=path_lengths, + 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): + failed = _assemble_batched_results( + [_empty_failure(dofs) for _ in range(batch)], dofs, resolved_device + ) + return PlanResult( + success=failed.success, + **{ + key: getattr(failed, key).to(waypoints.dtype) + for key in ("positions", "velocities", "accelerations", "dt") + }, + ) + return PlanResult( + **{key: value.to(resolved_device) for key, value in result.items()} + ) + + +__all__ = [ + "ToppraPlanner", + "ToppraPlannerCfg", + "ToppraPlanOptions", + "retime_joint_paths", +] @configclass class ToppraPlannerCfg(BasePlannerCfg): planner_type: str = "toppra" + backend: Literal["auto", "numpy", "warp"] = field(default="auto", kw_only=True) + """Auto selects Warp on CPU and CUDA; NumPy remains an explicit reference.""" + grid_size: int = field(default=100, kw_only=True) + """Minimum timing grid size; each spline segment is subdivided uniformly.""" max_workers: int | None = None - """Worker process count for the batched fan-out. None => min(cpu_count()//2, B).""" + """NumPy worker count. Ignored by Warp; None => min(cpu_count()//2, B).""" mp_context: str | None = None - """Multiprocessing start method for the batched fan-out. + """Multiprocessing start method for the NumPy batched fan-out; ignored by Warp. ``None`` (default) auto-selects based on the simulation device: ``'fork'`` on CPU and ``'spawn'`` on GPU. ``'fork'`` is faster — workers @@ -277,6 +391,19 @@ class ToppraPlannerCfg(BasePlannerCfg): """ +class _ToppraConstraints(dict): + """Copy config containers while preserving runtime tensor references.""" + + def __deepcopy__(self, memo: dict) -> _ToppraConstraints: + result = _ToppraConstraints() + memo[id(self)] = result + for key, value in self.items(): + result[deepcopy(key, memo)] = ( + value if isinstance(value, torch.Tensor) else deepcopy(value, memo) + ) + return result + + @configclass class ToppraPlanOptions(PlanOptions): @@ -293,6 +420,8 @@ class ToppraPlanOptions(PlanOptions): velocity resolves robot speed caps at planning time. Acceleration is a planning budget; ``None`` uses its compatibility default. Explicit speed requests are capped by available robot limits. + Tensor bounds retain their references through option construction/copying + so gradients reach the caller, including through non-leaf expressions. """ sample_method: TrajectorySampleMethod = TrajectorySampleMethod.TIME @@ -306,11 +435,22 @@ class ToppraPlanOptions(PlanOptions): If sample_method is 'time', this is the time interval in seconds. If sample_method is 'quantity', this is the total number of samples. + Both modes support first-order gradients of waypoints and tensor motion limits. """ + def __post_init__(self) -> None: + """Preserve tensor graphs before configclass copies mutable fields.""" + self.constraints = _ToppraConstraints(self.constraints) + class ToppraPlanner(BasePlanner): - """Time-optimal joint-space planner backed by TOPPRA.""" + """Rest-to-rest path timing with NumPy or batched Warp TOPPRA. + + Cubic path geometry uses uniform knots and not-a-knot boundaries. Timing + optimizes conservative interval constraints, so durations need not match + the former external TOPPRA package. CUDA auto mode keeps numerical work on + device; NumPy remains available explicitly for reference comparisons. + """ supported_move_types = frozenset({MoveType.JOINT_MOVE}) @@ -323,6 +463,9 @@ def __init__(self, cfg: ToppraPlannerCfg): Args: cfg: Configuration object containing ToppraPlanner settings """ + if cfg.backend not in {"auto", "numpy", "warp"}: + raise ValueError("TOPPRA backend must be auto, numpy, or warp") + _validate_grid_size(cfg.grid_size) super().__init__(cfg) self._pool = None @@ -460,11 +603,18 @@ def plan( b = _infer_batch_size(target_states) or 1 dofs = target_states[0].qpos.shape[-1] - # Build (B, N, DOF) numpy waypoints - waypoints = np.stack( - [s.qpos.detach().cpu().numpy() for s in target_states], axis=1 - ) # (B, N, DOF) - + backend = self.cfg.backend + if backend not in {"auto", "numpy", "warp"}: + raise ValueError("TOPPRA backend must be auto, numpy, or warp") + needs_grad = torch.is_grad_enabled() and ( + any(state.qpos.requires_grad for state in target_states) + or any( + isinstance(limit, torch.Tensor) and limit.requires_grad + for limit in options.constraints.values() + ) + ) + if needs_grad and backend == "numpy": + raise NotImplementedError("TOPPRA gradients require the Warp backend") constraints = _resolve_joint_constraints( self.robot, options.constraints, @@ -472,10 +622,65 @@ def plan( dof=dofs, batch_size=b, ) - vc = constraints["velocity"].numpy() - ac = constraints["acceleration"].numpy() + use_warp = backend != "numpy" + if use_warp: + from embodichain.compute.trajectory._toppra_warp import _retime_toppra_warp + + waypoints_tensor = torch.stack([s.qpos for s in target_states], dim=1).to( + self.device + ) + if waypoints_tensor.dtype not in (torch.float32, torch.float64): + waypoints_tensor = waypoints_tensor.to(torch.float64) + try: + result = _retime_toppra_warp( + waypoints_tensor, + constraints["velocity"], + constraints["acceleration"], + sample_count=( + int(options.sample_interval) + if options.sample_method == TrajectorySampleMethod.QUANTITY + else None + ), + sample_dt=( + options.sample_interval + if options.sample_method == TrajectorySampleMethod.TIME + else None + ), + grid_size=self.cfg.grid_size, + ) + except (ValueError, FloatingPointError, OverflowError): + return self._assemble_batched_result( + [_empty_failure(dofs) for _ in range(b)], dofs + ) + # The existing planner contract returns float32 for all floating + # fields, independently of input precision or selected backend. + return PlanResult( + **{ + key: value if key == "success" else value.to(torch.float32) + for key, value in result.items() + } + ) + + # Build (B, N, DOF) numpy waypoints + waypoints = np.stack( + [s.qpos.detach().cpu().numpy() for s in target_states], axis=1 + ) # (B, N, DOF) + + vc = constraints["velocity"] + ac = constraints["acceleration"] + if isinstance(vc, torch.Tensor): + vc = vc.detach().cpu().numpy() + if isinstance(ac, torch.Tensor): + ac = ac.detach().cpu().numpy() args_per_env = [ - (waypoints[i], vc, ac, options.sample_method, options.sample_interval) + ( + waypoints[i], + vc, + ac, + options.sample_method, + options.sample_interval, + self.cfg.grid_size, + ) for i in range(b) ] @@ -525,47 +730,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/pyproject.toml b/pyproject.toml index 37be9a05a..20930f944 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,7 +32,7 @@ dependencies = [ "huggingface_hub>=0.35.3,<1.0", "langchain", "langchain-openai", - "toppra==0.6.3", + "scipy", "pin", "pin-pink", "casadi", diff --git a/scripts/benchmark/motion_generation/planners/toppra.py b/scripts/benchmark/motion_generation/planners/toppra.py index 21374a937..2f9c87943 100644 --- a/scripts/benchmark/motion_generation/planners/toppra.py +++ b/scripts/benchmark/motion_generation/planners/toppra.py @@ -18,8 +18,6 @@ from __future__ import annotations -import importlib.util - from embodichain.lab.sim.motion.motion_generator import ( MotionGenCfg, MotionGenOptions, @@ -51,9 +49,7 @@ def __init__(self, spec: PlannerSpecCfg, context: PlannerContext) -> None: self.motion_generator: MotionGenerator | None = None def availability(self) -> tuple[bool, str | None]: - """Report whether the optional TOPPRA package is installed.""" - if importlib.util.find_spec("toppra") is None: - return False, "TOPPRA is not installed." + """The in-tree TOPPRA implementation needs no optional external package.""" return True, None def build(self) -> None: diff --git a/tests/benchmark/motion_generation/test_toppra_dependency.py b/tests/benchmark/motion_generation/test_toppra_dependency.py new file mode 100644 index 000000000..332f9c015 --- /dev/null +++ b/tests/benchmark/motion_generation/test_toppra_dependency.py @@ -0,0 +1,39 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""The benchmark must discover the in-tree TOPPRA backend on a clean install.""" + +from __future__ import annotations + +import importlib.util + +import pytest + + +@pytest.mark.no_sim +def test_toppra_adapter_is_available_without_external_package( + monkeypatch: pytest.MonkeyPatch, +) -> None: + original = importlib.util.find_spec + + def without_external(name, *args, **kwargs): + return None if name == "toppra" else original(name, *args, **kwargs) + + monkeypatch.setattr(importlib.util, "find_spec", without_external) + from scripts.benchmark.motion_generation.planners.toppra import ToppraAdapter + + adapter = object.__new__(ToppraAdapter) + assert adapter.availability() == (True, None) diff --git a/tests/compute/test_toppra.py b/tests/compute/test_toppra.py new file mode 100644 index 000000000..e66acc55d --- /dev/null +++ b/tests/compute/test_toppra.py @@ -0,0 +1,614 @@ +# ---------------------------------------------------------------------------- +# 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 and device contracts for the in-tree TOPPRA implementations.""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from embodichain.compute.trajectory._toppra import _NumpyToppra +from embodichain.compute.trajectory._toppra_warp import _retime_toppra_warp + + +@pytest.mark.parametrize("count", [2, 3, 4, 7, 20]) +def test_warp_matches_numpy_not_a_knot_geometry(count: int) -> None: + points = np.random.default_rng(count).normal(size=(3, count, 3)) + vlim = np.array([[-0.4, 1.0], [-1.2, 0.7], [-0.8, 1.4]]) + alim = np.array([[-1.2, 2.0], [-2.1, 0.8], [-0.6, 1.1]]) + result = _retime_toppra_warp(torch.tensor(points), vlim, alim, sample_count=128) + assert result["success"].all() + for row in range(len(points)): + reference = _NumpyToppra(points[row], vlim, alim) + duration = result["dt"][row].sum().item() + assert duration == pytest.approx(reference.duration, rel=1e-10) + expected = reference.sample(np.linspace(0, reference.duration, 128)) + for key, array in zip(("positions", "velocities", "accelerations"), expected): + np.testing.assert_allclose(result[key][row], array, rtol=1e-8, atol=1e-8) + + +@pytest.mark.parametrize("backend", ["numpy", "warp"]) +def test_straight_path_matches_analytic_trapezoid(backend: str) -> None: + # Distance 2, vmax 1, amax 2 => 0.5 + 1.5 + 0.5 = 2.5 seconds. + points = np.array([[0.0], [2.0]]) + if backend == "numpy": + trajectory = _NumpyToppra(points, 1.0, 2.0, grid_size=401) + duration = trajectory.duration + else: + result = _retime_toppra_warp( + torch.tensor(points[None]), 1.0, 2.0, grid_size=401, sample_count=128 + ) + assert result["success"].all() + duration = result["dt"].sum().item() + assert duration == pytest.approx(2.5, rel=1e-5) + + +@pytest.mark.parametrize("seed", [3, 12, 421]) +def test_continuous_asymmetric_limits_and_interval_dynamics(seed: int) -> None: + points = np.random.default_rng(seed).normal(size=(7, 3)) + vlim = np.array([[-0.4, 1.0], [-1.2, 0.7], [-0.8, 1.4]]) + alim = np.array([[-1.2, 2.0], [-2.1, 0.8], [-0.6, 1.1]]) + trajectory = _NumpyToppra(points, vlim, alim) + # Probe inside every timing interval, rather than only a uniform clock. + times = ( + trajectory.times[:-1, None] + + np.diff(trajectory.times)[:, None] * np.linspace(0.001, 0.999, 31) + ).ravel() + _, velocity, acceleration = trajectory.sample(times) + assert np.all(velocity >= vlim[:, 0] - 1e-9) + assert np.all(velocity <= vlim[:, 1] + 1e-9) + assert np.all(acceleration >= alim[:, 0] - 1e-9) + assert np.all(acceleration <= alim[:, 1] + 1e-9) + np.testing.assert_allclose( + np.diff(trajectory.speeds**2), + 2 * np.diff(trajectory.grid) * trajectory.accelerations, + atol=1e-12, + ) + np.testing.assert_allclose( + 0.5 + * (trajectory.speeds[:-1] + trajectory.speeds[1:]) + * np.diff(trajectory.times), + np.diff(trajectory.grid), + atol=1e-12, + ) + + +@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.gpu)]) +def test_mixed_batch_duplicates_stationary_failure_and_time_padding( + device: str, +) -> None: + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + points = torch.tensor( + [ + [[0.0, 2.0], [0.0, 2.0], [1.0, 2.0], [1.0, 2.0]], + [[3.0, 4.0], [3.0, 4.0], [3.0, 4.0], [3.0, 4.0]], + [[0.0, 0.0], [float("nan"), 0.0], [1.0, 0.0], [1.0, 0.0]], + [[0.0, 2.0], [0.0, 2.0], [0.1, 2.0], [0.1, 2.0]], + ], + dtype=torch.float64, + device=device, + ) + result = _retime_toppra_warp(points, 1.0, 2.0, sample_dt=0.03) + assert result["success"].tolist() == [True, True, False, True] + assert torch.equal( + result["positions"][1], points[1, :1].expand_as(result["positions"][1]) + ) + assert result["dt"][1:3].count_nonzero() == 0 + assert result["positions"][2].count_nonzero() == 0 + for row in (0, 3): + valid = 1 + int((result["dt"][row] > 0).sum()) + assert result["dt"][row].max() <= 0.03 + torch.testing.assert_close(result["positions"][row, 0], points[row, 0]) + torch.testing.assert_close( + result["positions"][row, valid - 1 :], + points[row, -1:].expand_as(result["positions"][row, valid - 1 :]), + ) + assert result["velocities"][row, [0, valid - 1]].count_nonzero() == 0 + assert result["velocities"][row, valid:].count_nonzero() == 0 + assert result["accelerations"][row, valid:].count_nonzero() == 0 + reference = _NumpyToppra(points[row].cpu().numpy(), 1.0, 2.0) + assert result["dt"][row].sum().item() == pytest.approx(reference.duration) + + +def test_path_lengths_exclude_rectangular_hold_padding_from_spline() -> None: + points = torch.tensor( + [ + [[0.0], [0.05], [0.1], [0.1]], + [[0.0], [0.03], [0.08], [0.12]], + ], + dtype=torch.float64, + ) + lengths = torch.tensor([3, 4]) + + result = _retime_toppra_warp( + points, 1.0, 2.0, path_lengths=lengths, sample_count=17 + ) + trimmed = _retime_toppra_warp(points[:1, :3], 1.0, 2.0, sample_count=17) + explicit_repeat = _retime_toppra_warp(points[:1], 1.0, 2.0, sample_count=17) + + torch.testing.assert_close(result["positions"][0], trimmed["positions"][0]) + torch.testing.assert_close(result["dt"][0], trimmed["dt"][0]) + assert result["dt"][0].sum() < explicit_repeat["dt"][0].sum() + assert result["positions"][0].max() <= points[0, 2, 0] + assert explicit_repeat["positions"][0].max() > points[0, 2, 0] + + +def test_path_lengths_leave_ignored_padding_with_zero_gradient() -> None: + points = torch.tensor( + [ + [[0.0], [0.05], [0.1], [0.1]], + [[0.0], [0.03], [0.08], [0.12]], + ], + dtype=torch.float64, + requires_grad=True, + ) + result = _retime_toppra_warp( + points, + 1.0, + 2.0, + path_lengths=torch.tensor([3, 4]), + sample_count=9, + ) + + result["dt"].sum().backward() + + assert points.grad[0, 3].count_nonzero() == 0 + assert points.grad[:, :3].count_nonzero() > 0 + + +@pytest.mark.parametrize( + "path_lengths", + [ + torch.tensor([[3]]), + torch.tensor([3.0]), + torch.tensor([1]), + torch.tensor([5]), + ], +) +def test_invalid_path_lengths_rejected(path_lengths: torch.Tensor) -> None: + with pytest.raises(ValueError, match="path_lengths"): + _retime_toppra_warp( + torch.zeros(1, 4, 1), + 1.0, + 2.0, + path_lengths=path_lengths, + sample_count=8, + ) + + +def test_one_sided_and_locked_joint_limits_isolate_infeasible_rows() -> None: + points = torch.tensor( + [[[0.0, 2.0], [1.0, 2.0]], [[0.0, 2.0], [-1.0, 2.0]], [[0.0, 2.0], [1.0, 3.0]]], + dtype=torch.float64, + ) + result = _retime_toppra_warp( + points, [[0.0, 1.0], [0.0, 0.0]], [[-2.0, 1.0], [0.0, 0.0]], sample_count=128 + ) + assert result["success"].tolist() == [True, False, False] + assert result["velocities"][0, :, 0].min() >= 0 + assert result["velocities"][0, :, 0].max() <= 1 + 1e-9 + assert result["accelerations"][0, :, 0].min() >= -2 - 1e-9 + assert result["accelerations"][0, :, 0].max() <= 1 + 1e-9 + + +@pytest.mark.parametrize("scale", [1e-4, 1.0, 1e4]) +def test_time_scaling_and_tiny_real_moves(scale: float) -> None: + points = np.array([[0.0], [1e-7]]) + result = _retime_toppra_warp( + torch.tensor(points[None]), scale, scale**2, sample_count=128 + ) + assert result["success"].all() + assert result["dt"].sum() > 0 + assert result["positions"][0, -1, 0] == points[-1, 0] + reference = _NumpyToppra(points, 1.0, 1.0) + assert result["dt"].sum().item() * scale == pytest.approx( + reference.duration, rel=1e-9 + ) + + +@pytest.mark.parametrize( + "kwargs", + [ + {"sample_count": 1}, + {"sample_count": 2.5}, + {"sample_dt": 0.0}, + {"sample_dt": float("nan")}, + {"sample_count": 10, "grid_size": 2}, + ], +) +def test_invalid_sampling_rejected(kwargs: dict) -> None: + with pytest.raises(ValueError): + _retime_toppra_warp(torch.zeros(1, 2, 1), 1.0, 2.0, **kwargs) + + +@pytest.mark.gpu +def test_cuda_matches_reference_on_nondefault_stream_without_host_waypoints( + monkeypatch: pytest.MonkeyPatch, +) -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + points = np.random.default_rng(47).normal(size=(8, 6, 6)) + expected = [_NumpyToppra(row, 1.0, 2.0) for row in points] + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + tensor = torch.tensor(points, device="cuda") + with monkeypatch.context() as patch: + + def forbidden(*args, **kwargs): + raise AssertionError("CUDA retiming must not copy waypoints to CPU") + + patch.setattr(torch.Tensor, "cpu", forbidden) + result = _retime_toppra_warp(tensor, 1.0, 2.0, sample_count=128) + # Queue consumption on the same stream before any explicit synchronize. + duration = result["dt"].sum(dim=1) + stream.synchronize() + assert result["success"].all() + assert all(value.is_cuda for value in result.values()) + for row, reference in enumerate(expected): + assert duration[row].item() == pytest.approx(reference.duration, rel=1e-10) + values = reference.sample(np.linspace(0, reference.duration, 128)) + for key, value in zip(("positions", "velocities", "accelerations"), values): + np.testing.assert_allclose( + result[key][row].cpu(), value, atol=1e-8, rtol=1e-8 + ) + + +@pytest.mark.parametrize( + ("points", "kept"), + [ + ([[0.0], [2e-7], [0.5]], [0, 2]), + ([[0.3], [0.3], [0.5], [0.5]], [0, 2, 3]), + ([[0.0], [6e-7], [1.2e-6], [0.5]], [0, 2, 3]), + ], +) +def test_legacy_waypoint_cleanup_preserves_spline_geometry( + points: list[list[float]], kept: list[int] +) -> None: + from scipy.interpolate import CubicSpline + + waypoints = np.array(points) + expected_points = waypoints[kept] + expected_path = CubicSpline(np.linspace(0.0, 1.0, len(kept)), expected_points) + reference = _NumpyToppra(waypoints, 1.0, 2.0) + phase = np.linspace(0.0, 1.0, 1001) + np.testing.assert_allclose(reference.path(phase), expected_path(phase), atol=1e-12) + result = _retime_toppra_warp( + torch.tensor(waypoints[None]), 1.0, 2.0, sample_count=128 + ) + assert result["success"].all() + assert result["dt"].sum().item() == pytest.approx(reference.duration, rel=1e-10) + expected = reference.sample(np.linspace(0.0, reference.duration, 128)) + for key, values in zip(("positions", "velocities", "accelerations"), expected): + np.testing.assert_allclose(result[key][0], values, rtol=1e-8, atol=1e-8) + + +@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.gpu)]) +def test_float32_cleanup_matches_original_threshold_arithmetic(device: str) -> None: + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + # float32(1e-6) is just below the double-precision threshold. Filtering must + # happen with the input arithmetic, before promotion for the spline solve. + points = np.array([[0.0], [1e-6], [0.5]], dtype=np.float32) + reference = _NumpyToppra(points, 1.0, 2.0) + # The original NumPy worker follows the installed NumPy scalar-promotion + # rules, which differ between 1.x and 2.x at this exact boundary. + expected_points = points if np.float32(1e-6) >= 1e-6 else points[[0, 2]] + np.testing.assert_array_equal(reference.points, expected_points) + result = _retime_toppra_warp( + torch.from_numpy(points[None]).to(device), 1.0, 2.0, sample_count=128 + ) + assert result["success"].all() + expected = reference.sample(np.linspace(0.0, reference.duration, 128)) + for key, value in zip(("positions", "velocities", "accelerations"), expected): + np.testing.assert_allclose(result[key][0].cpu(), value, rtol=1e-6, atol=1e-6) + + +@pytest.mark.parametrize("count", [2, 3, 4, 7]) +def test_waypoint_autograd_matches_finite_differences(count: int) -> None: + # Exercise linear, parabolic and general cubic paths, all output Jacobians, + # asymmetric bounds, and both velocity and acceleration active constraints. + generator = torch.Generator().manual_seed(123 + count) + points = torch.randn(1, count, 2, generator=generator, dtype=torch.float64) + points.requires_grad_() + + def evaluate(q: torch.Tensor) -> tuple[torch.Tensor, ...]: + result = _retime_toppra_warp( + q, + [[-0.7, 1.1], [-0.9, 0.8]], + [[-1.3, 2.0], [-1.6, 1.1]], + sample_count=9, + grid_size=25, + ) + assert result["success"].all() + return tuple( + result[key] for key in ("positions", "velocities", "accelerations", "dt") + ) + + tracked = evaluate(points) + with torch.no_grad(): + untracked = evaluate(points) + for actual, expected in zip(tracked, untracked): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert actual.requires_grad and not expected.requires_grad + assert torch.autograd.gradcheck( + evaluate, (points,), eps=1e-6, atol=2e-4, rtol=2e-3, fast_mode=True + ) + + +def test_duration_gradient_matches_acceleration_limited_scaling() -> None: + distance = 0.73 + points = torch.tensor( + [[[0.0], [distance]]], dtype=torch.float64, requires_grad=True + ) + result = _retime_toppra_warp(points, 100.0, [[-1.7, 2.3]], sample_count=17) + duration = result["dt"].sum() + # With inactive velocity bounds, T scales as sqrt(distance), including + # discretization. This probes the full solver rather than interpolation. + gradient = torch.autograd.grad(duration, points)[0] + derivative = duration.detach() / (2 * distance) + torch.testing.assert_close( + gradient[0, :, 0], torch.stack((-derivative, derivative)) + ) + + +def test_autograd_cleanup_holds_and_failures_keep_gradients_isolated() -> None: + points = torch.tensor( + [ + [[0.0, 0.0], [0.0, 0.0], [0.4, 0.2], [1.0, -0.1]], + [[2.0, 3.0], [2.0, 3.0], [2.0, 3.0], [2.0, 3.0]], + [[0.0, 0.0], [float("nan"), 0.0], [0.4, 0.2], [1.0, 0.0]], + # A finite but infeasible path with one-sided velocity bounds. + [[0.0, 0.0], [-0.3, 0.0], [-0.6, 0.0], [-1.0, 0.0]], + ], + dtype=torch.float64, + requires_grad=True, + ) + result = _retime_toppra_warp(points, [[0.0, 1.0], [-1.0, 1.0]], 2.0, sample_count=9) + assert result["success"].tolist() == [True, True, False, False] + sum( + result[key].sum() for key in ("positions", "velocities", "accelerations", "dt") + ).backward() + assert torch.isfinite(points.grad).all() + assert points.grad[0, 1].count_nonzero() == 0 # Removed duplicate. + assert points.grad[2:].count_nonzero() == 0 + torch.testing.assert_close( + points.grad[1, 0], torch.full((2,), 9.0, dtype=points.dtype) + ) + assert points.grad[1, 1:].count_nonzero() == 0 + + +def test_autograd_repeated_backward_and_independent_forwards() -> None: + points = torch.tensor( + [[[0.1], [0.7], [1.3]]], dtype=torch.float64, requires_grad=True + ) + first = _retime_toppra_warp(points, 1.0, 2.0, sample_count=9) + second = _retime_toppra_warp(points, 1.0, 2.0, sample_count=9) + seed = torch.linspace(0.2, 1.0, 9, dtype=points.dtype).view(1, 9, 1) + before = seed.clone() + one = torch.autograd.grad(first["positions"], points, seed, retain_graph=True)[0] + two = torch.autograd.grad(first["positions"], points, seed, retain_graph=True)[0] + both = torch.autograd.grad( + (first["positions"], second["positions"]), points, (seed, seed) + )[0] + torch.testing.assert_close(seed, before) + torch.testing.assert_close(one, two) + torch.testing.assert_close(both, 2 * one) + + +@pytest.mark.parametrize("sampling", [{"sample_count": 9}, {"sample_dt": 0.23}]) +@pytest.mark.parametrize( + ("velocity", "acceleration"), + [ + (0.9, 1.7), + ([0.9, 1.1], [1.7, 1.3]), + ([[-0.7, 1.1], [-0.9, 0.8]], [[-1.3, 2.0], [-1.6, 1.1]]), + ], +) +def test_waypoint_and_limit_gradients_match_both_sampling_modes( + sampling: dict, + velocity: object, + acceleration: object, +) -> None: + generator = torch.Generator().manual_seed(127) + points = torch.randn( + 1, 4, 2, generator=generator, dtype=torch.float64, requires_grad=True + ) + vlim = torch.tensor(velocity, dtype=torch.float64, requires_grad=True) + alim = torch.tensor(acceleration, dtype=torch.float64, requires_grad=True) + + def evaluate( + q: torch.Tensor, v: torch.Tensor, a: torch.Tensor + ) -> tuple[torch.Tensor, ...]: + result = _retime_toppra_warp(q, v, a, grid_size=25, **sampling) + assert result["success"].all() + return tuple( + result[key] for key in ("positions", "velocities", "accelerations", "dt") + ) + + tracked = evaluate(points, vlim, alim) + with torch.no_grad(): + untracked = evaluate(points, vlim, alim) + for actual, expected in zip(tracked, untracked): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert actual.requires_grad and not expected.requires_grad + assert torch.autograd.gradcheck( + evaluate, + (points, vlim, alim), + eps=1e-6, + atol=2e-4, + rtol=2e-3, + fast_mode=True, + ) + + +def test_time_sampling_gradient_preserves_padding_and_shared_limits() -> None: + points = torch.tensor( + [[[0.0], [0.23]], [[0.0], [0.73]], [[2.0], [2.0]], [[0.0], [float("nan")]]], + dtype=torch.float64, + requires_grad=True, + ) + acceleration = torch.tensor(1.7, dtype=torch.float64, requires_grad=True) + result = _retime_toppra_warp(points, 100.0, acceleration, sample_dt=0.17) + assert result["success"].tolist() == [True, True, True, False] + durations = result["dt"].sum(dim=1) + dq, da = torch.autograd.grad( + durations.sum(), (points, acceleration), retain_graph=True + ) + # Acceleration-limited duration scales as sqrt(distance / acceleration). + expected_q = durations[:2].detach() / (2 * points[:2, -1, 0].detach()) + torch.testing.assert_close(dq[:2, -1, 0], expected_q) + torch.testing.assert_close(dq[:2, 0, 0], -expected_q) + torch.testing.assert_close( + da, -durations.sum().detach() / (2 * acceleration.detach()) + ) + assert dq[2:].count_nonzero() == 0 + assert torch.isfinite(dq).all() + short_length = 1 + int((result["dt"][0] > 0).sum()) + assert short_length < result["positions"].shape[1] + padding = result["positions"][0, short_length:].sum() + padding_grad, limit_grad = torch.autograd.grad(padding, (points, acceleration)) + expected = torch.zeros_like(points) + expected[0, -1, 0] = result["positions"].shape[1] - short_length + torch.testing.assert_close(padding_grad, expected) + assert limit_grad == 0 + + +@pytest.mark.gpu +@pytest.mark.parametrize("sampling", [{"sample_count": 9}, {"sample_dt": 0.23}]) +def test_cuda_autograd_matches_cpu_and_finite_differences_on_current_stream( + monkeypatch: pytest.MonkeyPatch, + sampling: dict, +) -> None: + if not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + generator = torch.Generator().manual_seed(130) + inputs = ( + torch.randn(1, 7, 2, generator=generator, dtype=torch.float64), + torch.tensor([[-0.7, 1.1], [-0.9, 0.8]], dtype=torch.float64), + torch.tensor([[-1.3, 2.0], [-1.6, 1.1]], dtype=torch.float64), + ) + directions = tuple( + torch.randn(value.shape, generator=generator, dtype=value.dtype) + for value in inputs + ) + + def loss(q: torch.Tensor, v: torch.Tensor, a: torch.Tensor) -> torch.Tensor: + out = _retime_toppra_warp(q, v, a, grid_size=25, **sampling) + return sum( + out[key].square().sum() + for key in ("positions", "velocities", "accelerations", "dt") + ) + + cpu = tuple(value.clone().requires_grad_() for value in inputs) + expected = torch.autograd.grad(loss(*cpu), cpu) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + gpu = tuple(value.to("cuda").requires_grad_() for value in inputs) + perturbations = tuple(value.to("cuda") for value in directions) + with monkeypatch.context() as patch: + + def forbidden(*args, **kwargs): + raise AssertionError( + "Autograd must not copy waypoints or motion limits to CPU" + ) + + patch.setattr(torch.Tensor, "cpu", forbidden) + gradients = torch.autograd.grad(loss(*gpu), gpu) + eps = 1e-6 + with torch.no_grad(): + plus = tuple( + value + eps * direction + for value, direction in zip(gpu, perturbations) + ) + minus = tuple( + value - eps * direction + for value, direction in zip(gpu, perturbations) + ) + numerical = (loss(*plus) - loss(*minus)) / (2 * eps) + analytical = sum( + (grad * direction).sum() + for grad, direction in zip(gradients, perturbations) + ) + stream.synchronize() + for actual, reference in zip(gradients, expected): + torch.testing.assert_close(actual.cpu(), reference, atol=1e-7, rtol=1e-7) + torch.testing.assert_close(analytical, numerical, atol=2e-4, rtol=2e-3) + + +def test_autograd_batches_noncontiguous_inputs_without_cross_row_gradients() -> None: + generator = torch.Generator().manual_seed(901) + base = torch.randn( + 2, 3, 5, generator=generator, dtype=torch.float64, requires_grad=True + ) + points = base.transpose(1, 2) + assert not points.is_contiguous() + batched = _retime_toppra_warp(points, 1.0, 2.0, sample_count=11) + batched_gradient = torch.autograd.grad(batched["dt"][0].sum(), base)[0] + separate = points[0:1].detach().requires_grad_() + single = _retime_toppra_warp(separate, 1.0, 2.0, sample_count=11) + expected = torch.autograd.grad(single["dt"].sum(), separate)[0] + torch.testing.assert_close(batched_gradient[0], expected[0].T) + assert batched_gradient[1].count_nonzero() == 0 + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +@pytest.mark.parametrize("grad", [False, True]) +@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.gpu)]) +def test_frozen_joint_raw_excursions_fail_before_deduplication( + monkeypatch: pytest.MonkeyPatch, dtype: torch.dtype, grad: bool, device: str +) -> None: + from embodichain.compute.trajectory import _toppra_warp as adapter + + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + points = torch.tensor( + [ + [[0.0, 0.0], [0.0, 1e-7], [0.4, 0.0]], + [[0.0, 0.0], [0.4, 0.0], [0.4, 1e-7]], + [[0.0, 0.0], [0.0, 0.0], [0.4, 0.0]], + [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]], + ], + dtype=dtype, + device=device, + requires_grad=grad, + ) + # The second row's moved tail is outside its declared valid path prefix. + monkeypatch.setattr(adapter, "_cpu_workers", lambda *args: 2) + result = _retime_toppra_warp( + points, + [0.5, 0.0], + 0.5, + path_lengths=torch.tensor([3, 2, 3, 3], device=device), + sample_count=17, + ) + assert result["success"].tolist() == [False, True, True, True] + for key in ("positions", "velocities", "accelerations", "dt"): + assert result[key][0].count_nonzero() == 0 + assert result["positions"][1, :, 1].count_nonzero() == 0 + if grad: + result["dt"].sum().backward() + assert points.grad[0].count_nonzero() == 0 + assert points.grad[1, 2].count_nonzero() == 0 + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_numpy_frozen_joint_raw_excursion_is_not_dropped(dtype) -> None: + points = np.array([[0.0, 0.0], [0.0, 1e-7], [0.4, 0.0]], dtype=dtype) + with pytest.raises(ValueError, match="zero-speed joints"): + _NumpyToppra(points, [0.5, 0.0], 0.5) diff --git a/tests/compute/test_toppra_cpu.py b/tests/compute/test_toppra_cpu.py new file mode 100644 index 000000000..67fb4ba75 --- /dev/null +++ b/tests/compute/test_toppra_cpu.py @@ -0,0 +1,329 @@ +# ---------------------------------------------------------------------------- +# 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 TOPPRA projection and inference scheduling contracts.""" + +from __future__ import annotations + +import numpy as np +from concurrent.futures import ThreadPoolExecutor +from contextlib import nullcontext +import os +import subprocess +import sys +import textwrap +from threading import Event +import pytest +import torch +import warp as wp +from scipy.optimize import linprog + +from embodichain.compute.trajectory._toppra_warp import _retime_toppra_warp +from embodichain.compute.trajectory import _toppra_warp as adapter +from embodichain.compute.trajectory._warp import toppra as kernels + + +def _project(planes: np.ndarray) -> np.ndarray: + """Compute the largest feasible x for rest-feasible half-plane batches.""" + wp.init() + batch = len(planes) + arrays = [ + wp.from_numpy( + np.ascontiguousarray(planes[:, None, :, index]), + dtype=wp.float64, + device="cpu", + ) + for index in range(3) + ] + counts = wp.from_numpy(np.ones(batch, dtype=np.int32), device="cpu") + caps = wp.empty((batch, 1), dtype=wp.float64, device="cpu") + wp.launch( + kernels._project_caps_cpu, + dim=(batch, 1), + inputs=[*arrays, counts, caps], + device="cpu", + ) + return caps.numpy()[:, 0] + + +@pytest.mark.parametrize("scale", [1e-100, 1.0, 1e100]) +def test_cpu_projection_matches_analytic_polygons(scale: float) -> None: + # Every inequality has the form a*x + b*y <= c, with x,y >= 0. + planes = ( + np.array( + [ + [ + [1.0, 1.0, 1.0], + [1.0, -1.0, 0.0], + [1.0, 0.0, 2.0], + [0.0, 1.0, 2.0], + ], # x <= y and x+y <= 1 + [ + [1.0, 0.0, 2.0], + [0.0, 1.0, 3.0], + [0.0, -1.0, 0.0], + [0.0, 0.0, 0.0], + ], # independent box + [ + [1.0, 1.0, 1.0], + [1.0, -1.0, 1.0], + [1.0, 0.0, 2.0], + [0.0, 1.0, 2.0], + ], # optimum at y=0 + [ + [1.0, -1.0, 0.0], + [-1.0, 1.0, 0.0], + [1.0, 0.0, 3.0], + [0.0, 1.0, 2.0], + ], # line x=y + [ + [1.0, 0.0, 0.0], + [0.0, 1.0, 2.0], + [0.0, -1.0, 0.0], + [0.0, 0.0, 0.0], + ], # locked x + ] + ) + * scale + ) + np.testing.assert_allclose( + _project(planes), [0.5, 2.0, 1.0, 2.0, 0.0], rtol=1e-12, atol=1e-12 + ) + + +def test_cpu_projection_matches_independent_linear_programs() -> None: + rng = np.random.default_rng(712) + # This includes 14-DoF-sized plane sets, zero one-sided bounds and parallel + # planes. SciPy's HiGHS LP is independent of the TOPPRA projection algorithm. + planes = rng.normal(size=(96, 86, 3)) + planes[:, :, 2] = rng.uniform(0, 2, (96, 86)) + planes[::3, :4, 2] = 0 + planes[:, -2, :] = [1.0, 0.0, 3.0] + planes[:, -1, :] = [0.0, 1.0, 4.0] + expected = [] + for row in planes: + solution = linprog( + [-1.0, 0.0], + A_ub=row[:, :2], + b_ub=row[:, 2], + bounds=[(0, None), (0, None)], + method="highs", + ) + assert solution.success + expected.append(solution.x[0]) + np.testing.assert_allclose(_project(planes), expected, rtol=1e-10, atol=1e-12) + + +@pytest.mark.parametrize("sampling", [{"sample_count": 129}, {"sample_dt": 0.07}]) +def test_cpu_optimized_forward_matches_differentiable_solver(sampling: dict) -> None: + rng = np.random.default_rng(731) + points = torch.tensor(rng.normal(size=(8, 16, 7)), dtype=torch.float64) + velocity = np.column_stack((-rng.uniform(0.4, 1.2, 7), rng.uniform(0.5, 1.4, 7))) + acceleration = np.column_stack( + (-rng.uniform(0.8, 1.6, 7), rng.uniform(0.9, 1.8, 7)) + ) + optimized = _retime_toppra_warp(points, velocity, acceleration, **sampling) + tracked = _retime_toppra_warp( + points.requires_grad_(), velocity, acceleration, **sampling + ) + assert torch.equal(optimized["success"], tracked["success"]) + for key in ("positions", "velocities", "accelerations", "dt"): + torch.testing.assert_close(optimized[key], tracked[key], rtol=1e-10, atol=1e-10) + tracked["dt"].sum().backward() + assert torch.isfinite(points.grad).all() + + +@pytest.mark.parametrize("sampling", [{"sample_count": 33}, {"sample_dt": 0.07}]) +@pytest.mark.parametrize("mode", ["grad_enabled", "no_grad", "inference"]) +def test_cpu_chunks_preserve_padding_failure_and_torch_modes( + monkeypatch: pytest.MonkeyPatch, sampling: dict, mode: str +) -> None: + # Force two chunks on small inputs. The first chunk is entirely stationary + # or failed, so its local length is 2 even for quantity=33. The second + # includes different durations and ignored NaN tail padding. + points = torch.tensor( + [ + [[2.0, 3.0], [2.0, 3.0], [2.0, 3.0], [2.0, 3.0]], + [[0.0, 0.0], [float("nan"), 0.0], [0.2, 0.1], [0.5, 0.2]], + [[0.0, 0.0], [0.2, 0.1], [0.5, 0.2], [float("nan"), 0.0]], + [[0.0, 0.0], [0.3, 0.1], [0.6, 0.2], [1.0, 0.3]], + ], + dtype=torch.float64, + ) + lengths = torch.tensor([4, 4, 3, 4]) + if mode != "grad_enabled": + points.requires_grad_() + context = ( + torch.inference_mode() + if mode == "inference" + else torch.no_grad() if mode == "no_grad" else nullcontext() + ) + observed_modes = [] + original = adapter._cpu_chunk + + def record(*args, **kwargs): + result = original(*args, **kwargs) + observed_modes.append( + (result["positions"].requires_grad, result["positions"].is_inference()) + ) + return result + + monkeypatch.setattr(adapter, "_cpu_chunk", record) + with context: + with monkeypatch.context() as patch: + patch.setattr(adapter, "_cpu_workers", lambda *args: 1) + expected = _retime_toppra_warp( + points, 1.0, 2.0, path_lengths=lengths, **sampling + ) + monkeypatch.setattr(adapter, "_cpu_workers", lambda *args: 2) + actual = _retime_toppra_warp(points, 1.0, 2.0, path_lengths=lengths, **sampling) + assert actual["success"].tolist() == [True, False, True, True] + for key in actual: + torch.testing.assert_close(actual[key], expected[key], rtol=0, atol=0) + assert observed_modes == [(False, mode == "inference")] * 2 + + +def test_cpu_parallel_noncontiguous_rows_are_independent( + monkeypatch: pytest.MonkeyPatch, +) -> None: + points = torch.randn( + 12, 3, 7, generator=torch.Generator().manual_seed(491), dtype=torch.float64 + ).transpose(1, 2) + assert not points.is_contiguous() + monkeypatch.setattr(adapter, "_cpu_workers", lambda *args: 3) + batched = _retime_toppra_warp(points, 1.0, 2.0, sample_dt=0.09) + monkeypatch.setattr(adapter, "_cpu_workers", lambda *args: 1) + for row in range(len(points)): + expected = _retime_toppra_warp(points[row : row + 1], 1.0, 2.0, sample_dt=0.09) + valid = expected["positions"].shape[1] + for key in ("positions", "velocities", "accelerations", "dt"): + torch.testing.assert_close( + batched[key][row : row + 1, :valid], expected[key], rtol=0, atol=0 + ) + assert batched["dt"][row, valid:].count_nonzero() == 0 + + +def test_cpu_inference_does_not_record_into_foreign_warp_tape( + monkeypatch: pytest.MonkeyPatch, +) -> None: + points = torch.linspace(0, 1, 24, dtype=torch.float64).reshape(4, 3, 2) + monkeypatch.setattr(adapter, "_cpu_workers", lambda *args: 2) + # Compile/init before capture. Warp's tape is process-global; ordinary + # inference must not record launches into an unrelated caller's tape. + _retime_toppra_warp(points, 1.0, 2.0, sample_count=11) + with wp.Tape() as tape: + actual = _retime_toppra_warp(points, 1.0, 2.0, sample_count=11) + assert actual["success"].all() + assert not tape.launches + + +def test_cpu_concurrent_requests_use_independent_buffers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + points = torch.randn( + 8, 4, 3, generator=torch.Generator().manual_seed(292), dtype=torch.float64 + ) + monkeypatch.setattr(adapter, "_cpu_workers", lambda *args: 1) + expected = [ + _retime_toppra_warp(points * scale, 1.0, 2.0, sample_count=27) + for scale in (1.0, 2.0, 3.0) + ] + monkeypatch.setattr(adapter, "_cpu_workers", lambda *args: 2) + with ThreadPoolExecutor(max_workers=3) as requests: + futures = [ + requests.submit( + _retime_toppra_warp, points * scale, 1.0, 2.0, sample_count=27 + ) + for scale in (1.0, 2.0, 3.0) + ] + actual = [future.result() for future in futures] + for result, reference in zip(actual, expected): + for key in result: + torch.testing.assert_close(result[key], reference[key], rtol=0, atol=0) + + +def test_cpu_worker_error_drains_other_tasks_and_pool_remains_usable( + monkeypatch: pytest.MonkeyPatch, +) -> None: + points = torch.zeros(4, 3, 2, dtype=torch.float64) + points[:, :, 0] = torch.linspace(0.0, 1.0, 3) + points[2:, :, 1] = 1.0 + monkeypatch.setattr(adapter, "_cpu_workers", lambda *args: 2) + # Load before injecting an error so this tests request cleanup, not JIT. + _retime_toppra_warp(points, 1.0, 2.0, sample_count=11) + second_started, second_finished = Event(), Event() + original = adapter._cpu_chunk + + def faulty(chunk, *args, **kwargs): + if chunk[0, 0, 1] == 0: + assert second_started.wait(5) + raise RuntimeError("injected CPU worker failure") + second_started.set() + result = original(chunk, *args, **kwargs) + second_finished.set() + return result + + with monkeypatch.context() as patch: + patch.setattr(adapter, "_cpu_chunk", faulty) + with pytest.raises(RuntimeError, match="injected CPU worker failure"): + _retime_toppra_warp(points, 1.0, 2.0, sample_count=11) + assert second_finished.is_set() + assert _retime_toppra_warp(points, 1.0, 2.0, sample_count=11)["success"].all() + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="fork is unavailable") +def test_cpu_pool_is_recreated_after_fork() -> None: + # Isolate fork from any CUDA context or other native worker pools in pytest. + code = textwrap.dedent(""" + import os, signal, time + import torch + from embodichain.compute.trajectory import _toppra_warp as m + m._cpu_workers=lambda *args:2 + points=torch.linspace(0,1,24,dtype=torch.float64).reshape(4,3,2) + m._retime_toppra_warp(points,1.,2.,sample_count=11) + child=os.fork() + if child==0: + try: + result=m._retime_toppra_warp(points,1.,2.,sample_count=11) + os._exit(0 if result['success'].all() else 2) + except BaseException: + os._exit(3) + deadline=time.monotonic()+8 + while time.monotonic() None: planner = ToppraPlanner.__new__(ToppraPlanner) planner.robot = robot planner.device = robot.device - planner.cfg = ToppraPlannerCfg(robot_uid="robot", max_workers=2, mp_context="spawn") + planner.cfg = ToppraPlannerCfg( + robot_uid="robot", backend="numpy", max_workers=2, mp_context="spawn" + ) planner._pool = None planner._mp_context = "spawn" options = planner.default_plan_options() diff --git a/tests/sim/motion/planners/test_toppra_batched.py b/tests/sim/motion/planners/test_toppra_batched.py index e5dfc1386..eb5db4dcc 100644 --- a/tests/sim/motion/planners/test_toppra_batched.py +++ b/tests/sim/motion/planners/test_toppra_batched.py @@ -27,18 +27,31 @@ from embodichain.lab.sim.motion.planners.utils import TrajectorySampleMethod -# Note: TOPPRA's ParametrizeConstAccel robustly avoids returning None for smooth -# splines (verified empirically across tiny limits, huge displacements, zigzags, -# and plateaus), so the ``jnt_traj is None`` branch in ``_toppra_solve_one_env`` -# is defensive and not directly unit-tested here; infeasible inputs instead -# raise during path/constraint construction and are caught by the -# ``except Exception`` -> ``_empty_failure`` path. +class _UnboundedRobot: + """Supply real planner limit metadata without constructing a simulator.""" + + mimic_ids = () + mimic_parents = () + mimic_multipliers = () + control_parts = {} + + def __init__(self, dofs: int = 1, speed: float = 1000.0) -> None: + self.joint_names = [f"joint_{index}" for index in range(dofs)] + self.speed = speed + + def get_qvel_limits(self, *, joint_ids: list[int]) -> torch.Tensor: + return torch.full((256, len(joint_ids)), self.speed, dtype=torch.float64) + + +# Numerical contracts are also exercised without the simulation adapter in +# tests/compute/test_toppra.py. class TestToppraWorker: @pytest.mark.no_sim def test_backend_defaults_plan_without_an_explicit_sample_count(self): from embodichain.lab.sim.motion.planners.toppra_planner import ToppraPlanner planner = object.__new__(ToppraPlanner) + planner.robot = _UnboundedRobot() planner._pool = None options = planner.default_plan_options() assert options.sample_method is TrajectorySampleMethod.TIME @@ -71,7 +84,7 @@ def test_solve_one_env_quantity(self): assert out["dt"].shape == (20,) def test_solve_one_env_infeasible_exception(self): - # Single waypoint -> SplineInterpolator raises -> caught, returns failure + # A path needs at least two waypoints; the worker isolates this failure. wp = np.array([[0.0] * 6]) out = _toppra_solve_one_env( waypoints=wp, @@ -286,7 +299,9 @@ def test_plan_batched_pool_path(self, mp_context): ) sim.prepare() planner = ToppraPlanner( - ToppraPlannerCfg(robot_uid="p", max_workers=2, mp_context=mp_context) + ToppraPlannerCfg( + robot_uid="p", backend="numpy", max_workers=2, mp_context=mp_context + ) ) assert planner._mp_context == mp_context try: @@ -355,7 +370,10 @@ def test_workers_reaped_on_gc(self, mp_context): sim.prepare() planner = ToppraPlanner( ToppraPlannerCfg( - robot_uid="close_reap", max_workers=2, mp_context=mp_context + robot_uid="close_reap", + backend="numpy", + max_workers=2, + mp_context=mp_context, ) ) assert planner._mp_context == mp_context @@ -455,3 +473,553 @@ def test_batched_equals_inline_single(self): import embodichain.lab.sim as om om.SimulationManager.flush_cleanup_queue() + + +@pytest.mark.no_sim +@pytest.mark.parametrize("backend", ["numpy", "warp", "auto"]) +def test_in_tree_planner_preserves_batch_timing_and_failures(backend: str) -> None: + from embodichain.lab.sim.motion.planners.toppra_planner import ( + ToppraPlanner, + ToppraPlannerCfg, + ToppraPlanOptions, + ) + from embodichain.lab.sim.motion.planners.utils import PlanState + + planner = object.__new__(ToppraPlanner) + planner.robot = _UnboundedRobot() + planner.cfg = ToppraPlannerCfg(robot_uid="unused", backend=backend, max_workers=1) + planner.device = torch.device("cpu") + planner._pool = None + states = [ + PlanState.from_qpos(torch.tensor([[0.0], [0.0], [2.0]])), + PlanState.from_qpos(torch.tensor([[1.0], [float("nan")], [2.0]])), + ] + result = planner.plan( + states, + ToppraPlanOptions( + constraints={"velocity": [[-0.5, 1.0]], "acceleration": [[-2.0, 1.0]]}, + sample_method=TrajectorySampleMethod.TIME, + sample_interval=0.03, + ), + ) + assert result.success.tolist() == [True, False, True] + assert planner._pool is None + assert result.duration[0] > 0 + assert result.duration[1:].count_nonzero() == 0 + assert result.positions[2].eq(2.0).all() + assert result.velocities[2].count_nonzero() == 0 + assert result.dt.max() <= 0.03 + reference = _toppra_solve_one_env( + np.array([[0.0], [1.0]]), + [[-0.5, 1.0]], + [[-2.0, 1.0]], + TrajectorySampleMethod.TIME, + 0.03, + ) + np.testing.assert_allclose(result.positions[0], reference["positions"], atol=1e-6) + np.testing.assert_allclose(result.dt[0], reference["dt"], atol=1e-6) + + +@pytest.mark.no_sim +@pytest.mark.parametrize("backend", ["numpy", "warp"]) +def test_in_tree_planner_invalid_limits_rejected_before_dispatch(backend: str) -> None: + from embodichain.lab.sim.motion.planners.toppra_planner import ( + ToppraPlanner, + ToppraPlannerCfg, + ToppraPlanOptions, + ) + from embodichain.lab.sim.motion.planners.utils import PlanState + + planner = object.__new__(ToppraPlanner) + planner.robot = _UnboundedRobot() + planner.cfg = ToppraPlannerCfg(robot_uid="unused", backend=backend, max_workers=1) + planner.device = torch.device("cpu") + planner._pool = None + with pytest.raises(ValueError, match="positive"): + planner.plan( + [ + PlanState.from_qpos(torch.zeros(1, 1)), + PlanState.from_qpos(torch.ones(1, 1)), + ], + ToppraPlanOptions( + sample_method=TrajectorySampleMethod.QUANTITY, + constraints={"velocity": -1.0, "acceleration": 2.0}, + sample_interval=10, + ), + ) + + +@pytest.mark.no_sim +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float16, torch.float64]) +def test_cuda_auto_backend_keeps_plan_on_device( + monkeypatch: pytest.MonkeyPatch, + dtype: torch.dtype, +) -> None: + from embodichain.lab.sim.motion.planners.toppra_planner import ( + ToppraPlanner, + ToppraPlannerCfg, + ToppraPlanOptions, + ) + from embodichain.lab.sim.motion.planners.utils import PlanState + + if not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + planner = object.__new__(ToppraPlanner) + planner.robot = _UnboundedRobot() + planner.cfg = ToppraPlannerCfg(robot_uid="unused") + planner.device = torch.device("cuda") + planner._pool = None + points = torch.tensor( + [ + [[0.0], [0.4]], + [[1.0], [1.0]], + [[0.0], [float("nan")]], + ], + device="cuda", + dtype=dtype, + ) + states = [PlanState.from_qpos(points[:, i]) for i in range(2)] + with monkeypatch.context() as patch: + + def forbidden(*args, **kwargs): + raise AssertionError("auto CUDA planning must not copy waypoints to CPU") + + patch.setattr(torch.Tensor, "cpu", forbidden) + result = planner.plan( + states, + ToppraPlanOptions( + sample_method=TrajectorySampleMethod.QUANTITY, sample_interval=128.9 + ), + ) + assert result.success.tolist() == [True, True, False] + assert result.positions.is_cuda and result.dt.is_cuda + assert result.positions.shape == (3, 128, 1) + assert result.positions.dtype == result.dt.dtype == torch.float32 + assert planner._pool is None + assert result.duration[0] > 0 + assert result.duration[1:].count_nonzero() == 0 + torch.testing.assert_close(result.positions[0, -1], points[0, -1].float()) + + +@pytest.mark.no_sim +@pytest.mark.parametrize("backend", ["numpy", "warp"]) +@pytest.mark.parametrize( + "dtype", [torch.float16, torch.float32, torch.float64, torch.int64] +) +@pytest.mark.parametrize("quantity", [20, 20.0, 20.9]) +def test_legacy_planner_output_dtype_and_quantity_coercion( + backend: str, dtype: torch.dtype, quantity: float | int +) -> None: + from embodichain.lab.sim.motion.planners.toppra_planner import ( + ToppraPlanner, + ToppraPlannerCfg, + ToppraPlanOptions, + ) + from embodichain.lab.sim.motion.planners.utils import PlanState + + planner = object.__new__(ToppraPlanner) + planner.robot = _UnboundedRobot() + planner.cfg = ToppraPlannerCfg(robot_uid="unused", backend=backend, max_workers=1) + planner.device = torch.device("cpu") + planner._pool = None + result = planner.plan( + [ + PlanState.from_qpos(torch.zeros(1, 1, dtype=dtype)), + PlanState.from_qpos(torch.ones(1, 1, dtype=dtype)), + ], + ToppraPlanOptions( + sample_method=TrajectorySampleMethod.QUANTITY, sample_interval=quantity + ), + ) + assert result.success.tolist() == [True] + assert result.positions.shape == (1, 20, 1) + for values in ( + result.positions, + result.velocities, + result.accelerations, + result.dt, + ): + assert values.dtype == torch.float32 + + +@pytest.mark.no_sim +@pytest.mark.parametrize("backend", ["auto", "warp"]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.float32, torch.float64]) +def test_planner_preserves_waypoint_gradients_and_duration_contract( + backend: str, + dtype: torch.dtype, +) -> None: + from embodichain.lab.sim.motion.planners.toppra_planner import ( + ToppraPlanner, + ToppraPlannerCfg, + ToppraPlanOptions, + ) + from embodichain.lab.sim.motion.planners.utils import PlanState + + planner = object.__new__(ToppraPlanner) + planner.robot = _UnboundedRobot() + planner.cfg = ToppraPlannerCfg(robot_uid="unused", backend=backend) + planner.device = torch.device("cpu") + planner._pool = None + # Use a non-leaf input to verify gradients reach the caller's computation. + parameter = torch.tensor([[0.365]], dtype=dtype, requires_grad=True) + distance = 2 * parameter + result = planner.plan( + [ + PlanState.from_qpos(torch.zeros_like(distance)), + PlanState.from_qpos(distance), + ], + ToppraPlanOptions( + sample_method=TrajectorySampleMethod.QUANTITY, + constraints={"velocity": 100.0, "acceleration": 2.0}, + sample_interval=17, + ), + ) + assert result.success.all() + assert planner._pool is None + assert result.positions.dtype == torch.float32 + assert result.positions.requires_grad and result.duration.requires_grad + result.duration.sum().backward() + expected = result.duration.detach().to(dtype) / (2 * parameter.detach().flatten()) + torch.testing.assert_close(parameter.grad.flatten(), expected) + + +@pytest.mark.no_sim +def test_numpy_planner_rejects_autograd_instead_of_detaching() -> None: + from embodichain.lab.sim.motion.planners.toppra_planner import ( + ToppraPlanner, + ToppraPlannerCfg, + ToppraPlanOptions, + ) + from embodichain.lab.sim.motion.planners.utils import PlanState + + planner = object.__new__(ToppraPlanner) + planner.robot = _UnboundedRobot() + planner.cfg = ToppraPlannerCfg(robot_uid="unused", backend="numpy") + planner.device = torch.device("cpu") + planner._pool = None + states = [ + PlanState.from_qpos(torch.zeros(1, 1)), + PlanState.from_qpos(torch.ones(1, 1, requires_grad=True)), + ] + options = ToppraPlanOptions( + sample_method=TrajectorySampleMethod.TIME, + sample_interval=0.1, + ) + with pytest.raises(NotImplementedError, match="Warp backend"): + planner.plan(states, options) + + +@pytest.mark.no_sim +@pytest.mark.parametrize("backend", ["auto", "warp"]) +@pytest.mark.parametrize( + "sampling", [TrajectorySampleMethod.TIME, TrajectorySampleMethod.QUANTITY] +) +def test_planner_autograd_routes_trainable_limits_without_trainable_waypoints( + backend: str, + sampling: TrajectorySampleMethod, +) -> None: + from embodichain.lab.sim.motion.planners.toppra_planner import ( + ToppraPlanner, + ToppraPlannerCfg, + ToppraPlanOptions, + ) + from embodichain.lab.sim.motion.planners.utils import PlanState + + planner = object.__new__(ToppraPlanner) + planner.robot = _UnboundedRobot() + planner.cfg = ToppraPlannerCfg(robot_uid="unused", backend=backend) + planner.device = torch.device("cpu") + planner._pool = None + parameter = torch.tensor(0.6, dtype=torch.float64, requires_grad=True) + acceleration = parameter.exp() # Verify the caller's graph is preserved. + result = planner.plan( + [ + PlanState.from_qpos(torch.zeros(1, 1)), + PlanState.from_qpos(torch.full((1, 1), 0.73)), + ], + ToppraPlanOptions( + constraints={"velocity": 100.0, "acceleration": acceleration}, + sample_method=sampling, + sample_interval=0.13 if sampling == TrajectorySampleMethod.TIME else 17, + ), + ) + assert result.success.all() and result.duration.requires_grad + assert result.positions.dtype == torch.float32 and planner._pool is None + result.duration.sum().backward() + torch.testing.assert_close( + parameter.grad, -result.duration.sum().detach().double() / 2 + ) + + +@pytest.mark.no_sim +@pytest.mark.parametrize("backend", ["numpy", "warp", "auto"]) +def test_no_grad_planning_accepts_trainable_inputs_as_constants(backend: str) -> None: + from embodichain.lab.sim.motion.planners.toppra_planner import ( + ToppraPlanner, + ToppraPlannerCfg, + ToppraPlanOptions, + ) + from embodichain.lab.sim.motion.planners.utils import PlanState + + planner = object.__new__(ToppraPlanner) + planner.robot = _UnboundedRobot() + planner.cfg = ToppraPlannerCfg(robot_uid="unused", backend=backend) + planner.device = torch.device("cpu") + planner._pool = None + end = torch.ones(1, 1, requires_grad=True) + limit = torch.tensor(1.0, requires_grad=True) + with torch.no_grad(): + result = planner.plan( + [PlanState.from_qpos(torch.zeros_like(end)), PlanState.from_qpos(end)], + ToppraPlanOptions( + sample_method=TrajectorySampleMethod.QUANTITY, + constraints={"velocity": limit, "acceleration": 2 * limit}, + sample_interval=17, + ), + ) + assert result.success.all() and not result.duration.requires_grad + + +@pytest.mark.no_sim +def test_toppra_option_copies_preserve_tensor_graphs_and_isolate_containers() -> None: + from embodichain.lab.sim.motion.planners.toppra_planner import ToppraPlanOptions + + parameter = torch.tensor(0.6, requires_grad=True) + limit = parameter.exp() + options = ToppraPlanOptions(constraints={"velocity": [1.0], "acceleration": limit}) + copied = options.copy() + assert copied.constraints is not options.constraints + assert copied.constraints["acceleration"] is limit + copied.constraints["velocity"][0] = 2.0 + assert options.constraints["velocity"] == [1.0] + assert ToppraPlanOptions().constraints == {"velocity": None, "acceleration": None} + + +@pytest.mark.no_sim +def test_motion_generator_keeps_toppra_gradients_through_option_copy_and_resampling() -> ( + None +): + from embodichain.lab.sim.motion.motion_generator import ( + MotionGenerator, + MotionGenOptions, + ) + from embodichain.lab.sim.motion.planners.toppra_planner import ( + ToppraPlanner, + ToppraPlannerCfg, + ToppraPlanOptions, + ) + from embodichain.lab.sim.motion.planners.utils import PlanState + + planner = object.__new__(ToppraPlanner) + planner.robot = _UnboundedRobot() + planner.cfg = ToppraPlannerCfg(robot_uid="unused") + planner.device = torch.device("cpu") + planner._pool = None + generator = object.__new__(MotionGenerator) + generator.planner = planner + generator.device = planner.device + distance = torch.tensor([[0.73]], dtype=torch.float64, requires_grad=True) + parameter = torch.tensor(0.6, dtype=torch.float64, requires_grad=True) + result = generator.generate( + [ + PlanState.from_qpos(torch.zeros_like(distance)), + PlanState.from_qpos(distance), + ], + MotionGenOptions( + sample_count=29, + plan_opts=ToppraPlanOptions( + sample_method=TrajectorySampleMethod.QUANTITY, + constraints={"velocity": 100.0, "acceleration": parameter.exp()}, + sample_interval=17, + ), + ), + ) + assert result.success.all() and result.positions.shape == (1, 29, 1) + duration = result.duration.sum() + dq, da = torch.autograd.grad(duration, (distance, parameter)) + torch.testing.assert_close(dq, duration.detach().double() / (2 * distance.detach())) + torch.testing.assert_close(da, -duration.detach().double() / 2) + + +@pytest.mark.no_sim +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float16, torch.float32]) +def test_cuda_time_planner_backward_preserves_input_gradients( + dtype: torch.dtype, +) -> None: + from embodichain.lab.sim.motion.planners.toppra_planner import ( + ToppraPlanner, + ToppraPlannerCfg, + ToppraPlanOptions, + ) + from embodichain.lab.sim.motion.planners.utils import PlanState + + if not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + planner = object.__new__(ToppraPlanner) + planner.robot = _UnboundedRobot() + planner.cfg = ToppraPlannerCfg(robot_uid="unused") + planner.device = torch.device("cuda") + planner._pool = None + distance = torch.tensor([[0.73]], device="cuda", dtype=dtype, requires_grad=True) + acceleration = torch.tensor(1.7, device="cuda", requires_grad=True) + options = ToppraPlanOptions( + constraints={"velocity": 100.0, "acceleration": acceleration}, + sample_method=TrajectorySampleMethod.TIME, + sample_interval=0.13, + ) + result = planner.plan( + [ + PlanState.from_qpos(torch.zeros_like(distance)), + PlanState.from_qpos(distance), + ], + options, + ) + duration = result.duration.sum() + dq, da = torch.autograd.grad(duration, (distance, acceleration)) + assert result.success.all() and dq.is_cuda and da.is_cuda + expected_dq = duration.detach().double() / (2 * distance.detach().double()) + torch.testing.assert_close(dq, expected_dq.to(dtype)) + torch.testing.assert_close(da, -duration.detach() / (2 * acceleration.detach())) + + +@pytest.mark.no_sim +def test_cpu_auto_planner_uses_warp_without_numpy_workers(monkeypatch) -> None: + from embodichain.lab.sim.motion.planners import toppra_planner as module + from embodichain.compute.trajectory._toppra_warp import _retime_toppra_warp + from embodichain.lab.sim.motion.planners.utils import PlanState + + def forbidden(*args, **kwargs): + raise AssertionError("CPU auto planning must not use the NumPy process chain") + + monkeypatch.setattr(module, "_toppra_solve_one_env", forbidden) + monkeypatch.setattr(module.ToppraPlanner, "_get_pool", forbidden) + planner = object.__new__(module.ToppraPlanner) + planner.robot = _UnboundedRobot(7) + planner.cfg = module.ToppraPlannerCfg(robot_uid="unused") + planner.device = torch.device("cpu") + planner._pool = None + points = torch.zeros(32, 4, 7) + points[:, :, 0] = torch.linspace(0.0, 1.0, 4) + actual = planner.plan( + [PlanState.from_qpos(points[:, index]) for index in range(4)], + module.ToppraPlanOptions( + sample_method=TrajectorySampleMethod.QUANTITY, + constraints={"velocity": 1.0, "acceleration": 2.0}, + sample_interval=17, + ), + ) + reference = _retime_toppra_warp(points, 1.0, 2.0, sample_count=17) + for key, value in reference.items(): + torch.testing.assert_close(getattr(actual, key), value, rtol=0, atol=0) + assert actual.success.all() and planner._pool is None + + +@pytest.mark.no_sim +@pytest.mark.parametrize( + "dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64] +) +@pytest.mark.parametrize("grad", [False, True]) +def test_shared_cpu_retimer_preserves_cleanup_across_batch_sizes(dtype, grad) -> None: + from embodichain.lab.sim.motion.planners.toppra_planner import retime_joint_paths + + # float32(1e-6) crosses the NumPy cleanup boundary after double promotion. + # The same input must keep the same spline in single/batched/grad routes. + points = torch.tensor([[[0.0], [1e-6], [0.5]]], dtype=dtype, requires_grad=grad) + limits = {"velocity": 1.0, "acceleration": 2.0} + single = retime_joint_paths(points, constraints=limits, sample_interval=17) + batch = retime_joint_paths( + points.repeat(2, 1, 1), constraints=limits, sample_interval=17 + ) + assert single.success.all() and batch.success.all() + expected_dtype = dtype if dtype in (torch.float32, torch.float64) else torch.float32 + assert single.positions.dtype == batch.positions.dtype == expected_dtype + for key in ("positions", "velocities", "accelerations", "dt"): + torch.testing.assert_close( + getattr(single, key)[0], getattr(batch, key)[0], rtol=0, atol=0 + ) + + +@pytest.mark.no_sim +@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.gpu)]) +@pytest.mark.parametrize("requested", [0.3, 0.9]) +def test_planner_limit_gradients_survive_runtime_speed_caps( + device: str, requested: float +) -> None: + from embodichain.lab.sim.motion.planners.toppra_planner import ( + ToppraPlanner, + ToppraPlannerCfg, + ToppraPlanOptions, + ) + from embodichain.lab.sim.motion.planners.utils import PlanState + + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + planner = object.__new__(ToppraPlanner) + planner.robot = _UnboundedRobot(speed=0.5) + planner.cfg = ToppraPlannerCfg(robot_uid="unused") + planner.device = torch.device(device) + planner._pool = None + states = [ + PlanState.from_qpos(torch.tensor([[q]], device=device)) for q in (0.0, 2.0) + ] + parameter = torch.tensor( + requested, dtype=torch.float64, device=device, requires_grad=True + ) + + def solve(velocity: torch.Tensor): + return planner.plan( + states, + ToppraPlanOptions( + constraints={"velocity": velocity, "acceleration": 1000.0}, + sample_method=TrajectorySampleMethod.QUANTITY, + sample_interval=101, + ), + ) + + result = solve(parameter) + assert result.success.all() and result.duration.requires_grad + assert result.velocities.abs().max() <= 0.5 + 1e-6 + derivative = torch.autograd.grad(result.duration.sum(), parameter)[0] + if requested > 0.5: + assert derivative == 0 + else: + # An unsaturated speed request remains connected through the CPU + # runtime-limit intersection and the device-local Warp timing solve. + epsilon = 1e-3 + numerical = ( + solve(parameter.detach() + epsilon).duration.sum() + - solve(parameter.detach() - epsilon).duration.sum() + ) / (2 * epsilon) + assert derivative < 0 and derivative.device == parameter.device + torch.testing.assert_close(derivative, numerical.double(), rtol=3e-3, atol=3e-3) + + +@pytest.mark.no_sim +def test_toppra_cfg_keeps_legacy_positional_worker_fields() -> None: + from embodichain.lab.sim.motion.planners.toppra_planner import ToppraPlannerCfg + + cfg = ToppraPlannerCfg("robot", "toppra", 0, 2, "spawn") + assert cfg.max_workers == 2 and cfg.mp_context == "spawn" + assert cfg.backend == "auto" and cfg.grid_size == 100 + + +@pytest.mark.no_sim +@pytest.mark.parametrize( + "dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64] +) +def test_shared_retimer_failure_preserves_promoted_dtype(dtype: torch.dtype) -> None: + from embodichain.lab.sim.motion.planners.toppra_planner import retime_joint_paths + + result = retime_joint_paths( + torch.zeros(1, 2, 1, dtype=dtype), + constraints={"velocity": -1.0, "acceleration": 2.0}, + sample_interval=17, + ) + assert not result.success.any() + expected = dtype if dtype in (torch.float32, torch.float64) else torch.float32 + for key in ("positions", "velocities", "accelerations", "dt"): + value = getattr(result, key) + assert value.dtype == expected and value.count_nonzero() == 0 diff --git a/tests/sim/motion/test_motion_imports.py b/tests/sim/motion/test_motion_imports.py index 6ec0bed82..97530a91c 100644 --- a/tests/sim/motion/test_motion_imports.py +++ b/tests/sim/motion/test_motion_imports.py @@ -33,7 +33,14 @@ def test_motion_import_order_and_config_resolution(first_module: str) -> None: script = """ import importlib +import importlib.abc import sys + +class BlockExternalToppra(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname == "toppra" or fullname.startswith("toppra."): + raise AssertionError("External toppra must not be imported") +sys.meta_path.insert(0, BlockExternalToppra()) import dexsim importlib.import_module(sys.argv[1]) @@ -104,3 +111,35 @@ def test_motion_import_order_and_config_resolution(first_module: str) -> None: timeout=60, ) assert result.returncode == 0, result.stdout + result.stderr + + +def test_toppra_planning_and_backward_work_without_external_package() -> None: + script = """ +import importlib.abc +import sys +class BlockExternalToppra(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname == "toppra" or fullname.startswith("toppra."): + raise AssertionError("External toppra must not be imported") +sys.meta_path.insert(0, BlockExternalToppra()) +import numpy as np +import torch +from embodichain.compute.trajectory._toppra import _NumpyToppra +from embodichain.compute.trajectory._toppra_warp import _retime_toppra_warp +points = torch.tensor([[[0.0], [0.73]]], dtype=torch.float64, requires_grad=True) +acceleration = torch.tensor(1.7, dtype=torch.float64, requires_grad=True) +reference = _NumpyToppra(points.detach().numpy()[0], 1.0, 1.7) +result = _retime_toppra_warp(points, 1.0, acceleration, sample_dt=0.13) +assert result["success"].all() +assert np.isclose(result["dt"].sum().item(), reference.duration) +result["dt"].sum().backward() +assert torch.isfinite(points.grad).all() and torch.isfinite(acceleration.grad) +assert "toppra" not in sys.modules +""" + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + timeout=60, + ) + assert result.returncode == 0, result.stdout + result.stderr From c316cab197dd50512b773e82505ecfdb4d0e27a2 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 17 Oct 2026 09:14:09 +0000 Subject: [PATCH 2/3] refactor(motion): remove TOPPRA process backend --- .../topics/motion-planning/planner-details.md | 24 +- .../embodichain.lab.sim.motion.planners.rst | 9 +- .../overview/sim/motion/planners/index.rst | 2 +- .../sim/motion/planners/toppra_planner.md | 90 ++- embodichain/compute/trajectory/_toppra.py | 42 +- .../compute/trajectory/_toppra_common.py | 61 ++ .../compute/trajectory/_toppra_warp.py | 2 +- .../lab/sim/motion/planners/toppra_planner.py | 538 ++---------------- .../motion_generation/planners/toppra.py | 2 +- scripts/tutorials/sim/open_drawer.py | 2 - .../test_configured_integration.py | 4 +- .../test_simulation_environment.py | 6 +- .../sim/motion/planners/test_joint_limits.py | 41 -- .../motion/planners/test_toppra_batched.py | 504 ++++------------ .../motion/planners/test_toppra_planner.py | 1 - tests/sim/motion/test_motion_imports.py | 28 + tests/sim/utility/test_atom_action_utils.py | 28 +- 17 files changed, 378 insertions(+), 1006 deletions(-) create mode 100644 embodichain/compute/trajectory/_toppra_common.py diff --git a/agent_context/topics/motion-planning/planner-details.md b/agent_context/topics/motion-planning/planner-details.md index d0c00e834..736c554ba 100644 --- a/agent_context/topics/motion-planning/planner-details.md +++ b/agent_context/topics/motion-planning/planner-details.md @@ -77,12 +77,13 @@ retain their existing timing contracts and do not consume these backend limits. ### TOPPRA -`toppra_planner.py` selects Warp on CPU and CUDA in automatic mode. NumPy -remains an explicit reference backend. Numerical ownership is in -`compute/trajectory/_toppra.py`, `_toppra_warp.py`, and `_warp/toppra.py`. +`toppra_planner.py` uses Warp on CPU and CUDA. Numerical ownership is in +`compute/trajectory/_toppra_warp.py` and `_warp/toppra.py`; shared validation +lives in `_toppra_common.py`. `_toppra.py` is an offline NumPy oracle and must +not be imported by the runtime planner or Warp adapter. Robot-backed planning retains shared speed-cap and control-part resolution, including gradients of uncapped requested limits. -Both preserve uniform not-a-knot cubic geometry, the existing waypoint +The runtime and reference preserve uniform not-a-knot cubic geometry, the existing waypoint deduplication tolerance and final knot, and signed velocity/acceleration bounds containing rest. Planner outputs retain float32 and quantity sampling retains the existing integer coercion, independently of backend or input precision. @@ -97,18 +98,17 @@ their numerical buffers, inherit Torch grad/inference modes and merge using the global endpoint/zero-derivative/zero-dt padding contract. Ordinary inference launches must not record into Warp's process-global tape. CPU pool state is discarded after fork; it is not owned by an individual planner's `close()`. -Explicit NumPy selection fans batched solves into independent CPU processes, -which must remain picklable and avoid CUDA/Warp/simulation state. -Single-worker/single-row NumPy execution can run inline. The automatic process -context uses spawn with GPU simulation; do not force fork to reduce overhead. +Planning has no process backend, per-planner worker state or process cleanup. +The removed `backend`, `max_workers` and `mp_context` configuration fields must +be removed from callers; do not restore them as ignored compatibility options. Warp records first-order adjoints when Torch gradient tracking is enabled and waypoints or tensor velocity/acceleration bounds require gradients. Both time and quantity sampling propagate position, velocity, acceleration and duration losses through the timing solve. Scalar, per-joint and signed tensor bounds keep their graphs on either device; TOPPRA option copies preserve these tensor -references while copying mutable containers. Explicit NumPy selection rejects -gradient requests; `torch.no_grad()` retains ordinary forward planning. +references while copying mutable containers. `torch.no_grad()` retains +ordinary forward planning. Filtering, grid topology, sample counts and active-constraint choices are discrete: derivatives apply locally away from their switches. Time sampling keeps `ceil(duration / sample_dt) + 1` points uniformly spread across the full @@ -119,8 +119,8 @@ generator option copying/resampling and CUDA backward stream behavior in the TOPPRA tests. Variable-length outputs pad with the terminal position without adding duration. -Per-row failures preserve other rows; a broken worker pool is discarded and -rebuilt on a later call. Test these lifecycle boundaries in +Per-row failures preserve other rows. Test planner contracts and the absence +of process workers in `tests/sim/motion/planners/test_toppra_batched.py`; numerical limits, geometry, CPU/Warp agreement and CUDA stream behavior are covered in `tests/compute/test_toppra.py`; CPU projection, thread isolation, variable-length diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.planners.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.planners.rst index 0b166f09b..85444febe 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.planners.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.planners.rst @@ -148,8 +148,13 @@ Utilities Joint path retiming ------------------ -Re-parameterize existing joint paths with the differentiable Warp TOPPRA -backend. The helper preserves float32/float64 input precision, supports valid +TOPPRA uses a single Warp calculation chain on CPU and CUDA. Its configuration +accepts ``grid_size`` and the shared robot/manager fields; the removed +``backend``, ``max_workers`` and ``mp_context`` fields are no longer accepted. +CPU inference uses shared threads without creating planner-owned processes. + +Re-parameterize existing joint paths with the same differentiable timing +solve. The helper preserves float32/float64 input precision, supports valid per-row waypoint prefixes, and returns explicit timing with isolated failures. .. autosummary:: diff --git a/docs/source/overview/sim/motion/planners/index.rst b/docs/source/overview/sim/motion/planners/index.rst index a34ee47a3..17e6ad61b 100644 --- a/docs/source/overview/sim/motion/planners/index.rst +++ b/docs/source/overview/sim/motion/planners/index.rst @@ -17,7 +17,7 @@ Choose a planner - Native targets and requirements * - :doc:`toppra_planner` - Time-parameterizing joint paths under velocity and acceleration limits. - - Joint waypoints; requires TOPPRA. + - Joint waypoints; in-tree Warp on CPU or CUDA. * - :doc:`trapezoidal_planner` - Batched joint paths with trapezoidal or jerk-limited Double-S timing. - Joint waypoints; Torch or Warp. diff --git a/docs/source/overview/sim/motion/planners/toppra_planner.md b/docs/source/overview/sim/motion/planners/toppra_planner.md index 83b3c6a42..99915a561 100644 --- a/docs/source/overview/sim/motion/planners/toppra_planner.md +++ b/docs/source/overview/sim/motion/planners/toppra_planner.md @@ -1,8 +1,9 @@ # ToppraPlanner -`ToppraPlanner` time-parameterizes a joint path with -[TOPPRA](https://github.com/hungpham2511/toppra) under velocity and acceleration -limits. Install `toppra==0.6.3` before use. For Cartesian goals, use +`ToppraPlanner` time-parameterizes joint paths under velocity and acceleration +limits with the in-tree Warp implementation of +[TOPPRA](https://github.com/hungpham2511/toppra). CPU and CUDA use the same +calculation chain; the external `toppra` package is not required. For Cartesian goals, use [MotionGenerator](../motion_generator.md) to generate joint waypoints through IK. ## Direct planning @@ -30,6 +31,7 @@ goal_qpos[:, 0] += 0.1 # Choose a goal within this robot's limits. result = planner.plan( [PlanState.from_qpos(start_qpos), PlanState.from_qpos(goal_qpos)], ToppraPlanOptions( + control_part="arm", constraints={"velocity": 0.2, "acceleration": 0.5}, sample_method=TrajectorySampleMethod.TIME, sample_interval=0.01, @@ -41,17 +43,89 @@ planner.close() ## Constraints and sampling -- Set `constraints` on `ToppraPlanOptions`; a scalar applies to all controlled - joints and a list specifies per-joint limits. +- Set `constraints` on `ToppraPlanOptions`; positive scalars apply to all + controlled joints, vectors specify per-joint magnitudes, and `(DOF, 2)` + arrays specify signed bounds containing zero. Explicit velocity requests + are capped by the robot's runtime speed limits. `None` uses robot speed + limits and the default planning acceleration budget. +- `control_part` selects the trajectory's joint order. A partial-joint path + needs an explicit group when the mapping is ambiguous. Runtime zero-speed + joints must remain constant at every waypoint, including points that + would otherwise be removed by duplicate filtering. - `TIME` uses `sample_interval` as the requested time spacing in seconds; samples include both endpoints, so actual spacing can be smaller. - `QUANTITY` uses `sample_interval` as the output count, with at least two samples. Results follow the shared -[PlanResult contract](../motion_generator.md#inputs-and-results). Planning -runs on CPU; batched workloads can use the configured worker pool. Call -`close()` when a directly owned planner is no longer needed. +[PlanResult contract](../motion_generator.md#inputs-and-results). Planner +results use float32 on the robot's device. Internal spline fitting, interval +bounds, timing and sampling use float64. `grid_size` sets the minimum timing +grid size; segment alignment and path length can increase the actual size. +The interval bounds are conservative, so durations can differ from the +former external solver. Timing alone does not certify collision freedom. + +## CPU execution and configuration migration + +CPU inference automatically splits sufficiently large batches across a shared +pool of at most eight threads. Small requests and calls recording derivatives +run inline. CUDA numerical work remains on the input device. Planning does +not create a CPU process pool, and individual planners own no worker processes. +`close()` is an idempotent lifecycle hook; it does not close the shared CPU +thread executor. + +Remove `backend`, `max_workers`, and `mp_context` from existing +`ToppraPlannerCfg` constructors and configuration dictionaries. These options +are no longer accepted. The robot selects the device, and Warp is the only +runtime implementation: + +```python +cfg = ToppraPlannerCfg( + robot_uid=robot.uid, + sim_instance_id=sim.instance_id, + grid_size=100, +) +``` + +The NumPy implementation is an offline numerical reference used by tests; +it cannot be selected as a planner backend. + +## First-order gradients + +Waypoint tensors and tensor velocity/acceleration requests retain their graphs +through `ToppraPlanOptions` copying and the timing solve on CPU or CUDA. For +example, with a valid joint goal in the direct-planning setup above: + +```python +goal = goal_qpos.detach().requires_grad_() +result = planner.plan( + [PlanState.from_qpos(start_qpos), PlanState.from_qpos(goal)], + ToppraPlanOptions( + control_part="arm", + constraints={"velocity": 0.2, "acceleration": 0.5}, + sample_method=TrajectorySampleMethod.QUANTITY, + sample_interval=101, + ), +) +assert result.is_all_success() +result.duration.sum().backward() +assert goal.grad is not None +``` + +Gradients are local to fixed waypoint filtering, timing-grid topology, output +counts and active constraints. Higher-order derivatives are unsupported. +Failed rows contribute zero gradients; stationary rows differentiate only +their retained position. Robot speed-cap saturation stops gradients through +the capped request. Supply trainable boundary states through `PlanState.qpos`. + +## Retiming an existing path + +`retime_joint_paths()` exposes the same Warp timing solve for a `(B, N, DOF)` +tensor without constructing a planner. Provide explicit velocity and +acceleration limits; optional `path_lengths` exclude per-row tail padding +from spline fitting. Outputs preserve float32/float64 input precision, while +other input dtypes are promoted to float32. Recompute poses from the returned +positions because the helper resamples along the fitted spline. ## References diff --git a/embodichain/compute/trajectory/_toppra.py b/embodichain/compute/trajectory/_toppra.py index bddb731b0..3f6a4c359 100644 --- a/embodichain/compute/trajectory/_toppra.py +++ b/embodichain/compute/trajectory/_toppra.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Internal NumPy reference for rest-to-rest TOPPRA planning. +"""Offline NumPy reference for rest-to-rest TOPPRA planning. The path remains the planner's uniform, not-a-knot cubic spline. Squared path speeds at adjacent gridpoints obey half-planes a*x + b*y <= c. @@ -25,48 +25,12 @@ from __future__ import annotations -from numbers import Integral - import numpy as np from scipy.interpolate import CubicSpline -__all__ = [] - - -def _limits(value: object, dofs: int) -> np.ndarray: - """Canonicalize symmetric magnitudes or signed lower/upper bounds.""" - array = np.asarray(value, dtype=np.float64) - if array.ndim == 0: - array = np.full(dofs, float(array)) - if array.shape == (dofs,): - array = np.stack((-array, array), axis=-1) - if array.shape != (dofs, 2) or not np.isfinite(array).all(): - raise ValueError("limits must be finite scalar, (DOF,), or (DOF, 2)") - if np.any(array[:, 0] > 0) or np.any(array[:, 1] < 0): - raise ValueError("rest-to-rest limits must contain zero") - return array +from ._toppra_common import _limits, _validate_grid_size - -def _validate_sampling(sample_count: int | None, sample_dt: float | None) -> None: - if (sample_count is None) == (sample_dt is None): - raise ValueError("specify exactly one of sample_count and sample_dt") - if sample_count is not None and ( - isinstance(sample_count, bool) - or not isinstance(sample_count, Integral) - or sample_count < 2 - ): - raise ValueError("sample_count must be an integer >= 2") - if sample_dt is not None and (not np.isfinite(sample_dt) or sample_dt <= 0): - raise ValueError("sample_dt must be finite and positive") - - -def _validate_grid_size(grid_size: int) -> None: - if ( - isinstance(grid_size, bool) - or not isinstance(grid_size, Integral) - or grid_size < 3 - ): - raise ValueError("grid_size must be an integer >= 3") +__all__ = [] class _NumpyToppra: diff --git a/embodichain/compute/trajectory/_toppra_common.py b/embodichain/compute/trajectory/_toppra_common.py new file mode 100644 index 000000000..09a18637b --- /dev/null +++ b/embodichain/compute/trajectory/_toppra_common.py @@ -0,0 +1,61 @@ +# ---------------------------------------------------------------------------- +# 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 input validation for the Warp solver and offline numerical reference.""" + +from __future__ import annotations + +from numbers import Integral + +import numpy as np + +__all__ = [] + + +def _limits(value: object, dofs: int) -> np.ndarray: + """Canonicalize symmetric magnitudes or signed lower/upper bounds.""" + array = np.asarray(value, dtype=np.float64) + if array.ndim == 0: + array = np.full(dofs, float(array)) + if array.shape == (dofs,): + array = np.stack((-array, array), axis=-1) + if array.shape != (dofs, 2) or not np.isfinite(array).all(): + raise ValueError("limits must be finite scalar, (DOF,), or (DOF, 2)") + if np.any(array[:, 0] > 0) or np.any(array[:, 1] < 0): + raise ValueError("rest-to-rest limits must contain zero") + return array + + +def _validate_sampling(sample_count: int | None, sample_dt: float | None) -> None: + if (sample_count is None) == (sample_dt is None): + raise ValueError("specify exactly one of sample_count and sample_dt") + if sample_count is not None and ( + isinstance(sample_count, bool) + or not isinstance(sample_count, Integral) + or sample_count < 2 + ): + raise ValueError("sample_count must be an integer >= 2") + if sample_dt is not None and (not np.isfinite(sample_dt) or sample_dt <= 0): + raise ValueError("sample_dt must be finite and positive") + + +def _validate_grid_size(grid_size: int) -> None: + if ( + isinstance(grid_size, bool) + or not isinstance(grid_size, Integral) + or grid_size < 3 + ): + raise ValueError("grid_size must be an integer >= 3") diff --git a/embodichain/compute/trajectory/_toppra_warp.py b/embodichain/compute/trajectory/_toppra_warp.py index cb2d5dc30..f3cf825c9 100644 --- a/embodichain/compute/trajectory/_toppra_warp.py +++ b/embodichain/compute/trajectory/_toppra_warp.py @@ -33,7 +33,7 @@ import warp as wp from torch.autograd.function import once_differentiable -from ._toppra import _limits, _validate_grid_size, _validate_sampling +from ._toppra_common import _limits, _validate_grid_size, _validate_sampling from ._warp import toppra as kernels __all__ = [] diff --git a/embodichain/lab/sim/motion/planners/toppra_planner.py b/embodichain/lab/sim/motion/planners/toppra_planner.py index 5495668de..2514518d2 100644 --- a/embodichain/lab/sim/motion/planners/toppra_planner.py +++ b/embodichain/lab/sim/motion/planners/toppra_planner.py @@ -16,221 +16,44 @@ from __future__ import annotations -from dataclasses import field -from copy import deepcopy - -import os from collections.abc import Mapping, Sequence -from typing import Literal +from copy import deepcopy +from dataclasses import field import torch -import numpy as np -from concurrent.futures.process import BrokenProcessPool +from embodichain.compute.trajectory._toppra_common import _validate_grid_size +from embodichain.compute.trajectory._toppra_warp import _retime_toppra_warp from embodichain.utils import logger, configclass -from embodichain.lab.sim.motion.planners.utils import MoveType, TrajectorySampleMethod -from embodichain.lab.sim.motion.planners.base_planner import ( +from .base_planner import ( validate_plan_options, BasePlanner, BasePlannerCfg, PlanOptions, _infer_batch_size, ) -from .utils import PlanState, PlanResult -from ._joint_limits import ( - _DEFAULT_ACCELERATION, - _DEFAULT_VELOCITY, - _resolve_joint_constraints, -) - -from embodichain.compute.trajectory._toppra import ( - _NumpyToppra, - _validate_grid_size, - _validate_sampling, -) - - -def _toppra_solve_one_env( - waypoints: np.ndarray, - vel_constraint, - acc_constraint, - sample_method: "TrajectorySampleMethod", - sample_interval: float | int, - grid_size: int = 100, -) -> dict: - """Solve a single-env TOPPRA trajectory. Pure numpy/scipy — picklable, no torch/robot. +from ._joint_limits import _resolve_joint_constraints +from .utils import MoveType, TrajectorySampleMethod, PlanState, PlanResult - Args: - waypoints: ``(N, DOF)`` numpy array of joint waypoints. - vel_constraint / acc_constraint: Resolved scalar or per-DOF limits. - Standalone ``None`` inputs use compatibility defaults. Robot-backed - planning resolves effective limits before dispatching this worker. - sample_method: TIME or QUANTITY. - sample_interval: seconds (TIME) or sample count (QUANTITY). - - Returns: - dict with ``positions`` ``(N_b, DOF)``, ``velocities``, ``accelerations``, - ``dt`` ``(N_b,)``, ``success`` bool, and ``n`` int. - """ - dofs = waypoints.shape[1] - if vel_constraint is None: - vel_constraint = _DEFAULT_VELOCITY - if acc_constraint is None: - acc_constraint = _DEFAULT_ACCELERATION - try: - sample_dt = ( - sample_interval if sample_method == TrajectorySampleMethod.TIME else None - ) - sample_count = ( - int(sample_interval) - if sample_method == TrajectorySampleMethod.QUANTITY - else None - ) - _validate_sampling(sample_count, sample_dt) - trajectory = _NumpyToppra(waypoints, vel_constraint, acc_constraint, grid_size) - duration = trajectory.duration - if duration == 0: - count = 2 - elif sample_dt is not None: - count = max(2, int(np.ceil(duration / sample_dt)) + 1) - else: - count = int(sample_count) - ts = np.linspace(0.0, duration, count) - positions, velocities, accelerations = trajectory.sample(ts) - except (ValueError, RuntimeError, FloatingPointError, OverflowError): - return _empty_failure(dofs) - return { - "positions": positions, - "velocities": velocities, - "accelerations": accelerations, - "dt": np.diff(ts, prepend=0.0).astype(np.float32), - "success": True, - "n": len(ts), - } - - -def _empty_failure(dofs: int) -> dict: - z = np.zeros((2, dofs), dtype=np.float32) - return { - "positions": z, - "velocities": np.zeros_like(z), - "accelerations": np.zeros_like(z), - "dt": np.array([0.0, 0.0], dtype=np.float32), - "success": False, - "n": 2, - } - - -def _set_parent_death_signal() -> None: - r"""Best-effort: ask the kernel to SIGKILL this worker when its parent dies. - - This is the **only** cleanup mechanism that survives ``os._exit(0)`` in the - parent process, which is what :meth:`SimulationManager.destroy` does by - default (gated by ``EMBODICHAIN_SIM_EXIT_PROCESS``). ``os._exit`` skips - every Python-level finalizer — ``atexit`` handlers, ``__del__`` methods, - and ``concurrent.futures``' internal ``_python_exit`` that would otherwise - join/terminate daemon workers — so the worker processes would be orphaned - and reparented to init, leaving residual python processes behind. - - A kernel parent-death signal fires regardless of *how* the parent exits - (``os._exit``, a crash, or a normal return), so workers are reaped even - when no Python cleanup runs. No-op on non-Linux or if ``prctl`` is - unavailable; in that case workers still rely on ``__del__`` and the - daemon-process reaping that runs on a normal interpreter shutdown. - """ - try: - import ctypes - import signal - - libc = ctypes.CDLL(None) - PR_SET_PDEATHSIG = 1 - libc.prctl.argtypes = [ - ctypes.c_int, - ctypes.c_ulong, - ctypes.c_ulong, - ctypes.c_ulong, - ctypes.c_ulong, - ] - libc.prctl.restype = ctypes.c_int - libc.prctl(PR_SET_PDEATHSIG, signal.SIGKILL, 0, 0, 0) - except Exception: - pass - - -def _worker_init() -> None: - r"""Initializer run in each worker when the pool starts. - - Two responsibilities: - - 1. **Parent-death signal (Linux).** Install ``prctl(PR_SET_PDEATHSIG, - SIGKILL)`` so the kernel kills this worker the instant its parent - dies. This is what guarantees no residual worker processes when the - parent exits via ``os._exit(0)`` (see :func:`_set_parent_death_signal`). - Because of this, callers never need to explicitly shut the planner - down — workers self-terminate with the parent. - - 2. **atexit clearing (fork only).** When ``fork`` is used, the parent - has usually already initialized CUDA/GPU by the time the pool is - created, and ``fork`` copies the parent's atexit registry into the - child. Without clearing it the worker would try to run the parent's - cleanup routines on exit, which can deadlock or corrupt state. With - ``spawn`` this is not necessary, but the initializer still runs so - that ``fork`` remains usable for advanced callers who create the pool - before CUDA is initialized. - """ - import atexit - import multiprocessing as mp - - if mp.get_start_method() == "fork": - atexit._clear() - - _set_parent_death_signal() +__all__ = [ + "ToppraPlanner", + "ToppraPlannerCfg", + "ToppraPlanOptions", + "retime_joint_paths", +] -def _assemble_batched_results( - results: list[dict], - dofs: int, - device: torch.device, +def _failed_result( + batch: int, dofs: int, device: torch.device, dtype: torch.dtype ) -> 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 independent zero-filled rows in the caller's floating precision.""" + positions = torch.zeros((batch, 2, dofs), dtype=dtype, device=device) 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), + success=torch.zeros(batch, dtype=torch.bool, device=device), + positions=positions, + velocities=torch.zeros_like(positions), + accelerations=torch.zeros_like(positions), + dt=torch.zeros((batch, 2), dtype=dtype, device=device), ) @@ -321,7 +144,6 @@ def retime_joint_paths( path_lengths = path_lengths.contiguous() resolved_interval = count if sample_interval is None else sample_interval velocity, acceleration = constraints["velocity"], constraints["acceleration"] - from embodichain.compute.trajectory._toppra_warp import _retime_toppra_warp quantity = sample_method == TrajectorySampleMethod.QUANTITY waypoints = ( @@ -340,55 +162,19 @@ def retime_joint_paths( grid_size=grid_size, ) except (ValueError, FloatingPointError, OverflowError): - failed = _assemble_batched_results( - [_empty_failure(dofs) for _ in range(batch)], dofs, resolved_device - ) - return PlanResult( - success=failed.success, - **{ - key: getattr(failed, key).to(waypoints.dtype) - for key in ("positions", "velocities", "accelerations", "dt") - }, - ) + return _failed_result(batch, dofs, resolved_device, waypoints.dtype) return PlanResult( **{key: value.to(resolved_device) for key, value in result.items()} ) -__all__ = [ - "ToppraPlanner", - "ToppraPlannerCfg", - "ToppraPlanOptions", - "retime_joint_paths", -] - - @configclass class ToppraPlannerCfg(BasePlannerCfg): + """Configure the shared CPU/CUDA Warp time parameterizer.""" planner_type: str = "toppra" - backend: Literal["auto", "numpy", "warp"] = field(default="auto", kw_only=True) - """Auto selects Warp on CPU and CUDA; NumPy remains an explicit reference.""" grid_size: int = field(default=100, kw_only=True) """Minimum timing grid size; each spline segment is subdivided uniformly.""" - max_workers: int | None = None - """NumPy worker count. Ignored by Warp; None => min(cpu_count()//2, B).""" - mp_context: str | None = None - """Multiprocessing start method for the NumPy batched fan-out; ignored by Warp. - - ``None`` (default) auto-selects based on the simulation device: - ``'fork'`` on CPU and ``'spawn'`` on GPU. ``'fork'`` is faster — workers - inherit the parent's already-loaded modules, so pool startup is - near-instant — and is safe here because the TOPPRA worker - (:func:`_toppra_solve_one_env`) is pure numpy/scipy and never touches the - parent's Vulkan/Warp/CUDA context or render threads; ``_worker_init`` - clears the inherited atexit registry and installs ``prctl(PR_SET_PDEATHSIG)`` - so workers are reaped when the parent dies (incl. the ``os._exit`` path). - ``'spawn'`` is the safer choice when the parent has initialized CUDA - physics (``device='cuda'``) — fork-after-CUDA-init is the officially - unsupported case — or if fork deadlocks are observed, at the cost of - re-importing modules per worker. - """ class _ToppraConstraints(dict): @@ -444,138 +230,25 @@ def __post_init__(self) -> None: class ToppraPlanner(BasePlanner): - """Rest-to-rest path timing with NumPy or batched Warp TOPPRA. + """Rest-to-rest joint-path timing with differentiable CPU/CUDA Warp kernels. - Cubic path geometry uses uniform knots and not-a-knot boundaries. Timing - optimizes conservative interval constraints, so durations need not match - the former external TOPPRA package. CUDA auto mode keeps numerical work on - device; NumPy remains available explicitly for reference comparisons. + Uniform not-a-knot cubic geometry is timed under conservative interval + bounds. Results preserve the planner's float32 output contract. """ supported_move_types = frozenset({MoveType.JOINT_MOVE}) - def __init__(self, cfg: ToppraPlannerCfg): - r"""Initialize the TOPPRA trajectory planner. - - References: - - TOPPRA: Time-Optimal Path Parameterization for Robotic Systems (https://github.com/hungpham2511/toppra) - - Args: - cfg: Configuration object containing ToppraPlanner settings - """ - if cfg.backend not in {"auto", "numpy", "warp"}: - raise ValueError("TOPPRA backend must be auto, numpy, or warp") + def __init__(self, cfg: ToppraPlannerCfg) -> None: + """Validate the timing grid and resolve the configured robot.""" _validate_grid_size(cfg.grid_size) super().__init__(cfg) - self._pool = None - # Resolve the multiprocessing start method once, now that self.device - # (from BasePlanner) is known. None => auto: fork on CPU, spawn on GPU. - self._mp_context = self._resolve_mp_context(cfg.mp_context, self.device) - # No atexit / __del__-based shutdown is registered here: workers install - # prctl(PR_SET_PDEATHSIG) in _worker_init, so the kernel reaps them the - # moment the parent process dies — including the os._exit(0) path taken - # by SimulationManager.destroy(), which skips every Python finalizer. - # __del__ below only handles in-process GC of an abandoned planner. - def default_plan_options(self) -> ToppraPlanOptions: """Return a valid time-sampled default when no count is requested.""" return ToppraPlanOptions() - @staticmethod - def _resolve_mp_context(mp_context: str | None, device: torch.device) -> str: - """Return the multiprocessing start method to use for the worker pool. - - An explicit ``mp_context`` is honored as-is. ``None`` auto-selects: - ``'fork'`` on CPU (fast — workers inherit loaded modules; safe because - the worker is pure numpy/scipy), ``'spawn'`` everywhere else (safer - under an initialized CUDA context). - - Args: - mp_context: The cfg value; ``None`` means auto-select. - device: The physics device the planner's robot runs on. - - Returns: - One of ``'fork'`` / ``'spawn'``. - """ - if mp_context is not None: - return mp_context - return "fork" if device.type == "cpu" else "spawn" - - def _get_pool(self, batch_size: int): - if self._pool is not None: - return self._pool - import multiprocessing as mp - - max_workers = self.cfg.max_workers - if max_workers is None: - max_workers = max(1, min((os.cpu_count() or 2) // 2, batch_size)) - ctx = mp.get_context(self._mp_context) - from concurrent.futures import ProcessPoolExecutor - - self._pool = ProcessPoolExecutor( - max_workers=max_workers, - mp_context=ctx, - initializer=_worker_init, - ) - return self._pool - - def _shutdown_pool(self) -> None: - r"""Shut down the TOPPRA worker process pool (internal). - - We do **not** use ``shutdown(wait=True)``: with ``fork`` workers can - inherit the parent's CUDA/GPU context, causing them to deadlock inside - the driver at exit. ``wait=True`` would then hang the main process, - and if the user kills it the workers are left behind as residual - python processes. - - Instead we cancel pending work and then forcibly terminate/join/kill - every worker process so that this method returns only after all - workers are actually gone. - - This is **not** part of the public API. It exists for two internal - callers: the ``BrokenProcessPool`` recovery path in :meth:`plan`, and - ``__del__`` (so an abandoned planner in a long-running process does - not leak workers). Normal process exit does not rely on it — workers - install ``prctl(PR_SET_PDEATHSIG)`` in :func:`_worker_init` and are - reaped by the kernel when the parent dies, even under ``os._exit``. - """ - if self._pool is None: - return - - # Capture the worker process objects before shutdown clears them. - worker_processes = list(getattr(self._pool, "_processes", {}).values()) - - # Stop accepting new work and cancel any futures that have not started. - self._pool.shutdown(wait=False, cancel_futures=True) - - # Forcibly reap every worker. This is required for fork-based pools - # when the parent has initialized CUDA: the children inherit the - # context and may not exit cleanly on their own. - for proc in worker_processes: - if not proc.is_alive(): - continue - proc.terminate() - proc.join(timeout=5.0) - if proc.is_alive(): - proc.kill() - proc.join(timeout=1.0) - - self._pool = None - def close(self) -> None: - """Release TOPPRA worker processes owned by this planner.""" - self._shutdown_pool() - - def __del__(self): - # Only matters for in-process GC of an abandoned planner (and as a - # non-Linux fallback). Process-exit cleanup is handled by the kernel - # via PR_SET_PDEATHSIG installed in each worker, which survives the - # os._exit(0) path that SimulationManager.destroy() takes. - try: - self.close() - except Exception: - pass + """Complete the planner lifecycle; CPU workers belong to the shared runtime.""" @validate_plan_options(options_cls=ToppraPlanOptions) def plan( @@ -603,18 +276,6 @@ def plan( b = _infer_batch_size(target_states) or 1 dofs = target_states[0].qpos.shape[-1] - backend = self.cfg.backend - if backend not in {"auto", "numpy", "warp"}: - raise ValueError("TOPPRA backend must be auto, numpy, or warp") - needs_grad = torch.is_grad_enabled() and ( - any(state.qpos.requires_grad for state in target_states) - or any( - isinstance(limit, torch.Tensor) and limit.requires_grad - for limit in options.constraints.values() - ) - ) - if needs_grad and backend == "numpy": - raise NotImplementedError("TOPPRA gradients require the Warp backend") constraints = _resolve_joint_constraints( self.robot, options.constraints, @@ -622,113 +283,36 @@ def plan( dof=dofs, batch_size=b, ) - use_warp = backend != "numpy" - if use_warp: - from embodichain.compute.trajectory._toppra_warp import _retime_toppra_warp - waypoints_tensor = torch.stack([s.qpos for s in target_states], dim=1).to( - self.device - ) - if waypoints_tensor.dtype not in (torch.float32, torch.float64): - waypoints_tensor = waypoints_tensor.to(torch.float64) - try: - result = _retime_toppra_warp( - waypoints_tensor, - constraints["velocity"], - constraints["acceleration"], - sample_count=( - int(options.sample_interval) - if options.sample_method == TrajectorySampleMethod.QUANTITY - else None - ), - sample_dt=( - options.sample_interval - if options.sample_method == TrajectorySampleMethod.TIME - else None - ), - grid_size=self.cfg.grid_size, - ) - except (ValueError, FloatingPointError, OverflowError): - return self._assemble_batched_result( - [_empty_failure(dofs) for _ in range(b)], dofs - ) - # The existing planner contract returns float32 for all floating - # fields, independently of input precision or selected backend. - return PlanResult( - **{ - key: value if key == "success" else value.to(torch.float32) - for key, value in result.items() - } - ) - - # Build (B, N, DOF) numpy waypoints - waypoints = np.stack( - [s.qpos.detach().cpu().numpy() for s in target_states], axis=1 - ) # (B, N, DOF) - - vc = constraints["velocity"] - ac = constraints["acceleration"] - if isinstance(vc, torch.Tensor): - vc = vc.detach().cpu().numpy() - if isinstance(ac, torch.Tensor): - ac = ac.detach().cpu().numpy() - args_per_env = [ - ( - waypoints[i], - vc, - ac, - options.sample_method, - options.sample_interval, - self.cfg.grid_size, - ) - for i in range(b) - ] - - # Single-env planning never needs a process pool. - if b == 1: - results = [_toppra_solve_one_env(*a) for a in args_per_env] - else: - # Inline fallback for max_workers==1 or a single-core machine. - max_workers = self.cfg.max_workers - use_inline = (max_workers == 1) or ( - max_workers is None and ((os.cpu_count() or 2) // 2) <= 1 + waypoints_tensor = torch.stack([s.qpos for s in target_states], dim=1).to( + self.device + ) + if waypoints_tensor.dtype not in (torch.float32, torch.float64): + waypoints_tensor = waypoints_tensor.to(torch.float64) + try: + result = _retime_toppra_warp( + waypoints_tensor, + constraints["velocity"], + constraints["acceleration"], + sample_count=( + int(options.sample_interval) + if options.sample_method == TrajectorySampleMethod.QUANTITY + else None + ), + sample_dt=( + options.sample_interval + if options.sample_method == TrajectorySampleMethod.TIME + else None + ), + grid_size=self.cfg.grid_size, ) - if use_inline: - results = [_toppra_solve_one_env(*a) for a in args_per_env] - else: - pool = self._get_pool(b) - results = [None] * b - try: - futures = [ - pool.submit(_toppra_solve_one_env, *a) for a in args_per_env - ] - broken = False - for i, fut in enumerate(futures): - try: - results[i] = fut.result() - except BrokenProcessPool: - logger.log_warning( - "TOPPRA process pool broke; returning failure." - ) - self._shutdown_pool() - broken = True - break - except Exception: - results[i] = _empty_failure(dofs) - if broken: - for i in range(b): - if results[i] is None: - results[i] = _empty_failure(dofs) - except BrokenProcessPool: - # pool was already broken at submit time - logger.log_warning("TOPPRA process pool broke; returning failure.") - self._shutdown_pool() - for i in range(b): - if results[i] is None: - results[i] = _empty_failure(dofs) - - 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`.""" - return _assemble_batched_results(results, dofs, self.device) + except (ValueError, FloatingPointError, OverflowError): + return _failed_result(b, dofs, self.device, torch.float32) + # The existing planner contract returns float32 for all floating + # fields, independently of input precision. + return PlanResult( + **{ + key: value if key == "success" else value.to(torch.float32) + for key, value in result.items() + } + ) diff --git a/scripts/benchmark/motion_generation/planners/toppra.py b/scripts/benchmark/motion_generation/planners/toppra.py index 2f9c87943..3d4a40c7e 100644 --- a/scripts/benchmark/motion_generation/planners/toppra.py +++ b/scripts/benchmark/motion_generation/planners/toppra.py @@ -86,7 +86,7 @@ def plan(self, case: BenchmarkCase) -> PlanResult: ) def close(self) -> None: - """Release TOPPRA worker pools and drop the motion generator.""" + """Drop the motion generator through the common planner lifecycle.""" self._close_motion_generator() diff --git a/scripts/tutorials/sim/open_drawer.py b/scripts/tutorials/sim/open_drawer.py index e788cbe4e..04aa461c4 100644 --- a/scripts/tutorials/sim/open_drawer.py +++ b/scripts/tutorials/sim/open_drawer.py @@ -600,8 +600,6 @@ def main(args: argparse.Namespace | None = None) -> None: cfg=MotionGenCfg( planner_cfg=ToppraPlannerCfg( robot_uid=robot.uid, - # Keep this small tutorial deterministic across platforms. - max_workers=1, ), ) ) diff --git a/tests/gym/envs/task_program/test_configured_integration.py b/tests/gym/envs/task_program/test_configured_integration.py index 98d435cdb..190aa7c89 100644 --- a/tests/gym/envs/task_program/test_configured_integration.py +++ b/tests/gym/envs/task_program/test_configured_integration.py @@ -298,9 +298,9 @@ def test_execution_policy_planner_config_reaches_configured_adapter() -> None: def test_planner_policy_decoder_validates_backend_and_options() -> None: """Planner backends remain closed while typed options stay configurable.""" assert _decode_planner_config( - {"type": "toppra", "config": {"max_workers": 2}}, + {"type": "toppra", "config": {"grid_size": 127}}, path="execution policy.planner", - ) == {"type": "toppra", "config": {"max_workers": 2}} + ) == {"type": "toppra", "config": {"grid_size": 127}} with pytest.raises(ValueError, match="execution policy.planner.type"): _decode_planner_config( diff --git a/tests/gym/envs/task_program/test_simulation_environment.py b/tests/gym/envs/task_program/test_simulation_environment.py index 306547000..306ffa770 100644 --- a/tests/gym/envs/task_program/test_simulation_environment.py +++ b/tests/gym/envs/task_program/test_simulation_environment.py @@ -1600,7 +1600,7 @@ def test_adapter_factory_forwards_policy_planner_config( SimulationSceneBinding(registry_id="scene"), _profile_binding(), ) - planner_config = {"type": "toppra", "config": {"max_workers": 2}} + planner_config = {"type": "toppra", "config": {"grid_size": 127}} factory = SimulationTaskProgramAdapterFactory( registration, planner_config=planner_config, @@ -1626,13 +1626,13 @@ def _create_adapter(bound_environment: object, **kwargs: object): def test_planner_config_decodes_to_typed_runtime_cfg() -> None: """A serialized planner selection becomes the selected planner config type.""" planner_cfg = simulation_environment_module._planner_cfg_from_config( - {"type": "toppra", "config": {"max_workers": 2}}, + {"type": "toppra", "config": {"grid_size": 127}}, robot_uid="robot", ) assert isinstance(planner_cfg, simulation_environment_module.ToppraPlannerCfg) assert planner_cfg.robot_uid == "robot" - assert planner_cfg.max_workers == 2 + assert planner_cfg.grid_size == 127 def test_curobo_scene_binding_preserves_live_object_references() -> None: diff --git a/tests/sim/motion/planners/test_joint_limits.py b/tests/sim/motion/planners/test_joint_limits.py index 597d372bd..75b5d6b76 100644 --- a/tests/sim/motion/planners/test_joint_limits.py +++ b/tests/sim/motion/planners/test_joint_limits.py @@ -118,9 +118,6 @@ def _zero_speed_planner(planner_type: str) -> ToppraPlanner | TrapezoidalPlanner planner.robot.velocities[:, 2] = 0.0 planner.device = torch.device("cpu") planner.cfg = cfg_class(robot_uid="robot") - if planner_type == "toppra": - planner.cfg.max_workers = 1 - planner._pool = None return planner @@ -348,9 +345,6 @@ def test_facade_and_direct_plans_apply_robot_limits(planner_type: str) -> None: planner.robot = robot planner.device = robot.device planner.cfg = cfg_class(robot_uid="robot") - if planner_type == "toppra": - planner.cfg.max_workers = 1 - planner._pool = None generator = MotionGenerator.__new__(MotionGenerator) generator.planner = planner states = [ @@ -399,41 +393,6 @@ def test_curobo_export_uses_robot_speeds_and_planner_acceleration( assert exported["max_acceleration"] == [expected] * 3 -def test_toppra_spawn_workers_receive_resolved_joint_limits() -> None: - robot = _Robot() - planner = ToppraPlanner.__new__(ToppraPlanner) - planner.robot = robot - planner.device = robot.device - planner.cfg = ToppraPlannerCfg( - robot_uid="robot", backend="numpy", max_workers=2, mp_context="spawn" - ) - planner._pool = None - planner._mp_context = "spawn" - options = planner.default_plan_options() - options.control_part = "arm" - options.sample_method = TrajectorySampleMethod.QUANTITY - options.sample_interval = 100 - try: - result = planner.plan( - [ - PlanState.from_qpos(torch.zeros(2, 2)), - PlanState.from_qpos(torch.ones(2, 2)), - ], - options, - ) - assert result.is_all_success() - assert result.positions.shape == (2, 100, 2) - assert bool( - ( - result.velocities.abs().amax(dim=(0, 1)) - <= torch.tensor([1.2, 0.6]) * 1.03 - ).all() - ) - assert planner._pool is not None - finally: - planner.close() - - def test_mimic_follower_caps_constrain_curobo_active_joint() -> None: robot = _Robot() robot.velocities[:, 1] = 0.2 diff --git a/tests/sim/motion/planners/test_toppra_batched.py b/tests/sim/motion/planners/test_toppra_batched.py index eb5db4dcc..62c8c07c2 100644 --- a/tests/sim/motion/planners/test_toppra_batched.py +++ b/tests/sim/motion/planners/test_toppra_batched.py @@ -23,7 +23,6 @@ import pytest import torch -from embodichain.lab.sim.motion.planners.toppra_planner import _toppra_solve_one_env from embodichain.lab.sim.motion.planners.utils import TrajectorySampleMethod @@ -43,125 +42,75 @@ def get_qvel_limits(self, *, joint_ids: list[int]) -> torch.Tensor: return torch.full((256, len(joint_ids)), self.speed, dtype=torch.float64) -# Numerical contracts are also exercised without the simulation adapter in -# tests/compute/test_toppra.py. -class TestToppraWorker: - @pytest.mark.no_sim - def test_backend_defaults_plan_without_an_explicit_sample_count(self): - from embodichain.lab.sim.motion.planners.toppra_planner import ToppraPlanner - - planner = object.__new__(ToppraPlanner) - planner.robot = _UnboundedRobot() - planner._pool = None - options = planner.default_plan_options() - assert options.sample_method is TrajectorySampleMethod.TIME - assert options.sample_interval == 0.01 - out = _toppra_solve_one_env( - waypoints=np.array([[0.0] * 7, [0.2] * 7]), - vel_constraint=options.constraints["velocity"], - acc_constraint=options.constraints["acceleration"], - sample_method=options.sample_method, - sample_interval=options.sample_interval, - ) - assert out["success"] is True - assert out["dt"].sum() > 0.0 - assert np.max(out["dt"]) <= options.sample_interval + 1e-6 - np.testing.assert_allclose(out["positions"][-1], 0.2) - - def test_solve_one_env_quantity(self): - # 2-waypoint, 6-DOF - wp = np.array([[0.0] * 6, [0.5] * 6]) - out = _toppra_solve_one_env( - waypoints=wp, - vel_constraint=1.0, - acc_constraint=2.0, - sample_method=TrajectorySampleMethod.QUANTITY, - sample_interval=20, - ) - assert out["success"] is True - assert out["positions"].shape == (20, 6) - assert out["velocities"].shape == (20, 6) - assert out["dt"].shape == (20,) - - def test_solve_one_env_infeasible_exception(self): - # A path needs at least two waypoints; the worker isolates this failure. - wp = np.array([[0.0] * 6]) - out = _toppra_solve_one_env( - waypoints=wp, - vel_constraint=1.0, - acc_constraint=2.0, - sample_method=TrajectorySampleMethod.QUANTITY, - sample_interval=10, - ) - assert out["success"] is False - - def test_solve_one_env_time_sampling(self): - wp = np.array([[0.0] * 6, [0.5] * 6]) - out = _toppra_solve_one_env( - waypoints=wp, - vel_constraint=1.0, - acc_constraint=2.0, - sample_method=TrajectorySampleMethod.TIME, - sample_interval=0.05, - ) - assert out["success"] is True - assert out["positions"].shape[0] == out["n"] - assert out["n"] >= 2 - - def test_solve_one_env_same_waypoint_shortcut(self): - wp = np.array([[0.3] * 6, [0.3] * 6]) # identical - out = _toppra_solve_one_env( - waypoints=wp, - vel_constraint=1.0, - acc_constraint=2.0, - sample_method=TrajectorySampleMethod.QUANTITY, - sample_interval=20, - ) - assert out["success"] is True - assert out["n"] == 2 - assert out["dt"].sum() == 0.0 - - def test_solve_one_env_duplicate_plateau(self): - # Long plateaus of identical waypoints (e.g. from interpolating a - # segment where start_qpos equals the first target) must not make - # TOPPRA's controllable-set computation fail. - wp = np.array([[0.3] * 6] * 12 + [[0.5] * 6] * 12) - out = _toppra_solve_one_env( - waypoints=wp, - vel_constraint=1.0, - acc_constraint=2.0, - sample_method=TrajectorySampleMethod.QUANTITY, - sample_interval=20, - ) - assert out["success"] is True - assert out["positions"].shape == (20, 6) +@pytest.mark.no_sim +def test_planner_defaults_use_time_sampling_without_a_sample_count() -> None: + from embodichain.lab.sim.motion.planners.toppra_planner import ( + ToppraPlanner, + ToppraPlannerCfg, + ) + from embodichain.lab.sim.motion.planners.utils import PlanState + planner = object.__new__(ToppraPlanner) + planner.robot = _UnboundedRobot(7) + planner.device = torch.device("cpu") + planner.cfg = ToppraPlannerCfg(robot_uid="unused") + options = planner.default_plan_options() + assert options.sample_method is TrajectorySampleMethod.TIME + result = planner.plan( + [PlanState.from_qpos(torch.full((1, 7), q)) for q in (0.0, 0.2)], options + ) + assert result.success.all() and result.duration[0] > 0 + assert result.dt.max() <= options.sample_interval + 1e-6 -class TestToppraCfgFields: - def test_cfg_defaults(self): - from embodichain.lab.sim.motion.planners.toppra_planner import ToppraPlannerCfg - cfg = ToppraPlannerCfg(robot_uid="x") - assert cfg.max_workers is None - assert cfg.mp_context is None # None => auto-select by device +@pytest.mark.no_sim +@pytest.mark.parametrize("removed", ["backend", "max_workers", "mp_context"]) +def test_removed_process_configuration_is_rejected(removed: str) -> None: + from embodichain.lab.sim.motion.planners.toppra_planner import ToppraPlannerCfg - def test_mp_context_auto_resolution(self): - # _resolve_mp_context is a pure function of (cfg value, device) — no sim - # needed. None auto-selects: fork on CPU, spawn on GPU; explicit values - # are honored as-is. - import torch + with pytest.raises(TypeError, match=removed): + ToppraPlannerCfg(robot_uid="unused", **{removed: None}) - from embodichain.lab.sim.motion.planners.toppra_planner import ToppraPlanner - resolve = ToppraPlanner._resolve_mp_context - # Auto: CPU -> fork, CUDA -> spawn (creating a torch.device does NOT - # initialize CUDA, so this is safe to assert without a GPU). - assert resolve(None, torch.device("cpu")) == "fork" - assert resolve(None, torch.device("cuda", 0)) == "spawn" - assert resolve(None, torch.device("cuda")) == "spawn" - # Explicit override wins regardless of device. - assert resolve("spawn", torch.device("cpu")) == "spawn" - assert resolve("fork", torch.device("cuda", 0)) == "fork" +@pytest.mark.no_sim +def test_planner_batch_matches_offline_numpy_reference() -> None: + from embodichain.compute.trajectory._toppra import _NumpyToppra + from embodichain.lab.sim.motion.planners.toppra_planner import ( + ToppraPlanner, + ToppraPlannerCfg, + ToppraPlanOptions, + ) + from embodichain.lab.sim.motion.planners.utils import PlanState + + points = np.random.default_rng(712).normal(size=(3, 4, 2)).astype(np.float32) + tensor = torch.from_numpy(points) + planner = object.__new__(ToppraPlanner) + planner.robot = _UnboundedRobot(2) + planner.device = torch.device("cpu") + planner.cfg = ToppraPlannerCfg(robot_uid="unused") + result = planner.plan( + [PlanState.from_qpos(tensor[:, k]) for k in range(4)], + ToppraPlanOptions( + constraints={"velocity": 1.0, "acceleration": 2.0}, + sample_method=TrajectorySampleMethod.QUANTITY, + sample_interval=33, + ), + ) + assert result.success.all() + for row, values in enumerate(points): + reference = _NumpyToppra(values, 1.0, 2.0) + expected = reference.sample(np.linspace(0, reference.duration, 33)) + for key, value in zip(("positions", "velocities", "accelerations"), expected): + np.testing.assert_allclose( + getattr(result, key)[row], value, rtol=1e-6, atol=1e-6 + ) + assert result.duration[row].item() == pytest.approx( + reference.duration, rel=1e-6 + ) + + +# Numerical contracts are also exercised without the simulation adapter in +# tests/compute/test_toppra.py. class TestToppraPlanBatched: @@ -182,7 +131,7 @@ def _make_planner(self): ) ) sim.prepare() - planner = ToppraPlanner(ToppraPlannerCfg(robot_uid="t", max_workers=1)) + planner = ToppraPlanner(ToppraPlannerCfg(robot_uid="t")) return planner, sim def test_plan_batched_quantity_uniform_N(self): @@ -271,213 +220,9 @@ def test_plan_batched_time_tailpads(self): om.SimulationManager.flush_cleanup_queue() - @pytest.mark.slow - @pytest.mark.parametrize("mp_context", ["fork", "spawn"]) - def test_plan_batched_pool_path(self, mp_context): - # Exercise the real ProcessPoolExecutor branch (max_workers=2, B=3) - # under both start methods: 'fork' is the CPU default, 'spawn' is the - # GPU fallback. Both must plan correctly and reap workers on GC. - from embodichain.lab.sim.motion.planners.utils import ( - PlanState, - TrajectorySampleMethod, - ) - from embodichain.lab.sim.motion.planners.toppra_planner import ( - ToppraPlanner, - ToppraPlannerCfg, - ToppraPlanOptions, - ) - from embodichain.lab.sim import SimulationManager, SimulationManagerCfg - from embodichain_robots import CobotMagicCfg - - sim = SimulationManager( - SimulationManagerCfg(headless=True, device="cpu", num_envs=3) - ) - sim.add_robot( - cfg=CobotMagicCfg.from_dict( - {"uid": "p", "init_pos": [0, 0, 0.7775], "init_qpos": [0.0] * 16} - ) - ) - sim.prepare() - planner = ToppraPlanner( - ToppraPlannerCfg( - robot_uid="p", backend="numpy", max_workers=2, mp_context=mp_context - ) - ) - assert planner._mp_context == mp_context - try: - B, dofs = 3, 6 - wp = torch.zeros(B, dofs) - wp[:, 0] = torch.linspace(0.1, 0.5, B) - states = [ - PlanState.from_qpos(torch.zeros(B, dofs)), - PlanState.from_qpos(wp), - ] - opts = ToppraPlanOptions( - control_part="left_arm", - sample_method=TrajectorySampleMethod.QUANTITY, - sample_interval=12, - constraints={"velocity": 1.0, "acceleration": 2.0}, - ) - r = planner.plan(states, opts) - assert r.success.shape == (B,) - assert r.success.all().item() - assert r.positions.shape == (B, 12, dofs) - finally: - del planner - gc.collect() - sim.destroy() - import embodichain.lab.sim as om - - om.SimulationManager.flush_cleanup_queue() - - @pytest.mark.slow - @pytest.mark.parametrize("mp_context", ["fork", "spawn"]) - def test_workers_reaped_on_gc(self, mp_context): - # Regression: abandoning the planner must reap its worker processes, - # under both 'fork' (CPU default) and 'spawn' (GPU fallback). - # - # Workers install prctl(PR_SET_PDEATHSIG) in _worker_init, so the kernel - # kills them the moment the parent process dies — including the - # os._exit(0) path taken by SimulationManager.destroy(), which skips - # every Python finalizer. That covers process exit. For in-process GC - # of an abandoned planner (e.g. between tests in a long-running - # session), __del__ -> _shutdown_pool() must terminate workers - # synchronously so none leak. - from embodichain.lab.sim.motion.planners.utils import ( - PlanState, - TrajectorySampleMethod, - ) - from embodichain.lab.sim.motion.planners.toppra_planner import ( - ToppraPlanner, - ToppraPlannerCfg, - ToppraPlanOptions, - ) - from embodichain.lab.sim import SimulationManager, SimulationManagerCfg - from embodichain_robots import CobotMagicCfg - - sim = SimulationManager( - SimulationManagerCfg(headless=True, device="cpu", num_envs=3) - ) - sim.add_robot( - cfg=CobotMagicCfg.from_dict( - { - "uid": "close_reap", - "init_pos": [0, 0, 0.7775], - "init_qpos": [0.0] * 16, - } - ) - ) - sim.prepare() - planner = ToppraPlanner( - ToppraPlannerCfg( - robot_uid="close_reap", - backend="numpy", - max_workers=2, - mp_context=mp_context, - ) - ) - assert planner._mp_context == mp_context - worker_processes = [] - try: - B, dofs = 3, 6 - wp = torch.zeros(B, dofs) - wp[:, 0] = torch.linspace(0.1, 0.5, B) - states = [ - PlanState.from_qpos(torch.zeros(B, dofs)), - PlanState.from_qpos(wp), - ] - opts = ToppraPlanOptions( - control_part="left_arm", - sample_method=TrajectorySampleMethod.QUANTITY, - sample_interval=12, - constraints={"velocity": 1.0, "acceleration": 2.0}, - ) - r = planner.plan(states, opts) - assert r.success.all().item() - - # Capture the worker processes created by the executor. - worker_processes = list(planner._pool._processes.values()) - assert len(worker_processes) == 2 - finally: - # Abandon the planner; __del__ -> _shutdown_pool must reap workers. - del planner - gc.collect() - for proc in worker_processes: - assert not proc.is_alive(), ( - f"TOPPRA worker process {proc.pid} survived planner GC; " - "worker processes must be reaped when the planner is collected." - ) - sim.destroy() - import embodichain.lab.sim as om - - om.SimulationManager.flush_cleanup_queue() - - -@pytest.mark.slow -class TestToppraNumericalRegression: - def test_batched_equals_inline_single(self): - from embodichain.lab.sim.motion.planners.toppra_planner import ( - _toppra_solve_one_env, - ToppraPlanner, - ToppraPlannerCfg, - ToppraPlanOptions, - ) - from embodichain.lab.sim.motion.planners.utils import ( - PlanState, - TrajectorySampleMethod, - ) - from embodichain.lab.sim import SimulationManager, SimulationManagerCfg - from embodichain_robots import CobotMagicCfg - - sim = SimulationManager( - SimulationManagerCfg(headless=True, device="cpu", num_envs=4) - ) - sim.add_robot( - cfg=CobotMagicCfg.from_dict( - {"uid": "r", "init_pos": [0, 0, 0.7775], "init_qpos": [0.0] * 16} - ) - ) - sim.prepare() - planner = ToppraPlanner(ToppraPlannerCfg(robot_uid="r", max_workers=1)) - try: - B, dofs = 4, 6 - wp = torch.zeros(B, dofs) - wp[:, 0] = torch.linspace(0.1, 0.6, B) - states = [ - PlanState.from_qpos(torch.zeros(B, dofs)), - PlanState.from_qpos(wp), - ] - opts = ToppraPlanOptions( - control_part="left_arm", - sample_method=TrajectorySampleMethod.QUANTITY, - sample_interval=20, - constraints={"velocity": 1.0, "acceleration": 2.0}, - ) - r = planner.plan(states, opts) - # Compare each env to the inline single-env solve - for b in range(B): - single = _toppra_solve_one_env( - np.stack([np.zeros(dofs), wp[b].numpy()]), - 1.0, - 2.0, - TrajectorySampleMethod.QUANTITY, - 20, - ) - assert np.allclose( - r.positions[b].cpu().numpy(), single["positions"], atol=1e-5 - ) - finally: - del planner - gc.collect() - sim.destroy() - import embodichain.lab.sim as om - - om.SimulationManager.flush_cleanup_queue() - @pytest.mark.no_sim -@pytest.mark.parametrize("backend", ["numpy", "warp", "auto"]) -def test_in_tree_planner_preserves_batch_timing_and_failures(backend: str) -> None: +def test_planner_preserves_batch_timing_and_failures() -> None: from embodichain.lab.sim.motion.planners.toppra_planner import ( ToppraPlanner, ToppraPlannerCfg, @@ -487,9 +232,8 @@ def test_in_tree_planner_preserves_batch_timing_and_failures(backend: str) -> No planner = object.__new__(ToppraPlanner) planner.robot = _UnboundedRobot() - planner.cfg = ToppraPlannerCfg(robot_uid="unused", backend=backend, max_workers=1) + planner.cfg = ToppraPlannerCfg(robot_uid="unused") planner.device = torch.device("cpu") - planner._pool = None states = [ PlanState.from_qpos(torch.tensor([[0.0], [0.0], [2.0]])), PlanState.from_qpos(torch.tensor([[1.0], [float("nan")], [2.0]])), @@ -503,26 +247,23 @@ def test_in_tree_planner_preserves_batch_timing_and_failures(backend: str) -> No ), ) assert result.success.tolist() == [True, False, True] - assert planner._pool is None assert result.duration[0] > 0 assert result.duration[1:].count_nonzero() == 0 assert result.positions[2].eq(2.0).all() assert result.velocities[2].count_nonzero() == 0 assert result.dt.max() <= 0.03 - reference = _toppra_solve_one_env( - np.array([[0.0], [1.0]]), - [[-0.5, 1.0]], - [[-2.0, 1.0]], - TrajectorySampleMethod.TIME, - 0.03, + from embodichain.compute.trajectory._toppra import _NumpyToppra + + reference = _NumpyToppra(np.array([[0.0], [1.0]]), [[-0.5, 1.0]], [[-2.0, 1.0]]) + times = np.linspace(0, reference.duration, result.positions.shape[1]) + np.testing.assert_allclose( + result.positions[0], reference.sample(times)[0], atol=1e-6 ) - np.testing.assert_allclose(result.positions[0], reference["positions"], atol=1e-6) - np.testing.assert_allclose(result.dt[0], reference["dt"], atol=1e-6) + np.testing.assert_allclose(result.dt[0], np.diff(times, prepend=0), atol=1e-6) @pytest.mark.no_sim -@pytest.mark.parametrize("backend", ["numpy", "warp"]) -def test_in_tree_planner_invalid_limits_rejected_before_dispatch(backend: str) -> None: +def test_in_tree_planner_invalid_limits_rejected_before_dispatch() -> None: from embodichain.lab.sim.motion.planners.toppra_planner import ( ToppraPlanner, ToppraPlannerCfg, @@ -532,9 +273,8 @@ def test_in_tree_planner_invalid_limits_rejected_before_dispatch(backend: str) - planner = object.__new__(ToppraPlanner) planner.robot = _UnboundedRobot() - planner.cfg = ToppraPlannerCfg(robot_uid="unused", backend=backend, max_workers=1) + planner.cfg = ToppraPlannerCfg(robot_uid="unused") planner.device = torch.device("cpu") - planner._pool = None with pytest.raises(ValueError, match="positive"): planner.plan( [ @@ -552,7 +292,7 @@ def test_in_tree_planner_invalid_limits_rejected_before_dispatch(backend: str) - @pytest.mark.no_sim @pytest.mark.gpu @pytest.mark.parametrize("dtype", [torch.float16, torch.float64]) -def test_cuda_auto_backend_keeps_plan_on_device( +def test_cuda_planner_keeps_plan_on_device( monkeypatch: pytest.MonkeyPatch, dtype: torch.dtype, ) -> None: @@ -569,7 +309,6 @@ def test_cuda_auto_backend_keeps_plan_on_device( planner.robot = _UnboundedRobot() planner.cfg = ToppraPlannerCfg(robot_uid="unused") planner.device = torch.device("cuda") - planner._pool = None points = torch.tensor( [ [[0.0], [0.4]], @@ -583,7 +322,7 @@ def test_cuda_auto_backend_keeps_plan_on_device( with monkeypatch.context() as patch: def forbidden(*args, **kwargs): - raise AssertionError("auto CUDA planning must not copy waypoints to CPU") + raise AssertionError("CUDA planning must not copy waypoints to CPU") patch.setattr(torch.Tensor, "cpu", forbidden) result = planner.plan( @@ -596,20 +335,18 @@ def forbidden(*args, **kwargs): assert result.positions.is_cuda and result.dt.is_cuda assert result.positions.shape == (3, 128, 1) assert result.positions.dtype == result.dt.dtype == torch.float32 - assert planner._pool is None assert result.duration[0] > 0 assert result.duration[1:].count_nonzero() == 0 torch.testing.assert_close(result.positions[0, -1], points[0, -1].float()) @pytest.mark.no_sim -@pytest.mark.parametrize("backend", ["numpy", "warp"]) @pytest.mark.parametrize( "dtype", [torch.float16, torch.float32, torch.float64, torch.int64] ) @pytest.mark.parametrize("quantity", [20, 20.0, 20.9]) def test_legacy_planner_output_dtype_and_quantity_coercion( - backend: str, dtype: torch.dtype, quantity: float | int + dtype: torch.dtype, quantity: float | int ) -> None: from embodichain.lab.sim.motion.planners.toppra_planner import ( ToppraPlanner, @@ -620,9 +357,8 @@ def test_legacy_planner_output_dtype_and_quantity_coercion( planner = object.__new__(ToppraPlanner) planner.robot = _UnboundedRobot() - planner.cfg = ToppraPlannerCfg(robot_uid="unused", backend=backend, max_workers=1) + planner.cfg = ToppraPlannerCfg(robot_uid="unused") planner.device = torch.device("cpu") - planner._pool = None result = planner.plan( [ PlanState.from_qpos(torch.zeros(1, 1, dtype=dtype)), @@ -644,10 +380,8 @@ def test_legacy_planner_output_dtype_and_quantity_coercion( @pytest.mark.no_sim -@pytest.mark.parametrize("backend", ["auto", "warp"]) @pytest.mark.parametrize("dtype", [torch.float16, torch.float32, torch.float64]) def test_planner_preserves_waypoint_gradients_and_duration_contract( - backend: str, dtype: torch.dtype, ) -> None: from embodichain.lab.sim.motion.planners.toppra_planner import ( @@ -659,9 +393,8 @@ def test_planner_preserves_waypoint_gradients_and_duration_contract( planner = object.__new__(ToppraPlanner) planner.robot = _UnboundedRobot() - planner.cfg = ToppraPlannerCfg(robot_uid="unused", backend=backend) + planner.cfg = ToppraPlannerCfg(robot_uid="unused") planner.device = torch.device("cpu") - planner._pool = None # Use a non-leaf input to verify gradients reach the caller's computation. parameter = torch.tensor([[0.365]], dtype=dtype, requires_grad=True) distance = 2 * parameter @@ -677,7 +410,6 @@ def test_planner_preserves_waypoint_gradients_and_duration_contract( ), ) assert result.success.all() - assert planner._pool is None assert result.positions.dtype == torch.float32 assert result.positions.requires_grad and result.duration.requires_grad result.duration.sum().backward() @@ -686,38 +418,10 @@ def test_planner_preserves_waypoint_gradients_and_duration_contract( @pytest.mark.no_sim -def test_numpy_planner_rejects_autograd_instead_of_detaching() -> None: - from embodichain.lab.sim.motion.planners.toppra_planner import ( - ToppraPlanner, - ToppraPlannerCfg, - ToppraPlanOptions, - ) - from embodichain.lab.sim.motion.planners.utils import PlanState - - planner = object.__new__(ToppraPlanner) - planner.robot = _UnboundedRobot() - planner.cfg = ToppraPlannerCfg(robot_uid="unused", backend="numpy") - planner.device = torch.device("cpu") - planner._pool = None - states = [ - PlanState.from_qpos(torch.zeros(1, 1)), - PlanState.from_qpos(torch.ones(1, 1, requires_grad=True)), - ] - options = ToppraPlanOptions( - sample_method=TrajectorySampleMethod.TIME, - sample_interval=0.1, - ) - with pytest.raises(NotImplementedError, match="Warp backend"): - planner.plan(states, options) - - -@pytest.mark.no_sim -@pytest.mark.parametrize("backend", ["auto", "warp"]) @pytest.mark.parametrize( "sampling", [TrajectorySampleMethod.TIME, TrajectorySampleMethod.QUANTITY] ) def test_planner_autograd_routes_trainable_limits_without_trainable_waypoints( - backend: str, sampling: TrajectorySampleMethod, ) -> None: from embodichain.lab.sim.motion.planners.toppra_planner import ( @@ -729,9 +433,8 @@ def test_planner_autograd_routes_trainable_limits_without_trainable_waypoints( planner = object.__new__(ToppraPlanner) planner.robot = _UnboundedRobot() - planner.cfg = ToppraPlannerCfg(robot_uid="unused", backend=backend) + planner.cfg = ToppraPlannerCfg(robot_uid="unused") planner.device = torch.device("cpu") - planner._pool = None parameter = torch.tensor(0.6, dtype=torch.float64, requires_grad=True) acceleration = parameter.exp() # Verify the caller's graph is preserved. result = planner.plan( @@ -746,7 +449,7 @@ def test_planner_autograd_routes_trainable_limits_without_trainable_waypoints( ), ) assert result.success.all() and result.duration.requires_grad - assert result.positions.dtype == torch.float32 and planner._pool is None + assert result.positions.dtype == torch.float32 result.duration.sum().backward() torch.testing.assert_close( parameter.grad, -result.duration.sum().detach().double() / 2 @@ -754,8 +457,7 @@ def test_planner_autograd_routes_trainable_limits_without_trainable_waypoints( @pytest.mark.no_sim -@pytest.mark.parametrize("backend", ["numpy", "warp", "auto"]) -def test_no_grad_planning_accepts_trainable_inputs_as_constants(backend: str) -> None: +def test_no_grad_planning_accepts_trainable_inputs_as_constants() -> None: from embodichain.lab.sim.motion.planners.toppra_planner import ( ToppraPlanner, ToppraPlannerCfg, @@ -765,9 +467,8 @@ def test_no_grad_planning_accepts_trainable_inputs_as_constants(backend: str) -> planner = object.__new__(ToppraPlanner) planner.robot = _UnboundedRobot() - planner.cfg = ToppraPlannerCfg(robot_uid="unused", backend=backend) + planner.cfg = ToppraPlannerCfg(robot_uid="unused") planner.device = torch.device("cpu") - planner._pool = None end = torch.ones(1, 1, requires_grad=True) limit = torch.tensor(1.0, requires_grad=True) with torch.no_grad(): @@ -816,7 +517,6 @@ def test_motion_generator_keeps_toppra_gradients_through_option_copy_and_resampl planner.robot = _UnboundedRobot() planner.cfg = ToppraPlannerCfg(robot_uid="unused") planner.device = torch.device("cpu") - planner._pool = None generator = object.__new__(MotionGenerator) generator.planner = planner generator.device = planner.device @@ -862,7 +562,6 @@ def test_cuda_time_planner_backward_preserves_input_gradients( planner.robot = _UnboundedRobot() planner.cfg = ToppraPlannerCfg(robot_uid="unused") planner.device = torch.device("cuda") - planner._pool = None distance = torch.tensor([[0.73]], device="cuda", dtype=dtype, requires_grad=True) acceleration = torch.tensor(1.7, device="cuda", requires_grad=True) options = ToppraPlanOptions( @@ -886,26 +585,29 @@ def test_cuda_time_planner_backward_preserves_input_gradients( @pytest.mark.no_sim -def test_cpu_auto_planner_uses_warp_without_numpy_workers(monkeypatch) -> None: - from embodichain.lab.sim.motion.planners import toppra_planner as module +def test_cpu_planner_never_starts_process_workers(monkeypatch) -> None: + from multiprocessing.process import BaseProcess + from embodichain.lab.sim.motion.planners.toppra_planner import ( + ToppraPlanner, + ToppraPlannerCfg, + ToppraPlanOptions, + ) from embodichain.compute.trajectory._toppra_warp import _retime_toppra_warp from embodichain.lab.sim.motion.planners.utils import PlanState def forbidden(*args, **kwargs): - raise AssertionError("CPU auto planning must not use the NumPy process chain") + raise AssertionError("TOPPRA planning must not start a process") - monkeypatch.setattr(module, "_toppra_solve_one_env", forbidden) - monkeypatch.setattr(module.ToppraPlanner, "_get_pool", forbidden) - planner = object.__new__(module.ToppraPlanner) + monkeypatch.setattr(BaseProcess, "start", forbidden) + planner = object.__new__(ToppraPlanner) planner.robot = _UnboundedRobot(7) - planner.cfg = module.ToppraPlannerCfg(robot_uid="unused") + planner.cfg = ToppraPlannerCfg(robot_uid="unused") planner.device = torch.device("cpu") - planner._pool = None points = torch.zeros(32, 4, 7) points[:, :, 0] = torch.linspace(0.0, 1.0, 4) actual = planner.plan( [PlanState.from_qpos(points[:, index]) for index in range(4)], - module.ToppraPlanOptions( + ToppraPlanOptions( sample_method=TrajectorySampleMethod.QUANTITY, constraints={"velocity": 1.0, "acceleration": 2.0}, sample_interval=17, @@ -914,7 +616,9 @@ def forbidden(*args, **kwargs): reference = _retime_toppra_warp(points, 1.0, 2.0, sample_count=17) for key, value in reference.items(): torch.testing.assert_close(getattr(actual, key), value, rtol=0, atol=0) - assert actual.success.all() and planner._pool is None + planner.close() + planner.close() + assert actual.success.all() @pytest.mark.no_sim @@ -961,7 +665,6 @@ def test_planner_limit_gradients_survive_runtime_speed_caps( planner.robot = _UnboundedRobot(speed=0.5) planner.cfg = ToppraPlannerCfg(robot_uid="unused") planner.device = torch.device(device) - planner._pool = None states = [ PlanState.from_qpos(torch.tensor([[q]], device=device)) for q in (0.0, 2.0) ] @@ -997,15 +700,6 @@ def solve(velocity: torch.Tensor): torch.testing.assert_close(derivative, numerical.double(), rtol=3e-3, atol=3e-3) -@pytest.mark.no_sim -def test_toppra_cfg_keeps_legacy_positional_worker_fields() -> None: - from embodichain.lab.sim.motion.planners.toppra_planner import ToppraPlannerCfg - - cfg = ToppraPlannerCfg("robot", "toppra", 0, 2, "spawn") - assert cfg.max_workers == 2 and cfg.mp_context == "spawn" - assert cfg.backend == "auto" and cfg.grid_size == 100 - - @pytest.mark.no_sim @pytest.mark.parametrize( "dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64] diff --git a/tests/sim/motion/planners/test_toppra_planner.py b/tests/sim/motion/planners/test_toppra_planner.py index 99fd15cb6..af847325a 100644 --- a/tests/sim/motion/planners/test_toppra_planner.py +++ b/tests/sim/motion/planners/test_toppra_planner.py @@ -98,7 +98,6 @@ def test_trivial_trajectory(self) -> None: ) is True ) - assert self.planner._pool is None if __name__ == "__main__": diff --git a/tests/sim/motion/test_motion_imports.py b/tests/sim/motion/test_motion_imports.py index 97530a91c..e406e3104 100644 --- a/tests/sim/motion/test_motion_imports.py +++ b/tests/sim/motion/test_motion_imports.py @@ -143,3 +143,31 @@ def find_spec(self, fullname, path=None, target=None): timeout=60, ) assert result.returncode == 0, result.stdout + result.stderr + + +def test_warp_toppra_does_not_import_the_offline_numpy_solver() -> None: + script = """ +import importlib.abc +import sys +class BlockReference(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname == 'embodichain.compute.trajectory._toppra': + raise AssertionError('Production Warp must not load the offline solver') + if fullname == 'toppra' or fullname.startswith('toppra.'): + raise AssertionError('External toppra must not be imported') +sys.meta_path.insert(0, BlockReference()) +import torch +from embodichain.compute.trajectory._toppra_warp import _retime_toppra_warp +points = torch.linspace(0, .73, 4, dtype=torch.float64).view(1, 4, 1).expand(32, 4, 7).contiguous() +result = _retime_toppra_warp(points, 1., 1.7, sample_count=17) +assert result['success'].all() +points.requires_grad_() +result = _retime_toppra_warp(points, 1., 1.7, sample_dt=.13) +result['dt'].sum().backward() +assert torch.isfinite(points.grad).all() +assert 'embodichain.compute.trajectory._toppra' not in sys.modules +""" + result = subprocess.run( + [sys.executable, "-c", script], capture_output=True, text=True, timeout=60 + ) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/tests/sim/utility/test_atom_action_utils.py b/tests/sim/utility/test_atom_action_utils.py index 124554d49..f335bbfae 100644 --- a/tests/sim/utility/test_atom_action_utils.py +++ b/tests/sim/utility/test_atom_action_utils.py @@ -23,7 +23,7 @@ import torch from embodichain.lab.sim.utility import atom_action_utils -from embodichain.lab.sim.motion.planners.toppra_planner import _toppra_solve_one_env +from embodichain.lab.sim.motion.planners.toppra_planner import retime_joint_paths from embodichain.lab.sim.motion.planners.utils import TrajectorySampleMethod @@ -31,26 +31,32 @@ def test_legacy_plan_trajectory_keeps_requested_sample_count(monkeypatch, sample_count): """Count-based expert actions must not interpret their count as seconds.""" - class WorkerMotionGenerator: + class RetimeMotionGenerator: def __init__(self, cfg): self.cfg = cfg def generate(self, *, target_states, options): plan = options.plan_opts assert plan.sample_method is TrajectorySampleMethod.QUANTITY - solved = _toppra_solve_one_env( - waypoints=np.stack( - [state.qpos.reshape(-1).cpu().numpy() for state in target_states] - ), - vel_constraint=plan.constraints["velocity"], - acc_constraint=plan.constraints["acceleration"], + # This standalone fake has no robot from which to resolve caps. + limits = { + name: ( + fallback + if plan.constraints[name] is None + else plan.constraints[name] + ) + for name, fallback in (("velocity", 0.2), ("acceleration", 0.5)) + } + solved = retime_joint_paths( + torch.stack([state.qpos.reshape(-1) for state in target_states])[None], + constraints=limits, sample_method=plan.sample_method, sample_interval=plan.sample_interval, ) - assert solved["success"] - return SimpleNamespace(positions=torch.tensor(solved["positions"])[None]) + assert solved.is_all_success() + return solved - monkeypatch.setattr(atom_action_utils, "MotionGenerator", WorkerMotionGenerator) + monkeypatch.setattr(atom_action_utils, "MotionGenerator", RetimeMotionGenerator) env = SimpleNamespace( robot=SimpleNamespace(uid="robot"), sim=SimpleNamespace(instance_id="unused") ) From 3219ecb66a2f33ac5446ba2af480abda66d3d383 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Sat, 17 Oct 2026 12:14:10 +0000 Subject: [PATCH 3/3] perf(trajectory): accelerate TOPPRA CPU and CUDA inference --- .../topics/motion-planning/planner-details.md | 9 +- .../compute/trajectory/_toppra_warp.py | 36 ++++-- .../compute/trajectory/_warp/toppra.py | 89 ++++++++++++- .../compute/trajectory/_warp/toppra_cuda.py | 119 ++++++++++++++++++ tests/compute/test_toppra_cpu.py | 45 ++++--- tests/compute/test_toppra_cuda.py | 98 +++++++++++++++ 6 files changed, 371 insertions(+), 25 deletions(-) create mode 100644 embodichain/compute/trajectory/_warp/toppra_cuda.py create mode 100644 tests/compute/test_toppra_cuda.py diff --git a/agent_context/topics/motion-planning/planner-details.md b/agent_context/topics/motion-planning/planner-details.md index 736c554ba..8f0c14b40 100644 --- a/agent_context/topics/motion-planning/planner-details.md +++ b/agent_context/topics/motion-planning/planner-details.md @@ -78,7 +78,8 @@ retain their existing timing contracts and do not consume these backend limits. ### TOPPRA `toppra_planner.py` uses Warp on CPU and CUDA. Numerical ownership is in -`compute/trajectory/_toppra_warp.py` and `_warp/toppra.py`; shared validation +`compute/trajectory/_toppra_warp.py` and `_warp/toppra.py`, with cooperative +CUDA inference in `_warp/toppra_cuda.py`; shared validation lives in `_toppra_common.py`. `_toppra.py` is an offline NumPy oracle and must not be imported by the runtime planner or Warp adapter. Robot-backed planning retains shared speed-cap and control-part resolution, @@ -92,6 +93,12 @@ TOPPRA package. Tiny real moves must retain positive duration. Warp constructs splines, bounds, controllable sets and samples on the input device. It honors the current Torch stream and reads a scalar to size output. +CUDA inference reduces constraint planes across one block per path; all +early exits must remain uniform across the block. Inference may use fast +polygon projection, while recorded calls retain the selected-expression +projection and serial solver needed by the adjoint. Dense CPU inference sampling +shares scalar timing and interval searches across the joints of each sample; +recorded sampling retains its per-joint adjoint path. CPU inference amortizes sufficiently large batches over a shared bounded thread pool; small requests and adjoint-recording calls stay inline. Row blocks own their numerical buffers, inherit Torch grad/inference modes and merge using diff --git a/embodichain/compute/trajectory/_toppra_warp.py b/embodichain/compute/trajectory/_toppra_warp.py index f3cf825c9..d175f6c4b 100644 --- a/embodichain/compute/trajectory/_toppra_warp.py +++ b/embodichain/compute/trajectory/_toppra_warp.py @@ -507,8 +507,8 @@ def _launch_toppra( ) wp.launch( ( - kernels._project_caps_cpu - if device.is_cpu and not differentiable + kernels._project_caps_fast + if not differentiable and (device.is_cpu or batch >= 128) else kernels._project_caps ), dim=(batch, capacity), @@ -526,9 +526,18 @@ def _launch_toppra( durations = wp.zeros( batch, dtype=wp.float64, device=device, requires_grad=differentiable ) + solver, solve_dim, solve_block = kernels._solve, batch, 256 + if device.is_cuda and not differentiable: + from ._warp import toppra_cuda + + solver = toppra_cuda._solve + solve_dim, solve_block = ( + batch, + toppra_cuda._BLOCK_DIM, + ), toppra_cuda._BLOCK_DIM wp.launch( - kernels._solve, - dim=batch, + solver, + dim=solve_dim, inputs=[ a, b, @@ -543,6 +552,7 @@ def _launch_toppra( ], device=device, record_tape=differentiable, + block_dim=solve_block, ) duration_tensor = wp.to_torch(durations) ok = wp.to_torch(success).bool() @@ -560,9 +570,19 @@ def _launch_toppra( sample_lengths = desired.to(torch.int32).clamp_min(2) sample_lengths = torch.where(moving, sample_lengths, 2).contiguous() samples = int(sample_lengths.max().item()) + # Reuse scalar searches only when dense output amortizes the + # point-wise joint loop; small outputs retain the flat sampler. + cpu_sampling = device.is_cpu and not differentiable and samples >= 512 + # Narrow only on final stores. Avoid nested Torch cast work in + # CPU worker threads and an extra double-precision output buffer. + sample_dtype = ( + wp.float32 + if cpu_sampling and waypoints.dtype == torch.float32 + else wp.float64 + ) positions = wp.zeros( (batch, samples, dofs), - dtype=wp.float64, + dtype=sample_dtype, device=device, requires_grad=differentiable, ) @@ -571,13 +591,13 @@ def _launch_toppra( ) dt = wp.zeros( (batch, samples), - dtype=wp.float64, + dtype=sample_dtype, device=device, requires_grad=differentiable, ) wp.launch( - kernels._sample, - dim=(batch, samples, dofs), + kernels._sample_cpu if cpu_sampling else kernels._sample, + dim=(batch, samples) if cpu_sampling else (batch, samples, dofs), inputs=[ clean, coefficients, diff --git a/embodichain/compute/trajectory/_warp/toppra.py b/embodichain/compute/trajectory/_warp/toppra.py index 3f0c58379..1bf72fd02 100644 --- a/embodichain/compute/trajectory/_warp/toppra.py +++ b/embodichain/compute/trajectory/_warp/toppra.py @@ -23,6 +23,8 @@ from __future__ import annotations +from typing import Any + import warp as wp __all__ = [] @@ -402,7 +404,7 @@ def _project_caps( @wp.kernel(enable_backward=False) -def _project_caps_cpu( +def _project_caps_fast( a: wp.array3d(dtype=wp.float64), b: wp.array3d(dtype=wp.float64), c: wp.array3d(dtype=wp.float64), @@ -696,3 +698,88 @@ def _sample( velocity_value = wp.float64(0.0) positions[row, sample, joint] = position velocities[row, sample, joint] = velocity_value + + +@wp.kernel(enable_backward=False) +def _sample_cpu( + points: wp.array3d(dtype=wp.float64), + coefficients: wp.array4d(dtype=wp.float64), + counts: wp.array(dtype=int), + intervals: wp.array(dtype=int), + success: wp.array(dtype=int), + x: wp.array2d(dtype=wp.float64), + times: wp.array2d(dtype=wp.float64), + durations: wp.array(dtype=wp.float64), + sample_counts: wp.array(dtype=int), + positions: wp.array3d(dtype=Any), + velocities: wp.array3d(dtype=Any), + accelerations: wp.array3d(dtype=Any), + dt: wp.array2d(dtype=Any), +): + # CPU inference shares the time search and scalar state across joints. + # Keep the recorded per-joint sampler for its existing adjoint contract. + row, sample = wp.tid() + if success[row] == 0: + return + n, valid = counts[row], sample_counts[row] + if sample >= valid or n == 1: + for joint in range(points.shape[2]): + positions[row, sample, joint] = positions.dtype(points[row, n - 1, joint]) + return + step = durations[row] / wp.float64(valid - 1) + query = step * wp.float64(sample) + if sample > 0: + dt[row, sample] = dt.dtype(step) + total = intervals[row] + i = _sample_interval(times, row, total, query) + elapsed = query - times[row, i] + accel = (x[row, i + 1] - x[row, i]) * wp.float64(total) / wp.float64(2.0) + initial = _speed(x[row, i]) + speed = wp.max(wp.float64(0.0), initial + accel * elapsed) + s = wp.clamp( + wp.float64(i) / wp.float64(total) + + initial * elapsed + + wp.float64(0.5) * accel * elapsed * elapsed, + wp.float64(i) / wp.float64(total), + wp.float64(i + 1) / wp.float64(total), + ) + segment = wp.min(int(s * wp.float64(n - 1)), n - 2) + delta = s - wp.float64(segment) / wp.float64(n - 1) + for joint in range(points.shape[2]): + ca, cb = ( + coefficients[row, segment, joint, 0], + coefficients[row, segment, joint, 1], + ) + cc, cd = ( + coefficients[row, segment, joint, 2], + coefficients[row, segment, joint, 3], + ) + first = cb + delta * (wp.float64(2.0) * cc + wp.float64(3.0) * delta * cd) + second = wp.float64(2.0) * cc + wp.float64(6.0) * delta * cd + position = ca + delta * (cb + delta * (cc + delta * cd)) + velocity_value = first * speed + accelerations[row, sample, joint] = accelerations.dtype( + second * speed * speed + first * accel + ) + if sample == 0: + position = points[row, 0, joint] + velocity_value = wp.float64(0.0) + elif sample == valid - 1: + position = points[row, n - 1, joint] + velocity_value = wp.float64(0.0) + positions[row, sample, joint] = positions.dtype(position) + velocities[row, sample, joint] = velocities.dtype(velocity_value) + + +# Register both concrete variants before the submitting thread loads this +# module. CPU worker threads must not race to instantiate a generic kernel. +for _output_dtype in (wp.float32, wp.float64): + wp.overload( + _sample_cpu, + { + "positions": wp.array3d(dtype=_output_dtype), + "velocities": wp.array3d(dtype=_output_dtype), + "accelerations": wp.array3d(dtype=_output_dtype), + "dt": wp.array2d(dtype=_output_dtype), + }, + ) diff --git a/embodichain/compute/trajectory/_warp/toppra_cuda.py b/embodichain/compute/trajectory/_warp/toppra_cuda.py new file mode 100644 index 000000000..8ee7bad6b --- /dev/null +++ b/embodichain/compute/trajectory/_warp/toppra_cuda.py @@ -0,0 +1,119 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Cooperative float64 CUDA inference; one 64-lane block owns each path. + +The serial recorded solver remains responsible for autograd. Bounds and +feasibility tolerances match it; reductions only parallelize plane scans. +""" + +from __future__ import annotations + +import warp as wp + +from .toppra import _forward_bound, _speed + +__all__: list[str] = [] + +_BLOCK_DIM = 64 + + +@wp.kernel(enable_backward=False) +def _solve( + a: wp.array3d(dtype=wp.float64), + b: wp.array3d(dtype=wp.float64), + c: wp.array3d(dtype=wp.float64), + intervals: wp.array(dtype=int), + caps: wp.array2d(dtype=wp.float64), + controllable: wp.array2d(dtype=wp.float64), + x: wp.array2d(dtype=wp.float64), + times: wp.array2d(dtype=wp.float64), + durations: wp.array(dtype=wp.float64), + success: wp.array(dtype=int), +): + # Every CTA owns one complete row. All returns and reductions are + # uniform across its 64 lanes, including invalid/stationary rows. + row, lane = wp.tid() + total = intervals[row] + if success[row] == 0 or total == 0: + return + next_cap = wp.float64(0.0) + for reverse in range(total): + i = total - 1 - reverse + candidate = wp.float64(wp.inf) + for group in range((a.shape[2] + 63) // 64): + j = group * 64 + lane + if j < a.shape[2]: + aj, bj = a[row, i, j], b[row, i, j] + if aj > wp.float64(0.0) and bj < wp.float64(0.0): + bound = (c[row, i, j] - bj * next_cap) / aj + if bound < candidate: + candidate = bound + limit = wp.tile_extract(wp.tile_min(wp.tile(candidate)), 0) + cap = wp.min(caps[row, i], limit) + if not wp.isfinite(cap) or cap < wp.float64(0.0): + if lane == 0: + success[row] = 0 + return + if lane == 0: + controllable[row, i] = cap + next_cap = cap + # A cross-warp reduction synchronizes each preceding controllable write. + # Forward step i reads K[i+1], written before the final backward reduction. + current, arrival = wp.float64(0.0), wp.float64(0.0) + for i in range(total): + candidate = wp.float64(wp.inf) + for group in range((a.shape[2] + 63) // 64): + j = group * 64 + lane + if j < a.shape[2] and b[row, i, j] > wp.float64(0.0): + bound = _forward_bound( + a[row, i, j], b[row, i, j], c[row, i, j], current + ) + if bound < candidate: + candidate = bound + limit = wp.tile_extract(wp.tile_min(wp.tile(candidate)), 0) + high = wp.max(wp.float64(0.0), wp.min(controllable[row, i + 1], limit)) + invalid = int(0) + for group in range((a.shape[2] + 63) // 64): + j = group * 64 + lane + if j < a.shape[2]: + left, right, check_bound = ( + a[row, i, j] * current, + b[row, i, j] * high, + c[row, i, j], + ) + tolerance = wp.float64(1e-10) * wp.max( + wp.max(wp.abs(left), wp.abs(right)), wp.abs(check_bound) + ) + if left + right - check_bound > tolerance: + invalid = 1 + failed = wp.tile_extract(wp.tile_max(wp.tile(invalid)), 0) + denominator = _speed(current) + _speed(high) + if failed != 0 or denominator <= wp.float64(0.0): + if lane == 0: + success[row] = 0 + return + next_arrival = arrival + wp.float64(2.0) / (wp.float64(total) * denominator) + if not wp.isfinite(next_arrival) or next_arrival <= arrival: + if lane == 0: + success[row] = 0 + return + if lane == 0: + x[row, i + 1] = high + times[row, i + 1] = next_arrival + current, arrival = high, next_arrival + if lane == 0: + durations[row] = arrival diff --git a/tests/compute/test_toppra_cpu.py b/tests/compute/test_toppra_cpu.py index 67fb4ba75..4140054d6 100644 --- a/tests/compute/test_toppra_cpu.py +++ b/tests/compute/test_toppra_cpu.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""CPU TOPPRA projection and inference scheduling contracts.""" +"""Fast TOPPRA projection and CPU inference scheduling contracts.""" from __future__ import annotations @@ -36,7 +36,7 @@ from embodichain.compute.trajectory._warp import toppra as kernels -def _project(planes: np.ndarray) -> np.ndarray: +def _project(planes: np.ndarray, device: str = "cpu") -> np.ndarray: """Compute the largest feasible x for rest-feasible half-plane batches.""" wp.init() batch = len(planes) @@ -44,23 +44,26 @@ def _project(planes: np.ndarray) -> np.ndarray: wp.from_numpy( np.ascontiguousarray(planes[:, None, :, index]), dtype=wp.float64, - device="cpu", + device=device, ) for index in range(3) ] - counts = wp.from_numpy(np.ones(batch, dtype=np.int32), device="cpu") - caps = wp.empty((batch, 1), dtype=wp.float64, device="cpu") + counts = wp.from_numpy(np.ones(batch, dtype=np.int32), device=device) + caps = wp.empty((batch, 1), dtype=wp.float64, device=device) wp.launch( - kernels._project_caps_cpu, + kernels._project_caps_fast, dim=(batch, 1), inputs=[*arrays, counts, caps], - device="cpu", + device=device, ) return caps.numpy()[:, 0] @pytest.mark.parametrize("scale", [1e-100, 1.0, 1e100]) -def test_cpu_projection_matches_analytic_polygons(scale: float) -> None: +@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.gpu)]) +def test_fast_projection_matches_analytic_polygons(scale: float, device: str) -> None: + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") # Every inequality has the form a*x + b*y <= c, with x,y >= 0. planes = ( np.array( @@ -100,11 +103,14 @@ def test_cpu_projection_matches_analytic_polygons(scale: float) -> None: * scale ) np.testing.assert_allclose( - _project(planes), [0.5, 2.0, 1.0, 2.0, 0.0], rtol=1e-12, atol=1e-12 + _project(planes, device), [0.5, 2.0, 1.0, 2.0, 0.0], rtol=1e-12, atol=1e-12 ) -def test_cpu_projection_matches_independent_linear_programs() -> None: +@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.gpu)]) +def test_fast_projection_matches_independent_linear_programs(device: str) -> None: + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") rng = np.random.default_rng(712) # This includes 14-DoF-sized plane sets, zero one-sided bounds and parallel # planes. SciPy's HiGHS LP is independent of the TOPPRA projection algorithm. @@ -124,13 +130,18 @@ def test_cpu_projection_matches_independent_linear_programs() -> None: ) assert solution.success expected.append(solution.x[0]) - np.testing.assert_allclose(_project(planes), expected, rtol=1e-10, atol=1e-12) + np.testing.assert_allclose( + _project(planes, device), expected, rtol=1e-10, atol=1e-12 + ) -@pytest.mark.parametrize("sampling", [{"sample_count": 129}, {"sample_dt": 0.07}]) -def test_cpu_optimized_forward_matches_differentiable_solver(sampling: dict) -> None: +@pytest.mark.parametrize("sampling", [{"sample_count": 513}, {"sample_dt": 0.07}]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_cpu_optimized_forward_matches_differentiable_solver( + sampling: dict, dtype: torch.dtype +) -> None: rng = np.random.default_rng(731) - points = torch.tensor(rng.normal(size=(8, 16, 7)), dtype=torch.float64) + points = torch.tensor(rng.normal(size=(8, 16, 7)), dtype=dtype) velocity = np.column_stack((-rng.uniform(0.4, 1.2, 7), rng.uniform(0.5, 1.4, 7))) acceleration = np.column_stack( (-rng.uniform(0.8, 1.6, 7), rng.uniform(0.9, 1.8, 7)) @@ -140,8 +151,12 @@ def test_cpu_optimized_forward_matches_differentiable_solver(sampling: dict) -> points.requires_grad_(), velocity, acceleration, **sampling ) assert torch.equal(optimized["success"], tracked["success"]) + tolerance = 1e-6 if dtype == torch.float32 else 1e-10 for key in ("positions", "velocities", "accelerations", "dt"): - torch.testing.assert_close(optimized[key], tracked[key], rtol=1e-10, atol=1e-10) + assert optimized[key].dtype == dtype + torch.testing.assert_close( + optimized[key], tracked[key], rtol=tolerance, atol=tolerance + ) tracked["dt"].sum().backward() assert torch.isfinite(points.grad).all() diff --git a/tests/compute/test_toppra_cuda.py b/tests/compute/test_toppra_cuda.py new file mode 100644 index 000000000..803d3488f --- /dev/null +++ b/tests/compute/test_toppra_cuda.py @@ -0,0 +1,98 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Cooperative CUDA inference agrees with the serial differentiable solver.""" + +from __future__ import annotations + +import pytest +import torch + +from embodichain.compute.trajectory._toppra_warp import _retime_toppra_warp + +pytestmark = pytest.mark.gpu + + +@pytest.mark.parametrize("dofs", [1, 14]) +@pytest.mark.parametrize("batch", [4, 129]) +@pytest.mark.parametrize("sampling", [{"sample_count": 17}, {"sample_dt": 0.2}]) +def test_cuda_cooperative_inference_preserves_serial_solution( + dofs: int, batch: int, sampling: dict[str, int | float] +) -> None: + """Cover multi-warp planes, both projection routes, and isolated failure rows.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + generator = torch.Generator().manual_seed(804) + points = 0.03 * torch.randn( + batch, 6, dofs, generator=generator, dtype=torch.float64 + ).cumsum(dim=1) + points[1] = points[1, :1].clone() + points[2, 0, 0] = torch.nan + points[3, 0, 0] = torch.inf + points = points.cuda() + lengths = (4 + torch.arange(batch) % 3).to(device="cuda", dtype=torch.int32) + velocity = torch.stack( + (-torch.linspace(0.7, 1.1, dofs), torch.linspace(0.8, 1.2, dofs)), dim=-1 + ).to(device="cuda", dtype=torch.float64) + acceleration = velocity * 2.0 + # Requires-grad uses the independently retained serial projection/solve + # kernels; inference uses cooperative plane reductions and, for large + # batches, the fast projection. + reference = _retime_toppra_warp( + points.clone().requires_grad_(), + velocity, + acceleration, + path_lengths=lengths, + **sampling, + ) + with torch.no_grad(): + actual = _retime_toppra_warp( + points, velocity, acceleration, path_lengths=lengths, **sampling + ) + assert reference["success"][:2].all() + assert not reference["success"][2:4].any() + assert reference["success"][4:].all() + for name in reference: + torch.testing.assert_close(actual[name], reference[name], rtol=1e-9, atol=1e-9) + + +@pytest.mark.parametrize("scale", [1e-4, 1.0, 1e4]) +def test_cuda_cooperative_solver_rejects_finite_infeasible_rows(scale: float) -> None: + """One-sided motion fails inside solve while frozen motion fails in prepare.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + # A large mixed batch also exercises fast projection. The decreasing + # path contains finite inputs but cannot move under a nonnegative speed. + points = torch.tensor( + [ + [[0.0, 2.0], [1.0, 2.0]], + [[0.0, 2.0], [-1.0, 2.0]], + [[0.0, 2.0], [1.0, 3.0]], + [[0.5, 2.0], [0.5, 2.0]], + ], + dtype=torch.float64, + device="cuda", + ).repeat(32, 1, 1) + velocity = [[0.0, scale], [0.0, 0.0]] + acceleration = [[-2.0 * scale**2, scale**2], [0.0, 0.0]] + reference = _retime_toppra_warp( + points.clone().requires_grad_(), velocity, acceleration, sample_count=33 + ) + with torch.no_grad(): + actual = _retime_toppra_warp(points, velocity, acceleration, sample_count=33) + assert actual["success"].tolist() == [True, False, False, True] * 32 + for name in reference: + torch.testing.assert_close(actual[name], reference[name], rtol=1e-9, atol=1e-9)