diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 8fae060e0..e9714da82 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -407,6 +407,9 @@ topics: source_of_truth: - embodichain/lab/sim/planners/base_planner.py - embodichain/lab/sim/planners/toppra_planner.py + - embodichain/lab/sim/planners/trapezoidal_planner.py + - embodichain/lab/sim/planners/trapezoidal_warp.py + - scripts/benchmark/motion_generation/trapezoidal_planner.py - embodichain/lab/sim/planners/curobo/curobo_planner.py - embodichain/lab/sim/planners/curobo/curobo_yaml.py - embodichain/lab/sim/planners/motion_generator.py diff --git a/agent_context/topics/ik-solvers/ik-solvers.md b/agent_context/topics/ik-solvers/ik-solvers.md index ac0ddcbf0..998e6ca0f 100644 --- a/agent_context/topics/ik-solvers/ik-solvers.md +++ b/agent_context/topics/ik-solvers/ik-solvers.md @@ -205,6 +205,10 @@ and available CUDA backends in seeded and full redundancy-search modes. - `a1, a2, b, c1–c4, offsets, flip_axes, has_parallelogram`: OPW kinematic parameters. - `safe_margin`: joint-limit safety margin in radians. +- `get_ik(..., return_all_solutions=True)` is the candidate-generation path for + continuous batch IK. `Robot.compute_batch_ik(..., continuous=True)` performs + sequential branch selection through an internal OPW selector, using the + previous sample as the next seed and the configured nearest-solution weights. --- diff --git a/agent_context/topics/motion-planning/motion-planning.md b/agent_context/topics/motion-planning/motion-planning.md index e6028a532..7a34ad584 100644 --- a/agent_context/topics/motion-planning/motion-planning.md +++ b/agent_context/topics/motion-planning/motion-planning.md @@ -7,6 +7,8 @@ | Planner registry | `embodichain/lab/sim/planners/__init__.py` | | Base planner class & config | `embodichain/lab/sim/planners/base_planner.py` → `BasePlanner`, `BasePlannerCfg`, `CollisionWorldInfo`, `PlanOptions`, `validate_plan_options` | | TOPPRA planner | `embodichain/lab/sim/planners/toppra_planner.py` → `ToppraPlanner`, `ToppraPlannerCfg`, `ToppraPlanOptions` | +| Trapezoidal planner | `embodichain/lab/sim/planners/trapezoidal_planner.py` → `TrapezoidalPlanner`, `TrapezoidalPlannerCfg`, `TrapezoidalPlanOptions` | +| Trapezoidal Warp kernels | `embodichain/utils/warp/kinematics/trapezoidal_warp.py` → batched profile construction and sampling kernels | | Neural planner | `embodichain/lab/sim/planners/neural_planner.py` → `NeuralPlanner`, `NeuralPlannerCfg`, `NeuralPlanOptions` | | cuRobo planner | `embodichain/lab/sim/planners/curobo/curobo_planner.py` → `CuroboPlanner`, `CuroboPlannerCfg`, `CuroboWorldCfg`, `CuroboPlanOptions` | | Planner assets | `embodichain/data/assets/planner_assets.py` → `download_neural_planner_checkpoint()` | @@ -30,6 +32,7 @@ The entire stack is **env-batched** (`B = num_envs`). `PlanState` / `PlanResult` ``` BasePlanner (ABC) ├─ ToppraPlanner Time-optimal path parameterization (fork-pool fan-out) + ├─ TrapezoidalPlanner Batched trapezoidal or jerk-limited Double-S timing ├─ NeuralPlanner (experimental) APG waypoint rollout (native batching) └─ CuroboPlanner CUDA collision-aware planning (native batching) @@ -40,12 +43,14 @@ Config hierarchy: ``` BasePlannerCfg robot_uid (MISSING), planner_type ├─ ToppraPlannerCfg planner_type = "toppra", max_workers, mp_context + ├─ TrapezoidalPlannerCfg planner_type = "trapezoidal" └─ NeuralPlannerCfg planner_type = "neural", checkpoint_path (MISSING) MotionGenCfg planner_cfg (MISSING — must be a BasePlannerCfg subclass) PlanOptions (empty base) ├─ ToppraPlanOptions constraints, sample_method, sample_interval + ├─ TrapezoidalPlanOptions profile, constraints, sample_method, sample_interval └─ NeuralPlanOptions control_part, start_qpos, max_steps MotionGenOptions strategy, sample_count, velocity/acceleration limits, @@ -85,6 +90,94 @@ Worker details: - `TIME` sampling can produce per-env waypoint counts; shorter trajectories are tail-padded by repeating the final waypoint and `duration` records the real endpoint per env. - Per-env failures set `success[b] = False` and fill the env's trajectory with its start qpos; other envs continue. `BrokenProcessPool` tears the pool down and rebuilds it on the next call. +### TrapezoidalPlanner + +The trapezoidal planner is a dependency-free, batched Torch backend for +piecewise-linear joint paths. Each input waypoint is a rest point. The default +profile is acceleration-limited trapezoidal timing (with triangular fallback +for short moves); ``profile="double_s"`` selects the rest-to-rest linear-path +subset of the seven-phase Double-S time law. It preserves that +implementation's discrete ``amax *= 0.9`` feasibility search for moves without +a cruise phase and its 1% ``EnforceJointLimits`` margin. Scalar or per-joint +velocity, acceleration, and jerk limits are projected onto each linear path +segment. Golden tests cover durations and sampled position, velocity, and +acceleration against the reference trajectory implementation. It supports fixed +quantity and approximate fixed-time sampling and returns explicit ``dt``. +``minimum_duration`` applies one uniform per-environment time scale, preserving +the path while reducing velocity, acceleration, and jerk. The minimal +``scripts/tutorials/sim/planner/trapezoidal_profile.py`` example plots scalar position, +velocity, acceleration, and jerk without starting simulation. A runnable +batched robot example lives in ``scripts/tutorials/sim/planner/trapezoidal_planner.py``. +The tutorial names the two diagnostics ``velocity_trapezoidal`` (the +``trapezoidal`` backend profile) and ``acceleration_trapezoidal`` (the +jerk-limited ``double_s`` backend profile). It can save diagnostic plots with +``--plot-output`` and supports headless execution with ``--no-show-plot``. +Multi-path/profile runs finish planning before showing all figures together. +``--path joint`` generates synchronized, limit-clamped motion on every arm +joint. ``--path cartesian`` is Cartesian-first rather than joint-first: the +trapezoidal backend time-parameterizes metric line arclength ``s(t)`` under +``--cartesian-velocity``, ``--cartesian-acceleration``, and +``--cartesian-jerk``; every resulting sample becomes an exact +fixed-orientation point on the line before continuous-seed IK. The joint +trajectory is never resampled afterward, so its desired EEF path retains the +planned Cartesian geometry and time law. ``--path both`` runs both diagnostics +and reports the scalar Cartesian derivative peaks plus maximum FK line error. +``--cartesian-distance`` controls line length and ``--cartesian-step`` sets a +minimum IK sample density. For a six-axis arm the tutorial uses the same +non-singular seed as the TOPPRA demo before solving the default 0.10 m downward +line. Joint derivatives used by the diagnostic and replay come from +second-order differential kinematics. The solver Jacobian gives ``dq/ds`` for +the fixed-orientation Cartesian tangent, its path derivative gives +``d²q/ds²``, and the Double-S outputs are composed by the chain rule as +``dq/dt = dq/ds * ds/dt`` and +``d²q/dt² = d²q/ds² * (ds/dt)² + dq/ds * d²s/dt²``. The implementation does +not numerically differentiate sampled IK positions in time. Cartesian derivative constraints +apply to ``s(t)`` and do not imply identical joint-space jerk bounds after the +nonlinear IK mapping. No display filtering is applied. Before derivative +evaluation, +each analytic IK sample is replaced only by the joint-limit-valid ``2π`` +equivalent nearest to the previous seed, removing representation wrap jumps +without changing the physical configuration or Cartesian path. +For OPW robots the Cartesian tutorial submits the complete pose path through +``Robot.compute_batch_ik(..., continuous=True)``. This reuses the existing batch +IK boundary and asks OPW for all candidates in one solver call before its +internal temporally ordered branch selector runs; other solver types fail +explicitly rather than being silently treated as continuous path solvers. +Interactive replay consumes ``PlanResult.dt`` rather than submitting all +samples as fast as Python can loop: every command advances enough physics steps +for its scaled interval and windowed runs are wall-clock paced. The +``--replay-speed`` multiplier controls playback speed; headless runs retain +unthrottled wall-clock execution while preserving physics-step timing. Before +each run, both current joint state and drive target are reset to the planned +start pose. +By default every input waypoint remains a rest point. Set +``stop_at_waypoints=False`` to remove duplicate and straight, same-direction +interior points before timing; genuine direction changes remain explicit rest +points, and batch rows with fewer retained points are final-pose padded. +``backend="auto"`` uses Warp profile construction and sampling for CUDA float32 +inputs and the Torch reference path otherwise. ``backend="warp"`` also permits +explicit CPU Warp execution but requires float32. Path compression and limit +projection remain shared Torch tensor operations. Warp then constructs each +trapezoidal or Double-S scalar segment profile in parallel, including the +Double-S no-cruise acceleration reduction and phase integration. Shared +post-processing applies minimum-duration scaling and the Double-S duration +margin consistently across backends. +Torch uses batched ``searchsorted`` for segment lookup without materializing a +``(B, N, segments)`` comparison tensor. Warp uses binary segment search once +per ``(B, N)`` sample, then a separate ``(B, N, DOF)`` composition kernel so +joint dimensions do not repeat profile lookup work. +Torch phase lookup also uses ``searchsorted`` and expands only selected phase +durations. Position, velocity, acceleration, and jerk coefficients are gathered +directly by ``(batch, segment, phase)``, avoiding four additional +``(B, N, phases)`` temporary tensors. +An all-stationary batch takes a dedicated hold fast path: it generates only +sample times and zero-derivative outputs, skipping constraint projection, +profile construction, segment lookup, and Warp dispatch. +The reproducible microbenchmark at +``scripts/benchmark/motion_generation/trapezoidal_planner.py`` measures Torch +and Warp time, CPU/GPU memory, endpoint success, and cross-backend error, then +writes the required three-table Markdown report under ``outputs/benchmarks``. + ### NeuralPlanner (experimental) Learning-based EEF waypoint planner. Franka Panda only. @@ -179,7 +272,7 @@ import or branch on concrete planner option types. Unified interface for trajectory planning with optional pre-interpolation. - Wraps a `BasePlanner` instance (resolved from `planner_cfg.planner_type`). -- Supported planner types: TOPPRA, NeuralPlanner, and cuRobo. +- Supported planner types: TOPPRA, TrapezoidalPlanner, NeuralPlanner, and cuRobo. - `MotionGenCfg.planner_cfg` is **MISSING** — must be provided. - `generate()` and `interpolate_trajectory()` are env-batched (`B, N, DOF`). - `generate()` always returns a normalized `PlanResult`; failed rows hold the diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.planners.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.planners.rst index 3af9a9efc..73dc09e0e 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.planners.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.planners.rst @@ -10,6 +10,9 @@ embodichain.lab.sim.planners BasePlanner ToppraPlannerCfg ToppraPlanner + TrapezoidalPlanOptions + TrapezoidalPlannerCfg + TrapezoidalPlanner MotionGenCfg MotionGenerator TrajectorySampleMethod @@ -44,6 +47,74 @@ Toppra Planner :inherited-members: :show-inheritance: +Trapezoidal Planner +------------------- + +The trapezoidal planner applies either acceleration-limited trapezoidal timing +or jerk-limited Double-S timing to batched, piecewise-linear joint paths. Use +``TrapezoidalPlanOptions.minimum_duration`` to slow a trajectory without +changing its path or violating derivative limits. +For densely interpolated straight paths, set ``stop_at_waypoints=False`` to +remove redundant same-direction interior points while preserving real corners. +``backend="auto"`` selects Warp profile construction and sampling for CUDA +float32 inputs and retains the Torch reference implementation for CPU or +float64. The Warp path builds trapezoidal or Double-S phases in parallel per +batch segment, then evaluates scalar samples and composes all joints. +Both backends avoid a dense sample-by-segment lookup tensor: Torch uses batched +``searchsorted``, while Warp performs one binary lookup per batch sample before +parallel joint composition. +Torch also gathers profile coefficients directly by selected phase, avoiding +full sample-by-phase copies of position, velocity, acceleration, and jerk data. +All-stationary batches use a hold fast path and skip constraint projection, +profile construction, segment lookup, and backend dispatch. + +.. autoclass:: TrapezoidalPlanOptions + :members: + :exclude-members: __init__, copy, replace, to_dict, validate + +.. autoclass:: TrapezoidalPlannerCfg + :members: + :exclude-members: __init__, copy, replace, to_dict, validate + +.. autoclass:: TrapezoidalPlanner + :members: + :inherited-members: + :show-inheritance: + +Example +~~~~~~~ + +Select the Double-S profile when jerk continuity is required while retaining +the same ``MotionGenerator`` entry point:: + + generator = MotionGenerator( + MotionGenCfg( + planner_cfg=TrapezoidalPlannerCfg(robot_uid=robot.uid), + ) + ) + result = generator.generate( + [PlanState.from_qpos(start), PlanState.from_qpos(goal)], + MotionGenOptions( + plan_opts=TrapezoidalPlanOptions( + profile="double_s", + constraints={ + "velocity": 0.5, + "acceleration": 1.0, + "jerk": 3.0, + }, + sample_interval=200, + stop_at_waypoints=False, + backend="auto", + ) + ), + ) + +The complete batched simulation tutorial is +``scripts/tutorials/sim/planner/trapezoidal_planner.py``. +Pass ``--backend torch`` or ``--backend warp`` to compare implementations. +For repeatable timing and memory measurements, run +``scripts/benchmark/motion_generation/trapezoidal_planner.py``. + Motion Generator ---------------- diff --git a/docs/source/api_reference/public_api.rst b/docs/source/api_reference/public_api.rst index ea1aa59e5..6186e5a21 100644 --- a/docs/source/api_reference/public_api.rst +++ b/docs/source/api_reference/public_api.rst @@ -982,6 +982,31 @@ embodichain.lab.sim.planners.toppra_planner ToppraPlannerCfg ToppraPlanOptions +embodichain.lab.sim.planners.trapezoidal_planner +------------------------------------------------ + +.. currentmodule:: embodichain.lab.sim.planners.trapezoidal_planner + +.. autosummary:: + + TrapezoidalPlanOptions + TrapezoidalPlanner + TrapezoidalPlannerCfg + +embodichain.utils.warp.kinematics.trapezoidal_warp +-------------------------------------------------- + +.. currentmodule:: embodichain.utils.warp.kinematics.trapezoidal_warp + +Warp-accelerated helpers construct scalar trapezoidal or Double-S motion +profiles and compose their sampled path derivatives into batched joint-space +trajectories. + +.. autosummary:: + + build_profile_warp + compose_profile_samples_warp + embodichain.lab.sim.planners.utils ---------------------------------- @@ -1244,6 +1269,19 @@ embodichain.lab.sim.solvers.null_space_posture_task NullSpacePostureTask +embodichain.lab.sim.solvers.opw_solver +-------------------------------------- + +.. currentmodule:: embodichain.lab.sim.solvers.opw_solver + +Configuration and runtime solver for analytic OPW forward and inverse +kinematics of compatible six-axis manipulators. + +.. autosummary:: + + OPWSolver + OPWSolverCfg + embodichain.lab.sim.solvers.pink_solver --------------------------------------- diff --git a/docs/source/overview/sim/solvers/opw_solver.md b/docs/source/overview/sim/solvers/opw_solver.md index 8f7bff079..ae80a30b2 100644 --- a/docs/source/overview/sim/solvers/opw_solver.md +++ b/docs/source/overview/sim/solvers/opw_solver.md @@ -11,6 +11,7 @@ * Flexible configuration via `OPWSolverCfg` * Strict enforcement of joint limits * Forward kinematics (FK) and multiple IK solution branches +* Continuous whole-path IK through the existing batch-IK boundary ## Configuration @@ -104,6 +105,12 @@ solver = OPWSolver(cfg, device="cuda") ``` +Continuous selection is requested at the robot boundary with +`Robot.compute_batch_ik(..., continuous=True)`. The robot performs arena/root +frame conversion and calls `get_ik(..., return_all_solutions=True)` once for the +flattened pose batch. OPW then applies its internal sequential selector and +returns validity `(B, N)` plus continuous joint positions `(B, N, 6)`. + ## References * [OPW Kinematics Paper](https://doi.org/10.1109/TRO.2017.2776312) diff --git a/embodichain/lab/sim/objects/robot.py b/embodichain/lab/sim/objects/robot.py index 7b8a1340e..68dab7580 100644 --- a/embodichain/lab/sim/objects/robot.py +++ b/embodichain/lab/sim/objects/robot.py @@ -987,20 +987,35 @@ def compute_batch_ik( joint_seed: torch.Tensor | np.ndarray | None, name: str, env_ids: Sequence[int] | None = None, - ): - """Compute the inverse kinematics of the robot given joint positions and optionally a specific part name. - The input pose should be in the local arena frame. + *, + continuous: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor] | None: + """Compute batched inverse kinematics for arena-frame poses. + + When ``continuous`` is enabled, all analytic candidates are generated + in one solver call and the branch nearest to the previous path sample + is selected sequentially. The selected solver must support continuous + candidate selection. Args: - pose (torch.Tensor): The end effector pose of the robot, (num_envs, n_batch, 7) or (num_envs, n_batch, 4, 4). - joint_seed (torch.Tensor | None): The joint positions to use as a seed for the IK computation, (num_envs, n_batch, dof). If None, the zero joint positions will be used as the seed. - name (str | None): The name of the control part to compute the IK for. If None, the default part is used. - env_ids (Sequence[int] | None): Environment indices to apply the positions. Defaults to all environments. + pose: End-effector poses shaped ``(B, N, 7)`` or + ``(B, N, 4, 4)`` in the local arena frame. + joint_seed: Independent seeds shaped ``(B, N, DOF)``. In + continuous mode, the initial path seed shaped ``(B, DOF)``. + Defaults to zero independent seeds when omitted. + name: Control part whose solver should be used. + env_ids: Optional environment indices. Defaults to all + environments. + continuous: Preserve one temporally continuous IK branch across + the ``N`` path samples. Defaults to ``False``. Returns: - Tuple[torch.Tensor, torch.Tensor]: - Success Tensor with shape (num_envs, n_batch) - Qpos Tensor with shape (num_envs, n_batch, dof). + Per-sample success shaped ``(B, N)`` and joint positions shaped + ``(B, N, DOF)``, or ``None`` when no solver is configured. + + Raises: + ValueError: If an input shape is invalid or continuous selection + is unsupported by the configured solver. """ local_env_ids = self._all_indices if env_ids is None else env_ids @@ -1012,36 +1027,43 @@ def compute_batch_ik( return None pose = to_tensor(pose, device=self.device) - if pose.shape[0] != len(local_env_ids): - logger.log_error( - f"Pose batch size mismatch. Expected {len(local_env_ids)} but got {pose.shape[0]}." + batch_size = len(local_env_ids) + if pose.ndim not in {3, 4} or pose.shape[0] != batch_size: + raise ValueError( + f"pose must have shape ({batch_size}, N, 7) or " + f"({batch_size}, N, 4, 4)." ) n_batch = pose.shape[1] n_dof = solver.dof - if joint_seed is None: + if continuous: + if joint_seed is None: + joint_seed_tensor = torch.zeros( + (batch_size, n_dof), + dtype=torch.float32, + device=self.device, + ) + else: + joint_seed_tensor = to_tensor(joint_seed, device=self.device) + if joint_seed_tensor.shape != (batch_size, n_dof): + raise ValueError( + f"joint_seed must have shape ({batch_size}, {n_dof}) " + "when continuous=True." + ) + elif joint_seed is None: joint_seed = torch.zeros( - (len(local_env_ids), n_batch, n_dof), + (batch_size, n_batch, n_dof), dtype=torch.float32, device=self.device, ) - - if joint_seed.shape[0] != len(local_env_ids): - logger.log_error( - f"Joint seed env size mismatch. Expected {len(local_env_ids)} but got {joint_seed.shape[0]}." - ) - - if joint_seed.shape[1] != n_batch: - logger.log_error( - f"Joint seed batch size mismatch. Expected {n_batch} but got {joint_seed.shape[1]}." - ) - - if joint_seed.shape[-1] != n_dof: - logger.log_error( - f"Joint seed dof size mismatch. Expected {n_batch} but got {joint_seed.shape[-1]}." + else: + joint_seed = to_tensor(joint_seed, device=self.device) + if not continuous and joint_seed.shape != (batch_size, n_batch, n_dof): + raise ValueError( + f"joint_seed must have shape ({batch_size}, {n_batch}, {n_dof})." ) - if pose.shape[-1] == 7 and pose.dim() == 3: + if pose.shape[-1] == 7 and pose.ndim == 3: # Convert pose from (num_envs, n_batch, 7) to (num_envs * n_batch, 4, 4) pose_batch = pose.reshape(-1, 7) pos = pose_batch[:, :3] @@ -1054,9 +1076,14 @@ def compute_batch_ik( ) pose_batch[:, :3, :3] = rot pose_batch[:, :3, 3] = pos - else: + elif pose.shape[-2:] == (4, 4) and pose.ndim == 4: # Convert pose from (num_envs, n_batch, 4, 4) to (num_envs * n_batch, 4, 4) pose_batch = pose.reshape(-1, 4, 4) + else: + raise ValueError( + f"pose must have shape ({batch_size}, N, 7) or " + f"({batch_size}, N, 4, 4)." + ) # get xpos from link root base_xpos_n_envs = self.get_link_pose( @@ -1070,14 +1097,31 @@ def compute_batch_ik( ) pose_batch = torch.bmm(base_inv_xpos_batch, pose_batch) + if continuous: + candidate_valid, candidate_qpos = solver.get_ik( + target_xpos=pose_batch, + qpos_seed=None, + return_all_solutions=True, + ) + select_path = getattr(solver, "_select_continuous_ik_path", None) + if not callable(select_path): + raise ValueError( + f"Solver for {name!r} does not support continuous batch IK." + ) + return select_path( + candidate_qpos.reshape(batch_size, n_batch, -1, n_dof), + candidate_valid.reshape(batch_size, n_batch, -1), + joint_seed_tensor, + ) + joint_seed_batch = joint_seed.reshape(-1, n_dof) ret, qpos_batch = solver.get_ik( target_xpos=pose_batch, qpos_seed=joint_seed_batch, return_all_solutions=False, ) - ret = ret.reshape(len(local_env_ids), n_batch) - qpos = qpos_batch.reshape(len(local_env_ids), n_batch, n_dof) + ret = ret.reshape(batch_size, n_batch) + qpos = qpos_batch.reshape(batch_size, n_batch, n_dof) return ret, qpos def _init_control_parts(self, control_parts: Dict[str, List[str]]) -> None: diff --git a/embodichain/lab/sim/planners/__init__.py b/embodichain/lab/sim/planners/__init__.py index d0fd8ef09..07e75c205 100644 --- a/embodichain/lab/sim/planners/__init__.py +++ b/embodichain/lab/sim/planners/__init__.py @@ -22,6 +22,7 @@ from .utils import * from .base_planner import * from .toppra_planner import * +from .trapezoidal_planner import * from .neural_planner import * from .curobo.curobo_yaml import * from .curobo.curobo_planner import * diff --git a/embodichain/lab/sim/planners/motion_generator.py b/embodichain/lab/sim/planners/motion_generator.py index a15c224b8..4b4ed7091 100644 --- a/embodichain/lab/sim/planners/motion_generator.py +++ b/embodichain/lab/sim/planners/motion_generator.py @@ -34,6 +34,9 @@ ToppraPlanner, ToppraPlannerCfg, ToppraPlanOptions, + TrapezoidalPlanner, + TrapezoidalPlannerCfg, + TrapezoidalPlanOptions, NeuralPlanner, NeuralPlannerCfg, CuroboPlanner, @@ -63,7 +66,6 @@ @configclass class MotionGenCfg: - planner_cfg: BasePlannerCfg = MISSING """Configuration for the underlying planner. Must include 'planner_type' attribute to specify which planner to use, and any additional parameters required by that planner. @@ -74,7 +76,6 @@ class MotionGenCfg: @configclass class MotionGenOptions: - strategy: Literal["motion_gen", "ik_interp"] = "motion_gen" """Motion strategy: backend planning or deterministic IK interpolation.""" @@ -168,6 +169,7 @@ class MotionGenerator: _support_planner_dict = { "toppra": (ToppraPlanner, ToppraPlannerCfg), + "trapezoidal": (TrapezoidalPlanner, TrapezoidalPlannerCfg), "neural": (NeuralPlanner, NeuralPlannerCfg), "curobo": (CuroboPlanner, CuroboPlannerCfg), } @@ -237,8 +239,7 @@ def _validate_collision_pose_keys( for entity_id in entity_ids ): raise TypeError( - f"{field_name} keys must be non-empty strings without outer " - "whitespace." + f"{field_name} keys must be non-empty strings without outer whitespace." ) return set(entity_ids) @@ -414,17 +415,32 @@ def resolve_plan_options( raise ValueError("sample_count must be at least 2.") if plan_opts is not None: return deepcopy(plan_opts) - if sample_count is not None and self.planner.cfg.planner_type == "toppra": - return ToppraPlanOptions( - sample_method=TrajectorySampleMethod.QUANTITY, - sample_interval=sample_count, - constraints={ - "velocity": 0.2 if velocity_limit is None else velocity_limit, - "acceleration": ( - 0.5 if acceleration_limit is None else acceleration_limit - ), - }, + planner_type = getattr(getattr(self.planner, "cfg", None), "planner_type", None) + if planner_type in {"toppra", "trapezoidal"} and ( + sample_count is not None + or velocity_limit is not None + or acceleration_limit is not None + ): + options_type = ( + ToppraPlanOptions + if planner_type == "toppra" + else TrapezoidalPlanOptions ) + constraints: dict[str, float] = { + "velocity": 0.2 if velocity_limit is None else velocity_limit, + "acceleration": ( + 0.5 if acceleration_limit is None else acceleration_limit + ), + } + if planner_type == "trapezoidal": + constraints["jerk"] = 2.0 + options_kwargs: dict[str, object] = {"constraints": constraints} + if sample_count is not None: + options_kwargs.update( + sample_method=TrajectorySampleMethod.QUANTITY, + sample_interval=sample_count, + ) + return options_type(**options_kwargs) return self.planner.default_plan_options() @classmethod @@ -1085,7 +1101,7 @@ def ensure_batch_dim(tensor): alpha = 1.0 if batch_size == 1 else max(0.2, 1.0 / np.sqrt(batch_size)) for i in range(self.dofs): - label = f"Joint {i+1}" if b == 0 else "" + label = f"Joint {i + 1}" if b == 0 else "" axs[0].plot( time_steps, positions[b, :, i].numpy(), diff --git a/embodichain/lab/sim/planners/trapezoidal_planner.py b/embodichain/lab/sim/planners/trapezoidal_planner.py new file mode 100644 index 000000000..df2dc7ada --- /dev/null +++ b/embodichain/lab/sim/planners/trapezoidal_planner.py @@ -0,0 +1,807 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Batched trapezoidal and jerk-limited joint trajectory planning.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Literal + +import torch + +from embodichain.utils import configclass + +from .base_planner import ( + BasePlanner, + BasePlannerCfg, + PlanOptions, + validate_plan_options, +) +from .utils import MoveType, PlanResult, PlanState, TrajectorySampleMethod + +__all__ = [ + "TrapezoidalPlanOptions", + "TrapezoidalPlanner", + "TrapezoidalPlannerCfg", +] + + +@configclass +class TrapezoidalPlannerCfg(BasePlannerCfg): + """Configure the batched joint-space time-profile planner.""" + + planner_type: str = "trapezoidal" + + +@configclass +class TrapezoidalPlanOptions(PlanOptions): + """Configure trapezoidal or Double-S trajectory generation. + + Args: + profile: Scalar time profile used for every linear joint-path segment. + constraints: Positive scalar or per-joint ``velocity``, ``acceleration``, + and ``jerk`` limits. Jerk is required by the ``double_s`` profile. + sample_method: Fixed output quantity or approximately fixed time step. + sample_interval: Output count for ``QUANTITY`` or seconds for ``TIME``. + minimum_duration: Optional lower bound on each environment trajectory's + duration. Slower trajectories retain the same path and limits. + stop_at_waypoints: Whether every supplied waypoint is a rest point. When + false, redundant points on straight, same-direction runs are removed. + collinearity_tolerance: Cosine tolerance used by waypoint compression. + backend: Profile-construction and sampling backend. ``auto`` selects + Warp for CUDA float32 trajectories and Torch otherwise. + """ + + profile: Literal["trapezoidal", "double_s"] = "trapezoidal" + constraints: dict = { # noqa: RUF012 + "velocity": 0.2, + "acceleration": 0.5, + "jerk": 2.0, + } + sample_method: TrajectorySampleMethod = TrajectorySampleMethod.QUANTITY + sample_interval: float | int = 100 + minimum_duration: float | None = None + stop_at_waypoints: bool = True + collinearity_tolerance: float = 1e-5 + backend: Literal["auto", "torch", "warp"] = "auto" + + def __post_init__(self) -> None: + if self.profile not in {"trapezoidal", "double_s"}: + raise ValueError("profile must be 'trapezoidal' or 'double_s'.") + required = {"velocity", "acceleration"} + if self.profile == "double_s": + required.add("jerk") + missing = sorted(required.difference(self.constraints)) + if missing: + raise ValueError(f"constraints is missing required keys: {missing}.") + if self.sample_method is TrajectorySampleMethod.QUANTITY: + if ( + isinstance(self.sample_interval, bool) + or not isinstance(self.sample_interval, int) + or self.sample_interval < 2 + ): + raise ValueError( + "QUANTITY sample_interval must be an integer of at least 2." + ) + elif self.sample_method is TrajectorySampleMethod.TIME: + if ( + isinstance(self.sample_interval, bool) + or not isinstance(self.sample_interval, (int, float)) + or not math.isfinite(float(self.sample_interval)) + or float(self.sample_interval) <= 0.0 + ): + raise ValueError( + "TIME sample_interval must be finite and greater than zero." + ) + else: + raise ValueError(f"Unsupported sample method: {self.sample_method!r}.") + if self.minimum_duration is not None and ( + isinstance(self.minimum_duration, bool) + or not isinstance(self.minimum_duration, (int, float)) + or not math.isfinite(float(self.minimum_duration)) + or self.minimum_duration <= 0.0 + ): + raise ValueError( + "minimum_duration must be finite and greater than zero when set." + ) + if not isinstance(self.stop_at_waypoints, bool): + raise TypeError("stop_at_waypoints must be a bool.") + if ( + isinstance(self.collinearity_tolerance, bool) + or not isinstance(self.collinearity_tolerance, (int, float)) + or not math.isfinite(float(self.collinearity_tolerance)) + or not 0.0 <= self.collinearity_tolerance < 1.0 + ): + raise ValueError( + "collinearity_tolerance must be finite and in the range [0, 1)." + ) + if self.backend not in {"auto", "torch", "warp"}: + raise ValueError("backend must be 'auto', 'torch', or 'warp'.") + + +@dataclass(slots=True) +class _ProfileBatch: + """Internal fixed-shape phase representation.""" + + durations: torch.Tensor + positions: torch.Tensor + velocities: torch.Tensor + accelerations: torch.Tensor + jerks: torch.Tensor + + +def _compress_collinear_waypoints( + waypoints: torch.Tensor, + tolerance: float, +) -> torch.Tensor: + """Remove duplicate and straight-run interior points in a batched path. + + Rows may retain different numbers of points. Shorter rows are padded by + repeating their final point, which the profile builder treats as zero-time + segments. + """ + if waypoints.shape[1] <= 2: + return waypoints + edges = waypoints[:, 1:] - waypoints[:, :-1] + epsilon = max(1e-8, float(tolerance) * 1e-3) + deduplicate = torch.ones( + waypoints.shape[:2], dtype=torch.bool, device=waypoints.device + ) + deduplicate[:, 1:] = torch.linalg.vector_norm(edges, dim=-1) > epsilon + deduplicated_count = deduplicate.sum(dim=1) + deduplicated_size = max(2, int(deduplicated_count.max().item())) + deduplicated = waypoints[:, -1:].expand(-1, deduplicated_size, -1).clone() + deduplicated_index = deduplicate.cumsum(dim=1) - 1 + batch_index = torch.arange(waypoints.shape[0], device=waypoints.device)[:, None] + batch_index = batch_index.expand_as(deduplicate) + deduplicated[batch_index[deduplicate], deduplicated_index[deduplicate]] = waypoints[ + deduplicate + ] + + if deduplicated_size <= 2: + return deduplicated + edges = deduplicated[:, 1:] - deduplicated[:, :-1] + previous = edges[:, :-1] + following = edges[:, 1:] + previous_norm = torch.linalg.vector_norm(previous, dim=-1) + following_norm = torch.linalg.vector_norm(following, dim=-1) + active = (previous_norm > epsilon) & (following_norm > epsilon) + cosine = (previous * following).sum(dim=-1) / ( + previous_norm * following_norm + ).clamp_min(epsilon) + straight = active & (cosine >= 1.0 - tolerance) + point_ids = torch.arange(deduplicated_size, device=waypoints.device)[None] + last_point = (deduplicated_count - 1)[:, None] + keep = (point_ids == 0) | (point_ids == last_point) + real_interior = (point_ids[:, 1:-1] > 0) & (point_ids[:, 1:-1] < last_point) + keep[:, 1:-1] |= real_interior & ~straight + retained_count = keep.sum(dim=1) + output_count = max(2, int(retained_count.max().item())) + output = deduplicated[:, -1:].expand(-1, output_count, -1).clone() + output_index = keep.cumsum(dim=1) - 1 + batch_index = torch.arange(waypoints.shape[0], device=waypoints.device)[:, None] + batch_index = batch_index.expand_as(keep) + output[batch_index[keep], output_index[keep]] = deduplicated[keep] + return output + + +def _limit_tensor( + value: float | list | torch.Tensor, reference: torch.Tensor +) -> torch.Tensor: + """Return a positive per-joint limit tensor broadcastable to ``(B, S, D)``.""" + limit = torch.as_tensor(value, dtype=reference.dtype, device=reference.device) + if limit.ndim == 0: + limit = limit.expand(reference.shape[-1]) + if limit.shape != (reference.shape[-1],): + raise ValueError( + f"Joint limits must be scalar or shape ({reference.shape[-1]},), " + f"got {tuple(limit.shape)}." + ) + if not bool(torch.isfinite(limit).all().item()) or bool((limit <= 0).any().item()): + raise ValueError("Joint limits must contain finite positive values.") + return limit + + +def _scalar_path_limits( + delta: torch.Tensor, + constraints: dict, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Project per-joint derivative limits onto linear scalar path segments.""" + absolute_delta = delta.abs() + active = absolute_delta > 1e-8 + + def project(key: str) -> torch.Tensor: + joint_limit = _limit_tensor(constraints[key], delta) + ratios = joint_limit / absolute_delta.clamp_min(1e-8) + ratios.masked_fill_(~active, torch.inf) + return ratios.amin(dim=-1) + + velocity = project("velocity") + acceleration = project("acceleration") + jerk = ( + project("jerk") + if "jerk" in constraints + else torch.full_like(velocity, torch.inf) + ) + stationary = ~active.any(dim=-1) + zeros = torch.zeros_like(velocity) + return ( + torch.where(stationary, zeros, velocity), + torch.where(stationary, zeros, acceleration), + torch.where(stationary, zeros, jerk), + ) + + +def _integrate_phases(durations: torch.Tensor, jerks: torch.Tensor) -> _ProfileBatch: + """Integrate piecewise-constant jerk phases from rest.""" + positions = torch.zeros_like(durations) + velocities = torch.zeros_like(durations) + accelerations = torch.zeros_like(durations) + position = torch.zeros_like(durations[..., 0]) + velocity = torch.zeros_like(position) + acceleration = torch.zeros_like(position) + for phase in range(durations.shape[-1]): + positions[..., phase] = position + velocities[..., phase] = velocity + accelerations[..., phase] = acceleration + duration = durations[..., phase] + jerk = jerks[..., phase] + position = ( + position + + velocity * duration + + 0.5 * acceleration * duration.square() + + jerk * duration.pow(3) / 6.0 + ) + velocity = velocity + acceleration * duration + 0.5 * jerk * duration.square() + acceleration = acceleration + jerk * duration + return _ProfileBatch(durations, positions, velocities, accelerations, jerks) + + +def _apply_minimum_duration( + profile: _ProfileBatch, + minimum_duration: float | None, +) -> _ProfileBatch: + """Uniformly slow complete batch rows to a requested minimum duration.""" + if minimum_duration is None: + return profile + total_duration = profile.durations.sum(dim=(-2, -1)) + epsilon = torch.finfo(total_duration.dtype).eps + stationary = total_duration <= epsilon + requested_scale = minimum_duration / total_duration.clamp_min(epsilon) + scale = torch.maximum(torch.ones_like(total_duration), requested_scale) + scale = torch.where(stationary, torch.ones_like(scale), scale) + phase_scale = scale[:, None, None] + slowed = _ProfileBatch( + durations=profile.durations * phase_scale, + positions=profile.positions, + velocities=profile.velocities / phase_scale, + accelerations=profile.accelerations / phase_scale.square(), + jerks=profile.jerks / phase_scale.pow(3), + ) + if bool(stationary.any().item()): + slowed.durations[stationary] = 0.0 + slowed.durations[stationary, 0, slowed.durations.shape[-1] // 2] = ( + minimum_duration + ) + slowed.positions[stationary] = 0.0 + slowed.velocities[stationary] = 0.0 + slowed.accelerations[stationary] = 0.0 + slowed.jerks[stationary] = 0.0 + return slowed + + +def _scale_profile_time(profile: _ProfileBatch, scale: torch.Tensor) -> _ProfileBatch: + """Apply per-segment uniform time scaling without changing path position.""" + phase_scale = scale[..., None] + return _ProfileBatch( + durations=profile.durations * phase_scale, + positions=profile.positions, + velocities=profile.velocities / phase_scale, + accelerations=profile.accelerations / phase_scale.square(), + jerks=profile.jerks / phase_scale.pow(3), + ) + + +def _build_trapezoidal_profile( + velocity_limit: torch.Tensor, + acceleration_limit: torch.Tensor, +) -> _ProfileBatch: + """Build rest-to-rest unit-distance triangular or trapezoidal profiles.""" + stationary = velocity_limit <= 0.0 + acceleration_time = velocity_limit / acceleration_limit.clamp_min(1e-12) + acceleration_distance_twice = ( + velocity_limit.square() / acceleration_limit.clamp_min(1e-12) + ) + reaches_velocity = acceleration_distance_twice < 1.0 + peak_velocity = torch.where( + reaches_velocity, + velocity_limit, + torch.sqrt(acceleration_limit), + ) + acceleration_time = peak_velocity / acceleration_limit.clamp_min(1e-12) + cruise_time = torch.where( + reaches_velocity, + (1.0 - peak_velocity.square() / acceleration_limit) / peak_velocity, + torch.zeros_like(peak_velocity), + ) + acceleration_time = torch.where( + stationary, torch.zeros_like(acceleration_time), acceleration_time + ) + cruise_time = torch.where(stationary, torch.zeros_like(cruise_time), cruise_time) + durations = torch.stack([acceleration_time, cruise_time, acceleration_time], dim=-1) + # A trapezoidal profile has piecewise-constant acceleration. Represent it as + # zero-jerk phases and seed their accelerations explicitly below. + profile = _integrate_phases(durations, torch.zeros_like(durations)) + profile.accelerations[..., 0] = acceleration_limit + profile.accelerations[..., 1] = 0.0 + profile.accelerations[..., 2] = -acceleration_limit + profile.velocities[..., 0] = 0.0 + profile.velocities[..., 1] = peak_velocity + profile.velocities[..., 2] = peak_velocity + profile.positions[..., 0] = 0.0 + profile.positions[..., 1] = 0.5 * peak_velocity * acceleration_time + profile.positions[..., 2] = profile.positions[..., 1] + peak_velocity * cruise_time + return profile + + +def _build_double_s_profile( + velocity_limit: torch.Tensor, + acceleration_limit: torch.Tensor, + jerk_limit: torch.Tensor, +) -> _ProfileBatch: + """Build the rest-to-rest seven-phase Double-S profile. + + This follows ``TrajectoryDoubleS::_ComputeDoubleSProfile`` for zero path + boundary velocities. In particular, a move without a constant-velocity + phase repeatedly lowers the candidate acceleration by ``0.9`` until both + acceleration halves can contain their two jerk ramps. That deliberately + differs from the closed-form triangular-jerk fallback commonly used by + simplified Double-S implementations. + """ + stationary = velocity_limit <= 0.0 + epsilon = 1e-12 + safe_velocity = velocity_limit.clamp_min(epsilon) + safe_acceleration = acceleration_limit.clamp_min(epsilon) + safe_jerk = jerk_limit.clamp_min(epsilon) + + reaches_acceleration = safe_velocity * safe_jerk >= safe_acceleration.square() + tj = torch.where( + reaches_acceleration, + safe_acceleration / safe_jerk, + torch.sqrt(safe_velocity / safe_jerk), + ) + ta = torch.where( + reaches_acceleration, + tj + safe_velocity / safe_acceleration, + 2.0 * tj, + ) + cruise_time = 1.0 / safe_velocity - ta + no_cruise = (~stationary) & (cruise_time <= 0.0) + + # Use the intentionally discrete feasibility search required by the + # than substituting an analytic triangular-jerk solution. + candidate_acceleration = safe_acceleration.clone() + for _ in range(1001): + if not bool(no_cruise.any().item()): + break + candidate_tj = candidate_acceleration / safe_jerk + delta = torch.sqrt( + candidate_acceleration.pow(4) / safe_jerk.square() + + 4.0 * candidate_acceleration + ) + candidate_ta = (candidate_acceleration.square() / safe_jerk + delta) / ( + 2.0 * candidate_acceleration + ) + accepted = no_cruise & (candidate_ta >= 2.0 * candidate_tj) + tj = torch.where(accepted, candidate_tj, tj) + ta = torch.where(accepted, candidate_ta, ta) + no_cruise = no_cruise & ~accepted + candidate_acceleration = torch.where( + no_cruise, candidate_acceleration * 0.9, candidate_acceleration + ) + if bool(no_cruise.any().item()): + raise RuntimeError("Double-S acceleration search did not converge.") + cruise_time = torch.clamp_min(cruise_time, 0.0) + tj = torch.where(stationary, torch.zeros_like(tj), tj) + ta = torch.where(stationary, torch.zeros_like(ta), ta) + cruise_time = torch.where(stationary, torch.zeros_like(cruise_time), cruise_time) + constant_acceleration_time = torch.clamp_min(ta - 2.0 * tj, 0.0) + durations = torch.stack( + [ + tj, + constant_acceleration_time, + tj, + cruise_time, + tj, + constant_acceleration_time, + tj, + ], + dim=-1, + ) + signs = torch.tensor( + [1.0, 0.0, -1.0, 0.0, -1.0, 0.0, 1.0], + dtype=jerk_limit.dtype, + device=jerk_limit.device, + ) + jerks = jerk_limit[..., None] * signs + jerks = torch.where(stationary[..., None], torch.zeros_like(jerks), jerks) + return _integrate_phases(durations, jerks) + + +def _use_warp_backend( + reference: torch.Tensor, + backend: Literal["auto", "torch", "warp"], +) -> bool: + """Return whether a profile stage should execute through Warp.""" + return backend == "warp" or ( + backend == "auto" and reference.is_cuda and reference.dtype == torch.float32 + ) + + +def _build_scalar_profile( + *, + profile_name: Literal["trapezoidal", "double_s"], + velocity_limit: torch.Tensor, + acceleration_limit: torch.Tensor, + jerk_limit: torch.Tensor, + backend: Literal["auto", "torch", "warp"], +) -> _ProfileBatch: + """Build scalar phases with the selected backend and shared post-processing.""" + use_warp = _use_warp_backend(velocity_limit, backend) + if use_warp and velocity_limit.dtype != torch.float32: + if backend == "warp": + raise ValueError("The Warp trajectory backend requires float32 input.") + use_warp = False + if use_warp: + try: + from embodichain.utils.warp.kinematics.trapezoidal_warp import ( + build_profile_warp, + ) + except ImportError: + if backend == "warp": + raise ImportError( + "The Warp trajectory backend requires the 'warp' package." + ) from None + use_warp = False + else: + tensors = build_profile_warp( + profile=profile_name, + velocity_limits=velocity_limit, + acceleration_limits=acceleration_limit, + jerk_limits=jerk_limit, + ) + result = _ProfileBatch(*tensors) + if not use_warp: + result = ( + _build_double_s_profile(velocity_limit, acceleration_limit, jerk_limit) + if profile_name == "double_s" + else _build_trapezoidal_profile(velocity_limit, acceleration_limit) + ) + + if profile_name == "double_s": + # The joint-limit projection leaves a 1% margin whenever a + # sampled derivative reaches a limit. Linear Double-S segments always + # reach their projected jerk limit, so the resulting scale is 1.01. + margin = torch.where( + velocity_limit > 0.0, + torch.full_like(velocity_limit, 1.01), + torch.ones_like(velocity_limit), + ) + result = _scale_profile_time(result, margin) + return result + + +def _compose_profile_samples_torch( + *, + times: torch.Tensor, + cumulative_duration: torch.Tensor, + profile: _ProfileBatch, + segment_starts: torch.Tensor, + segment_deltas: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Evaluate scalar profiles and compose joint states with Torch.""" + segment = torch.searchsorted( + cumulative_duration.contiguous(), times.contiguous(), right=True + ) + segment = segment.clamp_max(segment_deltas.shape[1] - 1) + previous_duration = torch.cat( + [torch.zeros_like(cumulative_duration[:, :1]), cumulative_duration[:, :-1]], + dim=-1, + ) + local_time = times - torch.gather(previous_duration, 1, segment) + batch_ids = torch.arange(times.shape[0], device=times.device)[:, None].expand_as( + segment + ) + selected_durations = profile.durations[batch_ids, segment] + cumulative_phase_time = selected_durations.cumsum(dim=-1) + phase = torch.searchsorted( + cumulative_phase_time.contiguous(), + local_time.unsqueeze(-1).contiguous(), + right=True, + ).squeeze(-1) + phase = phase.clamp_max(profile.durations.shape[-1] - 1) + phase_start_time = torch.cat( + [ + torch.zeros_like(cumulative_phase_time[..., :1]), + cumulative_phase_time[..., :-1], + ], + dim=-1, + ) + gather = phase[..., None] + tau = ( + local_time - torch.gather(phase_start_time, -1, gather).squeeze(-1) + ).clamp_min(0.0) + p0 = profile.positions[batch_ids, segment, phase] + v0 = profile.velocities[batch_ids, segment, phase] + a0 = profile.accelerations[batch_ids, segment, phase] + jerk = profile.jerks[batch_ids, segment, phase] + path_position = p0 + v0 * tau + 0.5 * a0 * tau.square() + jerk * tau.pow(3) / 6.0 + path_velocity = v0 + a0 * tau + 0.5 * jerk * tau.square() + path_acceleration = a0 + jerk * tau + selected_delta = segment_deltas[batch_ids, segment] + selected_start = segment_starts[batch_ids, segment] + return ( + selected_start + selected_delta * path_position[..., None], + selected_delta * path_velocity[..., None], + selected_delta * path_acceleration[..., None], + ) + + +def _compose_profile_samples( + *, + times: torch.Tensor, + cumulative_duration: torch.Tensor, + profile: _ProfileBatch, + segment_starts: torch.Tensor, + segment_deltas: torch.Tensor, + backend: Literal["auto", "torch", "warp"], +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Dispatch sample composition to Torch or the optional Warp kernel.""" + use_warp = _use_warp_backend(times, backend) + if use_warp: + if times.dtype != torch.float32: + if backend == "warp": + raise ValueError("The Warp trajectory backend requires float32 input.") + else: + try: + from embodichain.utils.warp.kinematics.trapezoidal_warp import ( + compose_profile_samples_warp, + ) + except ImportError: + if backend == "warp": + raise ImportError( + "The Warp trajectory backend requires the 'warp' package." + ) from None + else: + return compose_profile_samples_warp( + times=times, + cumulative_segment_time=cumulative_duration, + phase_durations=profile.durations, + phase_positions=profile.positions, + phase_velocities=profile.velocities, + phase_accelerations=profile.accelerations, + phase_jerks=profile.jerks, + segment_starts=segment_starts, + segment_deltas=segment_deltas, + ) + return _compose_profile_samples_torch( + times=times, + cumulative_duration=cumulative_duration, + profile=profile, + segment_starts=segment_starts, + segment_deltas=segment_deltas, + ) + + +def _sample_times_from_intervals( + segment_duration: torch.Tensor, + intervals: torch.Tensor, + output_count: int, +) -> torch.Tensor: + """Compose sorted sample times whose segment endpoints are exact.""" + cumulative_intervals = intervals.cumsum(dim=1) + cumulative_duration = segment_duration.cumsum(dim=1) + sample_ids = torch.arange(output_count, device=segment_duration.device)[None] + last_sample = cumulative_intervals[:, -1:] + clamped_ids = sample_ids.expand(segment_duration.shape[0], -1).clamp_max( + last_sample + ) + segment = torch.searchsorted( + cumulative_intervals.contiguous(), clamped_ids.contiguous(), right=False + ).clamp_max(segment_duration.shape[1] - 1) + previous_intervals = torch.cat( + (torch.zeros_like(cumulative_intervals[:, :1]), cumulative_intervals[:, :-1]), + dim=1, + ) + previous_duration = torch.cat( + (torch.zeros_like(cumulative_duration[:, :1]), cumulative_duration[:, :-1]), + dim=1, + ) + segment_intervals = intervals.gather(1, segment).clamp_min(1) + local_interval = clamped_ids - previous_intervals.gather(1, segment) + times = previous_duration.gather(1, segment) + segment_duration.gather( + 1, segment + ) * (local_interval / segment_intervals) + times[:, 0] = 0.0 + return torch.where(last_sample > 0, times, torch.zeros_like(times)) + + +def _make_sample_times( + segment_duration: torch.Tensor, options: TrapezoidalPlanOptions +) -> torch.Tensor: + """Create samples per segment so every retained boundary is represented.""" + active = segment_duration > 1e-12 + if options.sample_method is TrajectorySampleMethod.QUANTITY: + count = int(options.sample_interval) + required = active.sum(dim=1) + 1 + if count < int(required.max().item()): + raise ValueError( + "Quantity sampling requires at least one sample per retained waypoint; " + f"received {count}, requires at least {int(required.max().item())}." + ) + base = active.to(torch.long) + extra = count - 1 - base.sum(dim=1) + total_duration = segment_duration.sum(dim=1, keepdim=True).clamp_min(1e-12) + raw_extra = segment_duration / total_duration * extra[:, None] + allocated_extra = torch.floor(raw_extra).to(torch.long) + remainder_count = extra - allocated_extra.sum(dim=1) + fractional = torch.where( + active, + raw_extra - allocated_extra, + torch.full_like(raw_extra, -1.0), + ) + order = fractional.argsort(dim=1, descending=True) + rank = torch.empty_like(order) + rank.scatter_( + 1, + order, + torch.arange(order.shape[1], device=order.device)[None].expand_as(order), + ) + allocated_extra += (rank < remainder_count[:, None]).to(torch.long) + intervals = base + allocated_extra + return _sample_times_from_intervals(segment_duration, intervals, count) + + intervals = torch.where( + active, + torch.ceil(segment_duration / float(options.sample_interval)).to(torch.long), + torch.zeros_like(segment_duration, dtype=torch.long), + ) + counts = torch.maximum( + intervals.sum(dim=1) + 1, + torch.full( + (segment_duration.shape[0],), + 2, + dtype=torch.long, + device=segment_duration.device, + ), + ) + return _sample_times_from_intervals( + segment_duration, intervals, int(counts.max().item()) + ) + + +def _plan_linear_profiles( + waypoints: torch.Tensor, + options: TrapezoidalPlanOptions, +) -> PlanResult: + """Plan batched piecewise-linear joint paths without simulation state.""" + if waypoints.ndim != 3 or waypoints.shape[1] < 2 or waypoints.shape[2] < 1: + raise ValueError("waypoints must have shape (B, K, DOF) with K >= 2.") + if not waypoints.is_floating_point() or not bool( + torch.isfinite(waypoints).all().item() + ): + raise ValueError("waypoints must be a finite floating-point tensor.") + if not options.stop_at_waypoints: + waypoints = _compress_collinear_waypoints( + waypoints, options.collinearity_tolerance + ) + delta = waypoints[:, 1:] - waypoints[:, :-1] + stationary_path = ~(delta.abs() > 1e-8).any(dim=(-2, -1)) + batch_size, _, dof = waypoints.shape + if bool(stationary_path.all().item()): + hold_duration = options.minimum_duration or 0.0 + duration = waypoints.new_full((batch_size,), hold_duration) + times = _make_sample_times(duration[:, None], options) + positions = waypoints[:, :1].expand(-1, times.shape[1], dof).clone() + velocities = torch.zeros_like(positions) + accelerations = torch.zeros_like(positions) + return PlanResult( + success=torch.ones(batch_size, dtype=torch.bool, device=waypoints.device), + positions=positions, + velocities=velocities, + accelerations=accelerations, + dt=torch.diff(times, dim=1, prepend=torch.zeros_like(times[:, :1])), + ) + + velocity_limit, acceleration_limit, jerk_limit = _scalar_path_limits( + delta, options.constraints + ) + profile = _build_scalar_profile( + profile_name=options.profile, + velocity_limit=velocity_limit, + acceleration_limit=acceleration_limit, + jerk_limit=jerk_limit, + backend=options.backend, + ) + profile = _apply_minimum_duration(profile, options.minimum_duration) + segment_duration = profile.durations.sum(dim=-1) + cumulative_duration = segment_duration.cumsum(dim=-1) + duration = cumulative_duration[:, -1] + times = _make_sample_times(segment_duration, options) + positions, velocities, accelerations = _compose_profile_samples( + times=times, + cumulative_duration=cumulative_duration, + profile=profile, + segment_starts=waypoints[:, :-1], + segment_deltas=delta, + backend=options.backend, + ) + positions[:, 0] = waypoints[:, 0] + positions[:, -1] = waypoints[:, -1] + velocities[:, -1] = 0.0 + accelerations[:, -1] = 0.0 + dt = torch.diff(times, dim=1, prepend=torch.zeros_like(times[:, :1])) + success = torch.ones(batch_size, dtype=torch.bool, device=waypoints.device) + if bool(stationary_path.any().item()): + positions[stationary_path] = waypoints[stationary_path, :1] + velocities[stationary_path] = 0.0 + accelerations[stationary_path] = 0.0 + if options.minimum_duration is None: + dt[stationary_path] = 0.0 + return PlanResult( + success=success, + positions=positions, + velocities=velocities, + accelerations=accelerations, + dt=dt, + ) + + +class TrapezoidalPlanner(BasePlanner): + """Plan batched linear joint paths with trapezoidal or Double-S timing.""" + + supported_move_types = frozenset({MoveType.JOINT_MOVE}) + + def default_plan_options(self) -> TrapezoidalPlanOptions: + """Return backend-default planning options.""" + return TrapezoidalPlanOptions() + + @validate_plan_options(options_cls=TrapezoidalPlanOptions) + def plan( + self, + target_states: list[PlanState], + options: TrapezoidalPlanOptions = TrapezoidalPlanOptions(), # noqa: B008 + ) -> PlanResult: + """Generate one batched joint trajectory through all target states. + + Args: + target_states: Joint waypoints with ``qpos`` shape ``(B, DOF)``. + options: Time-profile, derivative limits, and sampling configuration. + + Returns: + Batched positions, velocities, accelerations, and explicit timing. + + Raises: + ValueError: If fewer than two valid joint waypoints are supplied. + """ + if len(target_states) < 2 or any(state.qpos is None for state in target_states): + raise ValueError("TrapezoidalPlanner requires at least two qpos waypoints.") + waypoints = torch.stack([state.qpos for state in target_states], dim=1).to( + self.device + ) + return _plan_linear_profiles(waypoints, options) diff --git a/embodichain/lab/sim/solvers/opw_solver.py b/embodichain/lab/sim/solvers/opw_solver.py index e3202597e..8fefbae02 100644 --- a/embodichain/lab/sim/solvers/opw_solver.py +++ b/embodichain/lab/sim/solvers/opw_solver.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import torch import numpy as np import warp as wp @@ -29,6 +31,7 @@ OPWparam, opw_fk_kernel, opw_ik_kernel, + opw_ik_path_select_kernel, opw_ik_select_kernel, wp_vec6f, ) @@ -37,6 +40,8 @@ if TYPE_CHECKING: from typing import Self +__all__ = ["OPWSolver", "OPWSolverCfg"] + def normalize_to_pi(angle): angle = (angle + np.pi) % (2.0 * np.pi) - np.pi @@ -367,6 +372,55 @@ def get_ik_warp( best_ik_valid = wp.to_torch(best_ik_valid_wp).to(self.device) return best_ik_valid, best_ik_result + def _select_continuous_ik_path( + self, + candidate_qpos: torch.Tensor, + candidate_valid: torch.Tensor, + qpos_seed: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Select a continuous branch from precomputed OPW candidates.""" + if candidate_qpos.ndim != 4 or candidate_qpos.shape[-2:] != (8, 6): + raise ValueError("candidate_qpos must have shape (B, N, 8, 6).") + batch_size, sample_count = candidate_qpos.shape[:2] + if candidate_valid.shape != (batch_size, sample_count, 8): + raise ValueError( + f"candidate_valid must have shape ({batch_size}, {sample_count}, 8)." + ) + if qpos_seed.shape != (batch_size, 6): + raise ValueError(f"qpos_seed must have shape ({batch_size}, 6).") + kernel_device = standardize_device_string(self.device) + lower_limits = self.lower_qpos_limits.detach().cpu().tolist() + upper_limits = self.upper_qpos_limits.detach().cpu().tolist() + path_qpos = wp.zeros( + (batch_size, sample_count, 6), dtype=float, device=kernel_device + ) + path_valid = wp.zeros( + (batch_size, sample_count), dtype=int, device=kernel_device + ) + joint_weight = self.ik_nearest_weight.detach().cpu().tolist() + wp.launch( + kernel=opw_ik_path_select_kernel, + dim=batch_size, + inputs=[ + wp.from_torch(candidate_qpos.to(kernel_device).contiguous()), + wp.from_torch( + candidate_valid.to( + device=kernel_device, dtype=torch.int32 + ).contiguous() + ), + wp.from_torch( + qpos_seed.to(device=kernel_device, dtype=torch.float32).contiguous() + ), + wp_vec6f(*joint_weight), + wp_vec6f(*lower_limits), + wp_vec6f(*upper_limits), + self.cfg.safe_margin, + ], + outputs=[path_qpos, path_valid], + device=kernel_device, + ) + return wp.to_torch(path_valid).bool(), wp.to_torch(path_qpos) + def _calculate_dynamic_weights( self, current_joints, joint_limits, base_weights=None ) -> np.ndarray: diff --git a/embodichain/utils/warp/kinematics/opw_solver.py b/embodichain/utils/warp/kinematics/opw_solver.py index 877324d17..5bac08f70 100644 --- a/embodichain/utils/warp/kinematics/opw_solver.py +++ b/embodichain/utils/warp/kinematics/opw_solver.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import warp as wp import numpy as np from typing import Tuple @@ -46,6 +48,25 @@ def is_within_limit( return True +@wp.func +def nearest_equivalent_in_limit( + solution: float, + seed: float, + lower: float, + upper: float, + safe_margin: float, +) -> Tuple[float, bool]: + """Return the seed-nearest periodic equivalent inside the safe limits.""" + two_pi = 2.0 * wp.pi + lower_turn = wp.ceil((lower + safe_margin - solution) / two_pi) + upper_turn = wp.floor((upper - safe_margin - solution) / two_pi) + if lower_turn > upper_turn: + return solution, False + nearest_turn = wp.round((seed - solution) / two_pi) + selected_turn = wp.min(wp.max(nearest_turn, lower_turn), upper_turn) + return solution + selected_turn * two_pi, True + + @wp.func def safe_acos(x: float) -> float: return wp.acos(wp.clamp(x, -1.0, 1.0)) @@ -522,3 +543,71 @@ def opw_ik_select_kernel( else: # no valid solution best_ik_valid[i] = 0 + + +@wp.kernel +def opw_ik_path_select_kernel( + full_ik_result: wp.array(dtype=float, ndim=4), # [B, N, N_SOL, DOF] + full_ik_valid: wp.array(dtype=int, ndim=3), # [B, N, N_SOL] + initial_seed: wp.array(dtype=float, ndim=2), # [B, DOF] + joint_weights: wp_vec6f, + lower_limits: wp_vec6f, + upper_limits: wp_vec6f, + safe_margin: float, + path_result: wp.array(dtype=float, ndim=3), # [B, N, DOF] + path_valid: wp.array(dtype=int, ndim=2), # [B, N] +): + """Select a temporally continuous OPW branch for one environment.""" + batch = wp.tid() + sample_count = full_ik_result.shape[1] + for sample in range(sample_count): + best_distance = float(1.0e10) + best_solution = int(-1) + for candidate in range(8): + if full_ik_valid[batch, sample, candidate] == 0: + continue + distance = float(0.0) + is_continuous = bool(True) + for joint in range(6): + seed = initial_seed[batch, joint] + if sample > 0: + seed = path_result[batch, sample - 1, joint] + solution = full_ik_result[batch, sample, candidate, joint] + solution, equivalent_valid = nearest_equivalent_in_limit( + solution, + seed, + lower_limits[joint], + upper_limits[joint], + safe_margin, + ) + if not equivalent_valid: + is_continuous = False + error = (solution - seed) * joint_weights[joint] + distance += error * error + if is_continuous and distance < best_distance: + best_distance = distance + best_solution = candidate + if best_solution >= 0: + path_valid[batch, sample] = 1 + for joint in range(6): + seed = initial_seed[batch, joint] + if sample > 0: + seed = path_result[batch, sample - 1, joint] + solution = full_ik_result[batch, sample, best_solution, joint] + nearest, _ = nearest_equivalent_in_limit( + solution, + seed, + lower_limits[joint], + upper_limits[joint], + safe_margin, + ) + path_result[batch, sample, joint] = nearest + else: + path_valid[batch, sample] = 0 + for joint in range(6): + if sample == 0: + path_result[batch, sample, joint] = initial_seed[batch, joint] + else: + path_result[batch, sample, joint] = path_result[ + batch, sample - 1, joint + ] diff --git a/embodichain/utils/warp/kinematics/trapezoidal_warp.py b/embodichain/utils/warp/kinematics/trapezoidal_warp.py new file mode 100644 index 000000000..73e270933 --- /dev/null +++ b/embodichain/utils/warp/kinematics/trapezoidal_warp.py @@ -0,0 +1,399 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Warp profile construction and sampling for the trapezoidal planner.""" + +from __future__ import annotations + +import torch +import warp as wp + +__all__ = ["build_profile_warp", "compose_profile_samples_warp"] + + +@wp.kernel(enable_backward=False) +def _build_trapezoidal_profile_kernel( + velocity_limits: wp.array(dtype=wp.float32), + acceleration_limits: wp.array(dtype=wp.float32), + durations: wp.array(dtype=wp.float32), + positions: wp.array(dtype=wp.float32), + velocities: wp.array(dtype=wp.float32), + accelerations: wp.array(dtype=wp.float32), + jerks: wp.array(dtype=wp.float32), +) -> None: + profile = wp.tid() + velocity_limit = velocity_limits[profile] + requested_acceleration = acceleration_limits[profile] + acceleration_limit = wp.max(requested_acceleration, 1.0e-12) + stationary = velocity_limit <= 0.0 + peak_velocity = velocity_limit + reaches_velocity = velocity_limit * velocity_limit / acceleration_limit < 1.0 + if not reaches_velocity: + peak_velocity = wp.sqrt(acceleration_limit) + acceleration_time = peak_velocity / acceleration_limit + cruise_time = float(0.0) + if reaches_velocity and not stationary: + cruise_time = ( + 1.0 - peak_velocity * peak_velocity / acceleration_limit + ) / peak_velocity + if stationary: + acceleration_time = 0.0 + cruise_time = 0.0 + + offset = profile * 3 + durations[offset] = acceleration_time + durations[offset + 1] = cruise_time + durations[offset + 2] = acceleration_time + positions[offset] = 0.0 + positions[offset + 1] = 0.5 * peak_velocity * acceleration_time + positions[offset + 2] = positions[offset + 1] + peak_velocity * cruise_time + velocities[offset] = 0.0 + velocities[offset + 1] = peak_velocity + velocities[offset + 2] = peak_velocity + accelerations[offset] = requested_acceleration + accelerations[offset + 1] = 0.0 + accelerations[offset + 2] = -requested_acceleration + jerks[offset] = 0.0 + jerks[offset + 1] = 0.0 + jerks[offset + 2] = 0.0 + + +@wp.kernel(enable_backward=False) +def _build_double_s_profile_kernel( + velocity_limits: wp.array(dtype=wp.float32), + acceleration_limits: wp.array(dtype=wp.float32), + jerk_limits: wp.array(dtype=wp.float32), + durations: wp.array(dtype=wp.float32), + positions: wp.array(dtype=wp.float32), + velocities: wp.array(dtype=wp.float32), + accelerations: wp.array(dtype=wp.float32), + jerks: wp.array(dtype=wp.float32), +) -> None: + profile = wp.tid() + velocity_limit = velocity_limits[profile] + acceleration_limit = wp.max(acceleration_limits[profile], 1.0e-12) + jerk_limit = wp.max(jerk_limits[profile], 1.0e-12) + stationary = velocity_limit <= 0.0 + safe_velocity = wp.max(velocity_limit, 1.0e-12) + + tj = float(acceleration_limit / jerk_limit) + ta = float(tj + safe_velocity / acceleration_limit) + if safe_velocity * jerk_limit < acceleration_limit * acceleration_limit: + tj = wp.sqrt(safe_velocity / jerk_limit) + ta = 2.0 * tj + cruise_time = float(1.0 / safe_velocity - ta) + if not stationary and cruise_time <= 0.0: + candidate_acceleration = float(acceleration_limit) + accepted = bool(False) + reduction = int(0) + while not accepted and reduction <= 1000: + candidate_tj = candidate_acceleration / jerk_limit + acceleration_squared = candidate_acceleration * candidate_acceleration + delta = wp.sqrt( + acceleration_squared * acceleration_squared / (jerk_limit * jerk_limit) + + 4.0 * candidate_acceleration + ) + candidate_ta = (acceleration_squared / jerk_limit + delta) / ( + 2.0 * candidate_acceleration + ) + if candidate_ta >= 2.0 * candidate_tj: + tj = candidate_tj + ta = candidate_ta + accepted = True + else: + candidate_acceleration *= 0.9 + reduction += 1 + cruise_time = 0.0 + if stationary: + tj = 0.0 + ta = 0.0 + cruise_time = 0.0 + + constant_acceleration_time = wp.max(ta - 2.0 * tj, 0.0) + offset = profile * 7 + for phase in range(7): + duration = tj + jerk = jerk_limit + if phase == 1 or phase == 5: + duration = constant_acceleration_time + jerk = 0.0 + elif phase == 3: + duration = cruise_time + jerk = 0.0 + elif phase == 2 or phase == 4: + jerk = -jerk_limit + if stationary: + jerk = 0.0 + durations[offset + phase] = duration + jerks[offset + phase] = jerk + + position = float(0.0) + velocity = float(0.0) + acceleration = float(0.0) + for phase in range(7): + index = offset + phase + positions[index] = position + velocities[index] = velocity + accelerations[index] = acceleration + duration = durations[index] + jerk = jerks[index] + position += ( + velocity * duration + + 0.5 * acceleration * duration * duration + + jerk * duration * duration * duration / 6.0 + ) + velocity += acceleration * duration + 0.5 * jerk * duration * duration + acceleration += jerk * duration + + +@wp.kernel(enable_backward=False) +def _evaluate_profile_samples_kernel( + times: wp.array(dtype=wp.float32), + cumulative_segment_time: wp.array(dtype=wp.float32), + phase_durations: wp.array(dtype=wp.float32), + phase_positions: wp.array(dtype=wp.float32), + phase_velocities: wp.array(dtype=wp.float32), + phase_accelerations: wp.array(dtype=wp.float32), + phase_jerks: wp.array(dtype=wp.float32), + sample_count: int, + segment_count: int, + phase_count: int, + sample_segments: wp.array(dtype=wp.int32), + path_positions: wp.array(dtype=wp.float32), + path_velocities: wp.array(dtype=wp.float32), + path_accelerations: wp.array(dtype=wp.float32), +) -> None: + batch, sample = wp.tid() + sample_offset = batch * sample_count + sample + time = times[sample_offset] + low = int(0) # noqa: RUF046, UP018 - mutable Warp binary-search bound. + high = int(segment_count) + while low < high: + middle = (low + high) // 2 + end_time = cumulative_segment_time[batch * segment_count + middle] + if time >= end_time: + low = middle + 1 + else: + high = middle + segment = low + if segment >= segment_count: + segment = segment_count - 1 + + previous_segment_time = float(0.0) # noqa: UP018 - mutable Warp value. + if segment > 0: + previous_segment_time = cumulative_segment_time[ + batch * segment_count + segment - 1 + ] + local_time = time - previous_segment_time + profile_offset = (batch * segment_count + segment) * phase_count + phase = int(0) # noqa: RUF046, UP018 - Warp requires a mutable typed value. + phase_start_time = float(0.0) # noqa: UP018 - mutable Warp value. + cumulative_phase_time = float(0.0) # noqa: UP018 - mutable Warp value. + for candidate in range(phase_count): + duration = phase_durations[profile_offset + candidate] + cumulative_phase_time += duration + if local_time >= cumulative_phase_time: + phase = candidate + 1 + phase_start_time = cumulative_phase_time + if phase >= phase_count: + phase = phase_count - 1 + phase_start_time = ( + cumulative_phase_time - phase_durations[profile_offset + phase] + ) + + profile_index = profile_offset + phase + tau = wp.max(local_time - phase_start_time, 0.0) + p0 = phase_positions[profile_index] + v0 = phase_velocities[profile_index] + a0 = phase_accelerations[profile_index] + jerk = phase_jerks[profile_index] + path_position = p0 + v0 * tau + 0.5 * a0 * tau * tau + jerk * tau * tau * tau / 6.0 + path_velocity = v0 + a0 * tau + 0.5 * jerk * tau * tau + path_acceleration = a0 + jerk * tau + + sample_segments[sample_offset] = segment + path_positions[sample_offset] = path_position + path_velocities[sample_offset] = path_velocity + path_accelerations[sample_offset] = path_acceleration + + +@wp.kernel(enable_backward=False) +def _compose_joint_samples_kernel( + sample_segments: wp.array(dtype=wp.int32), + path_positions: wp.array(dtype=wp.float32), + path_velocities: wp.array(dtype=wp.float32), + path_accelerations: wp.array(dtype=wp.float32), + segment_starts: wp.array(dtype=wp.float32), + segment_deltas: wp.array(dtype=wp.float32), + sample_count: int, + segment_count: int, + dof: int, + positions: wp.array(dtype=wp.float32), + velocities: wp.array(dtype=wp.float32), + accelerations: wp.array(dtype=wp.float32), +) -> None: + batch, sample, joint = wp.tid() + sample_offset = batch * sample_count + sample + segment = sample_segments[sample_offset] + joint_offset = (batch * segment_count + segment) * dof + joint + output_offset = sample_offset * dof + joint + delta = segment_deltas[joint_offset] + positions[output_offset] = ( + segment_starts[joint_offset] + delta * path_positions[sample_offset] + ) + velocities[output_offset] = delta * path_velocities[sample_offset] + accelerations[output_offset] = delta * path_accelerations[sample_offset] + + +def build_profile_warp( + *, + profile: str, + velocity_limits: torch.Tensor, + acceleration_limits: torch.Tensor, + jerk_limits: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Construct all scalar segment profiles in one Warp launch. + + Args: + profile: ``"trapezoidal"`` or ``"double_s"``. + velocity_limits: Projected path velocity limits shaped ``(B, S)``. + acceleration_limits: Projected acceleration limits shaped ``(B, S)``. + jerk_limits: Projected jerk limits shaped ``(B, S)``. + + Returns: + Phase durations, positions, velocities, accelerations, and jerks, each + shaped ``(B, S, P)`` where ``P`` is three or seven. + + Raises: + ValueError: If the profile, shapes, or dtypes are unsupported. + """ + if profile not in {"trapezoidal", "double_s"}: + raise ValueError(f"Unsupported Warp trajectory profile: {profile!r}.") + limits = (velocity_limits, acceleration_limits, jerk_limits) + if any(limit.dtype != torch.float32 for limit in limits): + raise ValueError("The Warp trajectory backend requires float32 tensors.") + if any(limit.device != velocity_limits.device for limit in limits[1:]): + raise ValueError("Warp trajectory limit tensors must share one device.") + if any(limit.shape != velocity_limits.shape for limit in limits[1:]): + raise ValueError("Warp trajectory limit tensors must have matching shapes.") + if velocity_limits.ndim != 2: + raise ValueError("Warp trajectory limits must have shape (B, S).") + + wp.init() + phase_count = 3 if profile == "trapezoidal" else 7 + output_shape = (*velocity_limits.shape, phase_count) + outputs = [ + torch.empty(output_shape, dtype=torch.float32, device=velocity_limits.device) + for _ in range(5) + ] + warp_limits = [wp.from_torch(limit.contiguous().flatten()) for limit in limits] + warp_outputs = [wp.from_torch(output.flatten()) for output in outputs] + kernel = ( + _build_trapezoidal_profile_kernel + if profile == "trapezoidal" + else _build_double_s_profile_kernel + ) + kernel_inputs = warp_limits[:2] if profile == "trapezoidal" else warp_limits + wp.launch( + kernel=kernel, + dim=velocity_limits.numel(), + inputs=kernel_inputs, + outputs=warp_outputs, + device=str(velocity_limits.device), + ) + return outputs[0], outputs[1], outputs[2], outputs[3], outputs[4] + + +def compose_profile_samples_warp( + *, + times: torch.Tensor, + cumulative_segment_time: torch.Tensor, + phase_durations: torch.Tensor, + phase_positions: torch.Tensor, + phase_velocities: torch.Tensor, + phase_accelerations: torch.Tensor, + phase_jerks: torch.Tensor, + segment_starts: torch.Tensor, + segment_deltas: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Evaluate scalar profiles once per sample, then compose all joints.""" + wp.init() + tensors = ( + times, + cumulative_segment_time, + phase_durations, + phase_positions, + phase_velocities, + phase_accelerations, + phase_jerks, + segment_starts, + segment_deltas, + ) + if any(tensor.dtype != torch.float32 for tensor in tensors): + raise ValueError("The Warp trajectory backend requires float32 tensors.") + batch_size, sample_count = times.shape + segment_count = segment_deltas.shape[1] + dof = segment_deltas.shape[2] + phase_count = phase_durations.shape[2] + sample_segments = torch.empty( + (batch_size, sample_count), dtype=torch.int32, device=times.device + ) + path_positions = torch.empty_like(times) + path_velocities = torch.empty_like(times) + path_accelerations = torch.empty_like(times) + positions = torch.empty( + (batch_size, sample_count, dof), dtype=torch.float32, device=times.device + ) + velocities = torch.empty_like(positions) + accelerations = torch.empty_like(positions) + inputs = [wp.from_torch(tensor.contiguous().flatten()) for tensor in tensors] + scalar_outputs = [ + wp.from_torch(sample_segments.flatten()), + wp.from_torch(path_positions.flatten()), + wp.from_torch(path_velocities.flatten()), + wp.from_torch(path_accelerations.flatten()), + ] + joint_outputs = [ + wp.from_torch(positions.flatten()), + wp.from_torch(velocities.flatten()), + wp.from_torch(accelerations.flatten()), + ] + wp.launch( + kernel=_evaluate_profile_samples_kernel, + dim=(batch_size, sample_count), + inputs=[ + *inputs[:7], + sample_count, + segment_count, + phase_count, + ], + outputs=scalar_outputs, + device=str(times.device), + ) + wp.launch( + kernel=_compose_joint_samples_kernel, + dim=(batch_size, sample_count, dof), + inputs=[ + *scalar_outputs, + *inputs[7:], + sample_count, + segment_count, + dof, + ], + outputs=joint_outputs, + device=str(times.device), + ) + return positions, velocities, accelerations diff --git a/scripts/benchmark/motion_generation/trapezoidal_planner.py b/scripts/benchmark/motion_generation/trapezoidal_planner.py new file mode 100644 index 000000000..ae9a04bc0 --- /dev/null +++ b/scripts/benchmark/motion_generation/trapezoidal_planner.py @@ -0,0 +1,256 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Benchmark batched trapezoidal-planner Torch and Warp backends. + +Run: python scripts/benchmark/motion_generation/trapezoidal_planner.py +""" + +from __future__ import annotations + +import argparse +import os +import time +from datetime import datetime, timezone +from pathlib import Path + +import psutil +import torch + +from embodichain.lab.sim.planners.trapezoidal_planner import ( + TrapezoidalPlanOptions, + _plan_linear_profiles, +) + + +def parse_args() -> argparse.Namespace: + """Parse benchmark sizes and device selection.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--device", default="cpu", help="Torch device to benchmark.") + parser.add_argument("--segments", type=int, default=32) + parser.add_argument("--samples", type=int, default=512) + parser.add_argument("--repeats", type=int, default=10) + parser.add_argument("--batch-sizes", type=int, nargs="+", default=[1, 16, 64]) + return parser.parse_args() + + +def memory_snapshot() -> dict[str, float]: + """Return current process RSS and Torch CUDA allocation in MiB.""" + cpu_mb = psutil.Process(os.getpid()).memory_info().rss / 1024**2 + gpu_mb = ( + torch.cuda.memory_allocated() / 1024**2 if torch.cuda.is_available() else 0.0 + ) + return {"cpu_mb": cpu_mb, "gpu_mb": gpu_mb} + + +def synchronize(device: torch.device) -> None: + """Synchronize CUDA timing when the selected device is asynchronous.""" + if device.type == "cuda": + torch.cuda.synchronize(device) + + +def make_waypoints( + batch_size: int, + segment_count: int, + device: torch.device, +) -> torch.Tensor: + """Create deterministic non-collinear seven-DOF benchmark paths.""" + parameter = torch.linspace( + 0.0, 4.0, segment_count + 1, dtype=torch.float32, device=device + ) + base = torch.stack( + [ + parameter, + torch.sin(parameter), + torch.cos(parameter), + 0.5 * torch.sin(2.0 * parameter), + 0.4 * torch.cos(1.5 * parameter), + 0.2 * parameter, + -0.1 * parameter, + ], + dim=-1, + ) + offsets = torch.arange(batch_size, dtype=torch.float32, device=device) + offsets = offsets[:, None, None] * 1e-3 + return base[None].expand(batch_size, -1, -1) + offsets + + +def benchmark_case( + *, + waypoints: torch.Tensor, + profile: str, + backend: str, + sample_count: int, + repeats: int, + reference: torch.Tensor | None, +) -> tuple[dict[str, object], dict[str, object], torch.Tensor]: + """Measure one backend/profile case and return report rows.""" + options = TrapezoidalPlanOptions( + profile=profile, + constraints={"velocity": 0.7, "acceleration": 1.4, "jerk": 4.0}, + sample_interval=sample_count, + backend=backend, + ) + result = _plan_linear_profiles(waypoints, options) + synchronize(waypoints.device) + if waypoints.is_cuda: + torch.cuda.reset_peak_memory_stats(waypoints.device) + before = memory_snapshot() + started = time.perf_counter() + for _ in range(repeats): + result = _plan_linear_profiles(waypoints, options) + synchronize(waypoints.device) + elapsed_ms = (time.perf_counter() - started) * 1000.0 / repeats + after = memory_snapshot() + peak_gpu_mb = ( + torch.cuda.max_memory_allocated(waypoints.device) / 1024**2 + if waypoints.is_cuda + else 0.0 + ) + max_error = ( + 0.0 + if reference is None + else float((result.positions - reference).abs().max().item()) + ) + success = bool( + result.is_all_success() + and torch.isfinite(result.positions).all() + and torch.allclose(result.positions[:, 0], waypoints[:, 0]) + and torch.allclose(result.positions[:, -1], waypoints[:, -1]) + ) + algorithm = f"{backend}-{profile}" + perf = { + "batch_size": waypoints.shape[0], + "algorithm": algorithm, + "cost_time_ms": f"{elapsed_ms:.4f}", + "cpu_delta_mb": f"{after['cpu_mb'] - before['cpu_mb']:+.2f}", + "gpu_delta_mb": f"{after['gpu_mb'] - before['gpu_mb']:+.2f}", + "peak_gpu_mb": f"{peak_gpu_mb:.2f}", + } + metric = { + "batch_size": waypoints.shape[0], + "algorithm": algorithm, + "success_rate": "1.0" if success else "0.0", + "max_position_error": f"{max_error:.3e}", + "duration_mean_s": f"{result.duration.mean().item():.4f}", + } + return perf, metric, result.positions + + +def write_markdown_report( + perf_rows: list[dict[str, object]], + metric_rows: list[dict[str, object]], + leaderboard_rows: list[dict[str, object]], +) -> Path: + """Write exactly three Markdown result tables.""" + output_dir = Path("outputs/benchmarks") + output_dir.mkdir(parents=True, exist_ok=True) + timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S") + path = output_dir / f"trapezoidal_planner_{timestamp}.md" + lines = ["# Trapezoidal Planner Benchmark", ""] + for title, rows in ( + ("Time & Memory", perf_rows), + ("Success & Other Metrics", metric_rows), + ("Leaderboard", leaderboard_rows), + ): + lines.extend([f"## {title}", ""]) + headers = list(rows[0]) + lines.append("| " + " | ".join(headers) + " |") + lines.append("| " + " | ".join(["---"] * len(headers)) + " |") + for row in rows: + lines.append("| " + " | ".join(str(row[key]) for key in headers) + " |") + lines.append("") + path.write_text("\n".join(lines), encoding="utf-8") + return path + + +def run_all_benchmarks() -> None: + """Run requested cases and save one three-table Markdown report.""" + args = parse_args() + device = torch.device(args.device) + perf_rows: list[dict[str, object]] = [] + metric_rows: list[dict[str, object]] = [] + print("=" * 60) + print("Trapezoidal Planner Performance Benchmarks") + print("=" * 60) + for batch_size in args.batch_sizes: + waypoints = make_waypoints(batch_size, args.segments, device) + for profile in ("trapezoidal", "double_s"): + torch_perf, torch_metric, reference = benchmark_case( + waypoints=waypoints, + profile=profile, + backend="torch", + sample_count=args.samples, + repeats=args.repeats, + reference=None, + ) + perf_rows.append(torch_perf) + metric_rows.append(torch_metric) + print( + f"batch={batch_size:>4d} torch-{profile:<11s} " + f"{torch_perf['cost_time_ms']:>10s} ms" + ) + try: + warp_perf, warp_metric, _ = benchmark_case( + waypoints=waypoints, + profile=profile, + backend="warp", + sample_count=args.samples, + repeats=args.repeats, + reference=reference, + ) + except (ImportError, OSError, RuntimeError, ValueError) as error: + print(f" warp-{profile} skipped: {error}") + else: + perf_rows.append(warp_perf) + metric_rows.append(warp_metric) + print( + f"batch={batch_size:>4d} warp-{profile:<12s} " + f"{warp_perf['cost_time_ms']:>10s} ms" + ) + algorithms = sorted({str(row["algorithm"]) for row in metric_rows}) + leaderboard_rows = [] + for algorithm in algorithms: + selected_metrics = [row for row in metric_rows if row["algorithm"] == algorithm] + selected_perf = [row for row in perf_rows if row["algorithm"] == algorithm] + success_rate = sum( + float(row["success_rate"]) for row in selected_metrics + ) / len(selected_metrics) + mean_ms = sum(float(row["cost_time_ms"]) for row in selected_perf) / len( + selected_perf + ) + leaderboard_rows.append( + { + "rank": 0, + "algorithm": algorithm, + "overall_success_rate": f"{success_rate:.3f}", + "mean_cost_time_ms": f"{mean_ms:.4f}", + } + ) + leaderboard_rows.sort( + key=lambda row: ( + -float(row["overall_success_rate"]), + float(row["mean_cost_time_ms"]), + ) + ) + for rank, row in enumerate(leaderboard_rows, start=1): + row["rank"] = rank + report = write_markdown_report(perf_rows, metric_rows, leaderboard_rows) + print(f"Markdown report saved: {report}") + + +if __name__ == "__main__": + run_all_benchmarks() diff --git a/scripts/tutorials/sim/planner/trapezoidal_planner.py b/scripts/tutorials/sim/planner/trapezoidal_planner.py new file mode 100644 index 000000000..82ead30bb --- /dev/null +++ b/scripts/tutorials/sim/planner/trapezoidal_planner.py @@ -0,0 +1,859 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Compare velocity- and acceleration-trapezoidal joint trajectories.""" + +from __future__ import annotations + +import argparse +import math +import os +import time +from pathlib import Path + +os.environ.setdefault("MPLCONFIGDIR", "/tmp/embodichain-matplotlib") + +import matplotlib.pyplot as plt +import torch +from matplotlib.figure import Figure +from mpl_toolkits.mplot3d.axes3d import Axes3D + +from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.objects import Robot +from embodichain.lab.sim.planners import ( + MotionGenCfg, + MotionGenerator, + MotionGenOptions, + PlanResult, + PlanState, + TrapezoidalPlannerCfg, + TrapezoidalPlanOptions, +) +from embodichain.lab.sim.planners.trapezoidal_planner import _plan_linear_profiles +from embodichain.lab.sim.robots import CobotMagicCfg +from embodichain.lab.visualization import visualization_cfg_from_args +from embodichain.utils.math import euler_xyz_from_quat + +DEFAULT_SAMPLES = 200 +DEFAULT_CARTESIAN_DISTANCE = 0.10 +DEFAULT_CARTESIAN_STEP = 0.01 +DEFAULT_REPLAY_SPEED = 1.0 +DEFAULT_CARTESIAN_VELOCITY = 0.15 +DEFAULT_CARTESIAN_ACCELERATION = 0.30 +DEFAULT_CARTESIAN_JERK = 1.0 +PROFILE_SPECS = { + "velocity_trapezoidal": ("trapezoidal", "velocity_trapezoidal"), + "acceleration_trapezoidal": ("double_s", "acceleration_trapezoidal"), +} + + +def configure_plot_fonts() -> None: + """Apply lightweight readable font defaults for this tutorial.""" + plt.rcParams.update( + { + "font.sans-serif": ["Noto Sans CJK SC", "DejaVu Sans", "sans-serif"], + "font.size": 11.0, + "axes.titlesize": 13.0, + "legend.fontsize": 9.0, + "axes.unicode_minus": False, + } + ) + + +def positive_float(value: str) -> float: + """Parse a finite positive command-line float.""" + parsed = float(value) + if not torch.isfinite(torch.tensor(parsed)) or parsed <= 0.0: + raise argparse.ArgumentTypeError("value must be finite and greater than zero") + return parsed + + +def sample_count(value: str) -> int: + """Parse a trajectory sample count accepted by the planner.""" + parsed = int(value) + if parsed < 2: + raise argparse.ArgumentTypeError("sample count must be at least 2") + return parsed + + +def parse_args() -> argparse.Namespace: + """Parse tutorial arguments.""" + # The shared launcher owns a boolean ``--profile`` flag for environment + # timing. This focused tutorial intentionally reuses that concise name for + # its trajectory profile and therefore replaces the shared action. + parser = argparse.ArgumentParser(description=__doc__, conflict_handler="resolve") + add_env_launcher_args_to_parser(parser) + parser.add_argument( + "--profile", + choices=(*PROFILE_SPECS, "both"), + default="both", + help=( + "Diagnostic to run: trapezoidal velocity, jerk-limited " + "trapezoidal acceleration, or both." + ), + ) + parser.add_argument( + "--samples", + type=sample_count, + default=DEFAULT_SAMPLES, + help="Number of output trajectory samples.", + ) + parser.add_argument( + "--path", + choices=("joint", "cartesian", "both"), + default="both", + help="Plan a synchronized multi-joint path, a straight EEF path, or both.", + ) + parser.add_argument( + "--cartesian-distance", + type=positive_float, + default=DEFAULT_CARTESIAN_DISTANCE, + help="Length in metres of the diagonal straight EEF demo path.", + ) + parser.add_argument( + "--cartesian-step", + type=positive_float, + default=DEFAULT_CARTESIAN_STEP, + help="Cartesian interpolation spacing in metres before IK.", + ) + parser.add_argument( + "--cartesian-velocity", + type=positive_float, + default=DEFAULT_CARTESIAN_VELOCITY, + help="Maximum straight-line EEF speed in m/s.", + ) + parser.add_argument( + "--cartesian-acceleration", + type=positive_float, + default=DEFAULT_CARTESIAN_ACCELERATION, + help="Maximum straight-line EEF acceleration in m/s².", + ) + parser.add_argument( + "--cartesian-jerk", + type=positive_float, + default=DEFAULT_CARTESIAN_JERK, + help="Maximum straight-line EEF jerk in m/s³.", + ) + parser.add_argument( + "--backend", + choices=("auto", "torch", "warp"), + default="auto", + help="Backend used to compose sampled joint states.", + ) + parser.add_argument( + "--replay-speed", + type=positive_float, + default=DEFAULT_REPLAY_SPEED, + help="Trajectory playback speed multiplier in the simulation window.", + ) + parser.add_argument( + "--plot-env", + type=int, + default=0, + help="Batch environment index shown in the diagnostic plot.", + ) + parser.add_argument( + "--plot-output", + type=Path, + default=None, + help="Optional PNG path. By default the figure is not saved.", + ) + parser.add_argument( + "--show-plot", + action=argparse.BooleanOptionalAction, + default=True, + help="Display the plot interactively (default: enabled).", + ) + return parser.parse_args() + + +def build_demo_waypoints(robot: Robot, control_part: str) -> list[torch.Tensor]: + """Create a batched path where every arm joint moves synchronously.""" + start = robot.get_qpos(name=control_part).clone() + dof = start.shape[1] + if dof == 6: + # Same non-singular seed used by the TOPPRA motion-generator tutorial. + safe_seed = start.new_tensor( + (0.0, torch.pi / 4.0, -torch.pi / 4.0, 0.0, torch.pi / 4.0, 0.0) + ) + start = safe_seed.unsqueeze(0).expand_as(start).clone() + joint_index = torch.arange(dof, dtype=start.dtype, device=start.device) + direction = torch.where(joint_index.remainder(2) == 0, 1.0, -1.0) + magnitude = torch.linspace(0.12, 0.30, dof, dtype=start.dtype, device=start.device) + middle = start + direction * magnitude + goal = start - direction * magnitude.flip(0) * 0.75 + joint_ids = robot.get_joint_ids(name=control_part) + limits = robot.get_qpos_limits(joint_ids=joint_ids).to(start) + margin = torch.minimum( + torch.full_like(limits[..., 0], 0.05), + (limits[..., 1] - limits[..., 0]).clamp_min(0.0) * 0.1, + ) + lower = limits[..., 0] + margin + upper = limits[..., 1] - margin + start = torch.maximum(torch.minimum(start, upper), lower) + middle = torch.maximum(torch.minimum(middle, upper), lower) + goal = torch.maximum(torch.minimum(goal, upper), lower) + return [start, middle, goal] + + +def build_cartesian_line_poses( + robot: Robot, + control_part: str, + start_qpos: torch.Tensor, + distance: float, +) -> list[torch.Tensor]: + """Create two poses defining a fixed-orientation straight EEF path.""" + if not torch.isfinite(torch.tensor(distance)) or distance <= 0.0: + raise ValueError("cartesian distance must be finite and greater than zero.") + start_pose = robot.compute_fk(qpos=start_qpos, name=control_part, to_matrix=True) + if start_pose is None: + raise RuntimeError(f"Forward kinematics is unavailable for {control_part!r}.") + goal_pose = start_pose.clone() + # Match the known-reachable first Cartesian segment from the TOPPRA demo. + direction = start_pose.new_tensor((0.0, 0.0, -1.0)) + goal_pose[:, :3, 3] += distance * direction + return [start_pose, goal_pose] + + +def joint_derivatives_from_path_time_law( + jacobians: torch.Tensor, + path_positions: torch.Tensor, + path_velocities: torch.Tensor, + path_accelerations: torch.Tensor, + cartesian_tangent: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Convert scalar path derivatives into joint derivatives. + + For a Cartesian line ``x(s)`` the tangent is constant and + ``d²x/ds² = 0``. The joint path derivatives therefore satisfy + ``J q_s = x_s`` and ``J q_ss = -J_s q_s``. Double-S then supplies the + time law through ``q_dot = q_s s_dot`` and + ``q_ddot = q_ss s_dot² + q_s s_ddot``. + + Args: + jacobians: Geometric Jacobians shaped ``(B, N, 6, DOF)``. + path_positions: Scalar metric path positions shaped ``(B, N)``. + path_velocities: Scalar path velocities shaped ``(B, N)``. + path_accelerations: Scalar path accelerations shaped ``(B, N)``. + cartesian_tangent: Constant unit twists shaped ``(B, 6)``. + + Returns: + Joint velocity and acceleration tensors shaped ``(B, N, DOF)``. + """ + batch_size, sample_count, _, dof = jacobians.shape + expected_scalar_shape = (batch_size, sample_count) + if ( + path_positions.shape != expected_scalar_shape + or path_velocities.shape != expected_scalar_shape + or path_accelerations.shape != expected_scalar_shape + or cartesian_tangent.shape != (batch_size, 6) + ): + raise ValueError("Path derivative tensors have incompatible shapes.") + if sample_count < 3: + raise ValueError("At least three path samples are required.") + if bool((torch.diff(path_positions, dim=1) <= 0.0).any().item()): + raise ValueError("Path positions must be strictly increasing.") + + jacobian_pinv = torch.linalg.pinv(jacobians) + tangent = cartesian_tangent[:, None, :, None].expand(-1, sample_count, -1, -1) + q_s = torch.matmul(jacobian_pinv, tangent).squeeze(-1) + jacobian_s_rows: list[torch.Tensor] = [] + for batch_index in range(batch_size): + jacobian_s_rows.append( + torch.gradient( + jacobians[batch_index], + spacing=(path_positions[batch_index],), + dim=(0,), + edge_order=2, + )[0] + ) + jacobian_s = torch.stack(jacobian_s_rows) + curvature_twist = torch.matmul(jacobian_s, q_s.unsqueeze(-1)) + q_ss = -torch.matmul(jacobian_pinv, curvature_twist).squeeze(-1) + velocities = q_s * path_velocities[..., None] + accelerations = ( + q_ss * path_velocities.square()[..., None] + q_s * path_accelerations[..., None] + ) + if velocities.shape != (batch_size, sample_count, dof): + raise RuntimeError("Unexpected joint derivative shape.") + return velocities, accelerations + + +def nearest_equivalent_joint_solution( + solution: torch.Tensor, + seed: torch.Tensor, + limits: torch.Tensor, +) -> torch.Tensor: + """Select the limit-valid ``2π`` equivalent closest to the previous seed.""" + if solution.shape != seed.shape or limits.shape != (*solution.shape, 2): + raise ValueError("solution, seed, and limits have incompatible shapes.") + turns = torch.round((seed - solution) / (2.0 * torch.pi)) + nearest = solution + turns * (2.0 * torch.pi) + valid = (nearest >= limits[..., 0]) & (nearest <= limits[..., 1]) + return torch.where(valid, nearest, solution) + + +def plan_cartesian_line( + robot: Robot, + control_part: str, + start_qpos: torch.Tensor, + *, + distance: float, + profile: str, + sample_count: int, + velocity_limit: float, + acceleration_limit: float, + jerk_limit: float, + backend: str, +) -> tuple[PlanResult, torch.Tensor, PlanResult]: + """Time-parameterize Cartesian line distance, then solve continuous IK. + + Returns: + Joint trajectory, desired Cartesian pose samples, and the scalar + Cartesian path-parameter trajectory carrying metric derivatives. + """ + if sample_count < 3: + raise ValueError("Cartesian planning requires sample_count >= 3.") + start_pose, goal_pose = build_cartesian_line_poses( + robot, control_part, start_qpos, distance + ) + batch_size = start_qpos.shape[0] + scalar_waypoints = start_qpos.new_zeros((batch_size, 2, 1)) + scalar_waypoints[:, 1, 0] = distance + scalar_plan = _plan_linear_profiles( + scalar_waypoints, + TrapezoidalPlanOptions( + profile=profile, + constraints={ + "velocity": velocity_limit, + "acceleration": acceleration_limit, + "jerk": jerk_limit, + }, + sample_interval=sample_count, + backend=backend, + ), + ) + progress = scalar_plan.positions[..., 0] / distance + desired_poses = start_pose[:, None].expand(-1, sample_count, -1, -1).clone() + translation = torch.lerp( + start_pose[:, None, :3, 3], + goal_pose[:, None, :3, 3], + progress[..., None], + ) + desired_poses[:, :, :3, 3] = translation + + joint_ids = robot.get_joint_ids(name=control_part) + joint_limits = robot.get_qpos_limits(joint_ids=joint_ids).to(start_qpos) + ik_result = robot.compute_batch_ik( + desired_poses, + start_qpos, + control_part, + continuous=True, + ) + if ik_result is None: + raise RuntimeError("Cartesian line IK solver is unavailable.") + sample_success, positions = ik_result + sample_success = sample_success.bool() + if not bool(sample_success.all().item()): + failure = torch.nonzero(~sample_success, as_tuple=False)[0].tolist() + raise RuntimeError( + f"Cartesian line IK failed at env {failure[0]}, sample {failure[1]}." + ) + previous = torch.cat((start_qpos[:, None], positions[:, :-1]), dim=1) + positions = nearest_equivalent_joint_solution( + positions, + previous, + joint_limits[:, None].expand(-1, sample_count, -1, -1), + ) + success = sample_success.all(dim=1) + solver = robot.get_solver(control_part) + if solver is None: + raise RuntimeError(f"Kinematic solver is unavailable for {control_part!r}.") + jacobians = solver.get_jacobian( + positions.reshape(batch_size * sample_count, -1), jac_type="full" + ).reshape(batch_size, sample_count, 6, -1) + base_pose = robot.get_control_part_base_pose(control_part, to_matrix=True) + world_direction = goal_pose[:, :3, 3] - start_pose[:, :3, 3] + world_direction /= torch.linalg.vector_norm(world_direction, dim=-1, keepdim=True) + root_direction = torch.matmul( + base_pose[:, :3, :3].transpose(-1, -2), world_direction.unsqueeze(-1) + ).squeeze(-1) + cartesian_tangent = torch.cat( + (root_direction, torch.zeros_like(root_direction)), dim=-1 + ).to(jacobians) + velocities, accelerations = joint_derivatives_from_path_time_law( + jacobians, + scalar_plan.positions[..., 0], + scalar_plan.velocities[..., 0], + scalar_plan.accelerations[..., 0], + cartesian_tangent, + ) + return ( + PlanResult( + success=success, + xpos_list=desired_poses, + positions=positions, + velocities=velocities, + accelerations=accelerations, + dt=scalar_plan.dt, + ), + desired_poses, + scalar_plan, + ) + + +def replay_plan( + sim: SimulationManager, + robot: Robot, + control_part: str, + positions: torch.Tensor, + dt: torch.Tensor, + *, + replay_speed: float = DEFAULT_REPLAY_SPEED, + realtime: bool = True, +) -> None: + """Drive a joint plan according to its explicit trajectory timing.""" + if positions.ndim != 3 or dt.shape != positions.shape[:2]: + raise ValueError("positions and dt must have shapes (B, N, DOF) and (B, N).") + if replay_speed <= 0.0: + raise ValueError("replay_speed must be greater than zero.") + if positions.shape[0] > 1 and not torch.allclose(dt, dt[:1].expand_as(dt)): + raise ValueError( + "replay_plan requires one environment or identical dt rows across environments." + ) + physics_dt = float(sim.sim_config.physics_dt) + wall_start = time.perf_counter() + target_elapsed = 0.0 + for sample_index, command in enumerate(positions.transpose(0, 1)): + sample_duration = float(dt[:, sample_index].max().item()) / replay_speed + robot.set_qpos(command, name=control_part) + physics_steps = max(1, math.ceil(sample_duration / physics_dt)) + sim.update(step=physics_steps) + if realtime: + target_elapsed += sample_duration + remaining = wall_start + target_elapsed - time.perf_counter() + if remaining > 0.0: + time.sleep(remaining) + + +def compute_eef_trajectory( + robot: Robot, + control_part: str, + joint_positions: torch.Tensor, + env_index: int, +) -> torch.Tensor: + """Evaluate FK for one row of a trajectory shaped ``(B, N, DOF)``. + + Returns: + End-effector position and quaternion as ``(N, 7)`` in the local arena + frame, ordered as ``x, y, z, qw, qx, qy, qz``. + """ + if joint_positions.ndim != 3: + raise ValueError("joint_positions must have shape (B, N, DOF).") + if not 0 <= env_index < joint_positions.shape[0]: + raise ValueError( + f"env_index must be in [0, {joint_positions.shape[0]}), got {env_index}." + ) + sample_count = joint_positions.shape[1] + poses = robot.compute_fk( + qpos=joint_positions[env_index], + name=control_part, + env_ids=[env_index] * sample_count, + to_matrix=False, + ) + if poses is None: + raise RuntimeError(f"Forward kinematics is unavailable for {control_part!r}.") + return poses + + +def unwrap_angles(angles: torch.Tensor) -> torch.Tensor: + """Remove artificial ``2π`` jumps from a sequence of Euler angles.""" + if angles.ndim != 2 or angles.shape[-1] != 3: + raise ValueError("angles must have shape (N, 3).") + if angles.shape[0] < 2: + return angles.clone() + delta = torch.diff(angles, dim=0) + wrapped_delta = torch.remainder(delta + torch.pi, 2.0 * torch.pi) - torch.pi + wrapped_delta = torch.where( + (wrapped_delta == -torch.pi) & (delta > 0.0), torch.pi, wrapped_delta + ) + correction = wrapped_delta - delta + correction = torch.where(delta.abs() < torch.pi, 0.0, correction) + return torch.cat((angles[:1], angles[1:] + correction.cumsum(dim=0)), dim=0) + + +def maximum_line_deviation(xyz: torch.Tensor) -> torch.Tensor: + """Return the largest orthogonal deviation from the endpoint line.""" + if xyz.ndim != 2 or xyz.shape[0] < 2 or xyz.shape[1] != 3: + raise ValueError("xyz must have shape (N, 3) with N >= 2.") + line = xyz[-1] - xyz[0] + squared_length = torch.dot(line, line) + if squared_length <= torch.finfo(xyz.dtype).eps: + return torch.linalg.vector_norm(xyz - xyz[:1], dim=-1).max() + progress = ((xyz - xyz[:1]) * line).sum(dim=-1) / squared_length + closest = xyz[:1] + progress[:, None] * line + return torch.linalg.vector_norm(xyz - closest, dim=-1).max() + + +def set_equal_3d_limits(axis: Axes3D, points: torch.Tensor) -> None: + """Use one physical scale for all axes of a 3D trajectory plot.""" + if points.ndim != 2 or points.shape[0] < 2 or points.shape[1] != 3: + raise ValueError("points must have shape (N, 3) with N >= 2.") + lower = points.amin(dim=0) + upper = points.amax(dim=0) + center = 0.5 * (lower + upper) + largest_span = float((upper - lower).max().item()) + radius = max(0.5 * largest_span * 1.1, 1e-4) + for setter, coordinate in ( + (axis.set_xlim, center[0]), + (axis.set_ylim, center[1]), + (axis.set_zlim, center[2]), + ): + midpoint = float(coordinate.item()) + setter(midpoint - radius, midpoint + radius) + axis.set_box_aspect((1.0, 1.0, 1.0)) + + +def plot_trajectory_diagnostics( + *, + dt: torch.Tensor, + eef_poses: torch.Tensor, + joint_positions: torch.Tensor, + joint_velocities: torch.Tensor, + joint_accelerations: torch.Tensor, + desired_eef_poses: torch.Tensor | None = None, + env_index: int, + output_path: Path | None, + profile_label: str, + show: bool = False, +) -> Figure: + """Plot Cartesian pose and joint derivatives for one batch environment.""" + configure_plot_fonts() + batch_size = joint_positions.shape[0] + if not 0 <= env_index < batch_size: + raise ValueError(f"env_index must be in [0, {batch_size}), got {env_index}.") + expected_prefix = joint_positions.shape[:2] + if ( + dt.shape != expected_prefix + or eef_poses.shape != (expected_prefix[1], 7) + or joint_velocities.shape != joint_positions.shape + or joint_accelerations.shape != joint_positions.shape + ): + raise ValueError("Trajectory diagnostic tensors have incompatible shapes.") + + time = dt[env_index].cumsum(dim=0).detach().cpu() + pose = eef_poses.detach().cpu() + roll, pitch, yaw = euler_xyz_from_quat(pose[:, 3:]) + orientation = unwrap_angles(torch.stack((roll, pitch, yaw), dim=-1)) + orientation_deg = torch.rad2deg(orientation) + position = joint_positions[env_index].detach().cpu() + velocity = joint_velocities[env_index].detach().cpu() + acceleration = joint_accelerations[env_index].detach().cpu() + + xyz = pose[:, :3] + duration = time[-1].item() + peak_velocity = velocity.abs().max().item() + peak_acceleration = acceleration.abs().max().item() + line_deviation_mm = maximum_line_deviation(xyz).item() * 1000.0 + + figure = plt.figure(figsize=(16, 13), layout="constrained") + grid = figure.add_gridspec(3, 2) + path_axis = figure.add_subplot(grid[0, 0], projection="3d") + axes = [ + figure.add_subplot(grid[0, 1]), + figure.add_subplot(grid[1, 0]), + figure.add_subplot(grid[1, 1]), + figure.add_subplot(grid[2, 0]), + figure.add_subplot(grid[2, 1]), + ] + path_axis.plot( + xyz[:, 0].numpy(), + xyz[:, 1].numpy(), + xyz[:, 2].numpy(), + color="tab:blue", + linewidth=2.5, + label="FK path", + ) + reference_xyz = ( + desired_eef_poses[:, :3, 3].detach().cpu() + if desired_eef_poses is not None + else torch.stack((xyz[0], xyz[-1])) + ) + path_axis.plot( + reference_xyz[:, 0].numpy(), + reference_xyz[:, 1].numpy(), + reference_xyz[:, 2].numpy(), + color="black", + linestyle="--", + alpha=0.65, + label="planned Cartesian line", + ) + path_axis.scatter(*xyz[0].tolist(), color="tab:green", s=70, label="start") + path_axis.scatter(*xyz[-1].tolist(), color="tab:red", s=70, label="goal") + set_equal_3d_limits(path_axis, torch.cat((xyz, reference_xyz), dim=0)) + path_axis.set_title("EEF path in Cartesian space") + path_axis.set_xlabel("x [m]") + path_axis.set_ylabel("y [m]") + path_axis.set_zlabel("z [m]") + path_axis.legend() + + colors = plt.get_cmap("tab10").colors + for axis, values, labels, title, ylabel in ( + (axes[0], xyz, ("x", "y", "z"), "EEF XYZ", "position [m]"), + ( + axes[1], + orientation_deg, + ("roll", "pitch", "yaw"), + "EEF orientation (XYZ Euler)", + "angle [deg]", + ), + ( + axes[2], + position, + tuple(f"q{i}" for i in range(position.shape[1])), + "Joint position", + "angle [rad]", + ), + ( + axes[3], + velocity, + tuple(f"dq{i}" for i in range(velocity.shape[1])), + "Joint velocity", + "velocity [rad/s]", + ), + ( + axes[4], + acceleration, + tuple(f"ddq{i}" for i in range(acceleration.shape[1])), + "Joint acceleration", + "acceleration [rad/s²]", + ), + ): + for curve_index, (curve, label) in enumerate( + zip(values.T.numpy(), labels, strict=True) + ): + axis.plot( + time.numpy(), + curve, + color=colors[curve_index % len(colors)], + linewidth=1.8, + label=label, + ) + axis.set_title(title) + axis.set_ylabel(ylabel) + axis.set_xlabel("time [s]") + axis.grid(True, alpha=0.22) + axis.legend(ncol=min(3, len(labels)), frameon=False) + if desired_eef_poses is not None: + desired_xyz = desired_eef_poses[:, :3, 3].detach().cpu() + for coordinate, label, color in zip( + desired_xyz.T, + ("x desired", "y desired", "z desired"), + colors[:3], + strict=True, + ): + axes[0].plot( + time.numpy(), + coordinate.numpy(), + color=color, + linestyle="--", + linewidth=1.2, + alpha=0.8, + label=label, + ) + axes[0].legend(ncol=3, frameon=False) + figure.suptitle( + f"{profile_label.replace('_', ' ').title()} — env {env_index}\n" + f"duration {duration:.3f} s | max |dq| {peak_velocity:.3f} rad/s | " + f"max |ddq| {peak_acceleration:.3f} rad/s² | " + f"line error {line_deviation_mm:.3f} mm", + ) + if output_path is not None: + output_path.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(output_path, dpi=160, bbox_inches="tight") + print(f"[INFO] trajectory plot saved to {output_path.resolve()}") + if show: + plt.show() + return figure + + +def diagnostic_output_path( + base_path: Path | None, diagnostic_label: str, multiple: bool +) -> Path | None: + """Return a stable output path for one requested diagnostic profile.""" + if base_path is None or not multiple: + return base_path + suffix = base_path.suffix or ".png" + return base_path.with_name(f"{base_path.stem}_{diagnostic_label}{suffix}") + + +def main() -> None: + """Create a robot, plan a trajectory, and replay it.""" + args = parse_args() + figures: list[Figure] = [] + sim = SimulationManager( + SimulationManagerCfg( + headless=args.headless, + sim_device=args.device, + num_envs=args.num_envs, + visualization=visualization_cfg_from_args(args), + ) + ) + try: + robot = sim.add_robot(CobotMagicCfg.from_dict({"uid": "CobotMagic"})) + if sim.is_use_gpu_physics: + sim.init_gpu_physics() + if not args.headless: + sim.open_window() + + control_part = "left_arm" + joint_waypoints = build_demo_waypoints(robot, control_part) + start_qpos = joint_waypoints[0] + generator = MotionGenerator( + MotionGenCfg( + planner_cfg=TrapezoidalPlannerCfg(robot_uid=robot.uid), + ) + ) + requested_profiles = ( + tuple(PROFILE_SPECS.items()) + if args.profile == "both" + else ((args.profile, PROFILE_SPECS[args.profile]),) + ) + requested_paths = ( + ("joint", "cartesian") if args.path == "both" else (args.path,) + ) + diagnostic_count = len(requested_profiles) * len(requested_paths) + for path_name in requested_paths: + if path_name == "joint": + target_states = [PlanState.from_qpos(qpos) for qpos in joint_waypoints] + for profile_name, ( + planner_profile, + profile_label, + ) in requested_profiles: + scalar_plan = None + if path_name == "cartesian": + result, desired_poses, scalar_plan = plan_cartesian_line( + robot, + control_part, + start_qpos, + distance=args.cartesian_distance, + profile=planner_profile, + sample_count=max( + args.samples, + math.ceil(args.cartesian_distance / args.cartesian_step) + + 1, + ), + velocity_limit=args.cartesian_velocity, + acceleration_limit=args.cartesian_acceleration, + jerk_limit=args.cartesian_jerk, + backend=args.backend, + ) + else: + desired_poses = None + result = generator.generate( + target_states, + MotionGenOptions( + control_part=control_part, + start_qpos=start_qpos, + plan_opts=TrapezoidalPlanOptions( + profile=planner_profile, + constraints={ + "velocity": 0.5, + "acceleration": 1.0, + "jerk": 3.0, + }, + sample_interval=args.samples, + stop_at_waypoints=False, + backend=args.backend, + ), + ), + ) + if ( + not result.is_all_success() + or result.positions is None + or result.velocities is None + or result.accelerations is None + or result.dt is None + ): + raise RuntimeError( + f"{path_name}/{profile_name} trajectory planning failed." + ) + print( + f"[INFO] path={path_name}, profile={profile_name}, " + f"shape={tuple(result.positions.shape)}, " + f"duration={result.duration.tolist()}" + ) + eef_poses = compute_eef_trajectory( + robot, control_part, result.positions, args.plot_env + ) + if path_name == "cartesian": + line_error = maximum_line_deviation(eef_poses[:, :3]) + assert scalar_plan is not None + print( + "[INFO] cartesian max line deviation=" + f"{line_error.item() * 1000.0:.3f} mm" + ) + print( + "[INFO] cartesian peaks: " + f"speed={scalar_plan.velocities.abs().max().item():.3f} m/s, " + "acceleration=" + f"{scalar_plan.accelerations.abs().max().item():.3f} m/s²" + ) + diagnostic_label = f"{path_name}_{profile_label}" + figures.append( + plot_trajectory_diagnostics( + dt=result.dt, + eef_poses=eef_poses, + joint_positions=result.positions, + joint_velocities=result.velocities, + joint_accelerations=result.accelerations, + desired_eef_poses=( + desired_poses[args.plot_env] + if desired_poses is not None + else None + ), + env_index=args.plot_env, + output_path=diagnostic_output_path( + args.plot_output, + diagnostic_label, + diagnostic_count > 1, + ), + profile_label=diagnostic_label, + ) + ) + robot.set_qpos(start_qpos, name=control_part, target=False) + robot.set_qpos(start_qpos, name=control_part) + sim.update(step=5) + replay_plan( + sim, + robot, + control_part, + result.positions, + result.dt, + replay_speed=args.replay_speed, + realtime=not args.headless, + ) + if args.show_plot: + plt.show() + finally: + for figure in figures: + plt.close(figure) + sim.destroy() + + +if __name__ == "__main__": + main() diff --git a/scripts/tutorials/sim/planner/trapezoidal_profile.py b/scripts/tutorials/sim/planner/trapezoidal_profile.py new file mode 100644 index 000000000..95d214b76 --- /dev/null +++ b/scripts/tutorials/sim/planner/trapezoidal_profile.py @@ -0,0 +1,134 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Plot a minimal scalar trapezoidal or Double-S trajectory example.""" + +from __future__ import annotations + +import argparse +import os + +os.environ.setdefault("MPLCONFIGDIR", "/tmp/embodichain-matplotlib") + +import matplotlib.pyplot as plt +import torch + +from embodichain.lab.sim.planners.trapezoidal_planner import ( + TrapezoidalPlanOptions, + _plan_linear_profiles, +) + + +def configure_plot_fonts() -> None: + """Apply lightweight readable font defaults for this example.""" + plt.rcParams.update( + { + "font.sans-serif": ["Noto Sans CJK SC", "DejaVu Sans", "sans-serif"], + "font.size": 11.0, + "axes.titlesize": 13.0, + "legend.fontsize": 9.0, + "axes.unicode_minus": False, + } + ) + + +def positive_float(value: str) -> float: + """Parse a finite positive floating-point argument.""" + parsed = float(value) + if not torch.isfinite(torch.tensor(parsed)) or parsed <= 0.0: + raise argparse.ArgumentTypeError("value must be finite and positive") + return parsed + + +def parse_args() -> argparse.Namespace: + """Parse example options.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--profile", + choices=("velocity_trapezoidal", "acceleration_trapezoidal"), + default="acceleration_trapezoidal", + ) + parser.add_argument("--distance", type=positive_float, default=0.1) + parser.add_argument("--velocity", type=positive_float, default=0.15) + parser.add_argument("--acceleration", type=positive_float, default=0.3) + parser.add_argument("--jerk", type=positive_float, default=1.0) + parser.add_argument("--samples", type=int, default=501) + parser.add_argument( + "--show-plot", + action=argparse.BooleanOptionalAction, + default=True, + help="Display the diagnostic figure (default: enabled).", + ) + return parser.parse_args() + + +def main() -> None: + """Plan the scalar move and show its derivatives.""" + args = parse_args() + configure_plot_fonts() + if args.samples < 3: + raise SystemExit("--samples must be at least 3") + planner_profile = ( + "trapezoidal" if args.profile == "velocity_trapezoidal" else "double_s" + ) + waypoints = torch.tensor([[[0.0], [args.distance]]], dtype=torch.float64) + result = _plan_linear_profiles( + waypoints, + TrapezoidalPlanOptions( + profile=planner_profile, + constraints={ + "velocity": args.velocity, + "acceleration": args.acceleration, + "jerk": args.jerk, + }, + sample_interval=args.samples, + backend="torch", + ), + ) + time = result.dt[0].cumsum(dim=0) + position = result.positions[0, :, 0] + velocity = result.velocities[0, :, 0] + acceleration = result.accelerations[0, :, 0] + jerk = torch.gradient(acceleration, spacing=(time,), edge_order=2)[0] + + print( + f"[INFO] profile={args.profile}, duration={result.duration.item():.12f} s, " + f"max_velocity={velocity.abs().max().item():.6f}, " + f"max_acceleration={acceleration.abs().max().item():.6f}, " + f"max_sampled_jerk={jerk.abs().max().item():.6f}" + ) + + figure, axes = plt.subplots(4, 1, figsize=(11, 10), sharex=True) + for axis, values, title, ylabel in ( + (axes[0], position, "Position", "q"), + (axes[1], velocity, "Velocity", "dq/dt"), + (axes[2], acceleration, "Acceleration", "d²q/dt²"), + (axes[3], jerk, "Jerk", "d³q/dt³"), + ): + axis.plot(time.numpy(), values.numpy(), linewidth=2.0) + axis.set_title(title) + axis.set_ylabel(ylabel) + axis.grid(True, alpha=0.25) + axes[-1].set_xlabel("time [s]") + figure.suptitle(args.profile.replace("_", " ").title()) + figure.tight_layout() + if args.show_plot: + plt.show() + plt.close(figure) + + +if __name__ == "__main__": + main() diff --git a/tests/lab/scripts/test_trapezoidal_planner_tutorial.py b/tests/lab/scripts/test_trapezoidal_planner_tutorial.py new file mode 100644 index 000000000..7991c5531 --- /dev/null +++ b/tests/lab/scripts/test_trapezoidal_planner_tutorial.py @@ -0,0 +1,452 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import argparse +from collections.abc import Callable +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +import matplotlib + +matplotlib.use("Agg") +from matplotlib import pyplot as plt + +from scripts.tutorials.sim.planner.trapezoidal_planner import ( + build_cartesian_line_poses, + build_demo_waypoints, + compute_eef_trajectory, + configure_plot_fonts, + diagnostic_output_path, + joint_derivatives_from_path_time_law, + maximum_line_deviation, + nearest_equivalent_joint_solution, + plan_cartesian_line, + plot_trajectory_diagnostics, + positive_float, + replay_plan, + sample_count, + set_equal_3d_limits, +) + + +class _Solver: + """Identity kinematics used to validate path derivative propagation.""" + + def get_jacobian(self, qpos: torch.Tensor, jac_type: str = "full") -> torch.Tensor: + assert jac_type == "full" + dof = qpos.shape[-1] + return torch.eye(6, dof, dtype=qpos.dtype).expand(qpos.shape[0], -1, -1) + + +class _Robot: + """Minimal robot interface used by the tutorial waypoint builders.""" + + def __init__(self, batch_size: int = 2, dof: int = 6) -> None: + self.qpos = torch.zeros(batch_size, dof) + self.commands: list[torch.Tensor] = [] + self.solver = _Solver() + self.fk_call_count = 0 + self.path_ik_call_count = 0 + + def get_qpos(self, name: str) -> torch.Tensor: + assert name == "left_arm" + return self.qpos + + def get_joint_ids(self, name: str) -> list[int]: + assert name == "left_arm" + return list(range(self.qpos.shape[1])) + + def get_qpos_limits(self, joint_ids: list[int]) -> torch.Tensor: + assert joint_ids == list(range(self.qpos.shape[1])) + limits = torch.tensor((-1.0, 1.0)).repeat(self.qpos.shape[1], 1) + return limits.unsqueeze(0).repeat(self.qpos.shape[0], 1, 1) + + def compute_fk( + self, + qpos: torch.Tensor, + name: str, + to_matrix: bool, + env_ids: list[int] | None = None, + ) -> torch.Tensor: + assert name == "left_arm" + self.fk_call_count += 1 + if env_ids is not None: + assert len(env_ids) == qpos.shape[0] + if to_matrix: + pose = torch.eye(4).repeat(qpos.shape[0], 1, 1) + pose[:, :3, 3] = qpos[:, :3] + return pose + quaternion = torch.zeros(qpos.shape[0], 4) + quaternion[:, 0] = 1.0 + return torch.cat((qpos[:, :3], quaternion), dim=-1) + + def compute_ik( + self, + pose: torch.Tensor, + joint_seed: torch.Tensor, + name: str, + ) -> tuple[torch.Tensor, torch.Tensor]: + assert name == "left_arm" + qpos = joint_seed.clone() + qpos[:, :3] = pose[:, :3, 3] + success = torch.ones(pose.shape[0], dtype=torch.bool) + return success, qpos + + def compute_batch_ik( + self, + pose: torch.Tensor, + joint_seed: torch.Tensor, + name: str, + *, + continuous: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + assert name == "left_arm" + assert continuous + self.path_ik_call_count += 1 + qpos = joint_seed[:, None].expand(-1, pose.shape[1], -1).clone() + qpos[:, :, :3] = pose[:, :, :3, 3] + success = torch.ones(pose.shape[:2], dtype=torch.bool) + return success, qpos + + def get_solver(self, name: str) -> _Solver: + assert name == "left_arm" + return self.solver + + def get_control_part_base_pose(self, name: str, to_matrix: bool) -> torch.Tensor: + assert name == "left_arm" + assert to_matrix + return torch.eye(4).repeat(self.qpos.shape[0], 1, 1) + + def set_qpos(self, qpos: torch.Tensor, name: str) -> None: + assert name == "left_arm" + self.commands.append(qpos.clone()) + + +class _Simulation: + def __init__(self, physics_dt: float = 0.01) -> None: + self.sim_config = SimpleNamespace(physics_dt=physics_dt) + self.update_steps: list[int] = [] + + def update(self, step: int) -> None: + self.update_steps.append(step) + + +def test_plot_font_style_has_readable_cjk_fallback_and_hierarchy() -> None: + configure_plot_fonts() + + assert plt.rcParams["font.sans-serif"][0] == "Noto Sans CJK SC" + assert plt.rcParams["axes.titlesize"] > plt.rcParams["font.size"] + assert plt.rcParams["font.size"] > plt.rcParams["legend.fontsize"] + assert plt.rcParams["axes.unicode_minus"] is False + + +def test_joint_demo_moves_every_joint_within_limits() -> None: + robot = _Robot() + + start, middle, goal = build_demo_waypoints(robot, "left_arm") + + assert torch.all(middle != start) + assert torch.all(goal != start) + assert torch.all(middle.abs() <= 1.0) + assert torch.all(goal.abs() <= 1.0) + + +def test_six_axis_demo_uses_toppra_nonsingular_seed() -> None: + robot = _Robot() + + start = build_demo_waypoints(robot, "left_arm")[0] + + assert torch.allclose(start[:, 1], torch.full((2,), torch.pi / 4.0)) + assert torch.allclose(start[:, 2], torch.full((2,), -torch.pi / 4.0)) + assert torch.allclose(start[:, 4], torch.full((2,), torch.pi / 4.0)) + + +def test_cartesian_demo_preserves_orientation_and_requested_distance() -> None: + robot = _Robot() + requested_distance = 0.18 + + start, goal = build_cartesian_line_poses( + robot, + "left_arm", + robot.qpos, + requested_distance, + ) + + displacement = goal[:, :3, 3] - start[:, :3, 3] + assert torch.allclose( + torch.linalg.vector_norm(displacement, dim=-1), + torch.full((robot.qpos.shape[0],), requested_distance), + ) + assert torch.equal(goal[:, :3, :3], start[:, :3, :3]) + + +def test_cartesian_demo_rejects_nonpositive_distance() -> None: + robot = _Robot() + + with pytest.raises(ValueError, match="greater than zero"): + build_cartesian_line_poses(robot, "left_arm", robot.qpos, 0.0) + + +def test_cartesian_demo_requires_three_samples_for_path_derivatives() -> None: + robot = _Robot() + + with pytest.raises(ValueError, match="sample_count >= 3"): + plan_cartesian_line( + robot, + "left_arm", + robot.qpos, + distance=0.1, + profile="trapezoidal", + sample_count=2, + velocity_limit=0.15, + acceleration_limit=0.3, + jerk_limit=1.0, + backend="torch", + ) + + +def test_cartesian_time_law_is_applied_before_ik() -> None: + robot = _Robot() + distance = 0.10 + velocity_limit = 0.15 + acceleration_limit = 0.30 + + result, desired_poses, scalar_plan = plan_cartesian_line( + robot, + "left_arm", + robot.qpos, + distance=distance, + profile="trapezoidal", + sample_count=101, + velocity_limit=velocity_limit, + acceleration_limit=acceleration_limit, + jerk_limit=1.0, + backend="torch", + ) + + desired_xyz = desired_poses[0, :, :3, 3] + assert torch.all(result.duration > 0.0) + assert maximum_line_deviation(desired_xyz).item() < 1e-7 + assert scalar_plan.velocities.abs().max() <= velocity_limit + 1e-6 + assert scalar_plan.accelerations.abs().max() <= acceleration_limit + 1e-6 + assert torch.allclose( + desired_xyz[-1] - desired_xyz[0], torch.tensor([0, 0, -distance]) + ) + assert torch.count_nonzero(result.velocities[:, 0]) == 0 + assert robot.path_ik_call_count == 1 + + +def test_eef_trajectory_uses_one_batched_fk_call() -> None: + robot = _Robot(batch_size=2) + sample_count = 101 + joint_positions = torch.zeros(2, sample_count, 6) + joint_positions[1, :, 2] = torch.linspace(0.0, -0.1, sample_count) + + poses = compute_eef_trajectory(robot, "left_arm", joint_positions, env_index=1) + + assert poses.shape == (sample_count, 7) + assert torch.equal(poses[:, 2], joint_positions[1, :, 2]) + assert robot.fk_call_count == 1 + + +def test_path_time_law_produces_exact_joint_derivatives_for_identity_model() -> None: + sample_count = 101 + parameter = torch.linspace(0.0, torch.pi, sample_count, dtype=torch.float64) + path_position = torch.linspace( + 0.0, 0.1, sample_count, dtype=torch.float64 + ).unsqueeze(0) + path_velocity = torch.sin(parameter).unsqueeze(0) + path_acceleration = torch.cos(parameter).unsqueeze(0) + jacobians = ( + torch.eye(6, dtype=torch.float64) + .reshape(1, 1, 6, 6) + .expand(1, sample_count, -1, -1) + ) + tangent = torch.tensor([[0.0, 0.0, -1.0, 0.0, 0.0, 0.0]], dtype=torch.float64) + + velocity, acceleration = joint_derivatives_from_path_time_law( + jacobians, + path_position, + path_velocity, + path_acceleration, + tangent, + ) + + assert torch.allclose(velocity[..., 2], -path_velocity, atol=1e-7) + assert torch.allclose(acceleration[..., 2], -path_acceleration, atol=1e-7) + assert torch.count_nonzero(velocity[..., [0, 1, 3, 4, 5]]) == 0 + + +def test_path_time_law_includes_jacobian_curvature_acceleration() -> None: + sample_count = 51 + path_position = torch.linspace( + 0.0, 1.0, sample_count, dtype=torch.float64 + ).unsqueeze(0) + path_velocity = torch.full_like(path_position, 0.5) + path_acceleration = torch.zeros_like(path_position) + jacobians = ( + torch.eye(6, dtype=torch.float64) + .reshape(1, 1, 6, 6) + .repeat(1, sample_count, 1, 1) + ) + jacobians[0, :, 0, 0] = 1.0 + path_position[0] + tangent = torch.tensor([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype=torch.float64) + + velocity, acceleration = joint_derivatives_from_path_time_law( + jacobians, + path_position, + path_velocity, + path_acceleration, + tangent, + ) + + expected_q_s = 1.0 / (1.0 + path_position) + expected_q_ss = -1.0 / (1.0 + path_position).square() + assert torch.allclose(velocity[..., 0], expected_q_s * 0.5, atol=1e-12) + assert torch.allclose(acceleration[..., 0], expected_q_ss * 0.25, atol=1e-12) + + +def test_nearest_equivalent_ik_solution_removes_angle_wrap() -> None: + seed = torch.tensor([[3.13]]) + wrapped_solution = torch.tensor([[-3.13]]) + limits = torch.tensor([[[-2.0 * torch.pi, 2.0 * torch.pi]]]) + + continuous = nearest_equivalent_joint_solution(wrapped_solution, seed, limits) + + assert (continuous - seed).abs().item() < 0.03 + + +def test_line_deviation_detects_off_axis_samples() -> None: + xyz = torch.tensor([[0.0, 0.0, 0.0], [0.5, 0.1, 0.0], [1.0, 0.0, 0.0]]) + + deviation = maximum_line_deviation(xyz) + + assert deviation.item() == pytest.approx(0.1) + + +def test_equal_3d_limits_do_not_magnify_cross_axis_noise() -> None: + figure = plt.figure() + axis = figure.add_subplot(111, projection="3d") + points = torch.tensor( + [[1e-6, -2e-6, 0.0], [-1e-6, 2e-6, -0.1]], dtype=torch.float64 + ) + + try: + set_equal_3d_limits(axis, points) + spans = torch.tensor( + [ + axis.get_xlim()[1] - axis.get_xlim()[0], + axis.get_ylim()[1] - axis.get_ylim()[0], + axis.get_zlim()[1] - axis.get_zlim()[0], + ] + ) + assert torch.allclose(spans, torch.full_like(spans, spans[0])) + assert spans[0].item() > 0.1 + finally: + plt.close(figure) + + +def test_multiple_diagnostics_receive_distinct_output_names() -> None: + base = Path("outputs/trajectory.png") + + joint = diagnostic_output_path(base, "joint_velocity_trapezoidal", True) + cartesian = diagnostic_output_path(base, "cartesian_velocity_trapezoidal", True) + + assert joint.name == "trajectory_joint_velocity_trapezoidal.png" + assert cartesian.name == "trajectory_cartesian_velocity_trapezoidal.png" + + +def test_diagnostics_do_not_save_without_explicit_output_path() -> None: + assert diagnostic_output_path(None, "cartesian_velocity_trapezoidal", True) is None + + +def test_diagnostic_dashboard_contains_six_focused_panels() -> None: + samples = 8 + dof = 6 + dt = torch.full((1, samples), 0.02) + dt[:, 0] = 0.0 + poses = torch.zeros(samples, 7) + poses[:, 2] = torch.linspace(0.0, -0.1, samples) + poses[:, 3] = 1.0 + desired_poses = torch.eye(4).repeat(samples, 1, 1) + desired_poses[:, 2, 3] = poses[:, 2] + positions = torch.linspace(0.0, 0.2, samples).view(1, samples, 1) + positions = positions.repeat(1, 1, dof) + + figure = plot_trajectory_diagnostics( + dt=dt, + eef_poses=poses, + joint_positions=positions, + joint_velocities=torch.full_like(positions, 0.1), + joint_accelerations=torch.full_like(positions, 0.2), + desired_eef_poses=desired_poses, + env_index=0, + output_path=None, + profile_label="cartesian_acceleration_trapezoidal", + ) + + try: + assert len(figure.axes) == 6 + assert "duration" in figure._suptitle.get_text() + assert figure.axes[0].name == "3d" + assert len(figure.axes[1].lines) == 6 + finally: + plt.close(figure) + + +def test_replay_uses_explicit_dt_to_advance_physics() -> None: + robot = _Robot(batch_size=1) + simulation = _Simulation(physics_dt=0.01) + positions = torch.zeros(1, 3, 6) + dt = torch.tensor([[0.0, 0.02, 0.03]]) + + replay_plan( + simulation, + robot, + "left_arm", + positions, + dt, + realtime=False, + ) + + assert simulation.update_steps == [1, 2, 3] + assert len(robot.commands) == positions.shape[1] + + +def test_replay_rejects_per_environment_timing_mismatch() -> None: + """Replay must not silently slow shorter batch rows to the longest row.""" + robot = _Robot(batch_size=2) + simulation = _Simulation(physics_dt=0.01) + positions = torch.zeros(2, 2, 6) + dt = torch.tensor([[0.0, 0.02], [0.0, 0.03]]) + + with pytest.raises(ValueError, match="identical dt rows"): + replay_plan(simulation, robot, "left_arm", positions, dt, realtime=False) + + +@pytest.mark.parametrize( + ("parser", "value"), + [(positive_float, "0"), (positive_float, "nan"), (sample_count, "1")], +) +def test_cli_numeric_constraints_fail_before_simulation( + parser: Callable[[str], float | int], value: str +) -> None: + with pytest.raises(argparse.ArgumentTypeError, match="must be"): + parser(value) diff --git a/tests/sim/objects/test_robot.py b/tests/sim/objects/test_robot.py index e6e050f91..c3b9a559d 100644 --- a/tests/sim/objects/test_robot.py +++ b/tests/sim/objects/test_robot.py @@ -18,6 +18,7 @@ import os from types import SimpleNamespace +from unittest.mock import Mock import numpy as np import pytest @@ -65,6 +66,32 @@ def test_get_qf_selects_control_part_joint_efforts(): assert torch.equal(actual_qf, full_qf[:, [3, 1]]) +def test_compute_batch_ik_continuous_reuses_all_solver_candidates() -> None: + """Continuous batch IK uses the existing batch boundary and one candidate call.""" + robot = object.__new__(Robot) + robot.device = torch.device("cpu") + robot._all_indices = [0, 1] + solver = Mock(dof=2, root_link_name="base") + candidate_valid = torch.ones(6, 8, dtype=torch.bool) + candidate_qpos = torch.zeros(6, 8, 2) + selected = torch.ones(2, 3, dtype=torch.bool) + selected_qpos = torch.full((2, 3, 2), 0.25) + solver.get_ik.return_value = (candidate_valid, candidate_qpos) + solver._select_continuous_ik_path.return_value = (selected, selected_qpos) + robot._solvers = {"arm": solver} + robot.get_link_pose = Mock(return_value=torch.eye(4).repeat(2, 1, 1)) + poses = torch.eye(4).repeat(2, 3, 1, 1) + seed = torch.zeros(2, 2) + + success, qpos = robot.compute_batch_ik(poses, seed, "arm", continuous=True) + + assert torch.equal(success, selected) + assert torch.equal(qpos, selected_qpos) + solver.get_ik.assert_called_once() + assert solver.get_ik.call_args.kwargs["return_all_solutions"] is True + solver._select_continuous_ik_path.assert_called_once() + + # Base test class for CPU and CUDA class BaseRobotTest: @classmethod diff --git a/tests/sim/planners/test_motion_generator_batched.py b/tests/sim/planners/test_motion_generator_batched.py index 4b84350d9..58f92188c 100644 --- a/tests/sim/planners/test_motion_generator_batched.py +++ b/tests/sim/planners/test_motion_generator_batched.py @@ -27,6 +27,7 @@ MotionGenerator, MotionGenOptions, ) +from embodichain.lab.sim.planners.trapezoidal_planner import TrapezoidalPlanOptions from embodichain.lab.sim.planners.utils import PlanState, PlanResult, MoveType BATCH_SIZE = 2 @@ -370,6 +371,24 @@ def _mock_planner(b=3, n=15, dofs=6): return planner +def test_resolve_trapezoidal_limits_without_sample_count() -> None: + planner = Mock() + planner.cfg.planner_type = "trapezoidal" + generator = object.__new__(MotionGenerator) + generator.planner = planner + + options = generator.resolve_plan_options( + plan_opts=None, + sample_count=None, + velocity_limit=0.01, + acceleration_limit=0.02, + ) + + assert isinstance(options, TrapezoidalPlanOptions) + assert options.constraints["velocity"] == 0.01 + assert options.constraints["acceleration"] == 0.02 + + def _mock_generator( *, batch_size: int = 2, diff --git a/tests/sim/planners/test_trapezoidal_planner.py b/tests/sim/planners/test_trapezoidal_planner.py new file mode 100644 index 000000000..23e7681d8 --- /dev/null +++ b/tests/sim/planners/test_trapezoidal_planner.py @@ -0,0 +1,569 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +import embodichain.lab.sim.planners.trapezoidal_planner as trapezoidal_module +from embodichain.lab.sim.planners.trapezoidal_planner import ( + TrapezoidalPlanner, + TrapezoidalPlanOptions, + _build_double_s_profile, + _build_scalar_profile, + _compress_collinear_waypoints, + _plan_linear_profiles, +) +from embodichain.lab.sim.planners.utils import PlanState, TrajectorySampleMethod + + +@pytest.mark.parametrize("profile", ["trapezoidal", "double_s"]) +def test_profiles_reach_endpoints_and_limits(profile: str) -> None: + waypoints = torch.tensor( + [ + [[0.0, 0.0], [1.0, -0.5]], + [[0.2, -0.1], [0.4, 0.7]], + ], + dtype=torch.float64, + ) + velocity_limit = 0.6 + acceleration_limit = 1.2 + options = TrapezoidalPlanOptions( + profile=profile, + constraints={ + "velocity": velocity_limit, + "acceleration": acceleration_limit, + "jerk": 4.0, + }, + sample_method=TrajectorySampleMethod.QUANTITY, + sample_interval=401, + ) + + result = _plan_linear_profiles(waypoints, options) + + assert result.success.all() + assert torch.allclose(result.positions[:, 0], waypoints[:, 0]) + assert torch.allclose(result.positions[:, -1], waypoints[:, -1]) + assert float(result.velocities.abs().max()) <= velocity_limit + 1e-9 + assert float(result.accelerations.abs().max()) <= acceleration_limit + 1e-9 + assert torch.all(result.dt >= 0.0) + assert torch.allclose(result.duration, result.dt.sum(dim=1)) + + +def test_triangular_profile_handles_zero_duration_cruise_boundary() -> None: + waypoints = torch.tensor([[[0.0], [1.0]]], dtype=torch.float64) + options = TrapezoidalPlanOptions( + profile="trapezoidal", + constraints={"velocity": 1.0, "acceleration": 1.0, "jerk": 1.0}, + sample_interval=3, + backend="torch", + ) + + result = _plan_linear_profiles(waypoints, options) + + assert torch.allclose( + result.positions[0, :, 0], + torch.tensor([0.0, 0.5, 1.0], dtype=torch.float64), + ) + assert torch.allclose( + result.velocities[0, :, 0], + torch.tensor([0.0, 1.0, 0.0], dtype=torch.float64), + ) + + +def test_double_s_has_continuous_acceleration_samples() -> None: + waypoints = torch.tensor([[[0.0], [1.0]]], dtype=torch.float64) + options = TrapezoidalPlanOptions( + profile="double_s", + constraints={"velocity": 0.8, "acceleration": 1.0, "jerk": 2.0}, + sample_interval=2001, + ) + + result = _plan_linear_profiles(waypoints, options) + acceleration_jump = torch.diff(result.accelerations[0, :, 0]).abs().max() + + assert float(acceleration_jump) < 0.01 + + +def test_double_s_respects_sampled_jerk_limit() -> None: + jerk_limit = 2.0 + waypoints = torch.tensor([[[0.0], [1.0]]], dtype=torch.float64) + options = TrapezoidalPlanOptions( + profile="double_s", + constraints={ + "velocity": 0.8, + "acceleration": 1.0, + "jerk": jerk_limit, + }, + sample_interval=4001, + ) + + result = _plan_linear_profiles(waypoints, options) + sampled_jerk = torch.diff(result.accelerations[0, :, 0]) / result.dt[0, 1:] + + assert float(sampled_jerk.abs().max()) <= jerk_limit + 1e-9 + + +@pytest.mark.parametrize( + ("distance", "velocity", "acceleration", "jerk", "expected_duration"), + [ + (1.0, 0.2, 0.5, 2.0, 5.7065), + (0.01, 0.2, 0.5, 2.0, 0.5483947163846388), + (0.1, 1.0, 1.0, 1.0, 1.4898240646067618), + (1.0, 1.0, 1.0, 1.0, 3.214067124738025), + ], +) +def test_double_s_duration_matches_reference_golden_values( + distance: float, + velocity: float, + acceleration: float, + jerk: float, + expected_duration: float, +) -> None: + """Match values generated by the reference trajectory implementation.""" + waypoints = torch.tensor([[[0.0], [distance]]], dtype=torch.float64) + options = TrapezoidalPlanOptions( + profile="double_s", + constraints={ + "velocity": velocity, + "acceleration": acceleration, + "jerk": jerk, + }, + sample_interval=101, + backend="torch", + ) + + result = _plan_linear_profiles(waypoints, options) + + assert result.duration.item() == pytest.approx(expected_duration, abs=1e-12) + + +def test_double_s_multi_waypoint_samples_match_reference() -> None: + """Cover segment lookup and derivative scaling with C++ golden samples.""" + waypoints = torch.tensor( + [[[0.0, 0.0], [0.4, -0.2], [0.1, 0.3]]], dtype=torch.float64 + ) + options = TrapezoidalPlanOptions( + profile="double_s", + constraints={"velocity": 0.6, "acceleration": 1.2, "jerk": 3.0}, + sample_interval=5, + backend="torch", + ) + result = _plan_linear_profiles(waypoints, options) + + assert result.duration.item() == pytest.approx(3.407298139310801, abs=1e-12) + assert torch.allclose(result.positions[0, 0], waypoints[0, 0]) + boundary_matches = torch.isclose( + result.positions[0], waypoints[0, 1], atol=1e-8, rtol=0.0 + ).all(dim=1) + assert boundary_matches.any() + boundary_index = boundary_matches.nonzero()[0, 0] + assert torch.allclose(result.positions[0, boundary_index], waypoints[0, 1]) + assert torch.allclose(result.positions[0, -1], waypoints[0, -1]) + assert torch.allclose( + result.velocities[0, 0], torch.zeros(2, dtype=result.velocities.dtype) + ) + assert torch.allclose( + result.velocities[0, boundary_index], + torch.zeros(2, dtype=result.velocities.dtype), + atol=1e-8, + ) + assert torch.allclose( + result.velocities[0, -1], torch.zeros(2, dtype=result.velocities.dtype) + ) + assert torch.all(result.dt[0, 1:] >= 0.0) + + +def test_double_s_short_move_breakpoints_match_reference() -> None: + """Lock the discrete acceleration-reduction branch to C++ behavior.""" + distance = 0.01 + velocity_limit = torch.tensor([[0.2 / distance]], dtype=torch.float64) + acceleration_limit = torch.tensor([[0.5 / distance]], dtype=torch.float64) + jerk_limit = torch.tensor([[2.0 / distance]], dtype=torch.float64) + profile = _build_double_s_profile(velocity_limit, acceleration_limit, jerk_limit) + expected_unscaled_durations = torch.tensor( + [ + 0.13286025, + 0.0057620338537815, + 0.13286025, + 0.0, + 0.13286025, + 0.0057620338537815, + 0.13286025, + ], + dtype=torch.float64, + ) + + assert torch.allclose( + profile.durations[0, 0], expected_unscaled_durations, atol=1e-14 + ) + + +def test_per_joint_limits_are_projected_onto_path() -> None: + velocity_limits = torch.tensor([0.25, 0.8], dtype=torch.float64) + acceleration_limits = torch.tensor([0.5, 1.5], dtype=torch.float64) + waypoints = torch.tensor([[[0.0, 0.0], [0.5, 1.0]]], dtype=torch.float64) + options = TrapezoidalPlanOptions( + constraints={ + "velocity": velocity_limits, + "acceleration": acceleration_limits, + "jerk": torch.tensor([1.0, 3.0], dtype=torch.float64), + }, + sample_interval=1001, + ) + + result = _plan_linear_profiles(waypoints, options) + + assert torch.all(result.velocities.abs().amax(dim=(0, 1)) <= velocity_limits) + assert torch.all(result.accelerations.abs().amax(dim=(0, 1)) <= acceleration_limits) + + +@pytest.mark.parametrize("profile", ["trapezoidal", "double_s"]) +def test_warp_backend_matches_torch_on_cpu(profile: str) -> None: + pytest.importorskip("warp") + waypoints = torch.tensor( + [ + [[0.0, 0.0], [0.4, -0.2], [0.1, 0.3]], + [[0.1, -0.1], [0.2, 0.5], [0.8, 0.2]], + ], + dtype=torch.float32, + ) + common = { + "profile": profile, + "constraints": {"velocity": 0.6, "acceleration": 1.2, "jerk": 3.0}, + "sample_interval": 257, + } + + torch_result = _plan_linear_profiles( + waypoints, TrapezoidalPlanOptions(**common, backend="torch") + ) + warp_result = _plan_linear_profiles( + waypoints, TrapezoidalPlanOptions(**common, backend="warp") + ) + + assert torch.allclose(warp_result.positions, torch_result.positions, atol=1e-5) + assert torch.allclose(warp_result.velocities, torch_result.velocities, atol=1e-5) + assert torch.allclose( + warp_result.accelerations, torch_result.accelerations, atol=1e-5 + ) + assert torch.equal(warp_result.dt, torch_result.dt) + + +@pytest.mark.parametrize("profile", ["trapezoidal", "double_s"]) +def test_warp_profile_construction_matches_torch(profile: str) -> None: + pytest.importorskip("warp") + velocity = torch.tensor([[0.0, 0.2, 1.0], [0.6, 4.0, 0.1]]) + acceleration = torch.tensor([[1.0, 0.5, 1.0], [1.2, 2.0, 0.3]]) + jerk = torch.tensor([[2.0, 2.0, 1.0], [3.0, 8.0, 0.5]]) + + torch_profile = _build_scalar_profile( + profile_name=profile, + velocity_limit=velocity, + acceleration_limit=acceleration, + jerk_limit=jerk, + backend="torch", + ) + warp_profile = _build_scalar_profile( + profile_name=profile, + velocity_limit=velocity, + acceleration_limit=acceleration, + jerk_limit=jerk, + backend="warp", + ) + + for torch_value, warp_value in zip( + ( + torch_profile.durations, + torch_profile.positions, + torch_profile.velocities, + torch_profile.accelerations, + torch_profile.jerks, + ), + ( + warp_profile.durations, + warp_profile.positions, + warp_profile.velocities, + warp_profile.accelerations, + warp_profile.jerks, + ), + strict=True, + ): + assert torch.allclose(warp_value, torch_value, atol=1e-5, rtol=1e-5) + + +def test_many_segment_path_preserves_constraints_and_endpoints() -> None: + waypoint_count = 129 + parameter = torch.linspace(0.0, 4.0, waypoint_count, dtype=torch.float64) + path = torch.stack( + [parameter, torch.sin(parameter), 0.5 * torch.cos(parameter)], dim=-1 + ) + waypoints = torch.stack([path, path + torch.tensor([0.1, -0.2, 0.3])]) + options = TrapezoidalPlanOptions( + profile="double_s", + constraints={"velocity": 0.7, "acceleration": 1.4, "jerk": 4.0}, + sample_interval=1025, + backend="torch", + ) + + result = _plan_linear_profiles(waypoints, options) + + assert torch.allclose(result.positions[:, 0], waypoints[:, 0]) + assert torch.allclose(result.positions[:, -1], waypoints[:, -1]) + assert float(result.velocities.abs().max()) <= 0.7 + 1e-9 + assert float(result.accelerations.abs().max()) <= 1.4 + 1e-9 + + +def test_explicit_warp_backend_rejects_float64() -> None: + waypoints = torch.tensor([[[0.0], [1.0]]], dtype=torch.float64) + options = TrapezoidalPlanOptions(backend="warp") + + with pytest.raises(ValueError, match="requires float32"): + _plan_linear_profiles(waypoints, options) + + +def test_multiple_segments_stop_at_internal_waypoint() -> None: + waypoints = torch.tensor([[[0.0], [0.5], [-0.25]]], dtype=torch.float64) + options = TrapezoidalPlanOptions( + profile="trapezoidal", + constraints={"velocity": 0.4, "acceleration": 0.8, "jerk": 2.0}, + sample_interval=801, + ) + + result = _plan_linear_profiles(waypoints, options) + nearest = torch.argmin((result.positions[0, :, 0] - 0.5).abs()) + + assert result.positions[0, nearest, 0] == pytest.approx(0.5, abs=1e-5) + assert result.velocities[0, nearest, 0] == pytest.approx(0.0, abs=2e-3) + + +def test_collinear_compression_removes_only_same_direction_points() -> None: + waypoints = torch.tensor( + [ + [[0.0, 0.0], [0.25, 0.0], [0.5, 0.0], [1.0, 0.0]], + [[0.0, 0.0], [0.5, 0.0], [0.25, 0.0], [1.0, 0.0]], + ] + ) + + compressed = _compress_collinear_waypoints(waypoints, tolerance=1e-5) + + assert compressed.shape == (2, 4, 2) + assert torch.equal(compressed[0, 0], waypoints[0, 0]) + assert torch.equal(compressed[0, 1], waypoints[0, -1]) + assert torch.equal(compressed[1], waypoints[1]) + + +def test_disabling_waypoint_stops_shortens_dense_straight_path() -> None: + waypoints = torch.linspace(0.0, 1.0, 21, dtype=torch.float64).reshape(1, 21, 1) + stopped = TrapezoidalPlanOptions( + profile="double_s", + sample_interval=501, + stop_at_waypoints=True, + ) + continuous = TrapezoidalPlanOptions( + profile="double_s", + sample_interval=501, + stop_at_waypoints=False, + ) + + stopped_result = _plan_linear_profiles(waypoints, stopped) + continuous_result = _plan_linear_profiles(waypoints, continuous) + + assert continuous_result.duration < stopped_result.duration + assert torch.allclose(continuous_result.positions[:, 0], waypoints[:, 0]) + assert torch.allclose(continuous_result.positions[:, -1], waypoints[:, -1]) + + +def test_collinear_compression_pads_batch_rows_with_final_point() -> None: + waypoints = torch.tensor( + [ + [[0.0], [0.25], [0.5], [1.0]], + [[0.0], [0.5], [0.25], [1.0]], + ] + ) + + compressed = _compress_collinear_waypoints(waypoints, tolerance=1e-5) + + assert compressed.shape == waypoints.shape + assert torch.all(compressed[0, 1:] == 1.0) + assert torch.equal(compressed[1], waypoints[1]) + + +def test_duplicate_waypoint_run_preserves_adjacent_corner() -> None: + waypoints = torch.tensor([[[0.0, 0.0], [1.0, 0.0], [1.0, 0.0], [1.0, 1.0]]]) + + compressed = _compress_collinear_waypoints(waypoints, tolerance=1e-5) + + assert torch.equal(compressed[0, :3], waypoints[0, (0, 1, 3)]) + + +def test_quantity_sampling_includes_exact_internal_waypoint_and_stop() -> None: + waypoints = torch.tensor([[[0.0, 0.0], [1.0, 0.0], [1.0, 1.0]]]) + options = TrapezoidalPlanOptions(sample_interval=4) + + result = _plan_linear_profiles(waypoints, options) + + waypoint_matches = torch.isclose( + result.positions[0], waypoints[0, 1], atol=1e-6, rtol=0.0 + ).all(dim=1) + assert waypoint_matches.any() + waypoint_index = waypoint_matches.nonzero()[0, 0] + assert torch.allclose(result.velocities[0, waypoint_index], torch.zeros(2)) + + +def test_quantity_sampling_rejects_fewer_samples_than_waypoints() -> None: + waypoints = torch.tensor([[[0.0], [1.0], [2.0]]]) + + with pytest.raises(ValueError, match="at least one sample per retained waypoint"): + _plan_linear_profiles(waypoints, TrapezoidalPlanOptions(sample_interval=2)) + + +def test_stationary_path_returns_zero_duration_hold() -> None: + waypoints = torch.full((2, 3, 4), 0.25) + options = TrapezoidalPlanOptions(sample_interval=10) + + result = _plan_linear_profiles(waypoints, options) + + assert result.success.all() + assert torch.all(result.positions == 0.25) + assert torch.count_nonzero(result.velocities) == 0 + assert torch.count_nonzero(result.accelerations) == 0 + assert torch.count_nonzero(result.dt) == 0 + + +def test_stationary_fast_path_skips_profile_construction(monkeypatch) -> None: + def fail_if_called(*args, **kwargs) -> None: + raise AssertionError("stationary paths must not construct a profile") + + monkeypatch.setattr( + trapezoidal_module, "_build_trapezoidal_profile", fail_if_called + ) + waypoints = torch.full((2, 3, 4), 0.25) + + result = _plan_linear_profiles( + waypoints, TrapezoidalPlanOptions(sample_interval=10) + ) + + assert result.positions.shape == (2, 10, 4) + assert torch.count_nonzero(result.dt) == 0 + + +def test_minimum_duration_slows_trajectory_without_changing_path() -> None: + requested_duration = 5.0 + waypoints = torch.tensor([[[0.0, 0.0], [0.3, -0.2]]], dtype=torch.float64) + base = TrapezoidalPlanOptions(profile="double_s", sample_interval=301) + slowed = TrapezoidalPlanOptions( + profile="double_s", + sample_interval=301, + minimum_duration=requested_duration, + ) + + base_result = _plan_linear_profiles(waypoints, base) + slowed_result = _plan_linear_profiles(waypoints, slowed) + + assert slowed_result.duration.item() == pytest.approx(requested_duration) + assert torch.allclose(base_result.positions, slowed_result.positions, atol=1e-6) + assert slowed_result.velocities.abs().max() < base_result.velocities.abs().max() + + +def test_stationary_path_can_hold_for_minimum_duration() -> None: + requested_duration = 2.5 + waypoints = torch.full((1, 2, 3), 0.25) + options = TrapezoidalPlanOptions( + sample_interval=11, + minimum_duration=requested_duration, + ) + + result = _plan_linear_profiles(waypoints, options) + + assert result.duration.item() == pytest.approx(requested_duration) + assert torch.all(result.positions == 0.25) + assert torch.count_nonzero(result.velocities) == 0 + + +def test_stationary_time_sampling_preserves_requested_hold_duration() -> None: + waypoints = torch.full((2, 2, 3), 0.25) + options = TrapezoidalPlanOptions( + sample_method=TrajectorySampleMethod.TIME, + sample_interval=0.2, + minimum_duration=0.5, + ) + + result = _plan_linear_profiles(waypoints, options) + + assert result.positions.shape == (2, 4, 3) + assert torch.allclose(result.duration, torch.full((2,), 0.5)) + assert torch.allclose(result.dt[:, 1:], torch.full((2, 3), 0.5 / 3.0)) + + +def test_time_sampling_tail_pads_shorter_rows() -> None: + waypoints = torch.tensor([[[0.0], [1.0]], [[0.0], [0.1]]]) + options = TrapezoidalPlanOptions( + sample_method=TrajectorySampleMethod.TIME, + sample_interval=0.05, + ) + + result = _plan_linear_profiles(waypoints, options) + + assert result.positions.shape[0] == 2 + assert torch.allclose(result.positions[:, -1], waypoints[:, -1]) + assert result.duration[0] > result.duration[1] + assert torch.count_nonzero(result.dt[1] == 0.0) > 1 + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"profile": "unknown"}, "profile"), + ({"sample_interval": 1}, "at least 2"), + ({"sample_interval": 2.5}, "integer"), + ({"minimum_duration": float("inf")}, "minimum_duration"), + ({"collinearity_tolerance": 1.0}, "collinearity_tolerance"), + ({"backend": "unknown"}, "backend"), + ({"constraints": {"velocity": 1.0}}, "missing required keys"), + ], +) +def test_options_reject_invalid_configuration(kwargs: dict, message: str) -> None: + with pytest.raises(ValueError, match=message): + TrapezoidalPlanOptions(**kwargs) + + +def test_rejects_nonpositive_joint_limit() -> None: + waypoints = torch.tensor([[[0.0], [1.0]]]) + options = TrapezoidalPlanOptions( + constraints={"velocity": 0.0, "acceleration": 1.0, "jerk": 1.0} + ) + + with pytest.raises(ValueError, match="finite positive"): + _plan_linear_profiles(waypoints, options) + + +def test_planner_adapts_batched_plan_states() -> None: + planner = TrapezoidalPlanner.__new__(TrapezoidalPlanner) + planner.device = torch.device("cpu") + planner.cfg = SimpleNamespace(planner_type="trapezoidal") + states = [ + PlanState.from_qpos(torch.tensor([[0.0, 0.0]])), + PlanState.from_qpos(torch.tensor([[0.5, -0.25]])), + ] + + result = planner.plan(states, TrapezoidalPlanOptions(sample_interval=20)) + + assert result.positions.shape == (1, 20, 2) + assert result.dt.shape == (1, 20) + assert result.is_all_success() diff --git a/tests/sim/solvers/test_opw_solver.py b/tests/sim/solvers/test_opw_solver.py index 3db7e9fb2..ad225ee41 100644 --- a/tests/sim/solvers/test_opw_solver.py +++ b/tests/sim/solvers/test_opw_solver.py @@ -188,6 +188,31 @@ def test_ik(self, arm_name: str): assert res[0] == False assert ik_qpos.shape == (1, dof) + @pytest.mark.parametrize("arm_name", ["left_arm", "right_arm"]) + def test_continuous_batch_ik_reconstructs_fk_path(self, arm_name: str): + """Continuous batch IK preserves a branch across an OPW pose path.""" + qpos_limits = self.robot.get_qpos_limits(name=arm_name) + qpos = grid_sample_qpos_from_limits( + qpos_limits, steps_per_joint=2, device=self.robot.device, max_samples=1 + ) + qpos_path = qpos[None].expand(1, 3, -1).contiguous() + target_path = self.robot.compute_batch_fk( + qpos=qpos_path, name=arm_name, to_matrix=True + ) + + success, solved_path = self.robot.compute_batch_ik( + pose=target_path, + joint_seed=qpos, + name=arm_name, + continuous=True, + ) + + assert success.all() + reconstructed = self.robot.compute_batch_fk( + qpos=solved_path, name=arm_name, to_matrix=True + ) + assert torch.allclose(target_path, reconstructed, atol=5e-3, rtol=5e-3) + def teardown_method(self): """Clean up resources after each test method.""" self.sim.destroy() diff --git a/tests/utils/test_opw_path_kernel.py b/tests/utils/test_opw_path_kernel.py new file mode 100644 index 000000000..318aad711 --- /dev/null +++ b/tests/utils/test_opw_path_kernel.py @@ -0,0 +1,161 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import torch +import warp as wp +import pytest + +from embodichain.utils.warp.kinematics.opw_solver import ( + opw_ik_path_select_kernel, + wp_vec6f, +) + + +def test_opw_path_selector_preserves_temporal_branch_continuity() -> None: + """The path selector must seed each sample from its previous result.""" + wp.init() + candidates = torch.zeros(1, 3, 8, 6) + validity = torch.zeros(1, 3, 8, dtype=torch.int32) + candidates[0, :, 0, 0] = torch.tensor((0.1, 0.2, 0.3)) + candidates[0, :, 1, 0] = torch.tensor((2.0, 1.9, 1.8)) + validity[:, :, :2] = 1 + output = torch.empty(1, 3, 6) + success = torch.empty(1, 3, dtype=torch.int32) + lower = wp_vec6f(-3.14, -3.14, -3.14, -3.14, -3.14, -3.14) + upper = wp_vec6f(3.14, 3.14, 3.14, 3.14, 3.14, 3.14) + + wp.launch( + kernel=opw_ik_path_select_kernel, + dim=1, + inputs=[ + wp.from_torch(candidates), + wp.from_torch(validity), + wp.from_torch(torch.zeros(1, 6)), + wp_vec6f(1.0, 1.0, 1.0, 1.0, 1.0, 1.0), + lower, + upper, + 0.0, + ], + outputs=[wp.from_torch(output), wp.from_torch(success)], + device="cpu", + ) + + assert torch.equal(success, torch.ones_like(success)) + assert torch.allclose(output[0, :, 0], torch.tensor((0.1, 0.2, 0.3))) + + +def test_opw_path_selector_uses_raw_valid_equivalent() -> None: + """A valid raw representation remains available when the nearest turn is blocked.""" + wp.init() + candidates = torch.zeros(1, 1, 8, 6) + candidates[0, 0, 0, 0] = -3.0 + validity = torch.zeros(1, 1, 8, dtype=torch.int32) + validity[0, 0, 0] = 1 + initial_seed = torch.zeros(1, 6) + initial_seed[0, 0] = 3.0 + output = torch.empty(1, 1, 6) + success = torch.empty(1, 1, dtype=torch.int32) + lower = wp_vec6f(-3.1, -3.1, -3.1, -3.1, -3.1, -3.1) + upper = wp_vec6f(3.1, 3.1, 3.1, 3.1, 3.1, 3.1) + + wp.launch( + kernel=opw_ik_path_select_kernel, + dim=1, + inputs=[ + wp.from_torch(candidates), + wp.from_torch(validity), + wp.from_torch(initial_seed), + wp_vec6f(1.0, 1.0, 1.0, 1.0, 1.0, 1.0), + lower, + upper, + 0.0, + ], + outputs=[wp.from_torch(output), wp.from_torch(success)], + device="cpu", + ) + + assert success.item() == 1 + assert output[0, 0, 0].item() == pytest.approx(-3.0) + + +def test_opw_path_selector_uses_equivalent_outside_safety_margin() -> None: + """A valid equivalent is selected when the nearest turn is in the margin.""" + wp.init() + candidates = torch.zeros(1, 1, 8, 6) + candidates[0, 0, 0, 0] = -6.0 + validity = torch.zeros(1, 1, 8, dtype=torch.int32) + validity[0, 0, 0] = 1 + initial_seed = torch.zeros(1, 6) + initial_seed[0, 0] = 0.3 + output = torch.empty(1, 1, 6) + success = torch.empty(1, 1, dtype=torch.int32) + lower = wp_vec6f(-6.4, -6.4, -6.4, -6.4, -6.4, -6.4) + upper = wp_vec6f(0.4, 0.4, 0.4, 0.4, 0.4, 0.4) + safe_margin = 0.2 + + wp.launch( + kernel=opw_ik_path_select_kernel, + dim=1, + inputs=[ + wp.from_torch(candidates), + wp.from_torch(validity), + wp.from_torch(initial_seed), + wp_vec6f(1.0, 1.0, 1.0, 1.0, 1.0, 1.0), + lower, + upper, + safe_margin, + ], + outputs=[wp.from_torch(output), wp.from_torch(success)], + device="cpu", + ) + + assert success.item() == 1 + assert output[0, 0, 0].item() == pytest.approx(-6.0) + + +def test_opw_path_selector_uses_other_valid_periodic_equivalent() -> None: + """A distant seed can select a different valid 2π representation.""" + wp.init() + candidates = torch.zeros(1, 1, 8, 6) + validity = torch.zeros(1, 1, 8, dtype=torch.int32) + validity[0, 0, 0] = 1 + initial_seed = torch.zeros(1, 6) + initial_seed[0, 0] = 20.0 + output = torch.empty(1, 1, 6) + success = torch.empty(1, 1, dtype=torch.int32) + lower = wp_vec6f(-0.1, -3.1, -3.1, -3.1, -3.1, -3.1) + upper = wp_vec6f(6.4, 3.1, 3.1, 3.1, 3.1, 3.1) + + wp.launch( + kernel=opw_ik_path_select_kernel, + dim=1, + inputs=[ + wp.from_torch(candidates), + wp.from_torch(validity), + wp.from_torch(initial_seed), + wp_vec6f(1.0, 1.0, 1.0, 1.0, 1.0, 1.0), + lower, + upper, + 0.0, + ], + outputs=[wp.from_torch(output), wp.from_torch(success)], + device="cpu", + ) + + assert success.item() == 1 + assert output[0, 0, 0].item() == pytest.approx(2.0 * torch.pi, abs=1e-5)