diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index b670e0448..abb08a6c7 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -71,7 +71,7 @@ topics: title: IK Solver System aliases: [ik solver, inverse kinematics, IK 求解器, 逆运动学] keywords: [BaseSolver, SolverCfg, SRSSolver, OPWSolver, URSolver, PinkSolver, PytorchSolver, PinocchioSolver, - DifferentialSolver, NeuralIKSolver, QposSeedSampler, NullSpacePostureTask, compute kinematics, Warp + DifferentialSolver, NeuralIKSolver, QposSeedSampler, QposSeedSelSampler, NullSpacePostureTask, compute kinematics, Warp kinematics] paths: [topics/ik-solvers/ik-solvers.md] source_of_truth: diff --git a/agent_context/topics/ik-solvers/solver-details.md b/agent_context/topics/ik-solvers/solver-details.md index 6738894b9..77a0b4526 100644 --- a/agent_context/topics/ik-solvers/solver-details.md +++ b/agent_context/topics/ik-solvers/solver-details.md @@ -103,6 +103,23 @@ batches for IK multi-start: - First sample = provided seed; remaining are uniform-random within limits. - `repeat_target_xpos(target_xpos, num_samples)` — repeats target poses to match expanded seed batch. +### `QposSeedSelSampler` (`qpos_seed_sel_sampler.py`) + +Database-driven drop-in extension of `QposSeedSampler` (SELIK-style retrieval). +Slots after the caller seed come from a lazily built Sobol FK database: pose-space +kNN (position + Frobenius rotation, `rot_scale` metres/radian) re-ranked by the +predicted joint step `||J⁺·Δpose||` when a Jacobian provider is configured. + +- Opt-in through `PytorchSolverCfg.enable_seed_selection` with `seed_db_size` + and `seed_rot_scale`; default off preserves shipped behaviour exactly. +- The database stores flange poses (no TCP), so runtime `set_tcp()` never + invalidates it; `get_ik` queries with its TCP-stripped target. +- Joint-limit changes trigger an automatic rebuild; `sample()` without a + `target_xpos` falls back to the parent's uniform-random behaviour. +- Cost: build ~11 ms / 6 MB at 20k entries on GPU (one-time, lazy); query + adds <1 ms per `get_ik` call (~0.2% of a solve). +- Analytic solvers (SRS/OPW/UR) do not consume seeds and are unaffected. + ### `NullSpacePostureTask` (`null_space_posture_task.py`) A `pink.tasks.Task` subclass for posture control in the null space of diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.solvers.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.solvers.rst index bc70ae02e..c8e58bd27 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.solvers.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.motion.solvers.rst @@ -149,3 +149,12 @@ Neural IK Solver :members: :inherited-members: :show-inheritance: + +Seed Selection +-------------- + +.. currentmodule:: embodichain.lab.sim.motion.solvers.qpos_seed_sel_sampler + +.. autoclass:: QposSeedSelSampler + :members: + :show-inheritance: diff --git a/embodichain/lab/sim/motion/solvers/pytorch_solver.py b/embodichain/lab/sim/motion/solvers/pytorch_solver.py index d1d257351..aaa45b662 100644 --- a/embodichain/lab/sim/motion/solvers/pytorch_solver.py +++ b/embodichain/lab/sim/motion/solvers/pytorch_solver.py @@ -23,6 +23,9 @@ from embodichain.utils import configclass, logger from embodichain.lab.sim.motion.solvers import SolverCfg, BaseSolver from embodichain.lab.sim.motion.solvers.qpos_seed_sampler import QposSeedSampler +from embodichain.lab.sim.motion.solvers.qpos_seed_sel_sampler import ( + QposSeedSelSampler, +) from embodichain.lab.sim.utility.solver_utils import validate_iteration_params if TYPE_CHECKING: @@ -70,11 +73,28 @@ class PytorchSolverCfg(SolverCfg): ik_nearest_weight: list[float] | None = None """Weights for the inverse kinematics nearest calculation. - + The weights influence how the solver prioritizes closeness to the seed position when multiple solutions are available. """ + enable_seed_selection: bool = False + """Retrieve multi-start seeds from a precomputed FK database. + + When enabled, the random slots of the multi-start seed batch are replaced + by database configurations whose flange poses are nearest to each target + (re-ranked by the predicted joint-space correction). Slot ``0`` still + holds the caller-provided seed. The database is built lazily on the first + ``get_ik`` call and rebuilt automatically when joint limits change. + """ + + seed_db_size: int = 20000 + """Number of joint configurations stored in the seed-selection database.""" + + seed_rot_scale: float = 0.2 + """Metres-per-radian weight of the rotation block in the seed-retrieval + pose metric.""" + def init_solver( self, device: torch.device = torch.device("cpu"), **kwargs ) -> "PytorchSolver": @@ -175,6 +195,33 @@ def __init__( self.dof = self.pk_serial_chain.n_joints + # Optional database-driven seed selection. The database stores flange + # poses (no TCP), so runtime ``set_tcp`` calls never invalidate it; + # ``get_ik`` queries it with the TCP-stripped target accordingly. + self._seed_sampler: QposSeedSelSampler | None = None + if cfg.enable_seed_selection: + self._seed_sampler = QposSeedSelSampler( + num_samples=self._num_samples, + dof=self.dof, + device=self.device, + fk_fn=self._compute_flange_fk, + jacobian_fn=self.get_jacobian, + db_size=cfg.seed_db_size, + rot_scale=cfg.seed_rot_scale, + ) + + def _compute_flange_fk(self, qpos: torch.Tensor) -> torch.Tensor: + """Compute end-link poses without the TCP transform. + + Args: + qpos (torch.Tensor): Joint positions with shape (N, dof). + + Returns: + torch.Tensor: Flange poses with shape (N, 4, 4). + """ + qpos = torch.as_tensor(qpos, dtype=torch.float32, device=self.device) + return self.compiled_fk(qpos)[-1, :, :, :] + def get_iteration_params(self) -> dict: r"""Returns the current iteration parameters. @@ -391,15 +438,28 @@ def get_ik( batch_size = target_xpos.shape[0] - sampler = QposSeedSampler( - num_samples=self._num_samples, dof=self.dof, device=self.device - ) - random_qpos_seeds = sampler.sample( - qpos_seed, - self.lower_qpos_limits, - self.upper_qpos_limits, - batch_size, - ) + if self._seed_sampler is not None: + # Database retrieval: ``target_xpos`` is in the flange frame here + # (TCP stripped above), matching the frame the database stores. + sampler = self._seed_sampler + sampler.num_samples = self._num_samples + random_qpos_seeds = sampler.sample( + qpos_seed, + self.lower_qpos_limits, + self.upper_qpos_limits, + batch_size, + target_xpos=target_xpos, + ) + else: + sampler = QposSeedSampler( + num_samples=self._num_samples, dof=self.dof, device=self.device + ) + random_qpos_seeds = sampler.sample( + qpos_seed, + self.lower_qpos_limits, + self.upper_qpos_limits, + batch_size, + ) target_xpos_repeated = sampler.repeat_target_xpos( target_xpos, self._num_samples ) diff --git a/embodichain/lab/sim/motion/solvers/qpos_seed_sel_sampler.py b/embodichain/lab/sim/motion/solvers/qpos_seed_sel_sampler.py new file mode 100644 index 000000000..835b4455d --- /dev/null +++ b/embodichain/lab/sim/motion/solvers/qpos_seed_sel_sampler.py @@ -0,0 +1,334 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Database-driven joint seed selection for iterative IK multi-start. + +Instead of filling the multi-start seed batch with uniform random draws, +:class:`QposSeedSelSampler` retrieves seeds from a precomputed forward-kinematics +database: the joint space is sampled once with a low-discrepancy sequence, the +end-effector pose of every sample is stored, and at query time the nearest +database entries to each target pose are returned as seeds. When a Jacobian +provider is available, candidates are re-ranked by the predicted joint-space +step ``||J^+ (target - seed_pose)||`` so that seeds requiring the smallest +correction are tried first. + +The retrieval strategy follows the SELIK solver of the WRS framework +(Wan Weiwei, 2023, MIT licensed); this is an independent batched PyTorch +re-implementation adapted to the EmbodiChain seed-sampler contract. + +The class is a drop-in extension of :class:`QposSeedSampler`: calls without a +target pose fall back to the parent's uniform random behaviour, and slot ``0`` +of the returned batch preserves the caller-provided seed so warm-start chains +(for example sequential waypoint solving) keep their continuity. +""" + +from __future__ import annotations + +from typing import Callable + +import torch + +from embodichain.utils import logger +from embodichain.utils.math import axis_angle_from_quat, quat_from_matrix + +from .qpos_seed_sampler import QposSeedSampler + +__all__ = ["QposSeedSelSampler"] + +FkFn = Callable[[torch.Tensor], torch.Tensor] +"""Batched forward kinematics: ``(N, dof)`` joints to ``(N, 4, 4)`` poses.""" + +JacobianFn = Callable[[torch.Tensor], torch.Tensor] +"""Batched geometric Jacobian: ``(N, dof)`` joints to ``(N, 6, dof)``.""" + + +def _pose_error_twist(cur: torch.Tensor, tgt: torch.Tensor) -> torch.Tensor: + """Compute the ``[dp, rotvec]`` twist between pose batches. + + Args: + cur: Current poses with shape ``(N, 4, 4)``. + tgt: Target poses with shape ``(N, 4, 4)``. + + Returns: + Twist tensor with shape ``(N, 6)``. + """ + dp = tgt[:, :3, 3] - cur[:, :3, 3] + rel = tgt[:, :3, :3] @ cur[:, :3, :3].transpose(1, 2) + # Quaternion route: numerically stable over the whole rotation range, + # including relative rotations at and near 180 degrees where a direct + # skew-symmetric axis extraction degenerates to zero. + rotvec = axis_angle_from_quat(quat_from_matrix(rel.contiguous())) + return torch.cat([dp, rotvec], dim=-1) + + +class QposSeedSelSampler(QposSeedSampler): + """SELIK-style seed sampler backed by a forward-kinematics database. + + The database is built lazily on the first target-aware :meth:`sample` call + using the joint limits supplied by the caller, and rebuilt automatically + whenever those limits change. Building is a batched FK sweep and typically + takes well under a second on GPU for the default database size. + + Args: + num_samples: Number of seeds per target (including the caller seed). + dof: Degrees of freedom. + device: Target device. + fk_fn: Batched forward kinematics of the solved chain, in the same + frame and with the same TCP as the IK targets. + jacobian_fn: Optional batched Jacobian provider. When given, retrieved + candidates are re-ranked by predicted joint-space step length. + db_size: Number of joint configurations stored in the database. + rot_scale: Metres-per-radian weight applied to the rotation block of + the pose metric used for nearest-neighbour retrieval. + k_max: Number of nearest neighbours fetched before re-ranking. + use_caller_seed: If ``True`` (default), slot ``0`` of every returned + batch is the caller-provided seed; if ``False`` all slots come + from the database. + sobol_seed: Scramble seed of the low-discrepancy joint sampler. + """ + + _LIMITS_ATOL = 1e-6 + _FK_CHUNK = 10000 + _QUERY_CHUNK = 256 + + def __init__( + self, + num_samples: int, + dof: int, + device: torch.device, + *, + fk_fn: FkFn, + jacobian_fn: JacobianFn | None = None, + db_size: int = 20000, + rot_scale: float = 0.2, + k_max: int = 200, + use_caller_seed: bool = True, + sobol_seed: int = 0, + ) -> None: + super().__init__(num_samples=num_samples, dof=dof, device=device) + if db_size <= 0: + raise ValueError(f"db_size must be positive, got {db_size}.") + if rot_scale <= 0.0: + raise ValueError(f"rot_scale must be positive, got {rot_scale}.") + if k_max <= 0: + raise ValueError(f"k_max must be positive, got {k_max}.") + self._fk_fn = fk_fn + self._jacobian_fn = jacobian_fn + self._db_size = db_size + self._rot_scale = rot_scale + self._k_max = k_max + self._use_caller_seed = use_caller_seed + self._sobol_seed = sobol_seed + + self._db_qpos: torch.Tensor | None = None + self._db_pose: torch.Tensor | None = None + self._db_vec: torch.Tensor | None = None + self._db_jinv: torch.Tensor | None = None + self._db_limits: tuple[torch.Tensor, torch.Tensor] | None = None + + # ------------------------------------------------------------------ database + + @property + def database_size(self) -> int: + """Number of entries in the built database, or ``0`` before build.""" + return 0 if self._db_qpos is None else int(self._db_qpos.shape[0]) + + def _limits_changed( + self, lower_limits: torch.Tensor, upper_limits: torch.Tensor + ) -> bool: + if self._db_limits is None: + return True + lo, hi = self._db_limits + return not ( + torch.allclose(lo, lower_limits, atol=self._LIMITS_ATOL) + and torch.allclose(hi, upper_limits, atol=self._LIMITS_ATOL) + ) + + def _pose_vec(self, pose: torch.Tensor) -> torch.Tensor: + """Flatten poses into the retrieval metric space, shape ``(N, 12)``.""" + return torch.cat( + [pose[:, :3, 3], self._rot_scale * pose[:, :3, :3].reshape(-1, 9)], + dim=-1, + ) + + def _build_database( + self, lower_limits: torch.Tensor, upper_limits: torch.Tensor + ) -> None: + """Sample the joint space and store poses, metric vectors and J-pinv.""" + sobol = torch.quasirandom.SobolEngine( + dimension=self.dof, scramble=True, seed=self._sobol_seed + ) + unit = sobol.draw(self._db_size).to( + device=lower_limits.device, dtype=lower_limits.dtype + ) + qpos = lower_limits + unit * (upper_limits - lower_limits) + + poses, vecs, jinvs = [], [], [] + with torch.no_grad(): + for start in range(0, self._db_size, self._FK_CHUNK): + chunk = qpos[start : start + self._FK_CHUNK] + pose = self._fk_fn(chunk) + poses.append(pose) + vecs.append(self._pose_vec(pose)) + if self._jacobian_fn is not None: + jac = self._jacobian_fn(chunk) + jinvs.append(torch.linalg.pinv(jac, rcond=1e-4)) + + self._db_qpos = qpos + self._db_pose = torch.cat(poses) + self._db_vec = torch.cat(vecs) + self._db_jinv = torch.cat(jinvs) if jinvs else None + self._db_limits = (lower_limits.clone(), upper_limits.clone()) + + # ------------------------------------------------------------------ retrieval + + def _query(self, target_xpos: torch.Tensor, k: int) -> torch.Tensor: + """Retrieve the top-``k`` seeds per target, shape ``(B, k, dof)``.""" + assert self._db_qpos is not None and self._db_vec is not None + k = min(k, self._db_qpos.shape[0]) + k_max = min(self._k_max, self._db_qpos.shape[0]) + query_vec = self._pose_vec(target_xpos) + + index_chunks = [] + for start in range(0, query_vec.shape[0], self._QUERY_CHUNK): + dist = torch.cdist( + query_vec[start : start + self._QUERY_CHUNK], self._db_vec + ) + index_chunks.append(dist.topk(k_max, largest=False).indices) + indices = torch.cat(index_chunks) + + if self._db_jinv is None: + return self._db_qpos[indices[:, :k]] + + batch = indices.shape[0] + seed_pose = self._db_pose[indices].reshape(-1, 4, 4) + target_rep = ( + target_xpos.unsqueeze(1).expand(batch, k_max, 4, 4).reshape(-1, 4, 4) + ) + twist = _pose_error_twist(seed_pose, target_rep).view(batch, k_max, 6) + step = torch.einsum("bkij,bkj->bki", self._db_jinv[indices], twist) + order = step.pow(2).sum(dim=-1).argsort(dim=1) + return self._db_qpos[torch.gather(indices, 1, order[:, :k])] + + def _pad_with_random( + self, + seeds: torch.Tensor, + count: int, + lower_limits: torch.Tensor, + upper_limits: torch.Tensor, + ) -> torch.Tensor: + """Pad retrieved seeds up to ``count`` slots per target. + + Retrieval is capped by the database size and by ``k_max``, so it can + return fewer candidates than requested. Shortfall slots are filled + with uniform random draws within the limits — the parent sampler's + behaviour — so the ``batch_size * num_samples`` output contract holds + for every configuration. + + Args: + seeds: Retrieved seeds with shape ``(B, k_got, dof)``. + count: Required number of slots per target. + lower_limits: ``(dof,)`` lower joint limits. + upper_limits: ``(dof,)`` upper joint limits. + + Returns: + torch.Tensor: Seeds with shape ``(B, count, dof)``. + """ + shortfall = count - seeds.shape[1] + if shortfall <= 0: + return seeds[:, :count] + random_fill = torch.rand( + seeds.shape[0], + shortfall, + self.dof, + device=seeds.device, + dtype=seeds.dtype, + ) + random_fill = lower_limits + random_fill * (upper_limits - lower_limits) + return torch.cat([seeds, random_fill], dim=1) + + # ------------------------------------------------------------------ sampling + + def sample( + self, + qpos_seed: torch.Tensor, + lower_limits: torch.Tensor, + upper_limits: torch.Tensor, + batch_size: int, + target_xpos: torch.Tensor | None = None, + ) -> torch.Tensor: + """Generate joint seeds, retrieving from the database when possible. + + Args: + qpos_seed: ``(batch_size, dof)`` or ``(dof,)`` caller seed. + lower_limits: ``(dof,)`` lower joint limits. + upper_limits: ``(dof,)`` upper joint limits. + batch_size: Number of targets. + target_xpos: Optional ``(batch_size, 4, 4)`` target poses in the + same frame ``fk_fn`` produces. When omitted, behaviour is + identical to :class:`QposSeedSampler`. + + Returns: + torch.Tensor: ``(batch_size * num_samples, dof)`` joint seeds, + target-major, slot ``0`` holding the caller seed unless + ``use_caller_seed=False``. The shape contract holds for every + configuration: when retrieval returns fewer candidates than + requested (database smaller than the seed count, or + ``num_samples - 1 > k_max``), the shortfall is filled with + uniform random draws within the limits. + """ + if target_xpos is None: + return super().sample(qpos_seed, lower_limits, upper_limits, batch_size) + + if target_xpos.shape != (batch_size, 4, 4): + logger.log_error( + f"target_xpos must have shape ({batch_size}, 4, 4), " + f"got {tuple(target_xpos.shape)}.", + ValueError, + ) + if qpos_seed.shape == (batch_size, self.dof): + seed_head = qpos_seed[:, None, :] + elif qpos_seed.shape == (self.dof,): + seed_head = qpos_seed.unsqueeze(0).repeat(batch_size, 1)[:, None, :] + else: + logger.log_error( + f"Invalid qpos_seed shape {qpos_seed.shape} for batch_size " + f"{batch_size} and dof {self.dof}", + ValueError, + ) + + if self._limits_changed(lower_limits, upper_limits): + self._build_database(lower_limits, upper_limits) + + if self._use_caller_seed: + n_retrieved = self.num_samples - 1 + if n_retrieved == 0: + return seed_head.reshape(-1, self.dof) + retrieved = self._pad_with_random( + self._query(target_xpos, n_retrieved), + n_retrieved, + lower_limits, + upper_limits, + ) + joint_seeds = torch.cat([seed_head, retrieved], dim=1) + else: + joint_seeds = self._pad_with_random( + self._query(target_xpos, self.num_samples), + self.num_samples, + lower_limits, + upper_limits, + ) + return joint_seeds.reshape(-1, self.dof) diff --git a/tests/sim/motion/solvers/test_qpos_seed_sel_sampler.py b/tests/sim/motion/solvers/test_qpos_seed_sel_sampler.py new file mode 100644 index 000000000..bd28064fd --- /dev/null +++ b/tests/sim/motion/solvers/test_qpos_seed_sel_sampler.py @@ -0,0 +1,325 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +from __future__ import annotations + +import os + +import pytest +import torch + +from embodichain.data import get_data_path +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.motion.solvers import PytorchSolverCfg +from embodichain.lab.sim.cfg import RobotCfg +from embodichain.lab.sim.objects import Robot +from embodichain.lab.sim.motion.solvers.qpos_seed_sel_sampler import ( + QposSeedSelSampler, + _pose_error_twist, +) +from embodichain.utils.utility import reset_all_seeds + +_DB_SIZE = 5000 +_NUM_TARGETS = 24 + + +def _rot_x(angle: float) -> torch.Tensor: + pose = torch.eye(4) + c, s = torch.cos(torch.tensor(angle)), torch.sin(torch.tensor(angle)) + pose[1, 1], pose[1, 2], pose[2, 1], pose[2, 2] = c, -s, s, c + return pose + + +class TestPoseErrorTwist: + """Rotation-error stability over the whole rotation range (no sim needed).""" + + def test_exact_pi_rotation(self): + cur = torch.eye(4).unsqueeze(0) + tgt = _rot_x(torch.pi).unsqueeze(0) + twist = _pose_error_twist(cur, tgt) + assert abs(twist[0, 3:].norm().item() - torch.pi) < 1e-5 + + def test_near_pi_rotation(self): + angle = torch.pi - 1e-3 + cur = torch.eye(4).unsqueeze(0) + tgt = _rot_x(angle).unsqueeze(0) + twist = _pose_error_twist(cur, tgt) + assert abs(twist[0, 3:].norm().item() - angle) < 1e-4 + + def test_small_and_moderate_rotations(self): + for angle in (1e-4, 0.5, 2.0): + cur = torch.eye(4).unsqueeze(0) + tgt = _rot_x(angle).unsqueeze(0) + twist = _pose_error_twist(cur, tgt) + assert abs(twist[0, 3:].norm().item() - angle) < 1e-4 + + def test_flipped_candidate_not_assigned_zero_cost(self): + """A 180-degree-error candidate must rank behind an aligned one.""" + target = torch.eye(4).unsqueeze(0) + aligned = torch.eye(4).unsqueeze(0) + aligned[0, 0, 3] = 0.05 # small position offset, same orientation + flipped = _rot_x(torch.pi).unsqueeze(0) # same position, opposite pose + cost_aligned = _pose_error_twist(aligned, target).norm().item() + cost_flipped = _pose_error_twist(flipped, target).norm().item() + assert cost_flipped > 3.0 + assert cost_flipped > cost_aligned + + +class TestQposSeedSelSampler: + """Seed-selection sampler tests on the library DexforceW1 robot.""" + + def setup_method(self): + config = SimulationManagerCfg(headless=True, sim_device="cpu") + self.sim = SimulationManager(config) + + urdf = get_data_path("DexforceW1V021/DexforceW1_v02_1.urdf") + assert os.path.isfile(urdf) + + cfg_dict = { + "fpath": urdf, + "control_parts": { + "left_arm": [f"LEFT_J{i+1}" for i in range(7)], + }, + "solver_cfg": { + "left_arm": { + "class_type": "PytorchSolver", + "end_link_name": "left_ee", + "root_link_name": "left_arm_base", + "num_samples": 30, + }, + }, + } + self.robot: Robot = self.sim.add_robot(cfg=RobotCfg.from_dict(cfg_dict)) + self.solver = self.robot.get_solver("left_arm") + self.lower = self.solver.lower_qpos_limits + self.upper = self.solver.upper_qpos_limits + self.dof = self.solver.dof + + def teardown_method(self): + self.sim.destroy() + SimulationManager.flush_cleanup_queue() + + # ------------------------------------------------------------------ helpers + + def _make_sampler(self, num_samples: int, **kwargs) -> QposSeedSelSampler: + defaults = dict( + fk_fn=self.solver.get_fk, + jacobian_fn=self.solver.get_jacobian, + db_size=_DB_SIZE, + sobol_seed=0, + ) + defaults.update(kwargs) + return QposSeedSelSampler( + num_samples=num_samples, + dof=self.dof, + device=self.solver.device, + **defaults, + ) + + def _reachable_targets(self, n: int) -> tuple[torch.Tensor, torch.Tensor]: + """FK poses of random valid joint configurations, shape ``(n, 4, 4)``.""" + reset_all_seeds(0) + q_true = self.lower + torch.rand(n, self.dof) * (self.upper - self.lower) + with torch.no_grad(): + return self.solver.get_fk(q_true), q_true + + # ------------------------------------------------------------------ contract + + def test_fallback_and_target_aware_contract(self): + sampler = self._make_sampler(num_samples=6) + seed = torch.zeros(self.dof) + + # Without a target: parent behaviour, no database build. + out = sampler.sample(seed, self.lower, self.upper, batch_size=3) + assert out.shape == (18, self.dof) + grouped = out.view(3, 6, self.dof) + assert torch.allclose(grouped[:, 0], seed.expand(3, self.dof)) + assert sampler.database_size == 0 + + # With a target: same contract, slot 0 preserved, database built. + target, _ = self._reachable_targets(3) + out = sampler.sample( + seed, self.lower, self.upper, batch_size=3, target_xpos=target + ) + assert out.shape == (18, self.dof) + grouped = out.view(3, 6, self.dof) + assert torch.allclose(grouped[:, 0], seed.expand(3, self.dof)) + assert (grouped >= self.lower - 1e-6).all() + assert (grouped <= self.upper + 1e-6).all() + assert sampler.database_size == _DB_SIZE + + def test_retrieved_seeds_are_task_space_close(self): + """Retrieved seeds must reach far closer to the target than random.""" + sampler = self._make_sampler(num_samples=8) + target, _ = self._reachable_targets(_NUM_TARGETS) + out = sampler.sample( + torch.zeros(self.dof), + self.lower, + self.upper, + batch_size=_NUM_TARGETS, + target_xpos=target, + ) + retrieved = out.view(_NUM_TARGETS, 8, self.dof)[:, 1:] + + with torch.no_grad(): + seed_pos = self.solver.get_fk(retrieved.reshape(-1, self.dof))[:, :3, 3] + seed_pos = seed_pos.view(_NUM_TARGETS, 7, 3) + target_pos = target[:, :3, 3].unsqueeze(1) + db_dist = (seed_pos - target_pos).norm(dim=-1).mean() + + reset_all_seeds(1) + rand_q = self.lower + torch.rand(_NUM_TARGETS * 7, self.dof) * ( + self.upper - self.lower + ) + with torch.no_grad(): + rand_pos = self.solver.get_fk(rand_q)[:, :3, 3] + rand_dist = (rand_pos.view(_NUM_TARGETS, 7, 3) - target_pos).norm(dim=-1).mean() + + assert db_dist < 0.5 * rand_dist + + # ------------------------------------------------------------------ end to end + + def test_cfg_enabled_seed_selection_improves_get_ik(self): + """The configuration route must wire retrieval into the real get_ik.""" + target, _ = self._reachable_targets(_NUM_TARGETS) + k = 4 + + # Default-off solver: shipped random seeding, no sampler attached. + assert self.solver._seed_sampler is None + reset_all_seeds(0) + ok_shipped, _ = self.solver.get_ik(target.clone(), num_samples=k) + shipped = ok_shipped.float().mean().item() + + # Same chain with seed selection enabled through the config. + urdf = get_data_path("DexforceW1V021/DexforceW1_v02_1.urdf") + sel_solver = PytorchSolverCfg( + urdf_path=urdf, + end_link_name="left_ee", + root_link_name="left_arm_base", + # Standalone solvers do not get the robot-populated joint names + # and default nearest weights; provide both explicitly. + joint_names=[f"LEFT_J{i+1}" for i in range(7)], + ik_nearest_weight=[1.0] * 7, + num_samples=30, + enable_seed_selection=True, + seed_db_size=_DB_SIZE, + ).init_solver(device=self.solver.device) + assert sel_solver._seed_sampler is not None + + reset_all_seeds(0) + ok_sel, _ = sel_solver.get_ik(target.clone(), num_samples=k) + sel = ok_sel.float().mean().item() + + assert sel >= shipped + assert sel >= 0.8 + # The database was built lazily by the first get_ik call. + assert sel_solver._seed_sampler.database_size == _DB_SIZE + + # ------------------------------------------------------------------ capped retrieval + + def test_db_smaller_than_requested_seed_count(self): + """Reviewer case: db_size=2, num_samples=5, batch=3 -> (15, dof).""" + sampler = self._make_sampler(num_samples=5, db_size=2) + target, _ = self._reachable_targets(3) + seed = torch.zeros(self.dof) + out = sampler.sample( + seed, self.lower, self.upper, batch_size=3, target_xpos=target + ) + assert out.shape == (15, self.dof) + grouped = out.view(3, 5, self.dof) + assert torch.allclose(grouped[:, 0], seed.expand(3, self.dof)) + # Padded slots must still be valid configurations. + assert (grouped[:, 1:] >= self.lower - 1e-6).all() + assert (grouped[:, 1:] <= self.upper + 1e-6).all() + + def test_num_samples_exceeds_k_max(self): + """Retrieval capped by k_max must still honour the size contract.""" + sampler = self._make_sampler(num_samples=6, k_max=2) + target, _ = self._reachable_targets(4) + seed = torch.zeros(self.dof) + out = sampler.sample( + seed, self.lower, self.upper, batch_size=4, target_xpos=target + ) + assert out.shape == (24, self.dof) + grouped = out.view(4, 6, self.dof) + assert torch.allclose(grouped[:, 0], seed.expand(4, self.dof)) + assert (grouped[:, 1:] >= self.lower - 1e-6).all() + assert (grouped[:, 1:] <= self.upper + 1e-6).all() + + def test_capped_retrieval_without_caller_seed(self): + """use_caller_seed=False must also honour the size contract.""" + sampler = self._make_sampler(num_samples=5, db_size=2, use_caller_seed=False) + target, _ = self._reachable_targets(3) + out = sampler.sample( + torch.zeros(self.dof), + self.lower, + self.upper, + batch_size=3, + target_xpos=target, + ) + assert out.shape == (15, self.dof) + grouped = out.view(3, 5, self.dof) + assert (grouped >= self.lower - 1e-6).all() + assert (grouped <= self.upper + 1e-6).all() + + # ------------------------------------------------------------------ rebuild + + def test_limits_change_triggers_rebuild(self): + sampler = self._make_sampler(num_samples=4) + target, _ = self._reachable_targets(3) + seed = torch.zeros(self.dof) + sampler.sample(seed, self.lower, self.upper, batch_size=3, target_xpos=target) + first_db = sampler._db_qpos + + # Same limits: no rebuild. + sampler.sample(seed, self.lower, self.upper, batch_size=3, target_xpos=target) + assert sampler._db_qpos is first_db + + # Narrowed limits: rebuild, and every seed obeys the new bounds. + narrow_lo = self.lower * 0.5 + narrow_hi = self.upper * 0.5 + out = sampler.sample( + seed, narrow_lo, narrow_hi, batch_size=3, target_xpos=target + ) + assert sampler._db_qpos is not first_db + retrieved = out.view(3, 4, self.dof)[:, 1:] + assert (retrieved >= narrow_lo - 1e-6).all() + assert (retrieved <= narrow_hi + 1e-6).all() + + # ------------------------------------------------------------------ validation + + def test_invalid_inputs_raise(self): + sampler = self._make_sampler(num_samples=4) + target, _ = self._reachable_targets(3) + with pytest.raises(ValueError): + sampler.sample( + torch.zeros(self.dof), + self.lower, + self.upper, + batch_size=3, + target_xpos=torch.eye(4), + ) + with pytest.raises(ValueError): + sampler.sample( + torch.zeros(5, self.dof + 1), + self.lower, + self.upper, + batch_size=3, + target_xpos=target, + ) + with pytest.raises(ValueError): + self._make_sampler(num_samples=4, db_size=0) + with pytest.raises(ValueError): + self._make_sampler(num_samples=4, rot_scale=0.0)