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/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..0b9fa03d4 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,57 @@ | 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. | + +### 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. + 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..3ed3622bc 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, @@ -112,6 +116,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 +161,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 +210,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( @@ -221,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 @@ -345,44 +357,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 +627,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 +770,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 +848,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 +1030,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 +1069,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 +1099,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 +1107,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 +1185,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 +1241,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 +1255,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 +1358,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 +1373,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") @@ -1283,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() @@ -1438,9 +1576,22 @@ def analyze( logger.log_info("[3/3] Computing metrics...") metrics = self._compute_metrics() results["metrics"] = metrics + self._apply_manipulability(results) 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 +1914,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 @@ -1868,6 +2022,115 @@ 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_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 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))``. Condition + numbers (max/min singular value) are computed only when the metric + config enables isotropy. + + Returns: + 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, 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, None + + want_isotropy = self.config.metric.manipulability.compute_isotropy + chunk_size = 10000 + scores, conditions = [], [] + 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))) + 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. + """ + 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() + if scores is None: + # 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") + 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( + ( + 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.""" if self.workspace_points is None or len(self.workspace_points) == 0: @@ -1876,8 +2139,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"] = { @@ -1893,6 +2155,9 @@ def _compute_metrics(self) -> Dict[str, Any]: # Approximate volume (bounding box) metrics["bounding_box_volume"] = float(np.prod(dimensions)) + # Manipulability aggregates are attached by _apply_manipulability, + # which also runs on the cache-hit path. + logger.log_info(f"Computed {len(metrics)} metrics") return metrics @@ -1980,7 +2245,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 +2437,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,9 +2526,12 @@ 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") + 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/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/metrics/manipulability_metric.py b/embodichain/lab/sim/motion/workspace/metrics/manipulability_metric.py index b24b983a6..154118e27 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,8 @@ def compute( workspace_points: np.ndarray, 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. @@ -52,6 +57,11 @@ 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``. + condition_numbers: Precomputed per-point Jacobian condition + numbers, shape (N,). Used for isotropy statistics when + ``jacobians`` is not provided. **kwargs: Additional arguments. Returns: @@ -60,7 +70,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 +87,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()), @@ -107,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/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..895530e4d --- /dev/null +++ b/tests/sim/motion/workspace/test_analysis_batching.py @@ -0,0 +1,296 @@ +# ---------------------------------------------------------------------------- +# 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 = {} + # 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() + 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) diff --git a/tests/sim/motion/workspace/test_manipulability.py b/tests/sim/motion/workspace/test_manipulability.py new file mode 100644 index 000000000..b99478ccd --- /dev/null +++ b/tests/sim/motion/workspace/test_manipulability.py @@ -0,0 +1,281 @@ +# ---------------------------------------------------------------------------- +# 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 == {} + + 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.""" + + 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_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_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_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) + 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 + )