From 0196e3beeb775985bfbee77abf8154a5a9c43e3a Mon Sep 17 00:00:00 2001 From: yuecideng Date: Thu, 10 Sep 2026 08:07:50 +0000 Subject: [PATCH 1/5] perf(workspace): batch kinematics and reuse analytic IK buffers --- agent_context/topics/ik-solvers/ik-solvers.md | 17 + .../topics/robot-system/robot-system.md | 3 + .../robot-workspace/analysis-and-cache.md | 34 +- .../topics/robot-workspace/robot-workspace.md | 2 +- embodichain/lab/scripts/analyze_workspace.py | 19 +- .../lab/sim/motion/solvers/_buffers.py | 92 +++++ .../lab/sim/motion/solvers/base_solver.py | 55 +++ .../lab/sim/motion/solvers/opw_solver.py | 55 ++- .../lab/sim/motion/solvers/ur_solver.py | 15 +- .../lab/sim/motion/workspace/analyzer.py | 312 +++++++++++---- .../workspace/configs/sampling_config.py | 4 +- .../motion/workspace/samplers/base_sampler.py | 8 +- .../workspace/samplers/gaussian_sampler.py | 13 +- .../workspace/samplers/importance_sampler.py | 27 +- .../motion/workspace/samplers/lhs_sampler.py | 6 +- .../workspace/samplers/random_sampler.py | 6 +- .../workspace/samplers/sobol_sampler.py | 199 ++-------- embodichain/lab/sim/objects/robot.py | 31 +- .../benchmark_robot_workspace.py | 354 ++++++++++++++++++ .../workspace_optimization_report.md | 86 +++++ .../motion/solvers/test_analytic_batching.py | 225 +++++++++++ .../workspace/test_analysis_batching.py | 293 +++++++++++++++ 22 files changed, 1535 insertions(+), 321 deletions(-) create mode 100644 embodichain/lab/sim/motion/solvers/_buffers.py create mode 100644 scripts/benchmark/workspace_analyzer/benchmark_robot_workspace.py create mode 100644 scripts/benchmark/workspace_analyzer/workspace_optimization_report.md create mode 100644 tests/sim/motion/solvers/test_analytic_batching.py create mode 100644 tests/sim/motion/workspace/test_analysis_batching.py diff --git a/agent_context/topics/ik-solvers/ik-solvers.md b/agent_context/topics/ik-solvers/ik-solvers.md index bff1d9c56..7ea98be11 100644 --- a/agent_context/topics/ik-solvers/ik-solvers.md +++ b/agent_context/topics/ik-solvers/ik-solvers.md @@ -132,3 +132,20 @@ solver interfaces. The compute kernels do not import simulation modules. `utils/warp/kinematics/*_solver.py` are compatibility aliases. Validate kernel import/compilation with `tests/compute/test_imports.py` and solver behavior with the corresponding `tests/sim/motion/solvers/` tests. + +## Batch adapters and analytic scratch memory + +`BaseSolver.get_fk_batch()` and `get_ik_batch()` flatten and restore arbitrary +leading batch axes in the chain-root frame. Nearest IK returns a boolean success +mask with the leading shape and qpos with a final `dof` axis; existing concrete +`get_fk`/`get_ik` signatures remain available. Robot owns local-arena/root frame +conversion and broadcasts root transforms without materializing per-target copies. + +UR/OPW reuse internal candidate buffers through `solvers/_buffers.py`. +`prepare_buffers(max_batch)` reserves a high-water capacity; allocation is lazy +and later larger calls grow it. Public results remain independent of subsequent +calls, including `return_all_solutions=True`. Calls are serialized by a lock and +CUDA events; Warp kernels use the current Torch stream. Buffer storage is released +with the solver. UR retains all 512 periodic candidates and the existing nearest +selection; OPW retains eight candidates. OPW packs live joint limits in one host +transfer per call, so limit updates do not require cache invalidation. diff --git a/agent_context/topics/robot-system/robot-system.md b/agent_context/topics/robot-system/robot-system.md index dba892316..3da3f3810 100644 --- a/agent_context/topics/robot-system/robot-system.md +++ b/agent_context/topics/robot-system/robot-system.md @@ -40,6 +40,9 @@ A `Robot` is instantiated with a `RobotCfg` and a list of DexSim `Articulation` - `motion` subpackages load lazily, and workspace analyzer/visualization APIs retain a separate lazy export boundary. Do not introduce offline analysis imports into the Robot initialization path. +- `compute_batch_fk`/`compute_batch_ik` broadcast per-environment root transforms + and delegate batch shape handling to solver adapters; see + [solver batch contracts](../ik-solvers/ik-solvers.md#batch-adapters-and-analytic-scratch-memory). - Focused workspace coverage lives under `tests/sim/motion/workspace/`; solver tests live under `tests/sim/motion/solvers/`. diff --git a/agent_context/topics/robot-workspace/analysis-and-cache.md b/agent_context/topics/robot-workspace/analysis-and-cache.md index aaa733ef0..7fb17b223 100644 --- a/agent_context/topics/robot-workspace/analysis-and-cache.md +++ b/agent_context/topics/robot-workspace/analysis-and-cache.md @@ -4,7 +4,7 @@ ### Offline analysis, cache, and runtime path -1. `WorkspaceAnalysisConfig` chooses `JOINT_SPACE`, `CARTESIAN_SPACE`, or `PLANE_SAMPLING`, +1. `WorkspaceAnalyzerConfig` chooses `JOINT_SPACE`, `CARTESIAN_SPACE`, or `PLANE_SAMPLING`, plus sampler seed/count/batch size, cache, constraints, visualization, metrics, IK seeds/reference pose, control part, and plane settings. 2. Analysis is designed for one simulation environment; with multiple environments it warns @@ -14,7 +14,7 @@ 3. Joint-space mode samples within selected joint limits and applies FK; stored `workspace_points` and `joint_configurations` contain only valid aligned rows. 4. Cartesian mode samples requested XYZ bounds, or infers bounds from 1,000 random FK samples - plus margin/fallback, then runs seeded IK. It stores all sampled points and reachability + plus margin/fallback using one batched FK call, then runs seeded IK. It stores all sampled points and reachability outputs plus best joint configurations aligned to reachable points. 5. Plane mode creates samples on the configured plane and follows the same IK/result alignment path. 6. `WorkspaceAnalyzer.analyze()` checks `ResultsCache` before computation unless force is set, @@ -72,3 +72,33 @@ | Sampled pose is reachable but motion collides | Workspace is kinematic only; add motion planning/collision validation at the consuming layer. | | `workspace-cache list/clean` cannot see ResultsCache entries | It targets legacy session caches; use analyzer preview/cache directory semantics for result entries. | | Concurrent analyzers expose corrupt/incomplete cache | `ResultsCache.save()` writes result then metadata directly with no observed lock or temp-file replacement. Add atomic write/lock coverage before relying on shared concurrent writers. | + +### Sampling and allocation controls + +- `SamplingConfig` and `analyze-workspace --sampler` default to scrambled Sobol. + Explicit random/grid selections remain available. Sobol draws continue a + dimension-specific sequence; recreate the sampler to replay it. Samplers use + private RNGs and do not reset the application's Torch random state. +- `sample_within_constraints=True` (`--sample-within-constraints`) refills random, + Sobol or LHS proposals inside the permitted domain. Cartesian box/ground bounds + are intersected before generation; sphere proposals use a volume-uniform + transform; plane proposals are filtered after projection. Exclusion zones are + respected. Exhausting `max_sampling_rounds` raises instead of returning an + incomplete cache. Reachability is conditional on this permitted domain, so its + percentage is not interchangeable with the default full-domain statistic. +- `retain_diagnostics=False` (`--compact-results`) retains reachable positions, + aligned qpos and scores, while dropping all-target points and masks from the + returned result and archive. Counts and metrics remain available. Runtime and + preview accept both formats. This reduces retained/output storage; analysis + still materializes the original sampled domain before compaction. +- IK batches allocate target poses once across point/seed axes. One seed bypasses + seed-axis reductions and selection indexing. The analyzer never changes the + process-wide logging level. Analyzer FK/IK explicitly selects environment zero + and defaults to the robot device when no simulation manager is supplied. +- Cache keys include a sampling revision, compact/domain settings, geometry and + excluded zones. These changes intentionally invalidate older result entries. +- Focused logic tests: `tests/sim/motion/workspace/test_analysis_batching.py`. + Real analytic solver tests (including CUDA stream reuse) are in + `tests/sim/motion/solvers/test_analytic_batching.py`. The renderer-free benchmark + is `scripts/benchmark/workspace_analyzer/benchmark_robot_workspace.py`; it uses + real asset chains and Robot methods with fixed root snapshots and eager FK. diff --git a/agent_context/topics/robot-workspace/robot-workspace.md b/agent_context/topics/robot-workspace/robot-workspace.md index 34534d8a6..483613b29 100644 --- a/agent_context/topics/robot-workspace/robot-workspace.md +++ b/agent_context/topics/robot-workspace/robot-workspace.md @@ -10,7 +10,7 @@ Paths below are relative to `embodichain/lab/sim/` unless qualified. | Request | Owner | |---|---| -| Offline analysis | `workspace/analyzer.py`: `WorkspaceAnalyzer`, `WorkspaceAnalysisConfig` | +| Offline analysis | `workspace/analyzer.py`: `WorkspaceAnalyzer`, `WorkspaceAnalyzerConfig` | | Result-cache identity and persistence | `workspace/caches/results_cache.py`: `ResultsCache` | | Cache loading / point or voxel sampling | `workspace/runtime.py`: `RobotWorkspace`, `WorkspaceSample` | | Runtime binding config | `workspace/cfg.py`: `RobotWorkspaceCfg`; `cfg.py`: `RobotCfg.workspace_cfg` | diff --git a/embodichain/lab/scripts/analyze_workspace.py b/embodichain/lab/scripts/analyze_workspace.py index bc28c1b72..25e3a0605 100644 --- a/embodichain/lab/scripts/analyze_workspace.py +++ b/embodichain/lab/scripts/analyze_workspace.py @@ -419,6 +419,8 @@ def build_analyzer_config( constraint=constraint, ik_samples_per_point=args.ik_samples_per_point, control_part_name=control_part_name, + retain_diagnostics=not getattr(args, "compact_results", False), + sample_within_constraints=getattr(args, "sample_within_constraints", False), ) if mode == AnalysisMode.PLANE_SAMPLING: @@ -565,6 +567,8 @@ def _preview_points_and_colors( points = np.asarray(arrays["workspace_points"]) elif "all_points" in arrays: points = np.asarray(arrays["all_points"]) + elif "reachable_points" in arrays: + points = np.asarray(arrays["reachable_points"]) else: points = np.asarray(arrays[next(iter(arrays))]) @@ -904,8 +908,8 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: "--sampler", type=str, choices=["random", "sobol", "halton", "lhs", "uniform", "gaussian"], - default="random", - help="Sampling strategy (default: random).", + default="sobol", + help="Sampling strategy (default: sobol).", ) sampling.add_argument( "--seed", type=int, default=42, help="Random seed (default: 42)." @@ -914,6 +918,17 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: "--batch-size", type=int, default=1000, help="FK/IK batch size (default: 1000)." ) + sampling.add_argument( + "--sample-within-constraints", + action="store_true", + help="Refill samples inside the permitted domain; changes reachability denominator.", + ) + sampling.add_argument( + "--compact-results", + action="store_true", + help="Store reachable points, qpos and aligned scores without rejected-point diagnostics.", + ) + # --- Workspace / plane -------------------------------------------------- space = parser.add_argument_group("Workspace bounds & plane") space.add_argument( diff --git a/embodichain/lab/sim/motion/solvers/_buffers.py b/embodichain/lab/sim/motion/solvers/_buffers.py new file mode 100644 index 000000000..7d020904e --- /dev/null +++ b/embodichain/lab/sim/motion/solvers/_buffers.py @@ -0,0 +1,92 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Bounded analytic IK scratch storage with caller-owned results.""" + +from __future__ import annotations + +from contextlib import contextmanager, nullcontext +from functools import wraps +from threading import RLock + +import torch +import warp as wp + +__all__: list[str] = [] + + +class _IKBuffers: + """Reuse candidate storage, serializing callers across threads and streams.""" + + def __init__(self, device: torch.device) -> None: + self.device = torch.device(device) + self.capacity = 0 + self.arrays: dict[str, torch.Tensor] = {} + self.lock = RLock() + self.event: torch.cuda.Event | None = None + + def reserve(self, max_batch: int) -> None: + if max_batch < 1: + raise ValueError("max_batch must be positive") + with self.lock: + self.capacity = max(self.capacity, max_batch) + + @contextmanager + def borrow(self): + with self.lock: + stream = ( + torch.cuda.current_stream(self.device) + if self.device.type == "cuda" + else None + ) + if stream is not None and self.event is not None: + stream.wait_event(self.event) + scope = ( + wp.ScopedStream(wp.stream_from_torch(stream)) + if stream is not None + else nullcontext() + ) + with scope: + try: + yield + finally: + if stream is not None: + self.event = torch.cuda.Event() + self.event.record(stream) + + def zeros(self, name: str, batch: int, width: int, dtype: torch.dtype) -> wp.array: + self.capacity = max(self.capacity, batch) + size = self.capacity * width + tensor = self.arrays.get(name) + if tensor is None or tensor.numel() < size: + tensor = torch.empty(size, dtype=dtype, device=self.device) + self.arrays[name] = tensor + if self.device.type == "cuda": + tensor.record_stream(torch.cuda.current_stream(self.device)) + view = tensor[: batch * width] + view.zero_() + return wp.from_torch(view) + + +def _with_ik_buffers(method): + """Keep candidate buffers borrowed until result copies have been queued.""" + + @wraps(method) + def wrapped(self, *args, **kwargs): + with self._ik_buffers.borrow(): + return method(self, *args, **kwargs) + + return wrapped diff --git a/embodichain/lab/sim/motion/solvers/base_solver.py b/embodichain/lab/sim/motion/solvers/base_solver.py index 89ee5bb43..03ddf0a0f 100644 --- a/embodichain/lab/sim/motion/solvers/base_solver.py +++ b/embodichain/lab/sim/motion/solvers/base_solver.py @@ -29,6 +29,8 @@ from embodichain.lab.sim.utility.solver_utils import create_pk_serial_chain +from ._buffers import _IKBuffers + @configclass class SolverCfg: @@ -155,6 +157,8 @@ def __init__(self, cfg: SolverCfg = None, device: str = None, **kwargs): else: self.device = device + self._ik_buffers = _IKBuffers(self.device) + self.urdf_path = cfg.urdf_path self.joint_names = cfg.joint_names @@ -218,6 +222,57 @@ def __init__(self, cfg: SolverCfg = None, device: str = None, **kwargs): self._init_qpos_limits() + def prepare_buffers(self, max_batch: int) -> None: + """Reserve reusable analytic IK scratch capacity. + + Args: + max_batch: Maximum flattened target count, including IK seeds. + Storage is allocated lazily by analytic backends and grows when + necessary. Returned FK/IK tensors remain owned by the caller. + """ + self._ik_buffers.reserve(max_batch) + + def get_fk_batch(self, qpos: torch.Tensor) -> torch.Tensor: + """Compute chain-root FK while preserving arbitrary leading batch axes. + + Args: + qpos: Joint positions with shape ``(..., dof)`` on the solver device. + + Returns: + TCP transforms with shape ``(..., 4, 4)``. + """ + if qpos.ndim < 2 or qpos.shape[-1] != self.dof: + raise ValueError(f"Expected batched qpos ending in {self.dof} joints") + shape = qpos.shape[:-1] + return self.get_fk(qpos.reshape(-1, self.dof)).reshape(*shape, 4, 4) + + def get_ik_batch( + self, target_xpos: torch.Tensor, qpos_seed: torch.Tensor | None = None + ) -> tuple[torch.Tensor, torch.Tensor]: + """Compute nearest-solution IK with a uniform batch result contract. + + Args: + target_xpos: Chain-root TCP transforms with shape ``(..., 4, 4)``. + qpos_seed: Optional seeds with matching leading axes and ``dof`` joints. + + Returns: + Boolean success ``(...)`` and joint positions ``(..., dof)``. No + candidate axis is exposed; concrete ``get_ik`` APIs stay unchanged. + """ + if target_xpos.ndim < 3 or target_xpos.shape[-2:] != (4, 4): + raise ValueError("Expected batched target poses ending in (4, 4)") + shape = target_xpos.shape[:-2] + if qpos_seed is not None: + if qpos_seed.shape != (*shape, self.dof): + raise ValueError("Joint seed batch axes must match target poses") + qpos_seed = qpos_seed.reshape(-1, self.dof) + success, qpos = self.get_ik( + target_xpos=target_xpos.reshape(-1, 4, 4), + qpos_seed=qpos_seed, + return_all_solutions=False, + ) + return success.bool().reshape(shape), qpos.reshape(*shape, self.dof) + def set_ik_nearest_weight( self, ik_weight: np.ndarray, joint_ids: np.ndarray | None = None ) -> bool: diff --git a/embodichain/lab/sim/motion/solvers/opw_solver.py b/embodichain/lab/sim/motion/solvers/opw_solver.py index 498c47adc..04537a36d 100644 --- a/embodichain/lab/sim/motion/solvers/opw_solver.py +++ b/embodichain/lab/sim/motion/solvers/opw_solver.py @@ -20,6 +20,7 @@ import numpy as np import warp as wp import polars as pl +from ._buffers import _with_ik_buffers from itertools import product from typing import Union, Tuple, Any, Literal, TYPE_CHECKING @@ -228,6 +229,7 @@ def get_ik( """ return self.get_ik_warp(target_xpos, qpos_seed, return_all_solutions, **kwargs) + @_with_ik_buffers def get_ik_warp( self, target_xpos: torch.Tensor, @@ -250,45 +252,36 @@ def get_ik_warp( """ N_SOL = 8 DOF = 6 - n_sample = target_xpos.shape[0] kernel_device = standardize_device_string(self.device) if target_xpos.shape == (4, 4): target_xpos_batch = target_xpos[None, :, :].to(kernel_device) else: target_xpos_batch = target_xpos.to(kernel_device) + n_sample = target_xpos_batch.shape[0] target_xpos_wp = wp.from_torch(target_xpos_batch.reshape(-1)) - all_qpos_wp = wp.zeros( - n_sample * N_SOL * DOF, - dtype=float, - device=standardize_device_string(kernel_device), - ) - all_ik_valid_wp = wp.zeros( - n_sample * N_SOL, dtype=int, device=standardize_device_string(kernel_device) + all_qpos_wp = self._ik_buffers.zeros( + "qpos", n_sample, N_SOL * DOF, torch.float32 ) + all_ik_valid_wp = self._ik_buffers.zeros("valid", n_sample, N_SOL, torch.int32) # TODO: whether require gradient offsets_ = self.offsets.to(standardize_device_string(kernel_device)) sign_corrections_ = self.sign_corrections.to( standardize_device_string(kernel_device) ) - lower_limits_ = wp_vec6f( - self.lower_qpos_limits[0], - self.lower_qpos_limits[1], - self.lower_qpos_limits[2], - self.lower_qpos_limits[3], - self.lower_qpos_limits[4], - self.lower_qpos_limits[5], - ) - upper_limits_ = wp_vec6f( - self.upper_qpos_limits[0], - self.upper_qpos_limits[1], - self.upper_qpos_limits[2], - self.upper_qpos_limits[3], - self.upper_qpos_limits[4], - self.upper_qpos_limits[5], + # Pack both vectors in one host transfer instead of converting twelve + # CUDA scalars separately. Reading live limits preserves setter/in-place + # updates without an invalidation cache. + limits = ( + torch.stack((self.lower_qpos_limits, self.upper_qpos_limits)) + .detach() + .cpu() + .tolist() ) + lower_limits_ = wp_vec6f(*limits[0]) + upper_limits_ = wp_vec6f(*limits[1]) wp.launch( kernel=opw_ik_kernel, dim=(n_sample), @@ -309,7 +302,7 @@ def get_ik_warp( if return_all_solutions: all_qpos = wp.to_torch(all_qpos_wp).reshape(n_sample, N_SOL, DOF) all_ik_valid = wp.to_torch(all_ik_valid_wp).reshape(n_sample, N_SOL) - return all_ik_valid, all_qpos + return all_ik_valid.clone(), all_qpos.clone() if qpos_seed is not None: if qpos_seed.shape == ( n_sample, @@ -335,15 +328,13 @@ def get_ik_warp( qpos_seed_wp = wp.from_torch(qpos_seed) all_qpos_wp = all_qpos_wp.reshape((n_sample, N_SOL, DOF)) all_ik_valid_wp = all_ik_valid_wp.reshape((n_sample, N_SOL)) - joint_weight = kwargs.get("joint_weight", torch.ones(size=(DOF,), dtype=float)) - joint_weight_wp = wp_vec6f( - joint_weight[0], - joint_weight[1], - joint_weight[2], - joint_weight[3], - joint_weight[4], - joint_weight[5], + joint_weight = kwargs.get("joint_weight") + weights = ( + [1.0] * DOF + if joint_weight is None + else torch.as_tensor(joint_weight).detach().cpu().tolist() ) + joint_weight_wp = wp_vec6f(*weights) best_ik_result_wp = wp.zeros( (n_sample, 6), dtype=float, device=standardize_device_string(kernel_device) ) diff --git a/embodichain/lab/sim/motion/solvers/ur_solver.py b/embodichain/lab/sim/motion/solvers/ur_solver.py index 127e04074..64580b67a 100644 --- a/embodichain/lab/sim/motion/solvers/ur_solver.py +++ b/embodichain/lab/sim/motion/solvers/ur_solver.py @@ -27,6 +27,7 @@ ur_ik_kernel, ) import math +from ._buffers import _with_ik_buffers from embodichain.utils.device_utils import standardize_device_string @@ -133,7 +134,11 @@ def set_tcp(self, tcp: np.ndarray): self._tcp_inv = np.eye(4, dtype=float) self._tcp_inv[:3, :3] = self.tcp_xpos[:3, :3].T self._tcp_inv[:3, 3] = -self._tcp_inv[:3, :3] @ self.tcp_xpos[:3, 3] + self._tcp_inv_tensor = torch.tensor( + self._tcp_inv, dtype=torch.float32, device=self.device + ) + @_with_ik_buffers def get_ik( self, target_xpos: torch.Tensor, @@ -160,7 +165,7 @@ def get_ik( target_xpos_batch = target_xpos[None, :, :] else: target_xpos_batch = target_xpos - tcp_inv = torch.tensor(self._tcp_inv, dtype=torch.float32, device=self.device) + tcp_inv = self._tcp_inv_tensor target_xpos_batch = target_xpos_batch @ tcp_inv[None, :, :] n_sample = target_xpos_batch.shape[0] @@ -178,8 +183,10 @@ def get_ik( wp_device = standardize_device_string(self.device) # Flatten target poses to a 1-D float array for the Warp kernel. xpos_wp = wp.from_torch(target_xpos_batch.reshape(-1)) - all_qpos_wp = wp.zeros(n_sample * N_SOL * DOF, dtype=float, device=wp_device) - all_ik_valid_wp = wp.zeros(n_sample * N_SOL, dtype=int, device=wp_device) + all_qpos_wp = self._ik_buffers.zeros( + "qpos", n_sample, N_SOL * DOF, torch.float32 + ) + all_ik_valid_wp = self._ik_buffers.zeros("valid", n_sample, N_SOL, torch.int32) lower_qpos_limits_wp = wp.from_torch(self.lower_qpos_limits) upper_qpos_limits_wp = wp.from_torch(self.upper_qpos_limits) wp.launch( @@ -208,7 +215,7 @@ def get_ik( ) if return_all_solutions: - return all_solutions_validity, all_solutions + return all_solutions_validity.clone(), all_solutions.clone() # Select ik qpos based on the closest distance to the seed qpos qpos_seed_expanded = qpos_seed.unsqueeze(1).expand(-1, N_SOL, -1) distances = torch.norm( diff --git a/embodichain/lab/sim/motion/workspace/analyzer.py b/embodichain/lab/sim/motion/workspace/analyzer.py index 5478e7239..1819d3e08 100644 --- a/embodichain/lab/sim/motion/workspace/analyzer.py +++ b/embodichain/lab/sim/motion/workspace/analyzer.py @@ -112,6 +112,13 @@ class WorkspaceAnalyzerConfig: metric: MetricConfig = None """Metric configuration.""" + retain_diagnostics: bool = True + """Retain rejected points and per-target diagnostics; disable for runtime caches.""" + sample_within_constraints: bool = False + """Sample only the permitted domain. Reachability is then conditional on that domain.""" + max_sampling_rounds: int = 32 + """Bounded refill rounds for constrained sampling; exhaustion raises ValueError.""" + ik_samples_per_point: int = 1 """For Cartesian mode: number of random joint seeds to try for each Cartesian point.""" reference_pose: Any | None = None @@ -150,6 +157,8 @@ class WorkspaceAnalyzerConfig: def __post_init__(self): """Initialize sub-configs with defaults if not provided.""" + if self.ik_samples_per_point < 1 or self.max_sampling_rounds < 1: + raise ValueError("IK seed count and sampling rounds must be positive") if self.sampling is None: self.sampling = SamplingConfig() if self.cache is None: @@ -197,10 +206,8 @@ def __init__( # Check multi-environment compatibility and add protection self._check_num_envs_compatibility() - # Use sim_manager's device if available, otherwise default to CPU - self.device = ( - sim_manager.device if sim_manager is not None else torch.device("cpu") - ) + # Keep sampling and kinematics on the simulation or robot device. + self.device = sim_manager.device if sim_manager is not None else robot.device # Determine control part name from config self.control_part_name = self._determine_control_part( @@ -345,44 +352,29 @@ def _compute_dynamic_workspace_bounds(self) -> torch.Tensor: RandomSampler, ) - temp_sampler = RandomSampler(seed=self.config.sampling.seed) + temp_sampler = RandomSampler( + seed=self.config.sampling.seed, + device=self.device, + ) # Sample joint space to compute FK bounds joint_samples = temp_sampler.sample(num_samples=1000, bounds=self.qpos_limits) - # Compute FK for all samples with progress tracking - workspace_pts_list = [] - - pbar = self._create_optimized_tqdm( - range(len(joint_samples)), - desc="Computing Workspace Bounds (FK)", - unit="cfg", - color="cyan", - emoji="📏", - ) - - successful_fk = 0 - for i in pbar: - qpos = joint_samples[i : i + 1] # Keep batch dimension - try: - pose = self.robot.compute_fk( - qpos=qpos, - name=self.control_part_name, - to_matrix=True, - ) - position = pose[:, :3, 3] # Extract position - workspace_pts_list.append(position) - successful_fk += 1 - except Exception: - continue - - # Update progress bar with success rate - self._update_progress_with_stats( - pbar, i, successful_fk, metric_name="FK success", show_rate=True + # A single batched FK avoids one Python call and one solver dispatch per + # sample. Failed batches preserve the existing fallback bounds. + try: + poses = self.robot.compute_batch_fk( + qpos=joint_samples.unsqueeze(0), + name=self.control_part_name, + env_ids=[0], + to_matrix=True, ) + workspace_pts = poses[0, :, :3, 3] + except Exception as exc: + logger.log_warning(f"Batched FK failed while estimating bounds: {exc}") + workspace_pts = torch.empty((0, 3), device=self.device) - if workspace_pts_list: - workspace_pts = torch.cat(workspace_pts_list, dim=0) + if len(workspace_pts) > 0: # Compute min/max bounds for each dimension min_bounds = workspace_pts.min(dim=0).values max_bounds = workspace_pts.max(dim=0).values @@ -630,6 +622,119 @@ def sample_joint_space(self, num_samples: int | None = None) -> torch.Tensor: ) return joint_samples + def _sphere_parameters(self) -> tuple[torch.Tensor, float]: + """Resolve explicit sphere geometry without inferring it from sampled points.""" + cfg = self.config + bounds = cfg.constraint_bounds + if bounds is not None: + bounds = torch.as_tensor(bounds, dtype=torch.float32, device=self.device) + center = cfg.sphere_center + if center is None: + if bounds is None: + raise ValueError( + "Sphere sampling requires a center or constraint_bounds" + ) + center = bounds.mean(dim=1) + center = torch.as_tensor(center, dtype=torch.float32, device=self.device) + radius = cfg.sphere_radius + if radius is None: + if bounds is None: + raise ValueError( + "Sphere sampling requires a radius or constraint_bounds" + ) + half = (bounds[:, 1] - bounds[:, 0]) / 2 + if cfg.sphere_radius_mode == "inscribed": + radius = float(half.min()) + elif cfg.sphere_radius_mode == "circumscribed": + radius = float(torch.linalg.vector_norm(half)) + else: + raise ValueError("Unknown sphere_radius_mode") + if center.shape != (3,) or radius <= 0: + raise ValueError("Sphere requires a 3D center and a positive radius") + return center, float(radius) + + def _check_constraints(self, points: torch.Tensor) -> torch.Tensor: + """Apply workspace exclusions and the optional geometric domain.""" + valid = self.constraint_checker.check_constraints(points) + if self.config.constraint_type == "sphere": + center, radius = self._sphere_parameters() + valid &= ((points - center) ** 2).sum(dim=-1) <= radius**2 + elif self.config.constraint_type == "box": + bounds = torch.as_tensor( + self.config.constraint_bounds, device=points.device, dtype=points.dtype + ) + if bounds.shape != (3, 2): + raise ValueError("Box constraint_bounds must have shape (3, 2)") + valid &= ((points >= bounds[:, 0]) & (points <= bounds[:, 1])).all(-1) + elif self.config.constraint_type is not None: + raise ValueError(f"Unknown constraint_type: {self.config.constraint_type}") + return valid + + def _sample_in_domain( + self, bounds: torch.Tensor, num_samples: int, transform=None + ) -> torch.Tensor: + """Generate a bounded number of proposal batches, retaining only valid points.""" + if self.sampler.get_strategy_name() not in ("random", "sobol", "lhs"): + raise ValueError("Constrained sampling supports random, sobol, and lhs") + bounds = bounds.clone().to(dtype=torch.float32, device=self.device) + sphere = self.config.constraint_type == "sphere" and transform is None + if transform is None: + cfg = self.config.constraint + if cfg.min_bounds is not None: + bounds[:, 0] = torch.maximum( + bounds[:, 0], bounds.new_tensor(cfg.min_bounds) + ) + if cfg.max_bounds is not None: + bounds[:, 1] = torch.minimum( + bounds[:, 1], bounds.new_tensor(cfg.max_bounds) + ) + bounds[2, 0] = max(float(bounds[2, 0]), cfg.ground_height) + if self.config.constraint_type == "box": + box = torch.as_tensor( + self.config.constraint_bounds, + device=self.device, + dtype=bounds.dtype, + ) + bounds[:, 0] = torch.maximum(bounds[:, 0], box[:, 0]) + bounds[:, 1] = torch.minimum(bounds[:, 1], box[:, 1]) + if (bounds[:, 0] >= bounds[:, 1]).any(): + raise ValueError("Sampling bounds have no positive-volume intersection") + accepted = [] + remaining = num_samples + for _ in range(self.config.max_sampling_rounds): + count = min(max(remaining, 64), max(num_samples, 64)) + if sphere: + center, radius = self._sphere_parameters() + unit = self.sampler.sample( + num_samples=count, bounds=bounds.new_tensor([[0, 1]] * 3) + ) + z = 2 * unit[:, 0] - 1 + phi = 2 * torch.pi * unit[:, 1] + radial = torch.sqrt(torch.clamp(1 - z * z, min=0)) + direction = torch.stack( + (radial * torch.cos(phi), radial * torch.sin(phi), z), dim=1 + ) + points = center + radius * unit[:, 2:3].pow(1 / 3) * direction + in_bounds = ((points >= bounds[:, 0]) & (points <= bounds[:, 1])).all( + -1 + ) + else: + points = self.sampler.sample(num_samples=count, bounds=bounds) + if transform is not None: + points = transform(points) + in_bounds = torch.ones( + len(points), dtype=torch.bool, device=self.device + ) + points = points[in_bounds & self._check_constraints(points)][:remaining] + accepted.append(points) + remaining -= len(points) + if not remaining: + return torch.cat(accepted) + raise ValueError( + f"Constrained sampling exhausted {self.config.max_sampling_rounds} rounds; " + f"accepted {num_samples - remaining}/{num_samples} points" + ) + def sample_cartesian_space(self, num_samples: int | None = None) -> torch.Tensor: """Sample Cartesian positions within workspace bounds. @@ -660,10 +765,15 @@ def sample_cartesian_space(self, num_samples: int | None = None) -> torch.Tensor ) cartesian_bounds = self._compute_dynamic_workspace_bounds() - # Sample from Cartesian space using bounds - cartesian_samples = self.sampler.sample( - bounds=cartesian_bounds, num_samples=num_samples - ) + # Sample from Cartesian space using bounds. Sobol/LHS samplers perform + # their own vectorized generation; constraints are filtered in one pass + # before the considerably more expensive IK stage. + if self.config.sample_within_constraints: + cartesian_samples = self._sample_in_domain(cartesian_bounds, num_samples) + else: + cartesian_samples = self.sampler.sample( + bounds=cartesian_bounds, num_samples=num_samples + ) # Check how many samples pass workspace constraints valid_bounds = self.constraint_checker.check_bounds(cartesian_samples) @@ -733,10 +843,21 @@ def sample_plane( plane_bounds = plane_bounds.to(self.device) # Generate 2D samples and convert to 3D - plane_samples_2d = self.sampler.sample(num_samples, bounds=plane_bounds) - plane_samples_3d = self._plane_to_world_optimized( - plane_samples_2d, plane_normal, plane_point - ) + if self.config.sample_within_constraints: + plane_samples_3d = self._sample_in_domain( + plane_bounds, + num_samples, + transform=lambda uv: self._plane_to_world_optimized( + uv, plane_normal, plane_point + ), + ) + else: + plane_samples_2d = self.sampler.sample( + num_samples=num_samples, bounds=plane_bounds + ) + plane_samples_3d = self._plane_to_world_optimized( + plane_samples_2d, plane_normal, plane_point + ) logger.log_info( f"Generated {num_samples} plane samples using {self.sampler.get_strategy_name()}" @@ -904,6 +1025,7 @@ def _get_robot_base_position(self) -> torch.Tensor: current_pose = self.robot.compute_fk( qpos=self.robot.get_qpos()[None, :], # Add batch dimension name=self.control_part_name, + env_ids=[0], to_matrix=True, ) # Use current end-effector position projected to a reasonable height @@ -942,7 +1064,7 @@ def compute_workspace_points( num_samples = len(joint_configs) batch_size = batch_size or self.config.sampling.batch_size # Cap batch size to total samples - batch_size = min(batch_size, num_samples) + batch_size = max(1, min(batch_size, num_samples)) logger.log_info( f"Computing FK for {num_samples} samples (batch_size={batch_size})..." @@ -972,6 +1094,7 @@ def compute_workspace_points( poses = self.robot.compute_batch_fk( qpos=qpos_batch, name=self.control_part_name, + env_ids=[0], to_matrix=True, ) @@ -979,7 +1102,7 @@ def compute_workspace_points( positions = poses[0, :, :3, 3] # Vectorized constraint check for entire batch - valid_mask = self.constraint_checker.check_constraints(positions) + valid_mask = self._check_constraints(positions) if valid_mask.any(): workspace_points_list.append(positions[valid_mask]) @@ -1057,12 +1180,10 @@ def compute_reachability( num_samples = len(cartesian_points) ik_samples_per_point = self.config.ik_samples_per_point batch_size = batch_size or self.config.sampling.batch_size - batch_size = min(batch_size, num_samples) + batch_size = max(1, min(batch_size, num_samples)) # Pre-filter by workspace constraints (vectorized) - valid_cartesian_mask = self.constraint_checker.check_constraints( - cartesian_points - ) + valid_cartesian_mask = self._check_constraints(cartesian_points) logger.log_info( f"Pre-filtered Cartesian points: {valid_cartesian_mask.sum()}/{num_samples} " @@ -1115,9 +1236,13 @@ def compute_reachability( # Each position is repeated ik_samples_per_point times so that a single # compute_batch_ik call covers all (n_valid * K) targets at once. # Shape: (1, n_valid * K, 4, 4) - base_pose = current_ee_pose.unsqueeze(1).expand(1, n_valid, 4, 4).clone() - base_pose[0, :, :3, 3] = valid_positions - target_poses = base_pose.repeat_interleave(ik_samples_per_point, dim=1) + target_poses = ( + current_ee_pose.reshape(1, 1, 1, 4, 4) + .expand(1, n_valid, ik_samples_per_point, 4, 4) + .clone() + ) + target_poses[0, :, :, :3, 3] = valid_positions[:, None, :] + target_poses = target_poses.reshape(1, n_valid * ik_samples_per_point, 4, 4) # Generate all random seeds at once: (1, n_valid * K, num_joints) all_seeds = random_sampler.sample( @@ -1125,32 +1250,34 @@ def compute_reachability( ).unsqueeze(0) try: - logger.set_log_level("ERROR") success, qpos = self.robot.compute_batch_ik( pose=target_poses, joint_seed=all_seeds, name=self.control_part_name, - ) - logger.set_log_level("INFO") - - # Reshape results from flat batch to (n_valid, K) - success_2d = success[0].reshape(n_valid, ik_samples_per_point) - qpos_3d = qpos[0].reshape( - n_valid, ik_samples_per_point, self.num_joints + env_ids=[0], ) - # Success rate: fraction of seeds that solved IK for each point - success_rates_batch = success_2d.float().mean(dim=1) # (n_valid,) + if ik_samples_per_point == 1: + any_success = success[0].bool() + success_rates_batch = any_success.float() + best_qpos = qpos[0] + else: + # Reshape results from flat batch to (n_valid, K) + success_2d = success[0].reshape(n_valid, ik_samples_per_point) + qpos_3d = qpos[0].reshape( + n_valid, ik_samples_per_point, self.num_joints + ) - # Pick the joint config from the first successful seed per point - any_success = success_2d.any(dim=1) # (n_valid,) - first_success_idx = success_2d.float().argmax(dim=1) # (n_valid,) - best_qpos = qpos_3d[ - torch.arange(n_valid, device=self.device), first_success_idx - ] # (n_valid, num_joints) + # Success rate: fraction of seeds that solved IK for each point + success_rates_batch = success_2d.float().mean(dim=1) # (n_valid,) + # Pick the joint config from the first successful seed per point + any_success = success_2d.any(dim=1) # (n_valid,) + first_success_idx = success_2d.float().argmax(dim=1) # (n_valid,) + best_qpos = qpos_3d[ + torch.arange(n_valid, device=self.device), first_success_idx + ] # (n_valid, num_joints) except Exception as e: - logger.set_log_level("INFO") logger.log_warning( f"IK computation failed for batch [{batch_start}:{batch_end}]: {e}" ) @@ -1226,9 +1353,9 @@ def _get_reference_pose(self) -> torch.Tensor: hasattr(self.config, "reference_pose") and self.config.reference_pose is not None ): - reference_pose = self.config.reference_pose - if isinstance(reference_pose, np.ndarray): - reference_pose = torch.from_numpy(reference_pose).to(self.device) + reference_pose = torch.as_tensor( + self.config.reference_pose, dtype=torch.float32, device=self.device + ) if reference_pose.dim() == 2: reference_pose = reference_pose.unsqueeze(0) logger.log_info("Using provided reference pose for IK target orientation") @@ -1241,6 +1368,7 @@ def _get_reference_pose(self) -> torch.Tensor: current_ee_pose = self.robot.compute_fk( name=self.control_part_name, qpos=current_qpos.unsqueeze(0), + env_ids=[0], to_matrix=True, ) logger.log_info("Computing reference pose from current robot configuration") @@ -1441,6 +1569,18 @@ def analyze( results["config"] = self.config results["analysis_time"] = time.time() - start_time + if ( + not self.config.retain_diagnostics + and self.current_mode != AnalysisMode.JOINT_SPACE + ): + mask = results.pop("reachability_mask") + results["success_rates"] = results["success_rates"][mask] + results.pop("all_points") + results.pop("workspace_points") + self.workspace_points = results["reachable_points"] + self.success_rates = results["success_rates"] + self.reachability_mask = None + # Cache results (disk results cache; no-op when cache_dir is unset). if self._has_results_cache(): self._save_to_cache(results) @@ -1763,7 +1903,10 @@ def _generate_point_colors_and_sizes( if self.current_mode == AnalysisMode.CARTESIAN_SPACE else "Plane sampling" ) - if self.success_rates is not None and hasattr(self, "reachability_mask"): + if ( + self.success_rates is not None + and getattr(self, "reachability_mask", None) is not None + ): if filtered_to_reachable: # Points have been pre-filtered, but we still need to check IK reachability # Only color as green if we have verified IK solutions @@ -1980,7 +2123,7 @@ def visualize( self.current_mode in [AnalysisMode.CARTESIAN_SPACE, AnalysisMode.PLANE_SAMPLING] and not self.config.visualization.show_unreachable_points - and hasattr(self, "reachability_mask") + and getattr(self, "reachability_mask", None) is not None ): # Only show reachable points reachable_mask = self.reachability_mask.cpu().numpy() @@ -2172,8 +2315,23 @@ def serialize_parameter(value): "max_bounds": _tensor_to_list(constraint.max_bounds), "joint_limits_scale": constraint.joint_limits_scale, "ground_height": constraint.ground_height, + "exclude_zones": [ + [_tensor_to_list(lo), _tensor_to_list(hi)] + for lo, hi in constraint.exclude_zones + ], }, "ik_samples_per_point": cfg.ik_samples_per_point, + "sampling_revision": 2, + "retain_diagnostics": cfg.retain_diagnostics, + "sample_within_constraints": cfg.sample_within_constraints, + "max_sampling_rounds": cfg.max_sampling_rounds, + "geometry": { + "type": cfg.constraint_type, + "bounds": _tensor_to_list(cfg.constraint_bounds), + "center": _tensor_to_list(cfg.sphere_center), + "radius": cfg.sphere_radius, + "radius_mode": cfg.sphere_radius_mode, + }, } if cfg.reference_pose is not None: @@ -2246,7 +2404,9 @@ def _restore_analysis_state(self, results: Dict[str, Any]) -> None: self.current_mode = AnalysisMode(mode_str) if mode_str else None except ValueError: self.current_mode = None - self.workspace_points = results.get("workspace_points") + self.workspace_points = results.get( + "workspace_points", results.get("reachable_points") + ) self.joint_configurations = results.get("joint_configurations") self.success_rates = results.get("success_rates") if mode_str in ("cartesian_space", "plane_sampling"): diff --git a/embodichain/lab/sim/motion/workspace/configs/sampling_config.py b/embodichain/lab/sim/motion/workspace/configs/sampling_config.py index 96181480e..edb821089 100644 --- a/embodichain/lab/sim/motion/workspace/configs/sampling_config.py +++ b/embodichain/lab/sim/motion/workspace/configs/sampling_config.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- + +from __future__ import annotations from enum import Enum from dataclasses import dataclass from typing import Callable @@ -78,4 +80,4 @@ class SamplingConfig: def __post_init__(self): """Set default strategy after initialization.""" if self.strategy is None: - self.strategy = SamplingStrategy.UNIFORM + self.strategy = SamplingStrategy.SOBOL diff --git a/embodichain/lab/sim/motion/workspace/samplers/base_sampler.py b/embodichain/lab/sim/motion/workspace/samplers/base_sampler.py index 30a1bf97b..0b6d5dace 100644 --- a/embodichain/lab/sim/motion/workspace/samplers/base_sampler.py +++ b/embodichain/lab/sim/motion/workspace/samplers/base_sampler.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import numpy as np import torch from abc import ABC, abstractmethod @@ -83,10 +85,8 @@ def __init__( self.rng = np.random.RandomState(seed) self.device = device if device is not None else torch.device("cpu") - # Set torch seed - torch.manual_seed(seed) - if self.device.type == "cuda": - torch.cuda.manual_seed(seed) + self.device = torch.device(self.device) + self.generator = torch.Generator(device=self.device).manual_seed(seed) def sample( self, num_samples: int, bounds: torch.Tensor | np.ndarray | None = None diff --git a/embodichain/lab/sim/motion/workspace/samplers/gaussian_sampler.py b/embodichain/lab/sim/motion/workspace/samplers/gaussian_sampler.py index 9c32ffddb..d96a053a2 100644 --- a/embodichain/lab/sim/motion/workspace/samplers/gaussian_sampler.py +++ b/embodichain/lab/sim/motion/workspace/samplers/gaussian_sampler.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import numpy as np import torch from typing import Union @@ -172,7 +174,9 @@ def _generate_and_clip( Clipped samples (num_samples, n_dims). """ # Generate Gaussian samples - samples = torch.randn(num_samples, mean.shape[0], device=self.device) + samples = torch.randn( + num_samples, mean.shape[0], device=self.device, generator=self.generator + ) samples = mean + samples * std # Clip to bounds @@ -215,7 +219,12 @@ def _generate_with_rejection( num_generate = max(num_needed * 2, 100) # Generate 2x to reduce rejections # Generate Gaussian samples - samples = torch.randn(num_generate, mean.shape[0], device=self.device) + samples = torch.randn( + num_generate, + mean.shape[0], + device=self.device, + generator=self.generator, + ) samples = mean + samples * std # Check which samples are within bounds diff --git a/embodichain/lab/sim/motion/workspace/samplers/importance_sampler.py b/embodichain/lab/sim/motion/workspace/samplers/importance_sampler.py index 3eed34834..42576d1d6 100644 --- a/embodichain/lab/sim/motion/workspace/samplers/importance_sampler.py +++ b/embodichain/lab/sim/motion/workspace/samplers/importance_sampler.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import numpy as np import torch from typing import Callable, Union, TYPE_CHECKING @@ -91,7 +93,7 @@ def __init__( f"Invalid method '{method}'. Use 'rejection' or 'transform'." ) - def sample( + def _sample_from_bounds( self, bounds: torch.Tensor | np.ndarray, num_samples: int ) -> torch.Tensor: """Generate importance-weighted samples within the given bounds. @@ -157,7 +159,12 @@ def _rejection_sampling( num_needed = num_samples - len(accepted_samples) num_candidates_batch = max(num_needed * self.num_candidates, 100) - candidates = torch.rand(num_candidates_batch, n_dims, device=self.device) + candidates = torch.rand( + num_candidates_batch, + n_dims, + device=self.device, + generator=self.generator, + ) candidates = self._scale_samples(candidates, bounds) # Compute weights @@ -179,7 +186,9 @@ def _rejection_sampling( weights_normalized = torch.ones_like(weights) # Rejection sampling - accept_probs = torch.rand(num_candidates_batch, device=self.device) + accept_probs = torch.rand( + num_candidates_batch, device=self.device, generator=self.generator + ) accepted_mask = accept_probs < weights_normalized accepted_batch = candidates[accepted_mask] @@ -218,7 +227,9 @@ def _transform_sampling( # Generate candidate samples num_candidates_total = num_samples * self.num_candidates - candidates = torch.rand(num_candidates_total, n_dims, device=self.device) + candidates = torch.rand( + num_candidates_total, n_dims, device=self.device, generator=self.generator + ) candidates = self._scale_samples(candidates, bounds) # Compute weights @@ -244,12 +255,16 @@ def _transform_sampling( # Sample indices according to probabilities try: - indices = torch.multinomial(probabilities, num_samples, replacement=False) + indices = torch.multinomial( + probabilities, num_samples, replacement=False, generator=self.generator + ) except RuntimeError: # If probabilities are problematic, add small epsilon probabilities = probabilities + 1e-10 probabilities = probabilities / probabilities.sum() - indices = torch.multinomial(probabilities, num_samples, replacement=False) + indices = torch.multinomial( + probabilities, num_samples, replacement=False, generator=self.generator + ) selected_samples = candidates[indices] diff --git a/embodichain/lab/sim/motion/workspace/samplers/lhs_sampler.py b/embodichain/lab/sim/motion/workspace/samplers/lhs_sampler.py index f482506f1..c9535ba39 100644 --- a/embodichain/lab/sim/motion/workspace/samplers/lhs_sampler.py +++ b/embodichain/lab/sim/motion/workspace/samplers/lhs_sampler.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import numpy as np import torch from typing import Union @@ -97,7 +99,7 @@ def __init__( ) self.optimization = None - def sample( + def _sample_from_bounds( self, bounds: torch.Tensor | np.ndarray, num_samples: int ) -> torch.Tensor: """Generate Latin Hypercube samples within the given bounds. @@ -170,7 +172,7 @@ def _generate_lhs_scipy( """ # Create LHS engine lhs_engine = qmc.LatinHypercube( - d=n_dims, strength=strength, optimization=self.optimization, seed=self.seed + d=n_dims, strength=strength, optimization=self.optimization, seed=self.rng ) # Generate samples diff --git a/embodichain/lab/sim/motion/workspace/samplers/random_sampler.py b/embodichain/lab/sim/motion/workspace/samplers/random_sampler.py index e89b7a688..bece32aef 100644 --- a/embodichain/lab/sim/motion/workspace/samplers/random_sampler.py +++ b/embodichain/lab/sim/motion/workspace/samplers/random_sampler.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import numpy as np import torch from typing import Union, TYPE_CHECKING @@ -93,7 +95,9 @@ def _sample_from_bounds( n_dims = bounds.shape[0] # Generate random samples in [0, 1]^n_dims using torch - samples_unit = torch.rand(num_samples, n_dims, device=self.device) + samples_unit = torch.rand( + num_samples, n_dims, device=self.device, generator=self.generator + ) # Scale to the actual bounds samples = self._scale_samples(samples_unit, bounds) diff --git a/embodichain/lab/sim/motion/workspace/samplers/sobol_sampler.py b/embodichain/lab/sim/motion/workspace/samplers/sobol_sampler.py index f5ef64e3a..f619e96e8 100644 --- a/embodichain/lab/sim/motion/workspace/samplers/sobol_sampler.py +++ b/embodichain/lab/sim/motion/workspace/samplers/sobol_sampler.py @@ -14,52 +14,23 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import numpy as np import torch -from typing import Union - -try: - from scipy.stats import qmc - - SCIPY_AVAILABLE = True -except ImportError: - SCIPY_AVAILABLE = False -from embodichain.lab.sim.motion.workspace.configs.sampling_config import ( - SamplingStrategy, -) -from embodichain.lab.sim.motion.workspace.samplers.base_sampler import ( - BaseSampler, -) +from .base_sampler import BaseSampler +from ..configs.sampling_config import SamplingStrategy -from embodichain.utils import logger +__all__ = ["SobolSampler"] class SobolSampler(BaseSampler): - """Sobol sequence sampler using quasi-random low-discrepancy sequences. - - The Sobol sequence is a low-discrepancy sequence that provides excellent - uniformity in high-dimensional spaces. It's widely used in finance, - engineering, and scientific computing for Monte Carlo simulations. - - Advantages: - - Excellent uniformity in high dimensions (up to ~40 dimensions) - - Industry standard for quasi-Monte Carlo methods - - Better convergence than random sampling (O(1/n) vs O(1/√n)) - - Well-suited for integration and optimization + """Scrambled Sobol sampling with a persistent stream for each dimension. - Disadvantages: - - Requires scipy library - - Sequential generation (but can be scrambled for randomization) - - Initial points may not be well-distributed (use skip parameter) - - Attributes: - scramble: Whether to scramble the sequence for better randomization. - skip: Number of initial samples to skip. - - Notes: - This implementation uses scipy.stats.qmc.Sobol for efficient generation. - Falls back to a basic implementation if scipy is not available. + Successive draws continue the sequence. Power-of-two total sample counts + preserve Sobol's balance properties. Samples are drawn on CPU by Torch's + Sobol engine, then transferred as one tensor to the requested device. """ def __init__( @@ -68,149 +39,39 @@ def __init__( device: torch.device | None = None, scramble: bool = True, skip: int = 0, - ): - """Initialize the Sobol sampler. + ) -> None: + """Initialize the sequence. Args: - seed: Random seed for scrambling. Defaults to 42. - device: PyTorch device (cpu/cuda). Defaults to cpu. - scramble: Whether to scramble the sequence. Defaults to True. - Scrambling improves randomization while maintaining low discrepancy. - skip: Number of initial samples to skip. Defaults to 0. - Recommended: 0 for scrambled, >0 (e.g., 100) for unscrambled. - constraint: Optional geometric constraint for sampling (e.g., SphereConstraint). + seed: Scrambling seed. + device: Destination device for samples. + scramble: Whether to scramble the Sobol sequence. + skip: Number of leading sequence points to skip in each dimension. """ super().__init__(seed, device) + if skip < 0: + raise ValueError("skip must be non-negative") self.scramble = scramble self.skip = skip + self._engines: dict[int, torch.quasirandom.SobolEngine] = {} - if not SCIPY_AVAILABLE: - logger.log_warning( - "scipy is not available. Sobol sampler will use a basic fallback implementation. " - "For optimal performance, install scipy: pip install scipy" - ) - - def sample( + def _sample_from_bounds( self, bounds: torch.Tensor | np.ndarray, num_samples: int ) -> torch.Tensor: - """Generate Sobol sequence samples within the given bounds. - - Args: - bounds: Tensor/Array of shape (n_dims, 2) containing [lower, upper] bounds. - num_samples: Number of samples to generate. - - Returns: - Tensor of shape (num_samples, n_dims) containing the sampled points. - - Raises: - ValueError: If bounds are invalid or num_samples is non-positive. - - Examples: - >>> sampler = SobolSampler(scramble=True, seed=42) - >>> bounds = torch.tensor([[-1.0, 1.0], [-1.0, 1.0]], dtype=torch.float32) - >>> samples = sampler.sample(bounds, num_samples=100) - >>> samples.shape - torch.Size([100, 2]) - """ bounds = self._validate_bounds(bounds) - if num_samples <= 0: - raise ValueError(f"num_samples must be positive, got {num_samples}") - - n_dims = bounds.shape[0] - - if n_dims > 21201: # Maximum dimension for scipy's Sobol - raise ValueError( - f"Sobol sequence supports up to 21201 dimensions, got {n_dims}" + raise ValueError("num_samples must be positive") + dimension = bounds.shape[0] + if dimension not in self._engines: + engine = torch.quasirandom.SobolEngine( + dimension, scramble=self.scramble, seed=self.seed ) - - # Generate Sobol sequence - if SCIPY_AVAILABLE: - samples_unit = self._generate_sobol_scipy(n_dims, num_samples) - else: - samples_unit = self._generate_sobol_fallback(n_dims, num_samples) - - # Convert to tensor and scale to bounds - samples_unit_tensor = self._to_tensor(samples_unit) - samples = self._scale_samples(samples_unit_tensor, bounds) - - # Validate samples - self._validate_samples(samples, bounds) - - return samples - - def _generate_sobol_scipy(self, n_dims: int, num_samples: int) -> np.ndarray: - """Generate Sobol sequence using scipy. - - Args: - n_dims: Number of dimensions. - num_samples: Number of samples to generate. - - Returns: - Array of shape (num_samples, n_dims) with values in [0, 1]. - """ - # Create Sobol engine - sobol_engine = qmc.Sobol(d=n_dims, scramble=self.scramble, seed=self.seed) - - # Skip initial samples if requested - if self.skip > 0: - sobol_engine.fast_forward(self.skip) - - # Generate samples - samples = sobol_engine.random(n=num_samples) - - return samples.astype(np.float32) - - def _generate_sobol_fallback(self, n_dims: int, num_samples: int) -> np.ndarray: - """Fallback Sobol generator when scipy is not available. - - This is a basic implementation and may not match scipy's quality. - For production use, install scipy. - - Args: - n_dims: Number of dimensions. - num_samples: Number of samples to generate. - - Returns: - Array of shape (num_samples, n_dims) with values in [0, 1]. - """ - logger.log_warning( - "Using fallback Sobol generator. Results may not match scipy quality." - ) - - # Simple fallback: use stratified random sampling - # This is NOT a true Sobol sequence but provides reasonable coverage - samples = np.zeros((num_samples, n_dims), dtype=np.float32) - - for dim in range(n_dims): - # Divide [0, 1] into num_samples intervals - intervals = np.linspace(0, 1, num_samples + 1) - # Sample randomly within each interval - samples[:, dim] = intervals[:-1] + self.rng.rand(num_samples) / num_samples - - # Shuffle to reduce correlation - for dim in range(n_dims): - self.rng.shuffle(samples[:, dim]) - - return samples + if self.skip: + engine.fast_forward(self.skip) + self._engines[dimension] = engine + unit = self._engines[dimension].draw(num_samples).to(self.device) + return self._scale_samples(unit, bounds) def get_strategy_name(self) -> str: - """Get the name of the sampling strategy. - - Returns: - String identifier for the sampling strategy. - """ + """Return the registered strategy name.""" return SamplingStrategy.SOBOL.value - - def __repr__(self) -> str: - """String representation of the sampler.""" - return ( - f"{self.__class__.__name__}(" - f"strategy={self.get_strategy_name()}, " - f"scramble={self.scramble}, " - f"skip={self.skip}, " - f"seed={self.seed})" - ) - - -__all__ = ["SobolSampler"] diff --git a/embodichain/lab/sim/objects/robot.py b/embodichain/lab/sim/objects/robot.py index 37aef4063..d7c6b3cc3 100644 --- a/embodichain/lab/sim/objects/robot.py +++ b/embodichain/lab/sim/objects/robot.py @@ -969,17 +969,14 @@ def compute_batch_fk( qpos_ = qpos.to(self.device) n_batch = qpos_.shape[1] - qpos_batch = qpos_.reshape(-1, solver.dof) - xpos_batch = solver.get_fk(qpos=qpos_batch) + xpos_batch = solver.get_fk_batch(qpos_) # get xpos from link root base_xpos_n_envs = self.get_link_pose( link_name=solver.root_link_name, env_ids=local_env_ids, to_matrix=True ) - base_xpos_batch = ( - base_xpos_n_envs[:, None, :, :].repeat(1, n_batch, 1, 1).reshape(-1, 4, 4) - ) - result_matrix = torch.bmm(base_xpos_batch, xpos_batch) + result_matrix = torch.matmul(base_xpos_n_envs[:, None], xpos_batch) + result_matrix = result_matrix.reshape(-1, 4, 4) if to_matrix: result_matrix = result_matrix.reshape(len(local_env_ids), n_batch, 4, 4) @@ -1073,22 +1070,18 @@ def compute_batch_ik( link_name=solver.root_link_name, env_ids=local_env_ids, to_matrix=True ) base_inv_xpos_n_envs = torch.inverse(base_xpos_n_envs) - base_inv_xpos_batch = ( - base_inv_xpos_n_envs[:, None, :, :] - .repeat(1, n_batch, 1, 1) - .reshape(-1, 4, 4) + pose_batch = torch.matmul( + base_inv_xpos_n_envs[:, None], + pose_batch.reshape(len(local_env_ids), n_batch, 4, 4), ) - pose_batch = torch.bmm(base_inv_xpos_batch, pose_batch) - - joint_seed_batch = joint_seed.reshape(-1, n_dof) - ret, qpos_batch = solver.get_ik( + return solver.get_ik_batch( target_xpos=pose_batch, - qpos_seed=joint_seed_batch, - return_all_solutions=False, + qpos_seed=( + to_tensor(joint_seed, device=self.device) + if joint_seed is not None + else None + ), ) - ret = ret.reshape(len(local_env_ids), n_batch) - qpos = qpos_batch.reshape(len(local_env_ids), n_batch, n_dof) - return ret, qpos def _init_control_parts(self, control_parts: Dict[str, List[str]]) -> None: """Initialize the control parts of the robot. diff --git a/scripts/benchmark/workspace_analyzer/benchmark_robot_workspace.py b/scripts/benchmark/workspace_analyzer/benchmark_robot_workspace.py new file mode 100644 index 000000000..e43266724 --- /dev/null +++ b/scripts/benchmark/workspace_analyzer/benchmark_robot_workspace.py @@ -0,0 +1,354 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Measure UR5 and CobotMagic workspace paths without initializing a renderer. + +Run: python scripts/benchmark/workspace_analyzer/benchmark_robot_workspace.py --device cuda:0 +Run the same script with PYTHONPATH and cwd pointing at a baseline checkout to +compare revisions. Asset chains and analytic kernels are real; root poses are +fixed snapshots supplied by a minimal Robot adapter. FK uses eager asset-chain +code so compiler startup and specialization do not bias revision comparisons. +""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import subprocess +import time +from pathlib import Path +from types import MethodType, SimpleNamespace +from typing import Callable + +import psutil +import torch +import warp as wp + +from embodichain.lab.sim.objects.robot import Robot +from embodichain.lab.sim.robots import CobotMagicCfg, URRobotCfg +from embodichain.lab.sim.motion.workspace.analyzer import ( + AnalysisMode, + WorkspaceAnalyzer, + WorkspaceAnalyzerConfig, +) +from embodichain.lab.sim.motion.workspace.configs import ( + CacheConfig, + DimensionConstraint, + SamplingConfig, + SamplingStrategy, +) +from embodichain.utils import logger + + +def _robot(kind: str, device: torch.device) -> tuple[SimpleNamespace, str]: + """Bind real Robot batch methods to an asset-backed solver and fixed base.""" + preset = ( + URRobotCfg.from_dict({"robot_type": "ur5"}) + if kind == "ur5" + else CobotMagicCfg.from_dict({}) + ) + part = "arm" if kind == "ur5" else "left_arm" + chain = preset.build_pk_serial_chain(device)[part] + cfg = preset.solver_cfg[part] + cfg.joint_names = chain.get_joint_parameter_names() + solver = cfg.init_solver(device=device, pk_serial_chain=chain) + solver.compiled_fk = chain.forward_kinematics_tensor + base = torch.eye(4, device=device).unsqueeze(0) + if kind == "cobotmagic": + base[0, :3, 3] = base.new_tensor([0.233, 0.3, 0.0]) + joint_limits = torch.stack( + (solver.lower_qpos_limits, solver.upper_qpos_limits), dim=-1 + )[None] + robot = SimpleNamespace( + device=device, + cfg=preset, + num_envs=1, + _all_indices=[0], + _solvers={part: solver}, + control_parts={part: cfg.joint_names}, + joint_names=cfg.joint_names, + body_data=SimpleNamespace(qpos_limits=joint_limits), + get_joint_ids=lambda *args, **kw: list(range(6)), + get_qpos=lambda: solver.get_default_qpos_seed()[None], + get_link_pose=lambda **kw: base, + ) + for name in ("compute_fk", "compute_batch_fk", "compute_batch_ik"): + setattr(robot, name, MethodType(getattr(Robot, name), robot)) + return robot, part + + +class _Progress: + def __init__(self, iterable): + self.iterable = iterable + + def __iter__(self): + return iter(self.iterable) + + def close(self): + pass + + +def _measure( + call: Callable, device: torch.device, repeats: int +) -> tuple[float, float, float, float, object]: + """Return warmed median milliseconds, RSS/VRAM deltas, peak VRAM and output.""" + + def sync(): + if device.type == "cuda": + torch.cuda.synchronize(device) + wp.synchronize_device(str(device)) + + call() + sync() + process = psutil.Process(os.getpid()) + before_cpu = process.memory_info().rss + before_gpu = torch.cuda.memory_allocated(device) if device.type == "cuda" else 0 + if device.type == "cuda": + torch.cuda.reset_peak_memory_stats(device) + timings = [] + output = None + for _ in range(repeats): + output = None + sync() + start = time.perf_counter() + output = call() + sync() + timings.append((time.perf_counter() - start) * 1000) + cpu_delta = (process.memory_info().rss - before_cpu) / 1024**2 + gpu_delta = ( + (torch.cuda.memory_allocated(device) - before_gpu) / 1024**2 + if device.type == "cuda" + else 0.0 + ) + peak_gpu = ( + torch.cuda.max_memory_allocated(device) / 1024**2 + if device.type == "cuda" + else 0.0 + ) + return statistics.median(timings), cpu_delta, gpu_delta, peak_gpu, output + + +def _accuracy( + robot, part: str, mode: AnalysisMode, result: dict +) -> tuple[float, float]: + """Measure IK output against the default robot FK and target orientation.""" + qpos = result["joint_configurations"] + if mode == AnalysisMode.JOINT_SPACE or not len(qpos): + return 0.0, 0.0 + positions = result["reachable_points"] + target = robot.compute_fk(robot.get_qpos(), name=part, to_matrix=True) + actual = robot.compute_batch_fk(qpos[None], name=part, to_matrix=True)[0] + translation = ( + torch.linalg.vector_norm(actual[:, :3, 3] - positions, dim=-1).mean() * 1000 + ) + relative = actual[:, :3, :3].transpose(-1, -2) @ target[0, :3, :3] + cosine = ((relative.diagonal(dim1=-2, dim2=-1).sum(-1) - 1) / 2).clamp(-1, 1) + return float(translation), float(torch.acos(cosine).mean() * 180 / torch.pi) + + +def _table(rows: list[dict]) -> list[str]: + keys = list(rows[0]) + return [ + "| " + " | ".join(keys) + " |", + "| " + " | ".join(["---"] * len(keys)) + " |", + ] + ["| " + " | ".join(str(row[key]) for key in keys) + " |" for row in rows] + + +def run_all_benchmarks(args: argparse.Namespace) -> None: + """Measure three modes and dynamic bounds, then write one Markdown report. + + Args: + args: CLI device, sample counts, batch sizes, seeds and output settings. + """ + wp.init() + device = torch.device(args.device) + if device.type == "cuda" and not torch.cuda.is_available(): + print("Skipped: CUDA unavailable") + return + if device.type == "cuda": + torch.cuda.set_device(device) + torch.set_num_threads(args.threads) + logger.set_log_level("ERROR") + perf, quality = [], [] + for kind in ("ur5", "cobotmagic"): + robot, part = _robot(kind, device) + sim = SimpleNamespace(device=device, num_envs=1) + for n in args.samples: + for batch in args.batch_sizes: + for mode in AnalysisMode: + for seeds in ( + [1] if mode == AnalysisMode.JOINT_SPACE else args.seeds + ): + cfg = WorkspaceAnalyzerConfig( + mode=mode, + control_part_name=part, + sampling=SamplingConfig( + strategy=SamplingStrategy(args.sampler), + num_samples=n, + batch_size=batch, + seed=42, + ), + cache=CacheConfig(enabled=False), + constraint=DimensionConstraint( + min_bounds=( + [-1.3, -1.3, 0.0] + if kind == "ur5" + else [-0.5, -0.5, 0.0] + ), + max_bounds=( + [1.3, 1.3, 1.3] + if kind == "ur5" + else [1.0, 1.0, 1.0] + ), + ), + reference_pose=robot.compute_fk( + robot.get_qpos(), name=part, to_matrix=True + ), + ik_samples_per_point=seeds, + plane_normal=torch.tensor([0.0, 0.0, 1.0], device=device), + plane_point=torch.tensor([0.0, 0.0, 0.3], device=device), + plane_bounds=torch.tensor( + [[-1.0, 1.0], [-1.0, 1.0]], device=device + ), + ) + analyzer = WorkspaceAnalyzer(robot, cfg, sim) + analyzer._create_optimized_tqdm = lambda it, **kw: _Progress(it) + analyzer._update_progress_with_stats = lambda *a, **kw: None + analyzer._log_analysis_summary = lambda *a: None + + def call(): + # Each repeat uses the identical domain and random sequence. + analyzer.sampler = analyzer._create_sampler() + return analyzer.analyze( + force_recompute=True, visualize=False + ) + + ms, cpu, gpu, peak, result = _measure( + call, device, args.repeats + ) + label = f"{kind}/{mode.value}/K{seeds}/B{batch}/N{n}" + row = dict( + algorithm=label, + cost_time_ms=round(ms, 3), + cpu_delta_mb=round(cpu, 3), + gpu_delta_mb=round(gpu, 3), + peak_gpu_mb=round(peak, 3), + ) + perf.append(row) + rate = ( + result.get("num_reachable", result.get("num_valid", 0)) / n + ) + te, re = _accuracy(robot, part, mode, result) + quality.append( + dict( + algorithm=label, + success_rate=round(rate, 6), + translation_err_mm=round(te, 5), + rotation_err_deg=round(re, 5), + ) + ) + print(json.dumps(row), flush=True) + analyzer = WorkspaceAnalyzer( + robot, + WorkspaceAnalyzerConfig( + control_part_name=part, cache=CacheConfig(enabled=False) + ), + sim, + ) + analyzer._create_optimized_tqdm = lambda it, **kw: _Progress(it) + analyzer._update_progress_with_stats = lambda *a, **kw: None + ms, cpu, gpu, peak, bounds = _measure( + analyzer._compute_dynamic_workspace_bounds, device, args.repeats + ) + label = kind + "/dynamic_bounds/N1000" + perf.append( + dict( + algorithm=label, + cost_time_ms=round(ms, 3), + cpu_delta_mb=round(cpu, 3), + gpu_delta_mb=round(gpu, 3), + peak_gpu_mb=round(peak, 3), + ) + ) + quality.append( + dict( + algorithm=label, + success_rate=float(torch.isfinite(bounds).all()), + translation_err_mm="n/a", + rotation_err_deg="n/a", + ) + ) + rank = [ + dict(rank=i + 1, **row) + for i, row in enumerate( + sorted(quality, key=lambda row: row["success_rate"], reverse=True) + ) + ] + commit = subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip() + notes = [ + f"Revision: {commit}; label: {args.label}; device: {device}; threads: {args.threads}; repeats: {args.repeats}; sampler: {args.sampler}.", + "Times include sampling, Robot frame transforms, solver computation and basic metrics; exclude robot setup, visualization, persistent cache and initial warm-up.", + "No simulator is constructed. Fixed base snapshots exercise real Robot methods, real asset chains (eager FK) and real analytic IK kernels.", + "CPU memory is RSS delta. GPU columns count PyTorch allocator memory only; baseline Warp-owned allocations are excluded, so they cannot establish total VRAM savings.", + "IK residuals use default asset FK. Asset/analytic geometry differences and existing solver tolerances may produce nonzero residuals; compare both revisions before attributing them to these changes.", + "Success rates of joint, Cartesian and plane modes describe different domains and are not interchangeable quality rankings.", + ] + lines = [ + "# Robot workspace benchmark", + "", + *notes, + "", + "## Time & Memory", + "", + *_table(perf), + "", + "## Success & Other Metrics", + "", + *_table(quality), + "", + "## Leaderboard", + "", + *_table(rank), + ] + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text("\n".join(lines) + "\n", encoding="utf-8") + print(f"Markdown report saved: {args.output}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--device", default="cpu") + parser.add_argument("--samples", nargs="+", type=int, default=[1024, 8192]) + parser.add_argument("--batch-sizes", nargs="+", type=int, default=[1024]) + parser.add_argument("--seeds", nargs="+", type=int, default=[1, 4]) + parser.add_argument( + "--sampler", choices=["random", "sobol", "lhs"], default="random" + ) + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--threads", type=int, default=4) + parser.add_argument("--label", default="working-tree") + parser.add_argument( + "--output", type=Path, default=Path("outputs/benchmarks/robot_workspace.md") + ) + args = parser.parse_args() + if ( + min(args.samples + args.batch_sizes + args.seeds + [args.repeats, args.threads]) + <= 0 + ): + parser.error("counts must be positive") + run_all_benchmarks(args) diff --git a/scripts/benchmark/workspace_analyzer/workspace_optimization_report.md b/scripts/benchmark/workspace_analyzer/workspace_optimization_report.md new file mode 100644 index 000000000..4f416ca2a --- /dev/null +++ b/scripts/benchmark/workspace_analyzer/workspace_optimization_report.md @@ -0,0 +1,86 @@ +# Workspace optimization measurements + +Baseline: `797fe56ed51c0d74dc9337f35f71d358a1ff0e23`; implementation: `perf/workspace-batching`. +Environment: AMD EPYC 9K65, NVIDIA RTX PRO 5000 72GB Blackwell; Python 3.11, Torch 2.7.0+cu128, Warp 1.15.0. CPU runs use four Torch threads. Each row is the median of three warm repeats. + +The benchmark uses real UR5/CobotMagic analytic solvers, packaged asset chains and Robot batch methods, with fixed root-pose snapshots. It excludes renderer/simulation creation, persistent caching, compilation and visualization. Asset FK is eager; these timings are not full simulation deployment timings. Fixed-domain runs include sampling, FK/IK, frame conversion, filtering and basic metrics. Bounds rows measure the dynamic-bounds helper alone. + +Both revisions use explicit random sampling, seed 42 and identical fixed bounds. This isolates implementation changes from the new Sobol default. K denotes IK seeds per point, B denotes points per batch. Measurements on shared hardware have timing variability; small differences around 1x should be treated as unchanged. + +## CUDA timing + +| Robot/path | Before (ms) | After (ms) | Before / after | +| --- | ---: | ---: | ---: | +| ur5/joint_space/K1/B1024/N8192 | 15.009 | 14.302 | 1.05x | +| ur5/cartesian_space/K1/B1024/N8192 | 30.321 | 30.683 | 0.99x | +| ur5/cartesian_space/K4/B1024/N8192 | 38.410 | 28.280 | 1.36x | +| ur5/plane_sampling/K1/B1024/N8192 | 29.936 | 29.639 | 1.01x | +| ur5/plane_sampling/K4/B1024/N8192 | 37.896 | 28.291 | 1.34x | +| ur5/dynamic_bounds/N1000 | 879.574 | 1.470 | 598.35x | +| cobotmagic/joint_space/K1/B1024/N8192 | 9.898 | 9.581 | 1.03x | +| cobotmagic/cartesian_space/K1/B1024/N8192 | 21.215 | 20.737 | 1.02x | +| cobotmagic/cartesian_space/K4/B1024/N8192 | 22.145 | 19.960 | 1.11x | +| cobotmagic/plane_sampling/K1/B1024/N8192 | 25.148 | 19.519 | 1.29x | +| cobotmagic/plane_sampling/K4/B1024/N8192 | 20.747 | 21.156 | 0.98x | +| cobotmagic/dynamic_bounds/N1000 | 143.475 | 0.833 | 172.24x | + +## CPU timing + +| Robot/path | Before (ms) | After (ms) | Before / after | +| --- | ---: | ---: | ---: | +| ur5/joint_space/K1/B1024/N8192 | 18.229 | 12.636 | 1.44x | +| ur5/cartesian_space/K1/B1024/N8192 | 95.348 | 86.985 | 1.10x | +| ur5/cartesian_space/K4/B1024/N8192 | 433.853 | 383.950 | 1.13x | +| ur5/plane_sampling/K1/B1024/N8192 | 90.777 | 84.029 | 1.08x | +| ur5/plane_sampling/K4/B1024/N8192 | 556.121 | 385.301 | 1.44x | +| ur5/dynamic_bounds/N1000 | 209.597 | 1.603 | 130.75x | +| cobotmagic/joint_space/K1/B1024/N8192 | 19.188 | 14.113 | 1.36x | +| cobotmagic/cartesian_space/K1/B1024/N8192 | 17.676 | 15.849 | 1.12x | +| cobotmagic/cartesian_space/K4/B1024/N8192 | 46.012 | 42.893 | 1.07x | +| cobotmagic/plane_sampling/K1/B1024/N8192 | 12.510 | 11.069 | 1.13x | +| cobotmagic/plane_sampling/K4/B1024/N8192 | 28.233 | 25.973 | 1.09x | +| cobotmagic/dynamic_bounds/N1000 | 224.824 | 1.650 | 136.26x | + +## Accuracy comparison + +| Device | Cases compared | Identical success rates | Identical reported FK residuals | +| --- | ---: | ---: | ---: | +| cuda | 22 | 22/22 | 22/22 | +| cpu | 22 | 22/22 | 22/22 | + +UR keeps all 512 candidates (eight branches expanded across periodic joint shifts), its validity thresholds and nearest-seed selection. OPW keeps eight candidates. Existing default asset-FK versus analytic-geometry discrepancies are unchanged; exact candidate equality tests and separate analytic round trips validate allocation changes without silently weakening solver tolerances. + +The largest reliable gain is replacing 1,000 individual FK calls in bounds estimation with a batch. Fixed-domain FK/IK gains vary: some workloads improve, while several remain near parity. OPW packs live limits into one host transfer instead of twelve scalar conversions. Internal candidate arrays reuse storage; public returned arrays retain independent ownership. + +Memory columns in the raw reports are CPU RSS deltas and PyTorch allocator GPU deltas/peaks. Baseline Warp allocations are outside that GPU counter; these numbers must not be used to claim total VRAM reductions. Candidate capacity stays at the largest requested batch until solver destruction. Compact results reduce retained/cache arrays, not peak sampling allocation. + +## Reproduction + +From each checkout, run the same script with that checkout as `PYTHONPATH` and working directory: + +```bash +PYTHONPATH=. python /path/to/benchmark_robot_workspace.py --device cpu --output outputs/benchmarks/workspace_cpu.md +CUDA_VISIBLE_DEVICES=1 PYTHONPATH=. python /path/to/benchmark_robot_workspace.py --device cuda:0 --output outputs/benchmarks/workspace_cuda.md +``` + +The renderer-free benchmark supports CUDA remapping. Live DexSim tests must keep CUDA and Vulkan physical device selection aligned; the validation here ran live tests without `CUDA_VISIBLE_DEVICES` remapping. + +## New controls + +- `--sampler sobol` is now the CLI/API default; explicit `random` and `uniform` remain supported. +- `--sample-within-constraints` opts into bounded refill within box/ground/exclusion constraints (and configured sphere/plane domains through the analyzer API). Reachability then describes the permitted domain, not the original bounding box. +- `--compact-results` retains reachable points, qpos and aligned scores; dropped rejected-point diagnostics are unavailable for later analysis. +- `BaseSolver.prepare_buffers(max_batch)` optionally reserves scratch capacity; `get_fk_batch`/`get_ik_batch` preserve leading batch axes. + +Raw per-run reports are written under `outputs/benchmarks/workspace_{before,after}_{cpu,cuda}.md` in the implementation worktree. Each includes time/memory, success/pose metrics and a complete per-case ranking. Rankings across different analysis domains are descriptive, not an algorithm-quality comparison. + +## Validation + +Affected regression tests pass in two processes: 169 workspace/solver/config/motion-generator/randomization cases, plus 32 Robot CPU/CUDA cases. This includes real simulation tests, analytic candidate equality, scratch growth and CUDA stream ownership, and compact-cache visualization. The 23 project-context tests also pass. Black 26.3.1, `git diff --check`, context validation and API documentation coverage (1,853/1,853 exports) pass. + +The initial combined process passed 199 cases and failed two Robot CUDA FK cases after reaching Torch Dynamo's eight-recompilation limit for `forward_kinematics_tensor`. Running the Robot module separately passed all 32 cases; the other group passed all 169. No production compiler settings were changed to bypass that limit. + +```bash +PYTHONPATH=. python -m pytest -q tests/sim/motion/workspace tests/sim/motion/solvers/test_analytic_batching.py tests/sim/motion/solvers/test_base_solver.py tests/sim/objects/test_robot_cfg.py tests/sim/motion/test_motion_generator_batched.py tests/gym/envs/managers/test_workspace_randomization.py --run-gpu +PYTHONPATH=. python -m pytest -q tests/sim/objects/test_robot.py --run-gpu +``` diff --git a/tests/sim/motion/solvers/test_analytic_batching.py b/tests/sim/motion/solvers/test_analytic_batching.py new file mode 100644 index 000000000..00f465552 --- /dev/null +++ b/tests/sim/motion/solvers/test_analytic_batching.py @@ -0,0 +1,225 @@ +# ---------------------------------------------------------------------------- +# 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 warp as wp + +from embodichain.lab.sim.robots import CobotMagicCfg, URRobotCfg +from embodichain.lab.sim.objects.robot import Robot +from embodichain.lab.sim.motion.solvers._buffers import _IKBuffers + + +@pytest.fixture(params=["ur5", "cobotmagic"]) +def solver(request): + """Construct real analytic solvers and asset chains without a renderer.""" + wp.init() + preset = ( + URRobotCfg.from_dict({"robot_type": "ur5"}) + if request.param == "ur5" + else CobotMagicCfg.from_dict({}) + ) + part = "arm" if request.param == "ur5" else "left_arm" + device = torch.device(getattr(request, "param_device", "cpu")) + chain = preset.build_pk_serial_chain(device)[part] + cfg = preset.solver_cfg[part] + cfg.joint_names = chain.get_joint_parameter_names() + result = cfg.init_solver(device=device, pk_serial_chain=chain) + result.compiled_fk = chain.forward_kinematics_tensor + return result + + +def _analytic_fk(solver, qpos): + """Evaluate the analytic solver's geometry, independently of asset FK.""" + if hasattr(solver, "get_fk_warp"): + return solver.get_fk_warp(qpos) + cfg = solver.cfg + result = torch.eye(4, device=qpos.device).expand(len(qpos), 4, 4).clone() + params = [ + (cfg.d1, 0.0, torch.pi / 2), + (0.0, cfg.a2, 0.0), + (0.0, cfg.a3, 0.0), + (cfg.d4, 0.0, torch.pi / 2), + (cfg.d5, 0.0, -torch.pi / 2), + (cfg.d6, 0.0, 0.0), + ] + for joint, (d, a, alpha) in enumerate(params): + theta = qpos[:, joint] + ct, st = torch.cos(theta), torch.sin(theta) + ca, sa = torch.cos(theta.new_tensor(alpha)), torch.sin(theta.new_tensor(alpha)) + transform = torch.eye(4, device=qpos.device).expand(len(qpos), 4, 4).clone() + transform[:, 0, :] = torch.stack((ct, -st * ca, st * sa, a * ct), 1) + transform[:, 1, :] = torch.stack((st, ct * ca, -ct * sa, a * st), 1) + transform[:, 2, :] = theta.new_tensor([0.0, sa, ca, d]) + result = result @ transform + return result @ torch.as_tensor( + solver.tcp_xpos, dtype=qpos.dtype, device=qpos.device + ) + + +def _targets(solver, count): + generator = torch.Generator(device=solver.device).manual_seed(13) + lower, upper = solver.lower_qpos_limits, solver.upper_qpos_limits + qpos = lower + ( + 0.1 + 0.8 * torch.rand(count, 6, generator=generator, device=solver.device) + ) * (upper - lower) + # OPW's analytic geometry has a pre-existing difference from the packaged + # URDF chain; use its own FK to test analytic round trips, independently of sim. + fk = lambda q: _analytic_fk(solver, q) + return qpos, fk(qpos) + + +def test_candidate_reuse_and_old_results_remain_valid(solver, monkeypatch): + qpos, poses = _targets(solver, 17) + solver.prepare_buffers(32) + original_zeros = solver._ik_buffers.zeros + allocations = [] + for count in [9, 3, 17, 1]: + valid, result = solver.get_ik( + poses[:count], qpos[:count], return_all_solutions=True + ) + old_valid, old_result = valid.clone(), result.clone() + ptr = solver._ik_buffers.arrays["qpos"].data_ptr() + allocations.append(ptr) + bad_pose = poses[:1].clone() + bad_pose[:, :3, 3] = 100 + invalid, _ = solver.get_ik(bad_pose, qpos[:1]) + assert not invalid.bool().any() + torch.testing.assert_close(valid, old_valid) + torch.testing.assert_close(result, old_result) + + # Allocation-only reference: fresh zeroed buffers, identical unchanged kernels. + def fresh(name, batch, width, dtype): + return wp.from_torch( + torch.zeros(batch * width, dtype=dtype, device=solver.device) + ) + + with monkeypatch.context() as patch: + patch.setattr(solver._ik_buffers, "zeros", fresh) + expected_valid, expected_result = solver.get_ik( + poses[:count], qpos[:count], return_all_solutions=True + ) + torch.testing.assert_close(valid, expected_valid) + torch.testing.assert_close(result, expected_result) + assert len(set(allocations)) == 1 + solver._ik_buffers.zeros = original_zeros + + +def test_best_solution_round_trip_and_limits(solver): + qpos, poses = _targets(solver, 13) + success, result = solver.get_ik_batch( + poses.reshape(1, 13, 4, 4), qpos.reshape(1, 13, 6) + ) + assert success.dtype == torch.bool + assert success.shape == (1, 13) + assert result.shape == (1, 13, 6) + assert success.all() + old = result.clone() + solver.get_ik(poses[:2], qpos[:2]) + torch.testing.assert_close(result, old) + fk = lambda q: _analytic_fk(solver, q) + torch.testing.assert_close(fk(result.reshape(-1, 6)), poses, atol=3e-5, rtol=3e-5) + assert (result >= solver.lower_qpos_limits - 1e-6).all() + assert (result <= solver.upper_qpos_limits + 1e-6).all() + solver.set_qpos_limits(qpos[0] - 0.01, qpos[0] + 0.01) + valid, out = solver.get_ik(poses[:1], qpos[:1]) + assert valid.bool().all() + assert (out >= solver.lower_qpos_limits - 1e-6).all() + assert (out <= solver.upper_qpos_limits + 1e-6).all() + + +def test_robot_batch_transforms_broadcast_and_preserve_seed_shape(solver): + qpos, _ = _targets(solver, 6) + qpos = qpos.reshape(2, 3, 6) + bases = torch.eye(4).repeat(2, 1, 1) + bases[0, :3, 3] = torch.tensor([0.2, -0.3, 0.4]) + bases[1, :3, :3] = torch.tensor([[0.0, -1, 0], [1.0, 0, 0], [0, 0, 1]]) + bases[1, :3, 3] = torch.tensor([-0.4, 0.1, 0.5]) + robot = SimpleNamespace( + _all_indices=[0, 1], + _solvers={"arm": solver}, + device=solver.device, + get_link_pose=lambda **kwargs: bases[kwargs["env_ids"]], + ) + actual = Robot.compute_batch_fk(robot, qpos, "arm", to_matrix=True) + expected = torch.stack([bases[i] @ solver.get_fk(qpos[i]) for i in range(2)]) + torch.testing.assert_close(actual, expected) + # IK input uses the analytic FK, including the same per-environment roots. + fk = lambda q: _analytic_fk(solver, q) + targets = torch.stack([bases[i] @ fk(qpos[i]) for i in range(2)]) + valid, result = Robot.compute_batch_ik(robot, targets, qpos, "arm") + assert valid.shape == (2, 3) and valid.all() + torch.testing.assert_close( + fk(result.reshape(-1, 6)), fk(qpos.reshape(-1, 6)), atol=3e-5, rtol=3e-5 + ) + with pytest.raises(ValueError, match="batch axes"): + solver.get_ik_batch(targets, qpos[:, :1]) + + +def test_scratch_growth_zeroes_active_tail(): + wp.init() + buffers = _IKBuffers(torch.device("cpu")) + with buffers.borrow(): + first = buffers.zeros("q", 4, 6, torch.float32) + wp.to_torch(first).fill_(7) + with buffers.borrow(): + assert not wp.to_torch(buffers.zeros("q", 2, 6, torch.float32)).any() + with buffers.borrow(): + assert not wp.to_torch(buffers.zeros("q", 7, 6, torch.float32)).any() + + +@pytest.mark.gpu +@pytest.mark.parametrize("kind", ["ur5", "cobotmagic"]) +def test_cuda_stream_reuse(kind): + if not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + wp.init() + preset = ( + URRobotCfg.from_dict({"robot_type": "ur5"}) + if kind == "ur5" + else CobotMagicCfg.from_dict({}) + ) + part = "arm" if kind == "ur5" else "left_arm" + device = torch.device("cuda:0") + chain = preset.build_pk_serial_chain(device)[part] + cfg = preset.solver_cfg[part] + cfg.joint_names = chain.get_joint_parameter_names() + solver = cfg.init_solver(device=device, pk_serial_chain=chain) + solver.compiled_fk = chain.forward_kinematics_tensor + qpos, poses = _targets(solver, 19) + torch.cuda.synchronize() + outputs = [] + for count in [7, 19, 3, 11]: + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + outputs.append( + ( + count, + solver.get_ik( + poses[:count], qpos[:count], return_all_solutions=True + ), + ) + ) + torch.cuda.synchronize() + for count, (valid, result) in outputs: + expected_valid, expected_result = solver.get_ik( + poses[:count], qpos[:count], return_all_solutions=True + ) + torch.testing.assert_close(valid, expected_valid) + torch.testing.assert_close(result, expected_result) diff --git a/tests/sim/motion/workspace/test_analysis_batching.py b/tests/sim/motion/workspace/test_analysis_batching.py new file mode 100644 index 000000000..c5f5859cb --- /dev/null +++ b/tests/sim/motion/workspace/test_analysis_batching.py @@ -0,0 +1,293 @@ +# ---------------------------------------------------------------------------- +# 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 +from unittest.mock import Mock + +import numpy as np +import pytest +import torch + +from embodichain.lab.sim.motion.workspace.analyzer import ( + AnalysisMode, + WorkspaceAnalyzer, + WorkspaceAnalyzerConfig, +) +from embodichain.lab.sim.motion.workspace.configs import ( + CacheConfig, + DimensionConstraint, + SamplingConfig, + SamplingStrategy, +) +from embodichain.lab.sim.motion.workspace.samplers import ( + RandomSampler, + SobolSampler, + LatinHypercubeSampler, +) +from embodichain.lab.sim.motion.workspace.runtime import RobotWorkspace +from embodichain.lab.sim.motion.workspace.caches.results_cache import ResultsCache + + +def _analyzer(**kwargs) -> WorkspaceAnalyzer: + """Use deterministic kinematics to isolate analysis and sample alignment.""" + robot = Mock() + robot.device = torch.device("cpu") + robot.num_envs = 3 + robot.control_parts = {"arm": ["j0", "j1", "j2"]} + robot.get_joint_ids.return_value = [0, 1, 2] + robot.body_data = SimpleNamespace(qpos_limits=torch.tensor([[[-1.0, 1.0]] * 3])) + robot.get_qpos.return_value = torch.zeros(3, 3) + robot.cfg = SimpleNamespace(uid="test", fpath=None, solver_cfg={}) + robot._solvers = {} + + def fk(qpos, **kw): + pose = torch.eye(4).expand(*qpos.shape[:-1], 4, 4).clone() + pose[..., :3, 3] = qpos + return pose + + def ik(pose, joint_seed, **kw): + # Mix successful and unsuccessful points, independent of batch splitting. + success = pose[..., 0, 3] >= 0 + return success, pose[..., :3, 3].clone() + + robot.compute_batch_fk.side_effect = fk + robot.compute_fk.side_effect = fk + robot.compute_batch_ik.side_effect = ik + sampling = kwargs.pop("sampling", SamplingConfig(num_samples=64, batch_size=7)) + constraint = kwargs.pop( + "constraint", + DimensionConstraint( + min_bounds=np.array([-1.0, -1.0, 0.0]), + max_bounds=np.ones(3), + ), + ) + cfg = WorkspaceAnalyzerConfig( + sampling=sampling, + constraint=constraint, + cache=CacheConfig(enabled=False), + control_part_name="arm", + reference_pose=torch.eye(4), + **kwargs, + ) + result = WorkspaceAnalyzer(robot, cfg) + result._log_analysis_summary = Mock() + result._create_optimized_tqdm = lambda it, **kw: _Progress(it) + result._update_progress_with_stats = Mock() + return result + + +class _Progress: + def __init__(self, it): + self.it = it + + def __iter__(self): + return iter(self.it) + + def close(self): + pass + + +def test_dynamic_bounds_use_one_batch_and_env_zero(): + analyzer = _analyzer() + expected_q = RandomSampler(seed=42).sample(1000, analyzer.qpos_limits) + lo, hi = expected_q.amin(0), expected_q.amax(0) + expected = torch.stack((lo - 0.1 * (hi - lo), hi + 0.1 * (hi - lo)), 1) + torch.testing.assert_close(analyzer._compute_dynamic_workspace_bounds(), expected) + assert analyzer.robot.compute_batch_fk.call_count == 1 + assert analyzer.robot.compute_batch_fk.call_args.kwargs["env_ids"] == [0] + analyzer.robot.compute_fk.assert_not_called() + + +@pytest.mark.parametrize("seeds", [1, 4]) +def test_reachability_preserves_order_with_filtered_and_tail_batches(seeds): + analyzer = _analyzer(ik_samples_per_point=seeds) + points = torch.tensor( + [[0.2, 0.0, 0.1], [-0.2, 0.0, 0.1], [0.3, 0.0, -1.0], [0.4, 0.0, 0.2]] + ) + all_points, reachable, scores, mask, qpos = analyzer.compute_reachability( + points, batch_size=3 + ) + torch.testing.assert_close(all_points, points) + assert mask.tolist() == [True, False, False, True] + torch.testing.assert_close(scores, torch.tensor([1.0, 0.0, 0.0, 1.0])) + torch.testing.assert_close(reachable, points[[0, 3]]) + torch.testing.assert_close(qpos, reachable) + calls = analyzer.robot.compute_batch_ik.call_args_list + assert [c.kwargs["pose"].shape[1] for c in calls] == [2 * seeds, seeds] + assert all(c.kwargs["env_ids"] == [0] for c in calls) + + +def test_empty_analysis_entry_points(): + analyzer = _analyzer() + points, qpos = analyzer.compute_workspace_points(torch.empty(0, 3)) + assert points.shape == qpos.shape == (0, 3) + assert analyzer.compute_reachability(torch.empty(0, 3))[3].numel() == 0 + + +@pytest.mark.parametrize( + "sampler_cls", [RandomSampler, SobolSampler, LatinHypercubeSampler] +) +def test_sampler_is_reproducible_without_touching_global_rng(sampler_cls): + bounds = torch.tensor([[-1.0, 1.0]] * 3) + global_state = torch.random.get_rng_state().clone() + sampler = sampler_cls(seed=12) + first = sampler.sample(num_samples=16, bounds=bounds) + second = sampler.sample(num_samples=16, bounds=bounds) + torch.testing.assert_close(torch.random.get_rng_state(), global_state) + torch.testing.assert_close( + sampler_cls(seed=12).sample(num_samples=16, bounds=bounds), first + ) + assert not torch.equal(first, second) + assert ((first >= -1) & (first <= 1)).all() + + +def test_sobol_chunks_match_one_draw_and_plane_sampling_works(): + bounds = torch.tensor([[-1.0, 1.0]] * 3) + sampler = SobolSampler(seed=9) + chunks = torch.cat([sampler.sample(n, bounds) for n in [7, 9, 16]]) + torch.testing.assert_close(chunks, SobolSampler(seed=9).sample(32, bounds)) + points = _analyzer().sample_plane( + 32, plane_point=torch.tensor([0.0, 0.0, 0.5]), plane_bounds=bounds[:2] + ) + torch.testing.assert_close(points[:, 2], torch.full((32,), 0.5)) + + +def test_constrained_box_refills_excluded_zone(): + analyzer = _analyzer( + sample_within_constraints=True, + constraint_type="box", + constraint_bounds=torch.tensor([[0.1, 0.8], [-0.5, 0.5], [0.2, 0.8]]), + constraint=DimensionConstraint( + min_bounds=[-1, -1, 0], + max_bounds=[1, 1, 1], + exclude_zones=[([0.3, -1, 0], [0.6, 1, 1])], + ), + ) + points = analyzer.sample_cartesian_space(128) + assert points.shape == (128, 3) + assert analyzer._check_constraints(points).all() + assert ((points[:, 0] < 0.3) | (points[:, 0] > 0.6)).all() + + +def test_constrained_sphere_has_volume_distribution(): + analyzer = _analyzer( + sample_within_constraints=True, + constraint_type="sphere", + sphere_center=torch.tensor([0.0, 0.0, 0.5]), + sphere_radius=0.4, + ) + points = analyzer.sample_cartesian_space(4096) + radii = torch.linalg.vector_norm(points - torch.tensor([0.0, 0.0, 0.5]), dim=1) + assert radii.max() <= 0.4 + 1e-6 + assert float(radii.mean()) == pytest.approx(0.4 * 0.75, abs=0.005) + + +def test_constrained_plane_and_exhaustion(): + analyzer = _analyzer(sample_within_constraints=True) + bounds = torch.tensor([[-1.0, 1.0], [-1.0, 1.0]]) + points = analyzer.sample_plane( + 64, plane_point=torch.tensor([0.0, 0.0, 0.5]), plane_bounds=bounds + ) + assert analyzer._check_constraints(points).all() + torch.testing.assert_close(points[:, 2], torch.full((64,), 0.5)) + analyzer.config.max_sampling_rounds = 2 + with pytest.raises(ValueError, match="exhausted"): + analyzer.sample_plane( + 64, plane_point=torch.tensor([0.0, 0.0, -1.0]), plane_bounds=bounds + ) + + +@pytest.mark.parametrize( + "mode", [AnalysisMode.CARTESIAN_SPACE, AnalysisMode.PLANE_SAMPLING] +) +def test_compact_results_cache_and_visualization(tmp_path, mode): + kwargs = dict(mode=mode, retain_diagnostics=False) + if mode == AnalysisMode.PLANE_SAMPLING: + kwargs.update( + plane_point=torch.tensor([0.0, 0.0, 0.5]), + plane_bounds=torch.tensor([[-1.0, 1.0], [-1.0, 1.0]]), + ) + analyzer = _analyzer(**kwargs) + results = analyzer.analyze(visualize=False) + assert not {"workspace_points", "all_points", "reachability_mask"} & results.keys() + assert len(results["success_rates"]) == len(results["joint_configurations"]) + cache = ResultsCache(tmp_path) + path = cache.save("compact", results, {"mode": mode.value}) + runtime = RobotWorkspace.from_cache(path) + torch.testing.assert_close(runtime.qpos, results["joint_configurations"]) + restored = cache.load("compact") + analyzer._restore_analysis_state(restored) + colors = analyzer._generate_point_colors(analyzer.workspace_points.numpy()) + assert colors.shape == (len(runtime.qpos), 3) + analyzer.config.visualization.enabled = True + analyzer.config.visualization.show_unreachable_points = False + visualizer = Mock() + analyzer._create_visualizer_with_config = Mock(return_value=visualizer) + analyzer.visualize(show=False, backend="matplotlib") + assert visualizer.visualize.call_args.args[0].shape == (len(runtime.qpos), 3) + + +def test_cache_identity_tracks_new_options_and_exclusions(): + from embodichain.lab.sim.motion.workspace.caches.results_cache import ( + compute_cache_key, + ) + + analyzer = _analyzer() + key = lambda: compute_cache_key(analyzer._build_cache_key_metadata(64)) + initial = key() + analyzer.config.retain_diagnostics = False + assert key() != initial + compact = key() + analyzer.config.sample_within_constraints = True + assert key() != compact + before_exclusion = key() + analyzer.config.constraint.exclude_zones = [([0, 0, 0], [1, 1, 1])] + assert key() != before_exclusion + + +def test_reference_pose_uses_env_zero_and_accepts_serialized_pose(): + analyzer = _analyzer() + analyzer.config.reference_pose = None + pose = analyzer._get_reference_pose() + assert pose.shape == (1, 4, 4) + assert analyzer.robot.compute_fk.call_args.kwargs["env_ids"] == [0] + analyzer.config.reference_pose = torch.eye(4).tolist() + torch.testing.assert_close(analyzer._get_reference_pose(), torch.eye(4)[None]) + + +def test_cli_exposes_compact_constrained_sobol_defaults(): + from embodichain.lab.scripts.analyze_workspace import ( + parse_args, + build_analyzer_config, + _preview_points_and_colors, + ) + + args = parse_args( + ["--robot", "cobotmagic", "--compact-results", "--sample-within-constraints"] + ) + cfg = build_analyzer_config(args, "left_arm") + assert cfg.sampling.strategy == SamplingStrategy.SOBOL + assert cfg.sample_within_constraints and not cfg.retain_diagnostics + points = np.zeros((5, 3)) + # Array insertion order must not make preview choose qpos as XYZ. + xyz, colors = _preview_points_and_colors( + {"joint_configurations": np.zeros((5, 6)), "reachable_points": points}, + "cartesian_space", + False, + ) + assert xyz.shape == colors.shape == (5, 3) From 16e5a2ea4f2c9156fa88904baf514e643f043b86 Mon Sep 17 00:00:00 2001 From: Xinyi YUAN Date: Fri, 11 Sep 2026 21:00:26 +0900 Subject: [PATCH 2/5] feat(workspace): add per-point Yoshikawa manipulability to analysis Wire the metrics module into WorkspaceAnalyzer (closing the long-standing _compute_metrics TODO) and compute true per-configuration Yoshikawa manipulability w = sqrt(det(J J^T)) from the active solver's Jacobian after every analysis mode. Scores are row-aligned with joint_configurations (and with reachable points in Cartesian/plane modes), stored in results.npz, restored on cache hits, and aggregated under metrics["manipulability"]. Computation is gated on MetricConfig.enabled_metrics and costs ~10 ms per 470 configurations on GPU. Remove ManipulabilityMetric's centroid-distance placeholder: measured on Franka it is negatively correlated with true manipulability (corr = -0.37), so consumers ranking by it preferred worse configurations. Without Jacobians or precomputed scores the metric now warns and returns no statistics instead of fabricating them. The batching test's mock robot now returns None from get_solver, faithful to Robot.get_solver with no solvers attached. Also documented in the robot-workspace context: enabling the #599 seed-selection sampler speeds Cartesian reachability analysis 3.3x at unchanged num_samples=30 while detecting slightly more reachable points (measured on Franka, 4000 identical targets), with no analyzer changes. Covered by tests/sim/motion/workspace/test_manipulability.py; the full workspace suite passes (68 passed). Co-Authored-By: Claude Fable 5 --- agent_context/MAP.yaml | 3 +- .../robot-workspace/analysis-and-cache.md | 24 +++ .../lab/sim/motion/workspace/analyzer.py | 66 +++++++- .../motion/workspace/caches/results_cache.py | 1 + .../metrics/manipulability_metric.py | 56 +++---- .../workspace/test_analysis_batching.py | 3 + .../motion/workspace/test_manipulability.py | 142 ++++++++++++++++++ 7 files changed, 265 insertions(+), 30 deletions(-) create mode 100644 tests/sim/motion/workspace/test_manipulability.py diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 8e4b23cab..d7a3d8a4a 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -299,7 +299,8 @@ topics: title: Robot Workspace aliases: [robot workspace, workspace analyzer, reachability cache, 机器人工作空间, 可达空间, 工作空间缓存] keywords: [WorkspaceAnalyzer, AnalysisMode, ResultsCache, RobotWorkspace, RobotWorkspaceCfg, workspace_cfg, - sample_reachable_pose, analyze-workspace, workspace-cache, results.npz] + sample_reachable_pose, analyze-workspace, workspace-cache, results.npz, manipulability_scores, + ManipulabilityMetric] paths: [topics/robot-workspace/robot-workspace.md] source_of_truth: - embodichain/lab/sim/motion/workspace/analyzer.py diff --git a/agent_context/topics/robot-workspace/analysis-and-cache.md b/agent_context/topics/robot-workspace/analysis-and-cache.md index 7fb17b223..0b9fa03d4 100644 --- a/agent_context/topics/robot-workspace/analysis-and-cache.md +++ b/agent_context/topics/robot-workspace/analysis-and-cache.md @@ -73,6 +73,30 @@ | `workspace-cache list/clean` cannot see ResultsCache entries | It targets legacy session caches; use analyzer preview/cache directory semantics for result entries. | | Concurrent analyzers expose corrupt/incomplete cache | `ResultsCache.save()` writes result then metadata directly with no observed lock or temp-file replacement. Add atomic write/lock coverage before relying on shared concurrent writers. | +### Manipulability scores + +- When the metric config enables manipulability (default `ALL`), the analyzer + computes per-configuration Yoshikawa scores `w = sqrt(det(J J^T))` from the + active solver's Jacobian after analysis. `manipulability_scores` is + row-aligned with `joint_configurations` (and therefore with reachable points + in Cartesian/plane modes), stored in `results.npz`, and restored on cache + hits; aggregates land in `metrics["manipulability"]`. +- `ManipulabilityMetric` produces statistics only from Jacobians or + precomputed scores. The former centroid-distance placeholder was measured + to be negatively correlated with true manipulability (corr ≈ −0.37 on + Franka) and now yields a warning plus empty results instead. +- Cost is negligible: batched Jacobian + determinant is ~10 ms per 470 + configurations on GPU. + +### Seed selection for Cartesian/plane IK + +Cartesian and plane analysis verify reachability through the solver's +multi-start `get_ik`, so `PytorchSolverCfg.enable_seed_selection` applies +without analyzer changes. Measured on Franka (4000 identical points, warm): +enabling it at unchanged `num_samples=30` detected slightly more reachable +points at 3.3× lower wall time; `num_samples=8` reached 4.1× with <1% +detection loss. Analytic solvers (OPW/SRS/UR) are unaffected. + ### Sampling and allocation controls - `SamplingConfig` and `analyze-workspace --sampler` default to scrambled Sobol. diff --git a/embodichain/lab/sim/motion/workspace/analyzer.py b/embodichain/lab/sim/motion/workspace/analyzer.py index 1819d3e08..937ebb6ae 100644 --- a/embodichain/lab/sim/motion/workspace/analyzer.py +++ b/embodichain/lab/sim/motion/workspace/analyzer.py @@ -45,6 +45,10 @@ VisualizationType, VisualizationConfig, MetricConfig, + MetricType, +) +from embodichain.lab.sim.motion.workspace.metrics.manipulability_metric import ( + ManipulabilityMetric, ) from embodichain.lab.sim.motion.workspace.samplers import ( SamplerFactory, @@ -228,6 +232,7 @@ def __init__( self.metrics_results: Dict[str, Any] = {} self.current_mode: AnalysisMode | None = None self.success_rates: torch.Tensor | None = None + self.manipulability_scores: torch.Tensor | None = None # Path of the most recently written/read results cache entry (None until # a disk results cache is used). Exposed for CLI consumers. self._last_cache_path: Path | None = None @@ -1564,6 +1569,13 @@ def analyze( # Step 3: Compute metrics (common for both modes) logger.log_info("[3/3] Computing metrics...") + self.manipulability_scores = None + if self._manipulability_enabled(): + self.manipulability_scores = self._compute_manipulability_scores() + if self.manipulability_scores is not None: + # Row-aligned with joint_configurations (and therefore with + # reachable points in Cartesian/plane modes). + results["manipulability_scores"] = self.manipulability_scores metrics = self._compute_metrics() results["metrics"] = metrics results["config"] = self.config @@ -2011,6 +2023,47 @@ def _generate_point_colors(self, points: np.ndarray) -> np.ndarray: ) return colors + def _manipulability_enabled(self) -> bool: + """Whether manipulability computation is selected in the metric config.""" + enabled = self.config.metric.enabled_metrics or [] + return MetricType.ALL in enabled or MetricType.MANIPULABILITY in enabled + + def _compute_manipulability_scores(self) -> torch.Tensor | None: + """Compute per-configuration Yoshikawa manipulability scores. + + Uses the active control part's solver Jacobian on the stored + ``joint_configurations``, so every score row stays aligned with the + configuration (and, in Cartesian/plane modes, with the reachable + point) at the same index: ``w = sqrt(det(J @ J^T))``. + + Returns: + Scores with shape ``(N,)`` on the analysis device, or ``None`` + when no configurations or no solver Jacobian are available. + """ + qpos = self.joint_configurations + if qpos is None or len(qpos) == 0: + return None + solver = self.robot.get_solver(self.control_part_name) + if solver is None: + logger.log_warning( + "No solver available for manipulability computation; skipping." + ) + return None + + chunk_size = 10000 + scores = [] + with torch.no_grad(): + for start in range(0, len(qpos), chunk_size): + chunk = torch.as_tensor( + qpos[start : start + chunk_size], + dtype=torch.float32, + device=solver.device, + ) + jac = solver.get_jacobian(chunk) + jjt = jac @ jac.transpose(1, 2) + scores.append(torch.sqrt(torch.clamp(torch.det(jjt), min=0.0))) + return torch.cat(scores).to(self.device) + def _compute_metrics(self) -> Dict[str, Any]: """Compute workspace metrics based on configuration.""" if self.workspace_points is None or len(self.workspace_points) == 0: @@ -2019,8 +2072,7 @@ def _compute_metrics(self) -> Dict[str, Any]: metrics = {} - # TODO: Implement metric computation using metrics module - # For now, compute basic statistics + # Basic geometric statistics points_np = self.workspace_points.cpu().numpy() metrics["bounding_box"] = { @@ -2036,6 +2088,15 @@ def _compute_metrics(self) -> Dict[str, Any]: # Approximate volume (bounding box) metrics["bounding_box_volume"] = float(np.prod(dimensions)) + # True Yoshikawa manipulability aggregates from the per-point scores + # computed during analysis (see _compute_manipulability_scores). + if self.manipulability_scores is not None: + metric = ManipulabilityMetric(self.config.metric.manipulability) + metrics["manipulability"] = metric.compute( + points_np, + manipulability_scores=self.manipulability_scores.cpu().numpy(), + ) + logger.log_info(f"Computed {len(metrics)} metrics") return metrics @@ -2409,6 +2470,7 @@ def _restore_analysis_state(self, results: Dict[str, Any]) -> None: ) self.joint_configurations = results.get("joint_configurations") self.success_rates = results.get("success_rates") + self.manipulability_scores = results.get("manipulability_scores") if mode_str in ("cartesian_space", "plane_sampling"): self.reachable_points = results.get("reachable_points") self.reachability_mask = results.get("reachability_mask") diff --git a/embodichain/lab/sim/motion/workspace/caches/results_cache.py b/embodichain/lab/sim/motion/workspace/caches/results_cache.py index 8b02dc42c..a1fc58449 100644 --- a/embodichain/lab/sim/motion/workspace/caches/results_cache.py +++ b/embodichain/lab/sim/motion/workspace/caches/results_cache.py @@ -71,6 +71,7 @@ "joint_configurations", "success_rates", "reachability_mask", + "manipulability_scores", ) # Scalar/dict fields stored in meta.json (JSON-serializable form). diff --git a/embodichain/lab/sim/motion/workspace/metrics/manipulability_metric.py b/embodichain/lab/sim/motion/workspace/metrics/manipulability_metric.py index b24b983a6..f72f758f4 100644 --- a/embodichain/lab/sim/motion/workspace/metrics/manipulability_metric.py +++ b/embodichain/lab/sim/motion/workspace/metrics/manipulability_metric.py @@ -16,6 +16,7 @@ from typing import Dict, Any import numpy as np +from embodichain.utils import logger from embodichain.lab.sim.motion.workspace.metrics.base_metric import ( BaseMetric, ) @@ -27,8 +28,10 @@ class ManipulabilityMetric(BaseMetric): """Manipulability metric for workspace analysis. - Computes dexterity and manipulability measures throughout the workspace. - Note: Full implementation requires robot Jacobian computation. + Computes Yoshikawa manipulability statistics from robot Jacobians or from + precomputed per-point scores. Without either input no statistics are + produced: an earlier centroid-distance placeholder was measured to be + *negatively* correlated with true manipulability and has been removed. """ def __init__(self, config: ManipulabilityConfig | None = None): @@ -44,6 +47,7 @@ def compute( workspace_points: np.ndarray, joint_configurations: np.ndarray | None = None, jacobians: np.ndarray | None = None, + manipulability_scores: np.ndarray | None = None, **kwargs, ) -> Dict[str, Any]: """Compute manipulability metrics. @@ -52,6 +56,8 @@ def compute( workspace_points: Workspace points in Cartesian space, shape (N, 3). joint_configurations: Joint configurations, shape (N, num_joints). jacobians: Precomputed Jacobian matrices, shape (N, 6, num_joints). + manipulability_scores: Precomputed per-point Yoshikawa scores, + shape (N,). Takes precedence over ``jacobians``. **kwargs: Additional arguments. Returns: @@ -60,7 +66,12 @@ def compute( - std_manipulability: Standard deviation - min_manipulability: Minimum value - max_manipulability: Maximum value - - mean_condition: Average condition number (if isotropy enabled) + - mean_condition: Average condition number (if isotropy enabled + and Jacobians were provided) + + Without ``jacobians`` or ``manipulability_scores`` an empty dict is + returned: true manipulability cannot be derived from Cartesian + points alone, and fabricated statistics are worse than none. """ points = self._to_numpy(workspace_points) @@ -72,31 +83,22 @@ def compute( "max_manipulability": 0.0, } - # If Jacobians are not provided, we cannot compute true manipulability - # Return placeholder statistics - if jacobians is None: - # Estimate based on distance from centroid (simple heuristic) - centroid = points.mean(axis=0) - distances = np.linalg.norm(points - centroid, axis=1) - - # Normalize to [0, 1] range (higher manipulability near center) - max_dist = distances.max() if distances.max() > 0 else 1.0 - manipulability_scores = 1.0 - (distances / max_dist) - - # Filter by threshold - valid_mask = manipulability_scores >= self.config.jacobian_threshold - valid_scores = manipulability_scores[valid_mask] - - if len(valid_scores) == 0: - valid_scores = np.array([0.0]) - else: - # Compute true manipulability from Jacobians + if manipulability_scores is not None: + manipulability_scores = self._to_numpy(manipulability_scores) + elif jacobians is not None: manipulability_scores = self._compute_manipulability_index(jacobians) - valid_mask = manipulability_scores >= self.config.jacobian_threshold - valid_scores = manipulability_scores[valid_mask] - - if len(valid_scores) == 0: - valid_scores = np.array([0.0]) + else: + logger.log_warning( + "ManipulabilityMetric needs jacobians or precomputed scores; " + "skipping (no placeholder statistics are produced)." + ) + self.results = {} + return self.results + + valid_mask = manipulability_scores >= self.config.jacobian_threshold + valid_scores = manipulability_scores[valid_mask] + if len(valid_scores) == 0: + valid_scores = np.array([0.0]) self.results = { "mean_manipulability": float(valid_scores.mean()), diff --git a/tests/sim/motion/workspace/test_analysis_batching.py b/tests/sim/motion/workspace/test_analysis_batching.py index c5f5859cb..895530e4d 100644 --- a/tests/sim/motion/workspace/test_analysis_batching.py +++ b/tests/sim/motion/workspace/test_analysis_batching.py @@ -54,6 +54,9 @@ def _analyzer(**kwargs) -> WorkspaceAnalyzer: robot.get_qpos.return_value = torch.zeros(3, 3) robot.cfg = SimpleNamespace(uid="test", fpath=None, solver_cfg={}) robot._solvers = {} + # Faithful to Robot.get_solver with no solvers attached: manipulability + # computation then skips with a warning instead of touching a Mock. + robot.get_solver.return_value = None def fk(qpos, **kw): pose = torch.eye(4).expand(*qpos.shape[:-1], 4, 4).clone() diff --git a/tests/sim/motion/workspace/test_manipulability.py b/tests/sim/motion/workspace/test_manipulability.py new file mode 100644 index 000000000..b7c4183ba --- /dev/null +++ b/tests/sim/motion/workspace/test_manipulability.py @@ -0,0 +1,142 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- +"""Per-point Yoshikawa manipulability in workspace analysis.""" + +from __future__ import annotations + +import numpy as np +import torch + +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.objects import Robot +from embodichain.lab.sim.robots import CobotMagicCfg +from embodichain.lab.sim.motion.workspace.analyzer import ( + AnalysisMode, + WorkspaceAnalyzer, + WorkspaceAnalyzerConfig, +) +from embodichain.lab.sim.motion.workspace.caches.results_cache import ( + deserialize_results, + serialize_results, +) +from embodichain.lab.sim.motion.workspace.configs import ( + MetricConfig, + MetricType, + SamplingConfig, +) +from embodichain.lab.sim.motion.workspace.metrics.manipulability_metric import ( + ManipulabilityMetric, +) + + +class TestManipulabilityMetricUnit: + """Metric-class behaviour without a simulation.""" + + def test_exact_yoshikawa_from_jacobians(self): + # J = [I3 | 0] gives J @ J^T = I -> w = 1 exactly. + jac = np.zeros((5, 6, 7)) + for i in range(5): + jac[i, :6, :6] = np.eye(6) + metric = ManipulabilityMetric() + out = metric.compute(np.zeros((5, 3)), jacobians=jac) + assert abs(out["mean_manipulability"] - 1.0) < 1e-9 + assert out["num_valid_points"] == 5 + + def test_precomputed_scores_take_precedence(self): + metric = ManipulabilityMetric() + scores = np.array([0.2, 0.4, 0.6]) + out = metric.compute(np.zeros((3, 3)), manipulability_scores=scores) + assert abs(out["mean_manipulability"] - 0.4) < 1e-9 + assert abs(out["min_manipulability"] - 0.2) < 1e-9 + + def test_no_inputs_produces_no_fabricated_statistics(self): + """The centroid-distance placeholder was anti-correlated with truth + (measured corr = -0.37 on Franka) and must stay removed.""" + metric = ManipulabilityMetric() + out = metric.compute(np.random.rand(100, 3)) + assert out == {} + + +class TestAnalyzerManipulability: + """End-to-end score plumbing on the library CobotMagic robot.""" + + def setup_method(self): + config = SimulationManagerCfg(headless=True, sim_device="cpu") + self.sim = SimulationManager(config) + cfg_dict = { + "uid": "CobotMagic", + "solver_cfg": { + "left_arm": { + "class_type": "OPWSolver", + "end_link_name": "left_link6", + "root_link_name": "left_arm_base", + "tcp": [ + [1, 0, 0, 0], + [0, 1, 0, 0], + [0, 0, 1, 0.143], + [0, 0, 0, 1], + ], + }, + }, + } + self.robot: Robot = self.sim.add_robot(cfg=CobotMagicCfg.from_dict(cfg_dict)) + + def teardown_method(self): + self.sim.destroy() + SimulationManager.flush_cleanup_queue() + + def _analyze(self, metric_cfg: MetricConfig | None = None): + cfg = WorkspaceAnalyzerConfig( + mode=AnalysisMode.JOINT_SPACE, + sampling=SamplingConfig(num_samples=100), + control_part_name="left_arm", + metric=metric_cfg, + ) + analyzer = WorkspaceAnalyzer(robot=self.robot, config=cfg, sim_manager=self.sim) + return analyzer.analyze(num_samples=100, force_recompute=True), analyzer + + def test_scores_aligned_and_positive(self): + results, analyzer = self._analyze() + scores = results.get("manipulability_scores") + assert scores is not None + assert scores.shape == (len(results["joint_configurations"]),) + assert bool((scores >= 0).all()) + assert float(scores.mean()) > 0.0 + assert analyzer.manipulability_scores is not None + + def test_metrics_contain_true_aggregates(self): + results, _ = self._analyze() + manip = results["metrics"].get("manipulability") + assert manip is not None + assert manip["mean_manipulability"] > 0.0 + assert manip["max_manipulability"] >= manip["min_manipulability"] + + def test_disabled_metric_skips_scores(self): + results, analyzer = self._analyze( + MetricConfig(enabled_metrics=[MetricType.REACHABILITY]) + ) + assert "manipulability_scores" not in results + assert analyzer.manipulability_scores is None + + def test_scores_survive_cache_serialization(self): + results, _ = self._analyze() + arrays, meta = serialize_results(results) + assert "manipulability_scores" in arrays + restored = deserialize_results(arrays, meta) + original = results["manipulability_scores"].cpu().numpy() + np.testing.assert_allclose( + np.asarray(restored["manipulability_scores"]), original, rtol=1e-6 + ) From c9d59ce272749b1cf4956b637fa542a369babdc4 Mon Sep 17 00:00:00 2001 From: Xinyi YUAN Date: Fri, 11 Sep 2026 21:18:05 +0900 Subject: [PATCH 3/5] fix(workspace): repair cached metrics on load and honor isotropy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address two review findings on the manipulability integration: 1. Metric settings are deliberately not part of the results-cache key (a metric toggle must not invalidate the expensive sampling/IK work), so a cache entry written under a different metric configuration — or before manipulability existed — could be returned without scores. The cache-hit path now runs the same _apply_manipulability step as fresh analysis, recomputing scores and aggregates from the cached joint configurations in milliseconds and repairing such entries transparently. 2. The analyzer reduced Jacobians to Yoshikawa scalars and discarded them, so the default compute_isotropy=True could never produce its documented condition statistics. The chunked Jacobian sweep now also collects condition numbers (max/min singular value) when isotropy is enabled, and ManipulabilityMetric accepts them precomputed. Tests cover the repair path (an entry written with manipulability disabled is repaired by a later default-enabled hit on the same key), isotropy presence via the analyzer, and precomputed condition-number passthrough. Co-Authored-By: Claude Fable 5 --- .../lab/sim/motion/workspace/analyzer.py | 89 +++++++++++++------ .../metrics/manipulability_metric.py | 15 +++- .../motion/workspace/test_manipulability.py | 49 ++++++++++ 3 files changed, 124 insertions(+), 29 deletions(-) diff --git a/embodichain/lab/sim/motion/workspace/analyzer.py b/embodichain/lab/sim/motion/workspace/analyzer.py index 937ebb6ae..cd2ea89a9 100644 --- a/embodichain/lab/sim/motion/workspace/analyzer.py +++ b/embodichain/lab/sim/motion/workspace/analyzer.py @@ -1416,6 +1416,11 @@ def analyze( if cached_results is not None: logger.log_info("Loaded results from cache") self._restore_analysis_state(cached_results) + # Metric settings are not part of the cache key: entries + # written under a different metric configuration (or before + # manipulability existed) are repaired here from the cached + # joint configurations. + self._apply_manipulability(cached_results) self._log_analysis_summary(cached_results) if visualize: self._visualize_workspace() @@ -1569,15 +1574,9 @@ def analyze( # Step 3: Compute metrics (common for both modes) logger.log_info("[3/3] Computing metrics...") - self.manipulability_scores = None - if self._manipulability_enabled(): - self.manipulability_scores = self._compute_manipulability_scores() - if self.manipulability_scores is not None: - # Row-aligned with joint_configurations (and therefore with - # reachable points in Cartesian/plane modes). - results["manipulability_scores"] = self.manipulability_scores metrics = self._compute_metrics() results["metrics"] = metrics + self._apply_manipulability(results) results["config"] = self.config results["analysis_time"] = time.time() - start_time @@ -2028,30 +2027,36 @@ def _manipulability_enabled(self) -> bool: enabled = self.config.metric.enabled_metrics or [] return MetricType.ALL in enabled or MetricType.MANIPULABILITY in enabled - def _compute_manipulability_scores(self) -> torch.Tensor | None: - """Compute per-configuration Yoshikawa manipulability scores. + def _compute_manipulability_values( + self, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """Compute per-configuration Yoshikawa scores and condition numbers. Uses the active control part's solver Jacobian on the stored - ``joint_configurations``, so every score row stays aligned with the + ``joint_configurations``, so every row stays aligned with the configuration (and, in Cartesian/plane modes, with the reachable - point) at the same index: ``w = sqrt(det(J @ J^T))``. + point) at the same index: ``w = sqrt(det(J @ J^T))``. Condition + numbers (max/min singular value) are computed only when the metric + config enables isotropy. Returns: - Scores with shape ``(N,)`` on the analysis device, or ``None`` - when no configurations or no solver Jacobian are available. + Tuple of scores ``(N,)`` and condition numbers ``(N,)`` (or + ``None``) on the analysis device; ``(None, None)`` when no + configurations or no solver Jacobian are available. """ qpos = self.joint_configurations if qpos is None or len(qpos) == 0: - return None + return None, None solver = self.robot.get_solver(self.control_part_name) if solver is None: logger.log_warning( "No solver available for manipulability computation; skipping." ) - return None + return None, None + want_isotropy = self.config.metric.manipulability.compute_isotropy chunk_size = 10000 - scores = [] + scores, conditions = [], [] with torch.no_grad(): for start in range(0, len(qpos), chunk_size): chunk = torch.as_tensor( @@ -2062,7 +2067,47 @@ def _compute_manipulability_scores(self) -> torch.Tensor | None: jac = solver.get_jacobian(chunk) jjt = jac @ jac.transpose(1, 2) scores.append(torch.sqrt(torch.clamp(torch.det(jjt), min=0.0))) - return torch.cat(scores).to(self.device) + if want_isotropy: + singulars = torch.linalg.svdvals(jac) + conditions.append( + singulars[:, 0] / torch.clamp(singulars[:, -1], min=1e-15) + ) + return ( + torch.cat(scores).to(self.device), + torch.cat(conditions).to(self.device) if conditions else None, + ) + + def _apply_manipulability(self, results: Dict[str, Any]) -> None: + """Attach manipulability scores and aggregates to a results dict. + + Used on both the fresh-analysis path and the cache-hit path. Metric + settings are deliberately not part of the results-cache key, so a + cached entry may have been produced under a different metric + configuration (or before scores existed); recomputing from the cached + ``joint_configurations`` costs milliseconds and repairs such entries + transparently. + """ + self.manipulability_scores = None + if not self._manipulability_enabled(): + return + scores, conditions = self._compute_manipulability_values() + self.manipulability_scores = scores + if scores is None: + return + results["manipulability_scores"] = scores + metric = ManipulabilityMetric(self.config.metric.manipulability) + metrics = results.setdefault("metrics", {}) + metrics["manipulability"] = metric.compute( + ( + self.workspace_points.cpu().numpy() + if self.workspace_points is not None + else np.empty((0, 3)) + ), + manipulability_scores=scores.cpu().numpy(), + condition_numbers=( + conditions.cpu().numpy() if conditions is not None else None + ), + ) def _compute_metrics(self) -> Dict[str, Any]: """Compute workspace metrics based on configuration.""" @@ -2088,14 +2133,8 @@ def _compute_metrics(self) -> Dict[str, Any]: # Approximate volume (bounding box) metrics["bounding_box_volume"] = float(np.prod(dimensions)) - # True Yoshikawa manipulability aggregates from the per-point scores - # computed during analysis (see _compute_manipulability_scores). - if self.manipulability_scores is not None: - metric = ManipulabilityMetric(self.config.metric.manipulability) - metrics["manipulability"] = metric.compute( - points_np, - manipulability_scores=self.manipulability_scores.cpu().numpy(), - ) + # Manipulability aggregates are attached by _apply_manipulability, + # which also runs on the cache-hit path. logger.log_info(f"Computed {len(metrics)} metrics") diff --git a/embodichain/lab/sim/motion/workspace/metrics/manipulability_metric.py b/embodichain/lab/sim/motion/workspace/metrics/manipulability_metric.py index f72f758f4..154118e27 100644 --- a/embodichain/lab/sim/motion/workspace/metrics/manipulability_metric.py +++ b/embodichain/lab/sim/motion/workspace/metrics/manipulability_metric.py @@ -48,6 +48,7 @@ def compute( joint_configurations: np.ndarray | None = None, jacobians: np.ndarray | None = None, manipulability_scores: np.ndarray | None = None, + condition_numbers: np.ndarray | None = None, **kwargs, ) -> Dict[str, Any]: """Compute manipulability metrics. @@ -58,6 +59,9 @@ def compute( jacobians: Precomputed Jacobian matrices, shape (N, 6, num_joints). manipulability_scores: Precomputed per-point Yoshikawa scores, shape (N,). Takes precedence over ``jacobians``. + condition_numbers: Precomputed per-point Jacobian condition + numbers, shape (N,). Used for isotropy statistics when + ``jacobians`` is not provided. **kwargs: Additional arguments. Returns: @@ -109,10 +113,13 @@ def compute( } # Compute isotropy if requested - if self.config.compute_isotropy and jacobians is not None: - condition_numbers = self._compute_condition_numbers(jacobians) - self.results["mean_condition"] = float(condition_numbers.mean()) - self.results["std_condition"] = float(condition_numbers.std()) + if self.config.compute_isotropy: + if condition_numbers is None and jacobians is not None: + condition_numbers = self._compute_condition_numbers(jacobians) + if condition_numbers is not None: + condition_numbers = self._to_numpy(condition_numbers) + self.results["mean_condition"] = float(condition_numbers.mean()) + self.results["std_condition"] = float(condition_numbers.std()) return self.results diff --git a/tests/sim/motion/workspace/test_manipulability.py b/tests/sim/motion/workspace/test_manipulability.py index b7c4183ba..040fb72c3 100644 --- a/tests/sim/motion/workspace/test_manipulability.py +++ b/tests/sim/motion/workspace/test_manipulability.py @@ -69,6 +69,15 @@ def test_no_inputs_produces_no_fabricated_statistics(self): out = metric.compute(np.random.rand(100, 3)) assert out == {} + def test_precomputed_condition_numbers_feed_isotropy(self): + metric = ManipulabilityMetric() + out = metric.compute( + np.zeros((3, 3)), + manipulability_scores=np.array([0.5, 0.5, 0.5]), + condition_numbers=np.array([2.0, 4.0, 6.0]), + ) + assert abs(out["mean_condition"] - 4.0) < 1e-9 + class TestAnalyzerManipulability: """End-to-end score plumbing on the library CobotMagic robot.""" @@ -131,6 +140,46 @@ def test_disabled_metric_skips_scores(self): assert "manipulability_scores" not in results assert analyzer.manipulability_scores is None + def test_isotropy_condition_stats_honored(self): + results, _ = self._analyze() + manip = results["metrics"]["manipulability"] + # Default ManipulabilityConfig has compute_isotropy=True; the + # documented condition statistics must actually be produced. + assert manip["mean_condition"] >= 1.0 + assert manip["std_condition"] >= 0.0 + + def test_cache_hit_repairs_entry_from_different_metric_config(self, tmp_path): + from embodichain.lab.sim.motion.workspace.configs.cache_config import ( + CacheConfig, + ) + + def run(metric_cfg): + cfg = WorkspaceAnalyzerConfig( + mode=AnalysisMode.JOINT_SPACE, + sampling=SamplingConfig(num_samples=100), + control_part_name="left_arm", + metric=metric_cfg, + cache=CacheConfig(enabled=True, cache_dir=tmp_path), + ) + analyzer = WorkspaceAnalyzer( + robot=self.robot, config=cfg, sim_manager=self.sim + ) + return analyzer.analyze(num_samples=100) + + # First run writes a cache entry WITHOUT manipulability. + first = run(MetricConfig(enabled_metrics=[MetricType.REACHABILITY])) + assert "manipulability_scores" not in first + + # Second run (default metrics) hits the same key — metric settings are + # not part of the cache identity — and must repair the entry on load. + second = run(None) + assert "manipulability_scores" in second + assert second["manipulability_scores"].shape == ( + len(second["joint_configurations"]), + ) + assert second["metrics"]["manipulability"]["mean_manipulability"] > 0.0 + assert second["metrics"]["manipulability"]["mean_condition"] >= 1.0 + def test_scores_survive_cache_serialization(self): results, _ = self._analyze() arrays, meta = serialize_results(results) From 0c833d6c56b5e17cb2d8ce6a1c4028c82004957e Mon Sep 17 00:00:00 2001 From: Xinyi YUAN Date: Fri, 11 Sep 2026 21:27:29 +0900 Subject: [PATCH 4/5] fix(workspace): strip cached manipulability when the metric is disabled The repair-on-load path handled a disabled-producer entry hit by an enabled run, but not the symmetric direction: an enabled-producer entry hit by a disabled run leaked stale manipulability_scores and metrics["manipulability"] into the returned results, diverging from the fresh-analysis contract. _apply_manipulability now strips both fields when the metric is disabled, and keeps cached scores when the metric is enabled but locally not computable (they remain valid for the same joint configurations). Covered by test_cache_hit_strips_fields_when_metric_disabled, which writes a score-bearing entry first so the strip path is genuinely exercised on the hit. Co-Authored-By: Claude Fable 5 --- .../lab/sim/motion/workspace/analyzer.py | 17 +++++++++-- .../motion/workspace/test_manipulability.py | 30 +++++++++++++++++++ 2 files changed, 45 insertions(+), 2 deletions(-) diff --git a/embodichain/lab/sim/motion/workspace/analyzer.py b/embodichain/lab/sim/motion/workspace/analyzer.py index cd2ea89a9..ac6d6fa5d 100644 --- a/embodichain/lab/sim/motion/workspace/analyzer.py +++ b/embodichain/lab/sim/motion/workspace/analyzer.py @@ -2087,13 +2087,26 @@ def _apply_manipulability(self, results: Dict[str, Any]) -> None: ``joint_configurations`` costs milliseconds and repairs such entries transparently. """ - self.manipulability_scores = None if not self._manipulability_enabled(): + # Strip fields a differently-configured producer may have cached, + # so cached and fresh analyses expose the same result contract. + self.manipulability_scores = None + results.pop("manipulability_scores", None) + metrics = results.get("metrics") + if isinstance(metrics, dict): + metrics.pop("manipulability", None) return scores, conditions = self._compute_manipulability_values() - self.manipulability_scores = scores if scores is None: + # Enabled but not computable here (e.g. no solver): keep any + # cached scores — they were derived from the same joint + # configurations and remain valid. + cached = results.get("manipulability_scores") + self.manipulability_scores = ( + torch.as_tensor(cached) if cached is not None else None + ) return + self.manipulability_scores = scores results["manipulability_scores"] = scores metric = ManipulabilityMetric(self.config.metric.manipulability) metrics = results.setdefault("metrics", {}) diff --git a/tests/sim/motion/workspace/test_manipulability.py b/tests/sim/motion/workspace/test_manipulability.py index 040fb72c3..14cea99a7 100644 --- a/tests/sim/motion/workspace/test_manipulability.py +++ b/tests/sim/motion/workspace/test_manipulability.py @@ -180,6 +180,36 @@ def run(metric_cfg): assert second["metrics"]["manipulability"]["mean_manipulability"] > 0.0 assert second["metrics"]["manipulability"]["mean_condition"] >= 1.0 + def test_cache_hit_strips_fields_when_metric_disabled(self, tmp_path): + """Symmetric direction: an enabled run writes a score-bearing entry; + a later disabled run hitting the same key must strip the fields so + cached and fresh analyses expose the same result contract.""" + from embodichain.lab.sim.motion.workspace.configs.cache_config import ( + CacheConfig, + ) + + def run(metric_cfg): + cfg = WorkspaceAnalyzerConfig( + mode=AnalysisMode.JOINT_SPACE, + sampling=SamplingConfig(num_samples=100), + control_part_name="left_arm", + metric=metric_cfg, + cache=CacheConfig(enabled=True, cache_dir=tmp_path), + ) + analyzer = WorkspaceAnalyzer( + robot=self.robot, config=cfg, sim_manager=self.sim + ) + return analyzer.analyze(num_samples=100) + + # Enabled run computes scores and persists them in the entry. + first = run(None) + assert "manipulability_scores" in first + + # Disabled run hits the score-bearing entry: fields must be stripped. + second = run(MetricConfig(enabled_metrics=[MetricType.REACHABILITY])) + assert "manipulability_scores" not in second + assert "manipulability" not in second.get("metrics", {}) + def test_scores_survive_cache_serialization(self): results, _ = self._analyze() arrays, meta = serialize_results(results) From e94860ebd089d407c8b4d4dda20f278fe151d4c4 Mon Sep 17 00:00:00 2001 From: Xinyi YUAN Date: Fri, 11 Sep 2026 21:36:48 +0900 Subject: [PATCH 5/5] fix(workspace): recompute cached aggregates under the current metric config On a cache hit where the current robot has no solver, the previous branch kept both the cached scores and the producer's aggregate metrics, so the returned means/counts could reflect a different jacobian_threshold and condition statistics could be present with isotropy disabled (or stale when enabled). Cached scores are pure kinematics and stay valid, but aggregates now always go through ManipulabilityMetric under the CURRENT configuration; per-point condition numbers are not cached, so condition statistics are correctly absent on this path instead of leaking through. Covered by test_no_solver_hit_recomputes_aggregates_under_current_config: a mock no-solver robot hits an entry carrying producer aggregates from a different configuration, and the returned aggregates honour the current threshold while the stale mean_condition is removed. Co-Authored-By: Claude Fable 5 --- .../lab/sim/motion/workspace/analyzer.py | 23 ++++--- .../motion/workspace/test_manipulability.py | 60 +++++++++++++++++++ 2 files changed, 76 insertions(+), 7 deletions(-) diff --git a/embodichain/lab/sim/motion/workspace/analyzer.py b/embodichain/lab/sim/motion/workspace/analyzer.py index ac6d6fa5d..3ed3622bc 100644 --- a/embodichain/lab/sim/motion/workspace/analyzer.py +++ b/embodichain/lab/sim/motion/workspace/analyzer.py @@ -2098,16 +2098,25 @@ def _apply_manipulability(self, results: Dict[str, Any]) -> None: return scores, conditions = self._compute_manipulability_values() if scores is None: - # Enabled but not computable here (e.g. no solver): keep any - # cached scores — they were derived from the same joint - # configurations and remain valid. + # Enabled but not computable here (e.g. no solver): cached scores + # remain valid — they are pure kinematics of the same joint + # configurations. Per-point condition numbers are not cached, so + # condition statistics are unavailable on this path. cached = results.get("manipulability_scores") - self.manipulability_scores = ( - torch.as_tensor(cached) if cached is not None else None - ) - return + if cached is None: + self.manipulability_scores = None + metrics = results.get("metrics") + if isinstance(metrics, dict): + metrics.pop("manipulability", None) + return + scores = torch.as_tensor(cached).to(self.device) + conditions = None self.manipulability_scores = scores results["manipulability_scores"] = scores + # Aggregates are always recomputed under the CURRENT metric + # configuration: cached aggregates may reflect a different + # jacobian_threshold or isotropy setting, since metric settings are + # not part of the cache key. metric = ManipulabilityMetric(self.config.metric.manipulability) metrics = results.setdefault("metrics", {}) metrics["manipulability"] = metric.compute( diff --git a/tests/sim/motion/workspace/test_manipulability.py b/tests/sim/motion/workspace/test_manipulability.py index 14cea99a7..b99478ccd 100644 --- a/tests/sim/motion/workspace/test_manipulability.py +++ b/tests/sim/motion/workspace/test_manipulability.py @@ -210,6 +210,66 @@ def run(metric_cfg): assert "manipulability_scores" not in second assert "manipulability" not in second.get("metrics", {}) + def test_no_solver_hit_recomputes_aggregates_under_current_config(self): + """Cached scores stay valid without a solver, but aggregates must be + recomputed under the CURRENT metric configuration, and condition + statistics must not leak through when they cannot be derived.""" + from types import SimpleNamespace + from unittest.mock import Mock + + from embodichain.lab.sim.motion.workspace.configs.metric_config import ( + ManipulabilityConfig, + ) + + robot = Mock() + robot.device = torch.device("cpu") + robot.num_envs = 1 + robot.control_parts = {"arm": ["j0", "j1", "j2"]} + robot.get_joint_ids.return_value = [0, 1, 2] + robot.body_data = SimpleNamespace(qpos_limits=torch.tensor([[[-1.0, 1.0]] * 3])) + robot.get_qpos.return_value = torch.zeros(1, 3) + robot.cfg = SimpleNamespace(uid="t", fpath=None, solver_cfg={}) + robot._solvers = {} + robot.get_solver.return_value = None + + cfg = WorkspaceAnalyzerConfig( + mode=AnalysisMode.JOINT_SPACE, + sampling=SamplingConfig(num_samples=10), + control_part_name="arm", + metric=MetricConfig( + manipulability=ManipulabilityConfig( + jacobian_threshold=0.3, compute_isotropy=True + ) + ), + ) + analyzer = WorkspaceAnalyzer(robot, cfg) + analyzer.joint_configurations = torch.zeros(4, 3) + analyzer.workspace_points = torch.zeros(4, 3) + + cached_scores = torch.tensor([0.1, 0.2, 0.4, 0.6]) + results = { + "manipulability_scores": cached_scores.clone(), + "metrics": { + "manipulability": { + # Producer aggregates under a DIFFERENT configuration. + "mean_manipulability": 999.0, + "num_valid_points": 4, + "mean_condition": 123.0, + } + }, + } + analyzer._apply_manipulability(results) + + manip = results["metrics"]["manipulability"] + # Current jacobian_threshold=0.3 filters 0.1 and 0.2. + assert manip["num_valid_points"] == 2 + assert abs(manip["mean_manipulability"] - 0.5) < 1e-6 + # Isotropy is enabled but conditions cannot be derived without a + # solver: the stale producer value must not leak through. + assert "mean_condition" not in manip + # The kinematic scores themselves are kept unchanged. + assert torch.allclose(results["manipulability_scores"], cached_scores) + def test_scores_survive_cache_serialization(self): results, _ = self._analyze() arrays, meta = serialize_results(results)